Add feature for non inclusion proof

Author: wanseob

contracts/implementation.sol

@@ -10,6 +10,10 @@ contract PatriciaTreeImplementation {
     constructor () public {
     }
 
+    function insert(bytes key, bytes value) public {
+        tree.insert(key, value);
+    }
+
     function get(bytes key) public view returns (bytes) {
         return tree.get(key);
     }
@@ -42,12 +46,21 @@ contract PatriciaTreeImplementation {
         return tree.getProof(key);
     }
 
+    function getNonInclusionProof(bytes key) public view returns (
+        bytes32 leafLabel,
+        bytes32 leafNode,
+        uint branchMask,
+        bytes32[] _siblings
+    ) {
+        return tree.getNonInclusionProof(key);
+    }
+
     function verifyProof(bytes32 rootHash, bytes key, bytes value, uint branchMask, bytes32[] siblings) public pure {
         PatriciaTree.verifyProof(rootHash, key, value, branchMask, siblings);
     }
 
-    function insert(bytes key, bytes value) public {
-        tree.insert(key, value);
+    function verifyNonInclusionProof(bytes32 rootHash, bytes key, bytes32 leafLabel, bytes32 leafNode, uint branchMask, bytes32[] siblings) public pure {
+        PatriciaTree.verifyNonInclusionProof(rootHash, key, leafLabel, leafNode, branchMask, siblings);
     }
 }
 
@@ -134,7 +147,7 @@ contract PatriciaTreeMerkleProof {
     }
 
     function seal() public onlyFor(Status.OPENED) {
-//        require(_verifyEdge(originalRootEdge));
+        //        require(_verifyEdge(originalRootEdge));
         tree.rootEdge = originalRootEdge;
         tree.root = PatriciaTree.edgeHash(tree.rootEdge);
         _changeStatus(Status.ONGOING);

contracts/tree.sol

@@ -109,6 +109,54 @@ library PatriciaTree {
         }
     }
 
+    function getNonInclusionProof(Tree storage tree, bytes key) internal view returns (
+        bytes32 potentialSiblingLabel,
+        bytes32 potentialSiblingValue,
+        uint branchMask,
+        bytes32[] _siblings
+    ){
+        uint length;
+        uint numSiblings;
+
+        // Start from root edge
+        D.Label memory label = D.Label(keccak256(key), 256);
+        D.Edge memory e = tree.rootEdge;
+        bytes32[256] memory siblings;
+
+        while (true) {
+            // Find at edge
+            require(label.length >= e.label.length);
+            D.Label memory prefix;
+            D.Label memory suffix;
+            (prefix, suffix) = Utils.splitCommonPrefix(label, e.label);
+
+            // suffix.length == 0 means that the key exists. Thus the length of the suffix should be not zero
+            require(suffix.length != 0);
+
+            if (prefix.length >= e.label.length) {
+                // Partial matched, keep finding
+                length += prefix.length;
+                branchMask |= uint(1) << (255 - length);
+                length += 1;
+                uint head;
+                (head, label) = Utils.chopFirstBit(suffix);
+                siblings[numSiblings++] = edgeHash(tree.nodes[e.node].children[1 - head]);
+                e = tree.nodes[e.node].children[head];
+            } else {
+                // Found the potential sibling. Set data to return
+                potentialSiblingLabel = e.label.data;
+                potentialSiblingValue = e.node;
+                break;
+            }
+        }
+        if (numSiblings > 0)
+        {
+            _siblings = new bytes32[](numSiblings);
+            for (uint i = 0; i < numSiblings; i++)
+                _siblings[i] = siblings[i];
+        }
+    }
+
     function verifyProof(bytes32 rootHash, bytes key, bytes value, uint branchMask, bytes32[] siblings) public pure {
         D.Label memory k = D.Label(keccak256(key), 256);
         D.Edge memory e;
@@ -128,6 +176,29 @@ library PatriciaTree {
         require(rootHash == edgeHash(e));
     }
 
+    function verifyNonInclusionProof(bytes32 rootHash, bytes key, bytes32 potentialSiblingLabel, bytes32 potentialSiblingValue, uint branchMask, bytes32[] siblings) public pure {
+        D.Label memory k = D.Label(keccak256(key), 256);
+        D.Edge memory e;
+        for (uint i = 0; branchMask != 0; i++) {
+            uint bitSet = Utils.lowestBitSet(branchMask);
+            branchMask &= ~(uint(1) << bitSet);
+            (k, e.label) = Utils.splitAt(k, 255 - bitSet);
+            uint bit;
+            (bit, e.label) = Utils.chopFirstBit(e.label);
+            bytes32[2] memory edgeHashes;
+            if (i == 0) {
+                e.label.length = bitSet;
+                e.label.data = potentialSiblingLabel;
+                e.node = potentialSiblingValue;
+            }
+            edgeHashes[bit] = edgeHash(e);
+            edgeHashes[1 - bit] = siblings[siblings.length - i - 1];
+            e.node = keccak256(abi.encode(edgeHashes[0], edgeHashes[1]));
+        }
+        e.label = k;
+        require(rootHash == edgeHash(e));
+    }
+
     // TODO also return the proof
     function insert(Tree storage tree, bytes key, bytes value) internal {
         D.Label memory k = D.Label(keccak256(key), 256);

test/PatriciaTree.test.js

@@ -236,5 +236,45 @@ contract('PatriciaTree', async ([_, primary, nonPrimary]) => {
       })
     })
 
+    describe('getNonInclusionProof()', async () => {
+      let items = { key1: 'value1', key2: 'value2', key3: 'value3' }
+      it('should return proof data when the key does not exist', async () => {
+        for (const key of Object.keys(items)) {
+          await tree.insert(key, items[key], { from: primary })
+        }
+        await tree.getNonInclusionProof('key4')
+      })
+      it('should not return data when the key does exist', async () => {
+        for (const key of Object.keys(items)) {
+          await tree.insert(key, items[key], { from: primary })
+        }
+        try {
+          await tree.getNonInclusionProof('key1')
+          assert.fail('Did not reverted')
+        } catch (e) {
+          assert.ok('Reverted successfully')
+        }
+      })
+    })
+
+    describe('verifyNonInclusionProof()', async () => {
+      it('should be passed when we use correct proof data', async () => {
+        let items = { key1: 'value1', key2: 'value2', key3: 'value3' }
+        for (const key of Object.keys(items)) {
+          await tree.insert(key, items[key], { from: primary })
+        }
+        let rootHash = await tree.getRootHash()
+        let [potentialSiblingLabel, potentialSiblingValue, branchMask, siblings] = await tree.getNonInclusionProof('key4')
+        await tree.verifyNonInclusionProof(rootHash, 'key4', potentialSiblingLabel, potentialSiblingValue, branchMask, siblings)
+        for (const key of Object.keys(items)) {
+          try {
+            await tree.verifyNonInclusionProof(rootHash, key, potentialSiblingLabel, potentialSiblingValue, branchMask, siblings)
+            assert.fail('Did not reverted')
+          } catch (e) {
+            assert.ok('Reverted successfully')
+          }
+        }
+      })
+    })
   })
 })