Merge w/ feature/safe-get

Author: wanseob

contracts/implementation.sol

@@ -29,6 +29,14 @@ contract PartialMerkleTreeImplementation {
         return tree.get(key);
     }
 
+    function safeGet(bytes key) public view returns (bytes) {
+        return tree.safeGet(key);
+    }
+
+    function doesInclude(bytes key) public view returns (bool) {
+        return tree.doesInclude(key);
+    }
+
     function getValue(bytes32 hash) public view returns (bytes) {
         return tree.values[hash];
     }

contracts/tree.sol

@@ -138,6 +138,18 @@ library PartialMerkleTree {
         return getValue(tree, _findNode(tree, key));
     }
 
+    function safeGet(Tree storage tree, bytes key) internal view returns (bytes value) {
+        bytes32 valueHash = _findNode(tree, key);
+        require(valueHash != bytes32(0));
+        value = getValue(tree, valueHash);
+        require(valueHash == keccak256(value));
+    }
+
+    function doesInclude(Tree storage tree, bytes key) internal view returns (bool) {
+        bytes32 valueHash = _findNode(tree, key);
+        return (valueHash != bytes32(0));
+    }
+
     function getValue(Tree storage tree, bytes32 valueHash) internal view returns (bytes) {
         return tree.values[valueHash];
     }

test/PartialMerkleTree.Test.js

@@ -143,6 +143,30 @@ contract('PartialMerkleTree', async ([_, primary, nonPrimary]) => {
       })
     })
 
+    describe('safeGet()', async () => {
+      it('should return stored value for the given key', async () => {
+        await tree.insert('foo', 'bar', { from: primary })
+        assert.equal(web3.toUtf8(await tree.get('foo')), 'bar')
+      })
+      it('should throw if the given key is not included', async () => {
+        await tree.insert('foo', 'bar', { from: primary })
+        try {
+          await tree.get('fuz')
+          assert.fail('Did not reverted')
+        } catch (e) {
+          assert.ok('Reverted successfully')
+        }
+      })
+    })
+
+    describe('doesInclude()', async () => {
+      it('should return boolean whether the tree includes the given key or not', async () => {
+        await tree.insert('foo', 'bar', { from: primary })
+        assert.equal(await tree.doesInclude('foo'), true)
+        assert.equal(await tree.doesInclude('fuz'), false)
+      })
+    })
+
     describe('getNonInclusionProof()', async () => {
       let items = { key1: 'value1', key2: 'value2', key3: 'value3' }
       it('should return proof data when the key does not exist', async () => {