Feature for non-included keys

contracts/implementation.sol

@@ -9,9 +9,10 @@ contract PartialMerkleTreeImplementation {
     constructor () public {
     }
 
-    function initialize (bytes32 initialRoot) public {
+    function initialize(bytes32 initialRoot) public {
         tree.initialize(initialRoot);
     }
+
     function insert(bytes key, bytes value) public {
         tree.insert(key, value);
     }
@@ -20,6 +21,10 @@ contract PartialMerkleTreeImplementation {
         return tree.commitBranch(key, value, branchMask, siblings);
     }
 
+    function commitBranchOfNonInclusion(bytes key, bytes32 potentialSiblingLabel, bytes32 potentialSiblingValue, uint branchMask, bytes32[] siblings) public {
+        return tree.commitBranchOfNonInclusion(key, potentialSiblingLabel, potentialSiblingValue, branchMask, siblings);
+    }
+
     function get(bytes key) public view returns (bytes) {
         return tree.get(key);
     }
@@ -36,7 +41,20 @@ contract PartialMerkleTreeImplementation {
         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 {
         PartialMerkleTree.verifyProof(rootHash, key, value, branchMask, siblings);
     }
+
+    function verifyNonInclusionProof(bytes32 rootHash, bytes key, bytes32 leafLabel, bytes32 leafNode, uint branchMask, bytes32[] siblings) public pure {
+        PartialMerkleTree.verifyNonInclusionProof(rootHash, key, leafLabel, leafNode, branchMask, siblings);
+    }
 }

contracts/tree.sol

@@ -64,6 +64,56 @@ library PartialMerkleTree {
         tree.rootEdge = e;
     }
 
+    function commitBranchOfNonInclusion(
+        Tree storage tree,
+        bytes key,
+        bytes32 potentialSiblingLabel,
+        bytes32 potentialSiblingValue,
+        uint branchMask,
+        bytes32[] siblings
+    ) internal {
+        D.Label memory k = D.Label(keccak256(key), 256);
+        D.Edge memory e;
+        // e.node(0x083d)
+        for (uint i = 0; branchMask != 0; i++) {
+            // retrieve edge data with branch mask
+            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);
+
+            if (i == 0) {
+                e.label.length = bitSet;
+                e.label.data = potentialSiblingLabel;
+                e.node = potentialSiblingValue;
+            }
+
+            // find upper node with retrieved edge & sibling
+            bytes32[2] memory edgeHashes;
+            edgeHashes[bit] = edgeHash(e);
+            edgeHashes[1 - bit] = siblings[siblings.length - i - 1];
+            bytes32 upperNode = keccak256(abi.encode(edgeHashes[0], edgeHashes[1]));
+
+            // Update sibling information
+            D.Node storage parentNode = tree.nodes[upperNode];
+
+
+            // Put edge
+            parentNode.children[bit] = e;
+            // Put sibling edge if needed
+            if (parentNode.children[1 - bit].isEmpty()) {
+                parentNode.children[1 - bit].header = siblings[siblings.length - i - 1];
+            }
+            // go to upper edge
+            e.node = keccak256(abi.encode(edgeHashes[0], edgeHashes[1]));
+        }
+        e.label = k;
+        require(tree.root == edgeHash(e));
+        tree.root = edgeHash(e);
+        tree.rootEdge = e;
+    }
+
     function insert(Tree storage tree, bytes key, bytes value) internal {
         D.Label memory k = D.Label(keccak256(key), 256);
         bytes32 valueHash = keccak256(value);
@@ -152,6 +202,54 @@ library PartialMerkleTree {
         }
     }
 
+    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;
@@ -171,6 +269,29 @@ library PartialMerkleTree {
         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));
+    }
+
     function newEdge(bytes32 node, D.Label label) internal pure returns (D.Edge memory e){
         e.node = node;
         e.label = label;

test/PartialMerkleTree.Test.js

@@ -142,6 +142,47 @@ contract('PartialMerkleTree', async ([_, primary, nonPrimary]) => {
         assert.equal(web3.toUtf8(await tree.get('foo')), 'bar')
       })
     })
+
+    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')
+          }
+        }
+      })
+    })
   })
 
   context('We can reenact merkle tree transformation by submitting only referred siblings instead of submitting all nodes', async () => {
@@ -166,29 +207,54 @@ contract('PartialMerkleTree', async ([_, primary, nonPrimary]) => {
       siblingsForKey1 = proof[1]
     })
 
-    it('should start with same root hash by initialization', async()=> {
+    it('should start with same root hash by initialization', async () => {
       //initilaze with the first root hash
       await treeB.initialize(firstPhaseOfTreeA)
       assert.equal(await treeB.getRootHash(), firstPhaseOfTreeA)
     })
 
-    it('should not change root after committing branch data', async ()=> {
+    it('should not change root after committing branch data', async () => {
       // commit branch data
       await treeB.commitBranch('key1', referredValueForKey1, branchMaskForKey1, siblingsForKey1)
       assert.equal(await treeB.getRootHash(), firstPhaseOfTreeA)
     })
 
-    it('should be able to return proof data', async ()=> {
+    it('should be able to return proof data', async () => {
       // commit branch data
       await treeB.getProof('key1')
     })
 
+    let secondPhaseOfTreeA
+    let secondPhaseOfTreeB
     it('should have same root hash when we update key1', async () => {
       await treeA.insert('key1', 'val4')
       await treeB.insert('key1', 'val4')
-      let secondPhaseOfTreeA = await treeA.getRootHash()
-      let secondPhaseOfTreeB = await treeB.getRootHash()
+      secondPhaseOfTreeA = await treeA.getRootHash()
+      secondPhaseOfTreeB = await treeB.getRootHash()
       assert.equal(secondPhaseOfTreeA, secondPhaseOfTreeB)
     })
+
+    it('should revert before the branch data of non inclusion is committed', async () => {
+      try {
+        await treeB.insert('key4', 'val4')
+        assert.fail('Did not reverted')
+      } catch (e) {
+        assert.ok('Reverted successfully')
+      }
+    })
+
+    let thirdPhaseOfTreeA
+    let thirdPhaseOfTreeB
+    it('should be able to insert a non inclusion key-value pair after committting related branch data', async () => {
+      let [potentialSiblingLabel, potentialSiblingValue, branchMask, siblings] = await treeA.getNonInclusionProof('key4')
+      await treeB.commitBranchOfNonInclusion('key4', potentialSiblingLabel, potentialSiblingValue, branchMask, siblings)
+      assert.equal(await treeB.getRootHash(), secondPhaseOfTreeB)
+
+      await treeA.insert('key4', 'val4')
+      await treeB.insert('key4', 'val4')
+      thirdPhaseOfTreeA = await treeA.getRootHash()
+      thirdPhaseOfTreeB = await treeB.getRootHash()
+      assert.equal(thirdPhaseOfTreeA, thirdPhaseOfTreeB)
+    })
   })
 })