1package merkletree
  2
  3import (
  4	"fmt"
  5	"math/big"
  6)
  7
  8// NodeType defines the type of node in the MT.
  9type NodeType byte
 10
 11const (
 12	// NodeTypeMiddle indicates the type of middle Node that has children.
 13	NodeTypeMiddle NodeType = 0
 14	// NodeTypeLeaf indicates the type of a leaf Node that contains a key &
 15	// value.
 16	NodeTypeLeaf NodeType = 1
 17	// NodeTypeEmpty indicates the type of an empty Node.
 18	NodeTypeEmpty NodeType = 2
 19
 20	// DBEntryTypeRoot indicates the type of a DB entry that indicates the
 21	// current Root of a MerkleTree
 22	DBEntryTypeRoot NodeType = 3
 23)
 24
 25// Node is the struct that represents a node in the MT. The node should not be
 26// modified after creation because the cached key won't be updated.
 27type Node struct {
 28	// Type is the type of node in the tree.
 29	Type NodeType
 30	// ChildL is the left child of a middle node.
 31	ChildL *Hash
 32	// ChildR is the right child of a middle node.
 33	ChildR *Hash
 34	// Entry is the data stored in a leaf node.
 35	Entry [2]*Hash
 36	// key is a cache used to avoid recalculating key
 37	key *Hash
 38}
 39
 40// NewNodeLeaf creates a new leaf node.
 41func NewNodeLeaf(k, v *Hash) *Node {
 42	return &Node{Type: NodeTypeLeaf, Entry: [2]*Hash{k, v}}
 43}
 44
 45// NewNodeMiddle creates a new middle node.
 46func NewNodeMiddle(childL *Hash, childR *Hash) *Node {
 47	return &Node{Type: NodeTypeMiddle, ChildL: childL, ChildR: childR}
 48}
 49
 50// NewNodeEmpty creates a new empty node.
 51func NewNodeEmpty() *Node {
 52	return &Node{Type: NodeTypeEmpty}
 53}
 54
 55// NewNodeFromBytes creates a new node by parsing the input []byte.
 56func NewNodeFromBytes(b []byte) (*Node, error) {
 57	if len(b) < 1 {
 58		return nil, ErrNodeBytesBadSize
 59	}
 60	n := Node{Type: NodeType(b[0])}
 61	b = b[1:]
 62	switch n.Type {
 63	case NodeTypeMiddle:
 64		if len(b) != 2*ElemBytesLen {
 65			return nil, ErrNodeBytesBadSize
 66		}
 67		n.ChildL, n.ChildR = &Hash{}, &Hash{}
 68		copy(n.ChildL[:], b[:ElemBytesLen])
 69		copy(n.ChildR[:], b[ElemBytesLen:ElemBytesLen*2])
 70	case NodeTypeLeaf:
 71		if len(b) != 2*ElemBytesLen {
 72			return nil, ErrNodeBytesBadSize
 73		}
 74		n.Entry = [2]*Hash{{}, {}}
 75		copy(n.Entry[0][:], b[0:32])
 76		copy(n.Entry[1][:], b[32:64])
 77	case NodeTypeEmpty:
 78		break
 79	default:
 80		return nil, ErrInvalidNodeFound
 81	}
 82	return &n, nil
 83}
 84
 85// LeafKey computes the key of a leaf node given the hIndex and hValue of the
 86// entry of the leaf.
 87func LeafKey(k, v *Hash) (*Hash, error) {
 88	return HashElemsKey(big.NewInt(1), k.BigInt(), v.BigInt())
 89}
 90
 91// Key computes the key of the node by hashing the content in a specific way
 92// for each type of node.  This key is used as the hash of the merklee tree for
 93// each node.
 94func (n *Node) Key() (*Hash, error) {
 95	if n.key == nil { // Cache the key to avoid repeated hash computations.
 96		// NOTE: We are not using the type to calculate the hash!
 97		switch n.Type {
 98		case NodeTypeMiddle: // H(ChildL || ChildR)
 99			var err error
100			n.key, err = HashElems(n.ChildL.BigInt(), n.ChildR.BigInt())
101			if err != nil {
102				return nil, err
103			}
104		case NodeTypeLeaf:
105			var err error
106			n.key, err = LeafKey(n.Entry[0], n.Entry[1])
107			if err != nil {
108				return nil, err
109			}
110		case NodeTypeEmpty: // Zero
111			n.key = &HashZero
112		default:
113			n.key = &HashZero
114		}
115	}
116	return n.key, nil
117}
118
119// Value returns the value of the node.  This is the content that is stored in
120// the backend database.
121func (n *Node) Value() []byte {
122	switch n.Type {
123	case NodeTypeMiddle: // {Type || ChildL || ChildR}
124		return append([]byte{byte(n.Type)}, append(n.ChildL[:], n.ChildR[:]...)...)
125	case NodeTypeLeaf: // {Type || Data...}
126		return append([]byte{byte(n.Type)}, append(n.Entry[0][:], n.Entry[1][:]...)...)
127	case NodeTypeEmpty: // {}
128		return []byte{}
129	default:
130		return []byte{}
131	}
132}
133
134// String outputs a string representation of a node (different for each type).
135func (n *Node) String() string {
136	switch n.Type {
137	case NodeTypeMiddle: // {Type || ChildL || ChildR}
138		return fmt.Sprintf("Middle L:%s R:%s", n.ChildL, n.ChildR)
139	case NodeTypeLeaf: // {Type || Data...}
140		return fmt.Sprintf("Leaf I:%v D:%v", n.Entry[0], n.Entry[1])
141	case NodeTypeEmpty: // {}
142		return "Empty"
143	default:
144		return "Invalid Node"
145	}
146}