Refactor Trie

core/trie2/collector.go

@@ -0,0 +1,120 @@
+package trie2
+
+import (
+	"fmt"
+	"sync"
+
+	"github.com/NethermindEth/juno/core/felt"
+	"github.com/NethermindEth/juno/core/trie2/trienode"
+)
+
+// Used as a tool to collect all dirty nodes in a NodeSet
+type collector struct {
+	nodes *trienode.NodeSet
+}
+
+func newCollector(nodes *trienode.NodeSet) *collector {
+	return &collector{nodes: nodes}
+}
+
+// Collects the nodes in the node set and collapses a node into a hash node
+func (c *collector) Collect(n node, parallel bool) *hashNode {
+	return c.collect(new(Path), n, parallel).(*hashNode)
+}
+
+func (c *collector) collect(path *Path, n node, parallel bool) node {
+	// This path has not been modified, just return the cache
+	hash, dirty := n.cache()
+	if hash != nil && !dirty {
+		return hash
+	}
+
+	// Collect children and then parent
+	switch cn := n.(type) {
+	case *edgeNode:
+		collapsed := cn.copy()
+
+		// If the child is a binary node, recurse into it.
+		// Otherwise, it can only be a hashNode or valueNode.
+		// Combination of edge (parent) + edge (child) is not possible.
+		collapsed.child = c.collect(new(Path).Append(path, cn.path), cn.child, parallel)
+		return c.store(path, collapsed)
+	case *binaryNode:
+		collapsed := cn.copy()
+		collapsed.children = c.collectChildren(path, cn, parallel)
+		return c.store(path, collapsed)
+	case *hashNode:
+		return cn
+	case *valueNode: // each leaf node is stored as a single entry in the node set
+		return c.store(path, cn)
+	default:
+		panic(fmt.Sprintf("unknown node type: %T", cn))
+	}
+}
+
+// Collects the children of a binary node, may apply parallel processing if configured
+func (c *collector) collectChildren(path *Path, n *binaryNode, parallel bool) [2]node {
+	children := [2]node{}
+
+	// Helper function to process a single child
+	processChild := func(i int) {
+		child := n.children[i]
+		// Return early if it's already a hash node
+		if hn, ok := child.(*hashNode); ok {
+			children[i] = hn
+			return
+		}
+
+		// Create child path
+		childPath := new(Path).AppendBit(path, uint8(i))
+
+		if !parallel {
+			children[i] = c.collect(childPath, child, parallel)
+			return
+		}
+
+		// Parallel processing
+		childSet := trienode.NewNodeSet(c.nodes.Owner)
+		childCollector := newCollector(childSet)
+		children[i] = childCollector.collect(childPath, child, parallel)
+		c.nodes.MergeSet(childSet) //nolint:errcheck // guaranteed to succeed because same owner
+	}
+
+	if !parallel {
+		// Sequential processing
+		processChild(0)
+		processChild(1)
+		return children
+	}
+
+	// Parallel processing
+	var wg sync.WaitGroup
+	var mu sync.Mutex
+
+	for i := range 2 {
+		wg.Add(1)
+		go func(idx int) {
+			defer wg.Done()
+			mu.Lock()
+			processChild(idx)
+			mu.Unlock()
+		}(i)
+	}
+	wg.Wait()
+
+	return children
+}
+
+// Stores the node in the node set and returns the hash node
+func (c *collector) store(path *Path, n node) node {
+	hash, _ := n.cache()
+
+	blob := nodeToBytes(n)
+	if hash == nil { // this is a value node
+		c.nodes.Add(*path, trienode.NewNode(felt.Felt{}, blob))
+		return n
+	}
+
+	c.nodes.Add(*path, trienode.NewNode(hash.Felt, blob))
+	return hash
+}

core/trie2/errors.go

@@ -0,0 +1,8 @@
+package trie2
+
+import "errors"
+
+var (
+	ErrCommitted  = errors.New("trie is committed")
+	ErrEmptyRange = errors.New("empty range")
+)

core/trie2/hasher.go

