Merge pull request #2 from polymerdao/ian/fps_dev_v1.1.2

feat: Modify v1 IAVL to work with pure-preimage queries

benchmarks/bench_test.go

@@ -348,7 +348,6 @@ func runBenchmarks(b *testing.B, benchmarks []benchmark) {
 		)
 		if bb.dbType != "nodb" {
 			d, err = dbm.NewDB("test", bb.dbType, dirName)
-
 			if err != nil {
 				if strings.Contains(err.Error(), "unknown db_backend") {
 					// As an exception to run benchmarks: if the error is about cleveldb, or rocksdb,

import.go

@@ -68,15 +68,28 @@ func (i *Importer) writeNode(node *Node) error {
 	buf.Reset()
 	defer bufPool.Put(buf)
 
-	if err := node.writeBytes(buf); err != nil {
-		return err
+	if i.tree.useLegacyFormat {
+		if err := node.writeLegacyBytes(buf); err != nil {
+			return err
+		}
+	} else {
+		if err := node.writeBytes(buf); err != nil {
+			return err
+		}
 	}
 
 	bytesCopy := make([]byte, buf.Len())
 	copy(bytesCopy, buf.Bytes())
 
-	if err := i.batch.Set(i.tree.ndb.nodeKey(node.GetKey()), bytesCopy); err != nil {
-		return err
+	if i.tree.useLegacyFormat {
+		node.isLegacy = true
+		if err := i.batch.Set(i.tree.ndb.legacyNodeKey(node.GetKey()), bytesCopy); err != nil {
+			return err
+		}
+	} else {
+		if err := i.batch.Set(i.tree.ndb.nodeKey(node.GetKey()), bytesCopy); err != nil {
+			return err
+		}
 	}
 
 	i.batchSize++
@@ -193,19 +206,37 @@ func (i *Importer) Commit() error {
 		return ErrNoImport
 	}
 
+	var rootHash []byte
 	switch len(i.stack) {
 	case 0:
-		if err := i.batch.Set(i.tree.ndb.nodeKey(GetRootKey(i.version)), []byte{}); err != nil {
-			return err
+		if i.tree.useLegacyFormat {
+			rootHash = []byte{}
+			if err := i.batch.Set(i.tree.ndb.legacyNodeKey(GetRootKey(i.version)), []byte{}); err != nil {
+				return err
+			}
+		} else {
+			if err := i.batch.Set(i.tree.ndb.nodeKey(GetRootKey(i.version)), []byte{}); err != nil {
+				return err
+			}
 		}
 	case 1:
 		i.stack[0].nodeKey.nonce = 1
 		if err := i.writeNode(i.stack[0]); err != nil {
 			return err
 		}
 		if i.stack[0].nodeKey.version < i.version { // it means there is no update in the given version
-			if err := i.batch.Set(i.tree.ndb.nodeKey(GetRootKey(i.version)), i.tree.ndb.nodeKey(i.stack[0].nodeKey.GetKey())); err != nil {
-				return err
+			if i.tree.useLegacyFormat {
+				if len(i.stack[0].hash) == 0 {
+					i.stack[0]._hash(i.version)
+				}
+				rootHash = i.stack[0].hash
+				if err := i.batch.Set(i.tree.ndb.legacyNodeKey(GetRootKey(i.version)), i.tree.ndb.legacyNodeKey(rootHash)); err != nil {
+					return err
+				}
+			} else {
+				if err := i.batch.Set(i.tree.ndb.nodeKey(GetRootKey(i.version)), i.tree.ndb.nodeKey(i.stack[0].nodeKey.GetKey())); err != nil {
+					return err
+				}
 			}
 		}
 	default:
@@ -219,9 +250,16 @@ func (i *Importer) Commit() error {
 	}
 	i.tree.ndb.resetLatestVersion(i.version)
 
-	_, err = i.tree.LoadVersion(i.version)
-	if err != nil {
-		return err
+	if i.tree.useLegacyFormat {
+		_, err = i.tree.LoadVersionByRootHash(i.version, rootHash)
+		if err != nil {
+			return err
+		}
+	} else {
+		_, err = i.tree.LoadVersion(i.version)
+		if err != nil {
+			return err
+		}
 	}
 
 	i.Close()

mutable_tree.go

@@ -41,6 +41,8 @@ type MutableTree struct {
 	unsavedFastNodeRemovals  *sync.Map      // map[string]interface{} FastNodes that have not yet been removed from disk
 	ndb                      *nodeDB
 	skipFastStorageUpgrade   bool // If true, the tree will work like no fast storage and always not upgrade fast storage
+	useLegacyFormat          bool // If true, save nodes to the DB with the legacy format
+	rootHash                 []byte
 
 	mtx sync.Mutex
 }
@@ -66,6 +68,30 @@ func NewMutableTree(db dbm.DB, cacheSize int, skipFastStorageUpgrade bool, lg lo
 	}
 }
 
+func NewLegacyMutableTree(db dbm.DB, cacheSize int, skipFastStorageUpgrade, noStoreVersion bool,
+	rootHash []byte, lg log.Logger, options ...Option,
+) *MutableTree {
+	opts := DefaultOptions()
+	for _, opt := range options {
+		opt(&opts)
+	}
+
+	ndb := newLegacyNodeDB(db, cacheSize, opts, noStoreVersion, lg)
+	head := &ImmutableTree{ndb: ndb, skipFastStorageUpgrade: skipFastStorageUpgrade}
+
+	return &MutableTree{
+		logger:                   lg,
+		ImmutableTree:            head,
+		lastSaved:                head.clone(),
+		unsavedFastNodeAdditions: &sync.Map{},
+		unsavedFastNodeRemovals:  &sync.Map{},
+		ndb:                      ndb,
+		skipFastStorageUpgrade:   skipFastStorageUpgrade,
+		useLegacyFormat:          true,
+		rootHash:                 rootHash,
+	}
+}
+
 // IsEmpty returns whether or not the tree has any keys. Only trees that are
 // not empty can be saved.
 func (tree *MutableTree) IsEmpty() bool {
@@ -169,7 +195,15 @@ func (tree *MutableTree) Set(key, value []byte) (updated bool, err error) {
 // The returned value must not be modified, since it may point to data stored within IAVL.
 func (tree *MutableTree) Get(key []byte) ([]byte, error) {
 	if tree.root == nil {
-		return nil, nil
+		if tree.rootHash != nil {
+			root, err := tree.ndb.GetNode(tree.rootHash)
+			if err != nil {
+				return nil, err
+			}
+			tree.root = root
+		} else {
+			return nil, nil
+		}
 	}
 
 	if !tree.skipFastStorageUpgrade {
@@ -252,7 +286,7 @@ func (tree *MutableTree) set(key []byte, value []byte) (updated bool, err error)
 		if !tree.skipFastStorageUpgrade {
 			tree.addUnsavedAddition(key, fastnode.NewNode(key, value, tree.version+1))
 		}
-		tree.ImmutableTree.root = NewNode(key, value)
+		tree.ImmutableTree.root = NewNode(key, value, tree.useLegacyFormat)
 		return updated, nil
 	}
 
@@ -311,7 +345,7 @@ func (tree *MutableTree) recursiveSetLeaf(node *Node, key []byte, value []byte)
 			subtreeHeight: 1,
 			size:          2,
 			nodeKey:       nil,
-			leftNode:      NewNode(key, value),
+			leftNode:      NewNode(key, value, tree.useLegacyFormat),
 			rightNode:     node,
 		}, false, nil
 	case 1: // setKey > leafKey
@@ -321,10 +355,10 @@ func (tree *MutableTree) recursiveSetLeaf(node *Node, key []byte, value []byte)
 			size:          2,
 			nodeKey:       nil,
 			leftNode:      node,
-			rightNode:     NewNode(key, value),
+			rightNode:     NewNode(key, value, tree.useLegacyFormat),
 		}, false, nil
 	default:
-		return NewNode(key, value), true, nil
+		return NewNode(key, value, tree.useLegacyFormat), true, nil
 	}
 }
 
@@ -509,6 +543,41 @@ func (tree *MutableTree) LoadVersion(targetVersion int64) (int64, error) {
 	return latestVersion, nil
 }
 
+// LoadVersionByRootHash loads a tree using the provided version and roothash
+func (tree *MutableTree) LoadVersionByRootHash(version int64, rootHash []byte) (int64, error) {
+	if len(rootHash) == 0 {
+		return 0, errors.New("LoadVersionByRootHash must be provided a non-empty rootHash argument")
+	}
+
+	tree.mtx.Lock()
+	defer tree.mtx.Unlock()
+
+	t := &ImmutableTree{
+		ndb:                    tree.ndb,
+		version:                version,
+		skipFastStorageUpgrade: tree.skipFastStorageUpgrade,
+	}
+
+	var err error
+	t.root, err = tree.ndb.GetNode(rootHash)
+	if err != nil {
+		return 0, err
+	}
+
+	tree.ImmutableTree = t
+	tree.lastSaved = t.clone()
+	tree.rootHash = rootHash
+
+	if !tree.skipFastStorageUpgrade {
+		// Attempt to upgrade
+		if _, err := tree.enableFastStorageAndCommitIfNotEnabled(); err != nil {
+			return 0, err
+		}
+	}
+
+	return version, nil
+}
+
 // loadVersionForOverwriting attempts to load a tree at a previously committed
 // version, or the latest version below it. Any versions greater than targetVersion will be deleted.
 func (tree *MutableTree) LoadVersionForOverwriting(targetVersion int64) error {
@@ -737,17 +806,30 @@ func (tree *MutableTree) SaveVersion() ([]byte, int64, error) {
 	} else {
 		if tree.root.nodeKey != nil {
 			// it means there are no updated nodes
-			if err := tree.ndb.SaveRoot(version, tree.root.nodeKey); err != nil {
-				return nil, 0, err
-			}
-			// it means the reference node is a legacy node
-			if tree.root.isLegacy {
-				// it will update the legacy node to the new format
-				// which ensures the reference node is not a legacy node
-				tree.root.isLegacy = false
+			if tree.useLegacyFormat {
+				if len(tree.root.hash) == 0 {
+					tree.root._hash(version)
+				}
+				if err := tree.ndb.SaveLegacyRoot(version, tree.root.hash); err != nil {
+					return nil, 0, err
+				}
+				tree.root.isLegacy = true
 				if err := tree.ndb.SaveNode(tree.root); err != nil {
 					return nil, 0, fmt.Errorf("failed to save the reference legacy node: %w", err)
 				}
+			} else {
+				if err := tree.ndb.SaveRoot(version, tree.root.nodeKey); err != nil {
+					return nil, 0, err
+				}
+				// it means the reference node is a legacy node
+				if tree.root.isLegacy {
+					// it will update the legacy node to the new format
+					// which ensures the reference node is not a legacy node
+					tree.root.isLegacy = false
+					if err := tree.ndb.SaveNode(tree.root); err != nil {
+						return nil, 0, fmt.Errorf("failed to save the reference legacy node: %w", err)
+					}
+				}
 			}
 		} else {
 			if err := tree.saveNewNodes(version); err != nil {
@@ -771,7 +853,9 @@ func (tree *MutableTree) SaveVersion() ([]byte, int64, error) {
 		tree.unsavedFastNodeRemovals = &sync.Map{}
 	}
 
-	return tree.Hash(), version, nil
+	hash := tree.Hash()
+	tree.rootHash = hash
+	return hash, version, nil
 }
 
 func (tree *MutableTree) saveFastNodeVersion(latestVersion int64) error {
@@ -1009,9 +1093,11 @@ func (tree *MutableTree) saveNewNodes(version int64) error {
 	newNodes := make([]*Node, 0)
 	var recursiveAssignKey func(*Node) ([]byte, error)
 	recursiveAssignKey = func(node *Node) ([]byte, error) {
-		if node.nodeKey != nil {
+		node.isLegacy = tree.useLegacyFormat
+		if (!node.isLegacy && node.nodeKey != nil) || (node.isLegacy && node.hash != nil) {
 			return node.GetKey(), nil
 		}
+
 		nonce++
 		node.nodeKey = &NodeKey{
 			version: version,
@@ -1034,7 +1120,7 @@ func (tree *MutableTree) saveNewNodes(version int64) error {
 		node._hash(version)
 		newNodes = append(newNodes, node)
 
-		return node.nodeKey.GetKey(), nil
+		return node.GetKey(), nil
 	}
 
 	if _, err := recursiveAssignKey(tree.root); err != nil {

node.go

@@ -77,12 +77,13 @@ type Node struct {
 var _ cache.Node = (*Node)(nil)
 
 // NewNode returns a new node from a key, value and version.
-func NewNode(key []byte, value []byte) *Node {
+func NewNode(key []byte, value []byte, useLegacy bool) *Node {
 	return &Node{
 		key:           key,
 		value:         value,
 		subtreeHeight: 0,
 		size:          1,
+		isLegacy:      useLegacy,
 	}
 }
 
@@ -243,11 +244,9 @@ func MakeLegacyNode(hash, buf []byte) (*Node, error) {
 		nodeKey:       &NodeKey{version: ver},
 		key:           key,
 		hash:          hash,
-		isLegacy:      true,
 	}
 
 	// Read node body.
-
 	if node.isLeaf() {
 		val, _, err := encoding.DecodeBytes(buf)
 		if err != nil {
@@ -612,6 +611,9 @@ func (node *Node) writeBytes(w io.Writer) error {
 		if node.leftNodeKey == nil {
 			return ErrLeftNodeKeyEmpty
 		}
+		if node.rightNodeKey == nil {
+			return ErrRightNodeKeyEmpty
+		}
 		// check if children NodeKeys are legacy mode
 		if len(node.leftNodeKey) == hashSize {
 			mode += ModeLegacyLeftNode
@@ -662,6 +664,54 @@ func (node *Node) writeBytes(w io.Writer) error {
 	return nil
 }
 
+func (node *Node) writeLegacyBytes(w io.Writer) error {
+	if node == nil {
+		return errors.New("cannot write nil node")
+	}
+	err := encoding.EncodeVarint(w, int64(node.subtreeHeight))
+	if err != nil {
+		return fmt.Errorf("writing height, %w", err)
+	}
+	err = encoding.EncodeVarint(w, node.size)
+	if err != nil {
+		return fmt.Errorf("writing size, %w", err)
+	}
+	err = encoding.EncodeVarint(w, node.nodeKey.version)
+	if err != nil {
+		return fmt.Errorf("writing version, %w", err)
+	}
+
+	// Unlike writeHashBytes, key is written for inner nodes.
+	err = encoding.EncodeBytes(w, node.key)
+	if err != nil {
+		return fmt.Errorf("writing key, %w", err)
+	}
+
+	if node.isLeaf() {
+		err = encoding.EncodeBytes(w, node.value)
+		if err != nil {
+			return fmt.Errorf("writing value, %w", err)
+		}
+	} else {
+		if len(node.leftNodeKey) != hashSize {
+			return errors.New("node provided to writeLegacyBytes does not have a hash for leftNodeKey")
+		}
+		err = encoding.EncodeBytes(w, node.leftNodeKey)
+		if err != nil {
+			return fmt.Errorf("writing left hash, %w", err)
+		}
+
+		if len(node.leftNodeKey) != 32 {
+			return errors.New("node provided to writeLegacyBytes does not have a hash for rightNodeKey")
+		}
+		err = encoding.EncodeBytes(w, node.rightNodeKey)
+		if err != nil {
+			return fmt.Errorf("writing right hash, %w", err)
+		}
+	}
+	return nil
+}
+
 func (node *Node) getLeftNode(t *ImmutableTree) (*Node, error) {
 	if node.leftNode != nil {
 		return node.leftNode, nil