refactor(trie2): use more descriptive felt.Address type instead of felt.Felt (#3295)

* Commits from maksym/rawdb (squashed)

* address comments

* Commits from maksym/rawdb (squashed)

* Refactor the state to use more descriptive `felt.Address` type instead of `felt.Felt`

* unit tests, linter

* linter

* Refactor the state to use more descriptive `felt.Address` type instead of `felt.Felt`

* unit tests, linter

* comments

* address comments

* address comments

* linter

* address comments

* address comments

* linter

* address comments

Author: MaksymMalicki

core/felt/utils.go

@@ -0,0 +1,12 @@
+package felt
+
+func IsZero[F FeltLike](v *F) bool {
+	f := Felt(*v)
+	return f.IsZero()
+}
+
+func Equal[F FeltLike](a, b *F) bool {
+	fa := Felt(*a)
+	fb := Felt(*b)
+	return fa.Equal(&fb)
+}

core/state/database.go

@@ -38,5 +38,13 @@ func (s *StateDB) ContractTrie(stateComm *felt.Felt) (*trie2.Trie, error) {
 
 // Opens a contract storage trie for the given state root and contract address
 func (s *StateDB) ContractStorageTrie(stateComm, owner *felt.Felt) (*trie2.Trie, error) {
-	return trie2.New(trieutils.NewContractStorageTrieID(*stateComm, *owner), ContractStorageTrieHeight, crypto.Pedersen, s.triedb)
+	return trie2.New(
+		trieutils.NewContractStorageTrieID(
+			*stateComm,
+			felt.Address(*owner),
+		),
+		ContractStorageTrieHeight,
+		crypto.Pedersen,
+		s.triedb,
+	)
 }

core/state/object.go

@@ -51,7 +51,13 @@ func (s *stateObject) getStorage(key *felt.Felt) (felt.Felt, error) {
 	}
 
 	path := tr.FeltToPath(key)
-	v, err := trieutils.GetNodeByPath(s.state.db.disk, db.ContractTrieStorage, &s.addr, &path, true)
+	v, err := trieutils.GetNodeByPath(
+		s.state.db.disk,
+		db.ContractTrieStorage,
+		(*felt.Address)(&s.addr),
+		&path,
+		true,
+	)
 	if err != nil {
 		return felt.Zero, err
 	}

core/state/state.go

@@ -516,7 +516,8 @@ func (s *State) flush(
 				}
 
 				// TODO(weiihann): handle hash-based, and there should be better ways of doing this
-				if err := trieutils.DeleteStorageNodesByPath(batch, addr); err != nil {
+				err := trieutils.DeleteStorageNodesByPath(batch, (*felt.Address)(&addr))
+				if err != nil {
 					return err
 				}
 			} else { // updated

core/trie2/databasetest.go

@@ -26,7 +26,12 @@ func newTestNodeReader(id trieutils.TrieID, nodes []*trienode.MergeNodeSet, db d
 	return &testNodeReader{id: id, nodes: nodes, db: db, scheme: scheme}
 }
 
-func (n *testNodeReader) Node(owner *felt.Felt, path *trieutils.Path, hash *felt.Felt, isLeaf bool) ([]byte, error) {
+func (n *testNodeReader) Node(
+	owner *felt.Address,
+	path *trieutils.Path,
+	hash *felt.Felt,
+	isLeaf bool,
+) ([]byte, error) {
 	for _, nodes := range n.nodes {
 		var (
 			node trienode.TrieNode

core/trie2/errors.go

@@ -15,14 +15,14 @@ var (
 
 type MissingNodeError struct {
 	tt    trieutils.TrieType
-	owner felt.Felt
+	owner felt.Address
 	path  trieutils.Path
 	hash  felt.Felt
 	err   error
 }
 
 func (e *MissingNodeError) Error() string {
-	if e.owner.Equal(&felt.Zero) {
+	if felt.IsZero(&e.owner) {
 		return fmt.Sprintf("%s: missing trie node (path %v, hash %v) %v", e.tt, e.path, e.hash, e.err)
 	}
 	return fmt.Sprintf("%s: missing trie node (owner %v, path %v, hash %v) %v", e.tt, e.owner, e.path, e.hash, e.err)

core/trie2/trie.go

@@ -22,7 +22,7 @@ type Trie struct {
 	height uint8
 
 	// The owner of the trie, only used for contract trie. If not empty, this is a storage trie.
-	owner felt.Felt
+	owner felt.Address
 
 	// The root node of the trie
 	root trienode.Node

core/trie2/triedb/database/db.go

@@ -19,7 +19,7 @@ const (
 
 // Represents a reader for trie nodes
 type NodeReader interface {
-	Node(owner *felt.Felt, path *trieutils.Path, hash *felt.Felt, isLeaf bool) ([]byte, error)
+	Node(owner *felt.Address, path *trieutils.Path, hash *felt.Felt, isLeaf bool) ([]byte, error)
 }
 
 // Represents a database that produces a node reader for a given trie id

core/trie2/triedb/empty.go

@@ -23,6 +23,11 @@ func (EmptyNodeDatabase) NodeReader(id trieutils.TrieID) (database.NodeReader, e
 
 type EmptyNodeReader struct{}
 
-func (EmptyNodeReader) Node(owner *felt.Felt, path *trieutils.Path, hash *felt.Felt, isLeaf bool) ([]byte, error) {
+func (EmptyNodeReader) Node(
+	owner *felt.Address,
+	path *trieutils.Path,
+	hash *felt.Felt,
+	isLeaf bool,
+) ([]byte, error) {
 	return nil, nil
 }

core/trie2/triedb/hashdb/database.go

@@ -48,15 +48,27 @@ func New(disk db.KeyValueStore, config *Config) *Database {
 	}
 }
 
-func (d *Database) insert(owner *felt.Felt, path *trieutils.Path, hash *felt.Felt, isClass bool, node trienode.TrieNode) {
+func (d *Database) insert(
+	owner *felt.Address,
+	path *trieutils.Path,
+	hash *felt.Felt,
+	isClass bool,
+	node trienode.TrieNode,
+) {
 	_, found := d.dirtyCache.getNode(owner, path, hash, isClass)
 	if found {
 		return
 	}
 	d.dirtyCache.putNode(owner, path, hash, isClass, node)
 }
 
-func (d *Database) readNode(bucket db.Bucket, owner *felt.Felt, path *trieutils.Path, hash *felt.Felt, isLeaf bool) ([]byte, error) {
+func (d *Database) readNode(
+	bucket db.Bucket,
+	owner *felt.Address,
+	path *trieutils.Path,
+	hash *felt.Felt,
+	isLeaf bool,
+) ([]byte, error) {
 	if blob := d.cleanCache.getNode(path, hash); blob != nil {
 		return blob, nil
 	}
@@ -82,7 +94,7 @@ func (d *Database) readNode(bucket db.Bucket, owner *felt.Felt, path *trieutils.
 func (d *Database) NewIterator(id trieutils.TrieID) (db.Iterator, error) {
 	key := id.Bucket().Key()
 	owner := id.Owner()
-	if !owner.Equal(&felt.Zero) {
+	if !felt.IsZero(&owner) {
 		oBytes := owner.Bytes()
 		key = append(key, oBytes[:]...)
 	}
@@ -101,7 +113,16 @@ func (d *Database) Commit(_ *felt.Felt) error {
 		if err != nil {
 			return err
 		}
-		if err := trieutils.WriteNodeByHash(batch, db.ClassTrie, &felt.Zero, &path, &hash, node.IsLeaf(), node.Blob()); err != nil {
+		err = trieutils.WriteNodeByHash(
+			batch,
+			db.ClassTrie,
+			&felt.Address{},
+			&path,
+			&hash,
+			node.IsLeaf(),
+			node.Blob(),
+		)
+		if err != nil {
 			return err
 		}
 		d.cleanCache.putNode(&path, &hash, node.Blob())
@@ -112,7 +133,16 @@ func (d *Database) Commit(_ *felt.Felt) error {
 		if err != nil {
 			return err
 		}
-		if err := trieutils.WriteNodeByHash(batch, db.ContractTrieContract, &felt.Zero, &path, &hash, node.IsLeaf(), node.Blob()); err != nil {
+		err = trieutils.WriteNodeByHash(
+			batch,
+			db.ContractTrieContract,
+			&felt.Address{},
+			&path,
+			&hash,
+			node.IsLeaf(),
+			node.Blob(),
+		)
+		if err != nil {
 			return err
 		}
 		d.cleanCache.putNode(&path, &hash, node.Blob())
@@ -124,7 +154,16 @@ func (d *Database) Commit(_ *felt.Felt) error {
 			if err != nil {
 				return err
 			}
-			if err := trieutils.WriteNodeByHash(batch, db.ContractTrieStorage, &owner, &path, &hash, node.IsLeaf(), node.Blob()); err != nil {
+			err = trieutils.WriteNodeByHash(
+				batch,
+				db.ContractTrieStorage,
+				&owner,
+				&path,
+				&hash,
+				node.IsLeaf(),
+				node.Blob(),
+			)
+			if err != nil {
 				return err
 			}
 			d.cleanCache.putNode(&path, &hash, node.Blob())
@@ -158,7 +197,7 @@ func (d *Database) Update(
 
 	var classNodes map[trieutils.Path]trienode.TrieNode
 	var contractNodes map[trieutils.Path]trienode.TrieNode
-	var contractStorageNodes map[felt.Felt]map[trieutils.Path]trienode.TrieNode
+	var contractStorageNodes map[felt.Address]map[trieutils.Path]trienode.TrieNode
 
 	if mergedClassNodes != nil {
 		classNodes, _ = mergedClassNodes.Flatten()
@@ -170,15 +209,15 @@ func (d *Database) Update(
 		contractNodes, contractStorageNodes = mergedContractNodes.Flatten()
 	} else {
 		contractNodes = make(map[trieutils.Path]trienode.TrieNode)
-		contractStorageNodes = make(map[felt.Felt]map[trieutils.Path]trienode.TrieNode)
+		contractStorageNodes = make(map[felt.Address]map[trieutils.Path]trienode.TrieNode)
 	}
 
 	for path, node := range classNodes {
 		if _, ok := node.(*trienode.DeletedNode); ok {
 			continue // Since the hashdb is used for archive node only, there is no need to remove nodes
 		} else {
 			nodeHash := node.Hash()
-			d.insert(&felt.Zero, &path, &nodeHash, true, node)
+			d.insert(&felt.Address{}, &path, &nodeHash, true, node)
 		}
 	}
 
@@ -187,7 +226,7 @@ func (d *Database) Update(
 			continue
 		} else {
 			nodeHash := node.Hash()
-			d.insert(&felt.Zero, &path, &nodeHash, false, node)
+			d.insert(&felt.Address{}, &path, &nodeHash, false, node)
 		}
 	}
 
@@ -209,7 +248,12 @@ type reader struct {
 	d  *Database
 }
 
-func (r *reader) Node(owner *felt.Felt, path *trieutils.Path, hash *felt.Felt, isLeaf bool) ([]byte, error) {
+func (r *reader) Node(
+	owner *felt.Address,
+	path *trieutils.Path,
+	hash *felt.Felt,
+	isLeaf bool,
+) ([]byte, error) {
 	return r.d.readNode(r.id.Bucket(), owner, path, hash, isLeaf)
 }
 
@@ -228,15 +272,29 @@ func (d *Database) Close() error {
 func (d *Database) GetTrieRootNodes(classRootHash, contractRootHash *felt.Felt) (trienode.Node, trienode.Node, error) {
 	const contractClassTrieHeight = 251
 
-	classRootBlob, err := trieutils.GetNodeByHash(d.disk, db.ClassTrie, &felt.Zero, &trieutils.Path{}, classRootHash, false)
+	classRootBlob, err := trieutils.GetNodeByHash(
+		d.disk,
+		db.ClassTrie,
+		&felt.Address{},
+		&trieutils.Path{},
+		classRootHash,
+		false,
+	)
 	if err != nil {
 		return nil, nil, fmt.Errorf("class root node not found: %w", err)
 	}
 	if classRootBlob == nil {
 		return nil, nil, fmt.Errorf("class root node not found")
 	}
 
-	contractRootBlob, err := trieutils.GetNodeByHash(d.disk, db.ContractTrieContract, &felt.Zero, &trieutils.Path{}, contractRootHash, false)
+	contractRootBlob, err := trieutils.GetNodeByHash(
+		d.disk,
+		db.ContractTrieContract,
+		&felt.Address{},
+		&trieutils.Path{},
+		contractRootHash,
+		false,
+	)
 	if err != nil {
 		return nil, nil, fmt.Errorf("contract root node not found: %w", err)
 	}