[WIP] implement ClearAll for treeStore

Author: dylanlott

persistence/trees/trees.go

@@ -86,31 +86,32 @@ const (
 // of the underlying trees by utilizing the lazy loading
 // functionality provided by the underlying smt library.
 type treeStore struct {
-	merkleTrees map[merkleTree]*smt.SMT
-	nodeStores  map[merkleTree]kvstore.KVStore
+	treeStoreDir string
+	merkleTrees  map[merkleTree]*smt.SMT
+	nodeStores   map[merkleTree]kvstore.KVStore
 }
 
 // Update takes a transaction and a height and updates
 // all of the trees in the treeStore for that height.
 func (t *treeStore) Update(pgtx pgx.Tx, txi indexer.TxIndexer, height uint64) (string, error) {
-	// NB: Consider not wrapping this
 	return t.updateMerkleTrees(pgtx, txi, height)
 }
 
 // NewStateTrees takes a directory location and a TxIndexer  and returns a TreeStore
 // implementation that guarantees
-func NewStateTrees(treesStoreDir string, txi indexer.TxIndexer) (*treeStore, error) {
-	if treesStoreDir == "" {
+func NewStateTrees(dir string, txi indexer.TxIndexer) (*treeStore, error) {
+	if dir == "" {
 		return newMemStateTrees()
 	}
 
 	stateTrees := &treeStore{
-		merkleTrees: make(map[merkleTree]*smt.SMT, int(numMerkleTrees)),
-		nodeStores:  make(map[merkleTree]kvstore.KVStore, int(numMerkleTrees)),
+		treeStoreDir: dir,
+		merkleTrees:  make(map[merkleTree]*smt.SMT, int(numMerkleTrees)),
+		nodeStores:   make(map[merkleTree]kvstore.KVStore, int(numMerkleTrees)),
 	}
 
 	for tree := merkleTree(0); tree < numMerkleTrees; tree++ {
-		nodeStore, err := kvstore.NewKVStore(fmt.Sprintf("%s/%s_nodes", treesStoreDir, merkleTreeToString[tree]))
+		nodeStore, err := kvstore.NewKVStore(fmt.Sprintf("%s/%s_nodes", dir, merkleTreeToString[tree]))
 		if err != nil {
 			return nil, err
 		}
@@ -136,7 +137,6 @@ func newMemStateTrees() (*treeStore, error) {
 // updateMerkleTrees updates all of the merkle trees that TreeStore manages.
 // * it returns an hash of the output or an error.
 func (t *treeStore) updateMerkleTrees(pgtx pgx.Tx, txi indexer.TxIndexer, height uint64) (string, error) {
-	// TODO: loop through the trees, update each & commit; Trigger rollback if any fail to update.
 	for treeType := merkleTree(0); treeType < numMerkleTrees; treeType++ {
 		switch treeType {
 		// Actor Merkle Trees
@@ -369,7 +369,7 @@ func (t *treeStore) getActorsUpdated(
 
 	actors := make([]*coreTypes.Actor, len(addrs))
 	for i, addr := range addrs {
-		// TODO same goes here for int64 height cast here
+		// TECHDEBT #XXX: Avoid this cast to int64
 		actor, err := t.getActor(pgtx, actorSchema, addr, int64(height))
 		if err != nil {
 			return nil, err
@@ -466,7 +466,7 @@ func (t *treeStore) getParams(pgtx pgx.Tx, height uint64) ([]*coreTypes.Param, e
 	}
 	defer rows.Close()
 
-	var paramSlice []*coreTypes.Param // Store returned rows
+	var paramSlice []*coreTypes.Param
 	for rows.Next() {
 		param := new(coreTypes.Param)
 		if err := rows.Scan(&param.Name, &param.Value); err != nil {
@@ -475,6 +475,7 @@ func (t *treeStore) getParams(pgtx pgx.Tx, height uint64) ([]*coreTypes.Param, e
 		param.Height = int64(height)
 		paramSlice = append(paramSlice, param)
 	}
+
 	return paramSlice, nil
 }
 
@@ -539,20 +540,12 @@ func (t *treeStore) getChainsForActor(
 // be handled here in the TreeStore instead of the level above.
 func (t *treeStore) ClearAll() error {
 	for treeType := merkleTree(0); treeType < numMerkleTrees; treeType++ {
-		// nodeStore := t.nodeStores[treeType]
-		// tree := t.merkleTrees[treeType]
-		// TODO: implement the below
-		// 	if err := valueStore.ClearAll(); err != nil {
-		// 		return err
-		// 	}
-		// 	if err := tree.ClearAll(); err != nil {
-		// 		return err
-		// 	}
-
-		// p.stateTrees.merkleTrees[treeType] = smt.NewSparseMerkleTree(valueStore, nodeStore, sha256.New())
-		// TODO: like the line above, create a new set of state tree and assign it to the merkleTrees property
-		// // This is Needed in order to make sure the root is re-set correctly after clearing
-		// t.merkleTrees[treeType] = smt.NewSparseMerkleTree()
-	}
-	return fmt.Errorf("not impl")
+		nodeStore := t.nodeStores[treeType]
+		if err := nodeStore.ClearAll(); err != nil {
+			return fmt.Errorf("failed to clear %s node store: %w", merkleTreeToString[treeType], err)
+		}
+		t.merkleTrees[treeType] = smt.NewSparseMerkleTree(nodeStore, sha256.New())
+	}
+
+	return nil
 }