@@ -0,0 +1,117 @@
+package trie2
+
+import (
+	"fmt"
+	"sync"
+
+	"github.com/NethermindEth/juno/core/crypto"
+)
+
+// A tool for shashing nodes in the trie. It supports both sequential and parallel
+// hashing modes.
+type hasher struct {
+	hashFn   crypto.HashFn // The hash function to use
+	parallel bool          // Whether to hash binary node children in parallel
+}
+
+func newHasher(hash crypto.HashFn, parallel bool) hasher {
+	return hasher{
+		hashFn:   hash,
+		parallel: parallel,
+	}
+}
+
+// Computes the hash of a node and returns both the hash node and a cached
+// version of the original node. If the node already has a cached hash, returns
+// that instead of recomputing.
+func (h *hasher) hash(n node) (node, node) {
+	if hash, _ := n.cache(); hash != nil {
+		return hash, n
+	}
+
+	switch n := n.(type) {
+	case *edgeNode:
+		collapsed, cached := h.hashEdgeChild(n)
+		hn := &hashNode{Felt: *collapsed.hash(h.hashFn)}
+		cached.flags.hash = hn
+		return hn, cached
+	case *binaryNode:
+		collapsed, cached := h.hashBinaryChildren(n)
+		hn := &hashNode{Felt: *collapsed.hash(h.hashFn)}
+		cached.flags.hash = hn
+		return hn, cached
+	case *valueNode, *hashNode:
+		return n, n
+	default:
+		panic(fmt.Sprintf("unknown node type: %T", n))
+	}
+}
+
+func (h *hasher) hashEdgeChild(n *edgeNode) (collapsed, cached *edgeNode) {
+	collapsed, cached = n.copy(), n.copy()
+
+	switch n.child.(type) {
+	case *edgeNode, *binaryNode:
+		collapsed.child, cached.child = h.hash(n.child)
+	}
+
+	return collapsed, cached
+}
+
+func (h *hasher) hashBinaryChildren(n *binaryNode) (collapsed, cached *binaryNode) {
+	collapsed, cached = n.copy(), n.copy()
+
+	if h.parallel { // TODO(weiihann): double check this parallel strategy
+		var wg sync.WaitGroup
+		wg.Add(2)
+
+		go func() {
+			defer wg.Done()
+			if n.children[0] != nil {
+				collapsed.children[0], cached.children[0] = h.hash(n.children[0])
+			} else {
+				collapsed.children[0], cached.children[0] = nilValueNode, nilValueNode
+			}
+		}()
+
+		go func() {
+			defer wg.Done()
+			if n.children[1] != nil {
+				collapsed.children[1], cached.children[1] = h.hash(n.children[1])
+			} else {
+				collapsed.children[1], cached.children[1] = nilValueNode, nilValueNode
+			}
+		}()
+
+		wg.Wait()
+	} else {
+		if n.children[0] != nil {
+			collapsed.children[0], cached.children[0] = h.hash(n.children[0])
+		} else {
+			collapsed.children[0], cached.children[0] = nilValueNode, nilValueNode
+		}
+
+		if n.children[1] != nil {
+			collapsed.children[1], cached.children[1] = h.hash(n.children[1])
+		} else {
+			collapsed.children[1], cached.children[1] = nilValueNode, nilValueNode
+		}
+	}
+
+	return collapsed, cached
+}
+
+// Construct trie proofs and returns the collapsed node (i.e. nodes with hash children)
+// and the hashed node.
+func (h *hasher) proofHash(original node) (collapsed, hashed node) {
+	switch n := original.(type) {
+	case *edgeNode:
+		en, _ := h.hashEdgeChild(n)
+		return en, &hashNode{Felt: *en.hash(h.hashFn)}
+	case *binaryNode:
+		bn, _ := h.hashBinaryChildren(n)
+		return bn, &hashNode{Felt: *bn.hash(h.hashFn)}
+	default:
+		return n, n
+	}
+}

core/trie2/id.go

@@ -0,0 +1,61 @@
+package trie2
+
+import (
+	"github.com/NethermindEth/juno/core/felt"
+	"github.com/NethermindEth/juno/db"
+)
+
+type TrieType uint8
+
+const (
+	Empty TrieType = iota
+	ClassTrie
+	ContractTrie
+)
+
+// Represents the identifier for uniquely identifying a trie.
+type ID struct {
+	TrieType TrieType
+	Owner    felt.Felt // The contract address which the trie belongs to
+}
+
+// Returns the corresponding DB bucket for the trie
+func (id *ID) Bucket() db.Bucket {
+	switch id.TrieType {
+	case ClassTrie:
+		return db.ClassTrie
+	case ContractTrie:
+		if id.Owner == (felt.Felt{}) {
+			return db.ContractTrieContract
+		}
+		return db.ContractTrieStorage
+	case Empty:
+		return db.Bucket(0)
+	default:
+		panic("invalid trie type")
+	}
+}
+
+// Constructs an identifier for a class trie with the provided class trie root hash
+func ClassTrieID() *ID {
+	return &ID{
+		TrieType: ClassTrie,
+		Owner:    felt.Zero, // class trie does not have an owner
+	}
+}
+
+// Constructs an identifier for a contract trie or a contract's storage trie
+func ContractTrieID(owner felt.Felt) *ID {
+	return &ID{
+		TrieType: ContractTrie,
+		Owner:    owner,
+	}
+}
+
+// A general identifier, typically used for temporary trie
+func TrieID() *ID {
+	return &ID{
+		TrieType: Empty,
+		Owner:    felt.Zero,
+	}
+}