Add docs

deepsubtree.go

@@ -43,18 +43,22 @@ func NewDeepSubTree(db dbm.DB, cacheSize int, skipFastStorageUpgrade bool, versi
 	return &DeepSubTree{MutableTree: mutableTree, initialRootHash: nil, witnessData: nil, operationCounter: 0}
 }
 
+// Setter for witness data. Also, resets the operation counter back to 0.
 func (dst *DeepSubTree) SetWitnessData(witnessData []WitnessData) {
 	dst.witnessData = witnessData
 	dst.operationCounter = 0
 }
 
+// Returns the initial root hash if it is initialized and Deep Subtree root is nil.
+// Otherwise, returns the Deep Subtree working hash is considered the initial root hash.
 func (dst *DeepSubTree) GetInitialRootHash() ([]byte, error) {
 	if dst.root == nil && dst.initialRootHash != nil {
 		return dst.initialRootHash, nil
 	}
 	return dst.WorkingHash()
 }
 
+// Setter for initial root hash
 func (dst *DeepSubTree) SetInitialRootHash(initialRootHash []byte) {
 	dst.initialRootHash = initialRootHash
 }
@@ -174,7 +178,7 @@ func (dst *DeepSubTree) verifyOperation(operation Operation, key []byte, value [
 	return nil
 }
 
-// Set sets a key in the working tree with the given value.
+// Verifies the Set operation with witness data and perform the given write operation
 func (dst *DeepSubTree) Set(key []byte, value []byte) (updated bool, err error) {
 	err = dst.verifyOperation("write", key, value)
 	if err != nil {
@@ -183,7 +187,7 @@ func (dst *DeepSubTree) Set(key []byte, value []byte) (updated bool, err error)
 	return dst.set(key, value)
 }
 
-// Set sets a key in the working tree with the given value.
+// Sets a key in the working tree with the given value.
 func (dst *DeepSubTree) set(key []byte, value []byte) (updated bool, err error) {
 	if value == nil {
 		return updated, fmt.Errorf("attempt to store nil value at key '%s'", key)
@@ -285,8 +289,7 @@ func (dst *DeepSubTree) recursiveSet(node *Node, key []byte, value []byte) (
 	return newNode, updated, err
 }
 
-// Get returns the value of the specified key if it exists, or nil otherwise.
-// The returned value must not be modified, since it may point to data stored within IAVL.
+// Verifies the Get operation with witness data and perform the given read operation
 func (dst *DeepSubTree) Get(key []byte) (value []byte, err error) {
 	err = dst.verifyOperation("read", key, nil)
 	if err != nil {
@@ -305,8 +308,7 @@ func (dst *DeepSubTree) get(key []byte) ([]byte, error) {
 	return dst.ImmutableTree.Get(key)
 }
 
-// Remove tries to remove a key from the tree and if removed, returns its
-// value, and 'true'.
+// Verifies the Remove operation with witness data and perform the given delete operation
 func (dst *DeepSubTree) Remove(key []byte) (value []byte, removed bool, err error) {
 	err = dst.verifyOperation("delete", key, nil)
 	if err != nil {

deepsubtree_test.go

@@ -253,7 +253,8 @@ type testContext struct {
 	keys set.Set[string]
 }
 
-// Generates random new key half times and an existing key for the other half times.
+// Returns random new key half times if genRandom is true.
+// Otherwise, returns a randomly picked existing key.
 func (tc *testContext) getKey(genRandom bool, addsNewKey bool) (key []byte, err error) {
 	tree, r, keys := tc.tree, tc.r, tc.keys
 	if genRandom && readByte(r) < math.MaxUint8/2 {
@@ -276,6 +277,8 @@ func (tc *testContext) getKey(genRandom bool, addsNewKey bool) (key []byte, err
 	return []byte(kString), nil
 }
 
+// Performs the Set operation on full IAVL tree first, gets the witness data generated from
+// the operatio, and uses that witness data to peform the same operation on the Deep Subtree
 func (tc *testContext) setInDST(key []byte, value []byte) error {
 	if key == nil {
 		return nil
@@ -308,6 +311,8 @@ func (tc *testContext) setInDST(key []byte, value []byte) error {
 	return nil
 }
 
+// Performs the Remove operation on full IAVL tree first, gets the witness data generated from
+// the operatio, and uses that witness data to peform the same operation on the Deep Subtree
 func (tc *testContext) removeInDST(key []byte) error {
 	if key == nil {
 		return nil
@@ -345,6 +350,8 @@ func (tc *testContext) removeInDST(key []byte) error {
 	return nil
 }
 
+// Performs the Get operation on full IAVL tree first, gets the witness data generated from
+// the operation, and uses that witness data to peform the same operation on the Deep Subtree
 func (tc *testContext) getInDST(key []byte) error {
 	if key == nil {
 		return nil
@@ -373,6 +380,8 @@ func (tc *testContext) getInDST(key []byte) error {
 	return nil
 }
 
+// Fuzz tests different combinations of Get, Remove, Set operations generated in
+// a random order with keys related to operations chosen randomly
 func FuzzBatchAddReverse(f *testing.F) {
 	f.Fuzz(func(t *testing.T, input []byte) {
 		require := require.New(t)

mutable_tree.go

@@ -70,15 +70,19 @@ func NewMutableTreeWithOpts(db dbm.DB, cacheSize int, opts *Options, skipFastSto
 	}, nil
 }
 
+// Sets tracingEnabled to given boolean and also resets any existing witness data
 func (tree *MutableTree) SetTracingEnabled(tracingEnabled bool) {
 	tree.tracingEnabled = tracingEnabled
+	tree.ndb.setTracingEnabled(tracingEnabled)
 	tree.resetWitnessData()
 }
 
+// Resets witness data inside tree
 func (tree *MutableTree) resetWitnessData() {
 	tree.witnessData = make([]WitnessData, 0)
 }
 
+// Getter for witness data
 func (tree *MutableTree) GetWitnessData() []WitnessData {
 	return tree.witnessData
 }
@@ -156,10 +160,8 @@ func (tree *MutableTree) reapExistenceProofs(keysAccessed []string) ([]*ics23.Ex
 	return existenceProofs, nil
 }
 
-// Set sets a key in the working tree. Nil values are invalid. The given
-// key/value byte slices must not be modified after this call, since they point
-// to slices stored within IAVL. It returns true when an existing value was
-// updated, while false means it was a new key.
+// Wrapper around SetOp to add operation related data to the tree's witness data
+// when traicng is enabled
 func (tree *MutableTree) Set(key, value []byte) (updated bool, err error) {
 	if !tree.tracingEnabled {
 		return tree.SetOp(key, value)
@@ -210,8 +212,8 @@ func (tree *MutableTree) SetOp(key, value []byte) (updated bool, err error) {
 	return updated, nil
 }
 
-// Get returns the value of the specified key if it exists, or nil otherwise.
-// The returned value must not be modified, since it may point to data stored within IAVL.
+// Wrapper around GetOp to add operation related data to the tree's witness data
+// when traicng is enabled
 func (tree *MutableTree) Get(key []byte) ([]byte, error) {
 	if !tree.tracingEnabled {
 		return tree.GetOp(key)
@@ -411,8 +413,8 @@ func (tree *MutableTree) recursiveSet(node *Node, key []byte, value []byte, orph
 	}
 }
 
-// Remove removes a key from the working tree. The given key byte slice should not be modified
-// after this call, since it may point to data stored inside IAVL.
+// Wrapper around RemoveOp to add operation related data to the tree's witness data
+// when traicng is enabled
 func (tree *MutableTree) Remove(key []byte) ([]byte, bool, error) {
 	if !tree.tracingEnabled {
 		return tree.RemoveOp(key)

nodedb.go

@@ -80,7 +80,8 @@ type nodeDB struct {
 	latestVersion  int64            // Latest version of nodeDB.
 	nodeCache      cache.Cache      // Cache for nodes in the regular tree that consists of key-value pairs at any version.
 	fastNodeCache  cache.Cache      // Cache for nodes in the fast index that represents only key-value pairs at the latest version.
-	keysAccessed   set.Set[string]
+	keysAccessed   set.Set[string]  // Set of keys accessed so far, used when tracing is enabled
+	tracingEnabled bool
 }
 
 func newNodeDB(db dbm.DB, cacheSize int, opts *Options) *nodeDB {
@@ -105,17 +106,24 @@ func newNodeDB(db dbm.DB, cacheSize int, opts *Options) *nodeDB {
 		versionReaders: make(map[int64]uint32, 8),
 		storageVersion: string(storeVersion),
 		keysAccessed:   make(set.Set[string]),
+		tracingEnabled: false,
 	}
 }
 
-// Adds the given into a set of keys accessed
+// Adds the given into a set of keys accessed when tracing is enabled
 // Note: Used by Deep Subtrees to know which keys to add existence proofs for
 func (ndb *nodeDB) addTrace(key []byte) {
-	if ndb.keysAccessed != nil {
+	if ndb.tracingEnabled && ndb.keysAccessed != nil {
 		ndb.keysAccessed.Add(string(key))
 	}
 }
 
+// Sets tracingEnabled to given boolean and also resets keysAccessed
+func (ndb *nodeDB) setTracingEnabled(tracingEnabled bool) {
+	ndb.tracingEnabled = tracingEnabled
+	ndb.keysAccessed = make(set.Set[string])
+}
+
 // GetNode gets a node from memory or disk. If it is an inner node, it does not
 // load its children.
 func (ndb *nodeDB) GetNode(hash []byte) (*Node, error) {