feat(trie): implement new trie architecture with separated node types

core/trie2/collector.go

@@ -0,0 +1,133 @@
+package trie2
+
+import (
+	"fmt"
+	"sync"
+
+	"github.com/NethermindEth/juno/core/trie2/trienode"
+)
+
+const (
+	parallelThreshold = 8 // TODO: arbitrary number, configure this based on monitoring
+)
+
+// 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, 0).(*HashNode)
+}
+
+// Collects all dirty nodes in the trie and converts them to hash nodes.
+// Traverses the trie recursively, processing each node type:
+// - EdgeNode: Collects its child before processing the edge node itself
+// - BinaryNode: Collects both children (potentially in parallel) before processing the node
+// - HashNode: Already processed, returns as is
+// - ValueNode: Stores as a leaf in the node set
+// Returns a HashNode representing the processed node.
+func (c *collector) collect(path *Path, n Node, parallel bool, depth int) 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, depth+1)
+		return c.store(path, collapsed)
+	case *BinaryNode:
+		collapsed := cn.copy()
+		collapsed.Children = c.collectChildren(path, cn, parallel, depth+1)
+		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, depth int) [2]Node {
+	children := [2]Node{}
+
+	var mu sync.Mutex
+
+	// 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, depth)
+			return
+		}
+
+		// Parallel processing
+		childSet := trienode.NewNodeSet(c.nodes.Owner)
+		childCollector := newCollector(childSet)
+		children[i] = childCollector.collect(childPath, child, depth < parallelThreshold, depth)
+
+		// Merge the child set into the parent set
+		// Must be done under the mutex because node set is not thread safe
+		mu.Lock()
+		c.nodes.MergeSet(childSet) //nolint:errcheck // guaranteed to succeed because same owner
+		mu.Unlock()
+	}
+
+	if !parallel {
+		// Sequential processing
+		processChild(0)
+		processChild(1)
+		return children
+	}
+
+	// Parallel processing
+	var wg sync.WaitGroup
+	for i := range 2 {
+		wg.Add(1)
+		go func(idx int) {
+			defer wg.Done()
+			processChild(idx)
+		}(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.NewLeaf(blob))
+		return n
+	}
+
+	c.nodes.Add(*path, trienode.NewNonLeaf(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 hashing 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,64 @@
+package trie2
+
+import (
+	"github.com/NethermindEth/juno/core/felt"
+	"github.com/NethermindEth/juno/db"
+)
+
+var (
+	_ TrieID = (*ClassTrieID)(nil)
+	_ TrieID = (*ContractTrieID)(nil)
+	_ TrieID = (*ContractStorageTrieID)(nil)
+	_ TrieID = (*EmptyTrieID)(nil)
+)
+
+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
+}
+
+// Identifier for a class trie
+type ClassTrieID struct{}
+
+func NewClassTrieID() *ClassTrieID    { return &ClassTrieID{} }
+func (ClassTrieID) Bucket() db.Bucket { return db.ClassTrie }
+func (ClassTrieID) Owner() felt.Felt  { return felt.Zero }
+
+// Identifier for a contract trie
+type ContractTrieID struct{}
+
+func NewContractTrieID() *ContractTrieID    { return &ContractTrieID{} }
+func (id ContractTrieID) Bucket() db.Bucket { return db.ContractTrieContract }
+func (id ContractTrieID) Owner() felt.Felt  { return felt.Zero }
+
+// Identifier for a contract storage trie
+type ContractStorageTrieID struct {
+	owner felt.Felt
+}
+
+func NewContractStorageTrieID(owner felt.Felt) *ContractStorageTrieID {
+	return &ContractStorageTrieID{owner: owner}
+}
+
+func NewContractStorageTrieIDFromFelt(owner felt.Felt) *ContractStorageTrieID {
+	return &ContractStorageTrieID{owner: owner}
+}
+
+func (id ContractStorageTrieID) Bucket() db.Bucket { return db.ContractTrieStorage }
+func (id ContractStorageTrieID) Owner() felt.Felt  { return id.owner }
+
+// Identifier for an empty trie, only used for temporary purposes
+type EmptyTrieID struct{}
+
+func NewEmptyTrieID() *EmptyTrieID    { return &EmptyTrieID{} }
+func (EmptyTrieID) Bucket() db.Bucket { return db.Bucket(0) }
+func (EmptyTrieID) Owner() felt.Felt  { return felt.Zero }