fix get

Author: mmsqe

iavl/cli.py

@@ -13,7 +13,8 @@
                     decode_fast_node, diff_iterators, encode_stdint,
                     fast_node_key, get_node, get_root_node,
                     iavl_latest_version, iter_fast_nodes, iter_iavl_tree,
-                    load_commit_infos, root_key, store_prefix)
+                    legacy_root_key, load_commit_infos, node_key, root_key,
+                    store_prefix)
 
 
 @click.group
@@ -43,7 +44,9 @@ def root_hash(db, store: List[str], version: Optional[int]):
     for s in store:
         if version is None:
             version = iavl_latest_version(db, s)
-        bz = db.get(store_prefix(s) + root_key(version))
+        bz = db.get(store_prefix(s) + node_key(root_key(version)))
+        if not bz:
+            bz = db.get(store_prefix(s) + legacy_root_key(version))
         print(f"{s}: {binascii.hexlify(bz or b'').decode()}")
 
 

iavl/iavl.py

@@ -10,7 +10,7 @@
 import rocksdb
 
 from .utils import (ROOT_KEY_PREFIX, GetNode, PersistedNode, encode_bytes,
-                    node_key, root_key, visit_iavl_nodes)
+                    legacy_node_key, root_key, visit_iavl_nodes)
 
 NodeRef = Union[bytes, "Node"]
 
@@ -35,7 +35,7 @@ def get(self, hash: bytes) -> Optional[PersistedNode]:
         try:
             return self.cache[hash]
         except KeyError:
-            bz = self.db.get(self.prefix + node_key(hash))
+            bz = self.db.get(self.prefix + legacy_node_key(hash))
             if bz is None:
                 return
             node = PersistedNode.decode(bz, hash)
@@ -52,7 +52,7 @@ def batch_remove_node(self, hash: bytes):
         "remove node"
         if self.batch is None:
             self.batch = rocksdb.WriteBatch()
-        self.batch.delete(node_key(hash))
+        self.batch.delete(legacy_node_key(hash))
         self.cache.pop(hash, None)
 
     def batch_remove_root_hash(self, version: int):
@@ -64,7 +64,7 @@ def batch_set_node(self, hash: bytes, node: PersistedNode):
         if self.batch is None:
             self.batch = rocksdb.WriteBatch()
         self.cache[hash] = node
-        self.batch.put(node_key(hash), node.encode())
+        self.batch.put(legacy_node_key(hash), node.encode())
 
     def batch_set_root_hash(self, version: int, hash: bytes):
         if self.batch is None:

iavl/utils.py

@@ -1,4 +1,5 @@
 import hashlib
+import struct
 import itertools
 from collections.abc import Iterator
 from typing import Callable, List, NamedTuple, Optional, Tuple
@@ -11,9 +12,10 @@
 EMPTY_HASH = hashlib.sha256().digest()
 FAST_KEY_PREFIX = b"f"
 METADATA_KEY_PREFIX = b"m"
-NODE_KEY_PREFIX = b"n"
+NODE_KEY_PREFIX = b"s"
 ORPHAN_KEY_PREFIX = b"o"
 ROOT_KEY_PREFIX = b"r"
+LEGACY_NODE_KEY_PREFIX = b"n"
 GetNode = Callable[bytes, Optional["PersistedNode"]]
 
 
@@ -131,13 +133,33 @@ def prefix_iteritems(
 
 
 def root_key(v: int) -> bytes:
+    return struct.pack(">qI", v, 1)
+
+
+def legacy_root_key(v: int) -> bytes:
     return ROOT_KEY_PREFIX + v.to_bytes(8, "big")
 
 
 def node_key(hash: bytes) -> bytes:
+    if len(hash) != 12:
+        raise ValueError("Key must be 12 bytes")
     return NODE_KEY_PREFIX + hash
 
 
+def legacy_node_key(hash: bytes) -> bytes:
+    return LEGACY_NODE_KEY_PREFIX + hash
+
+
+def parse_node_key(key: bytes) -> tuple[int, int]:
+    if not key.startswith(NODE_KEY_PREFIX):
+        raise ValueError("Key must start with root prefix 's'")
+    key = key[len(NODE_KEY_PREFIX):]
+    if len(key) != 12:
+        raise ValueError("Key must be 12 bytes after prefix")
+    version, nonce = struct.unpack(">qI", key)
+    return version, nonce
+
+
 def fast_node_key(key: bytes) -> bytes:
     return FAST_KEY_PREFIX + key
 
@@ -246,10 +268,11 @@ def decode_fast_node(bz: bytes) -> (int, bytes, int):
 
 
 def get_node(
-    db: DBM, hash: bytes, store: Optional[str] = None
+    db: DBM, hash: bytes, store: Optional[str] = None, legacy=False,
 ) -> Optional[PersistedNode]:
-    prefix = store_prefix(store) if store is not None else b""
-    bz = db.get(prefix + node_key(hash))
+    prefix = store_prefix(store) if store else b""
+    key = legacy_node_key(hash) if legacy else node_key(hash)
+    bz = db.get(prefix + key)
     if not bz:
         return
     node, _ = decode_node(bz, hash)
@@ -259,10 +282,13 @@ def get_node(
 def get_root_node(
     db: DBM, version: int, store: Optional[str] = None
 ) -> Optional[PersistedNode]:
-    prefix = store_prefix(store) if store is not None else b""
-    hash = db.get(prefix + root_key(version))
+    prefix = store_prefix(store) if store else b""
+    hash = db.get(prefix + node_key(root_key(version)))
     if not hash:
-        return
+        hash = db.get(prefix + legacy_root_key(version))
+        if not hash:
+            return
+        return get_node(db, hash, store, legacy=True)
     return get_node(db, hash, store)
 
 
@@ -312,7 +338,7 @@ def iter_iavl_tree(
     prefix = store_prefix(store) if store is not None else b""
 
     def get_node(hash: bytes) -> PersistedNode:
-        n, _ = decode_node(db.get(prefix + node_key(hash)), hash)
+        n, _ = decode_node(db.get(prefix + legacy_node_key(hash)), hash)
         return n
 
     def prune_check(node: PersistedNode) -> (bool, bool):

iavl/visualize.py

@@ -5,7 +5,7 @@
 from hexbytes import HexBytes
 
 from .iavl import NodeDB
-from .utils import PersistedNode, decode_node, node_key
+from .utils import PersistedNode, decode_node, legacy_node_key
 
 
 def label(node: PersistedNode):
@@ -26,7 +26,7 @@ def visualize_iavl(
     g = Digraph(comment="IAVL Tree")
 
     def get_node(hash: bytes) -> PersistedNode:
-        n, _ = decode_node(db.get(prefix + node_key(hash)), hash)
+        n, _ = decode_node(db.get(prefix + legacy_node_key(hash)), hash)
         return n
 
     def vis_node(hash: bytes, n: PersistedNode):