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