inclusion proof

Author: turbocrime

src/TreeCbor.sol

@@ -43,15 +43,25 @@ library TreeCbor {
     function getCid(Tree memory tree, CidCbor.CidBytes32 indexCid, uint startIdx)
         internal
         pure
-        returns (TreeNodeCbor.TreeNode memory, uint index)
+        returns (TreeNodeCbor.TreeNode memory, uint)
     {
         bytes32 indexBytes = CidCbor.CidBytes32.unwrap(indexCid);
         for (uint i = startIdx; i < tree.cids.length; i++) {
             if (CidCbor.CidBytes32.unwrap(tree.cids[i]) == indexBytes) {
                 return (tree.nodes[i], i);
             }
         }
-        revert("node not found");
+        revert("couldn't get cid, not found");
+    }
+
+    function hasCid(Tree memory tree, CidCbor.CidBytes32 indexCid, uint startIdx) internal pure returns (bool, uint) {
+        bytes32 indexBytes = CidCbor.CidBytes32.unwrap(indexCid);
+        for (uint i = startIdx; i < tree.cids.length; i++) {
+            if (CidCbor.CidBytes32.unwrap(tree.cids[i]) == indexBytes) {
+                return (true, i);
+            }
+        }
+        return (false, 0);
     }
 
     function validTree(Tree memory tree) internal pure returns (Tree memory) {
@@ -64,4 +74,68 @@ library TreeCbor {
         }
         return tree;
     }
+
+    function verifyInclusion(
+        TreeCbor.Tree memory tree,
+        bytes[] memory treeCbor,
+        CidCbor.CidBytes32 entryCid,
+        bytes memory targetRecord,
+        string memory targetKey
+    ) internal pure returns (bool) {
+        CidCbor.CidBytes32 targetCid = CidCbor.CidBytes32.wrap(sha256(targetRecord));
+        CidCbor.CidBytes32[] memory queue = new CidCbor.CidBytes32[](1);
+        queue[0] = entryCid;
+
+        CidCbor.CidBytes32 leftWalk;
+        CidCbor.CidBytes32[] memory rightWalk;
+        CidCbor.CidBytes32 currentCid;
+        TreeNodeCbor.TreeNode memory currentNode;
+        uint currentIndex;
+
+        bool found = false;
+        bool hasCurrent = false;
+
+        CidCbor.CidBytes32[] memory newQueue;
+
+        while (queue.length > 0) {
+            currentCid = queue[0];
+            (hasCurrent, currentIndex) = TreeCbor.hasCid(tree, currentCid, currentIndex);
+            if (!hasCurrent) {
+                newQueue = new CidCbor.CidBytes32[](queue.length - 1);
+                for (uint i = 0; i < newQueue.length; i++) {
+                    newQueue[i] = queue[i + 1];
+                }
+                queue = newQueue;
+                continue;
+            }
+            (currentNode, currentIndex) = TreeCbor.getCid(tree, currentCid, currentIndex);
+
+            leftWalk = CidCbor.readCidBytes32(treeCbor[currentIndex], currentNode.left);
+            rightWalk = new CidCbor.CidBytes32[](currentNode.entries.length);
+
+            for (uint i = 0; i < currentNode.entries.length; i++) {
+                rightWalk[i] = CidCbor.readCidBytes32(treeCbor[currentIndex], currentNode.entries[i].tree);
+                if (keccak256(abi.encode(targetKey)) == keccak256(abi.encode(currentNode.entries[i].key))) {
+                    CidCbor.CidBytes32 valueCid =
+                        CidCbor.readCidBytes32(treeCbor[currentIndex], currentNode.entries[i].value);
+                    if (CidCbor.CidBytes32.unwrap(valueCid) == CidCbor.CidBytes32.unwrap(targetCid)) {
+                        require(!found, "duplicate entry");
+                        found = true;
+                    }
+                }
+            }
+
+            newQueue = new CidCbor.CidBytes32[](queue.length + rightWalk.length);
+            newQueue[0] = leftWalk;
+            for (uint i = 1; i < queue.length; i++) {
+                newQueue[i] = queue[i];
+            }
+            for (uint i = 0; i < rightWalk.length; i++) {
+                newQueue[queue.length + i] = rightWalk[i];
+            }
+            queue = newQueue;
+        }
+
+        return found;
+    }
 }

test/TreeCbor.t.sol

@@ -14,6 +14,9 @@ contract TreeTest is Test {
         hex"A56364696478206469643A706C633A6D74713365346D67743777796A6868616E69657A656A3637637265766D336C61796B6C746F73703232716464617461D82A5825000171122066DA6655BF8DA79B69A87299CF170FED8497FA3059379DC4A8BFE1E28CAB5D936470726576F66776657273696F6E03";
     bytes32 private constant correctLeftCidBytes = hex"6E7335ED248EDAE3ED49D47B88A5FCAD2985E15F416F8AE23A49DFC1231AEB91";
 
+    bytes private constant targetRecord =
+        hex"a4647465787478196361722066696c65732063616e6e6f74206875727420796f75652474797065726170702e62736b792e666565642e706f7374656c616e67738162656e696372656174656441747818323032342d31312d31355431323a31303a33322e3031345a";
+
     //bytes32 private constant recordCidBytes = hex"6d51f125727763752e073d9301892fbddaa4cc6090d5fff9af7d49106b92d457";
     //CidCbor.CidBytes32 private constant recordIdx = CidCbor.CidBytes32.wrap(recordCidBytes);
 
@@ -72,4 +75,8 @@ contract TreeTest is Test {
         console.log("node.left index", CidCbor.CidIndex.unwrap(rootNode.left));
         console.log("node.entries length", rootNode.entries.length);
     }
+
+    function test_verifyInclusion_only() public view {
+        TreeCbor.verifyInclusion(tree, nodeCbors, rootCid, targetRecord, "app.bsky.feed.post/3laydu3mgac2");
+    }
 }