Skip to content

Commit 98bd130

Browse files
wanseobwanseob
authored
Merge pull request #4 from commitground/feature/non-inclusion
Feature for non-included keys
2 parents 7320fe6 + d868d7c commit 98bd130

File tree

5 files changed

+210
-5
lines changed

5 files changed

+210
-5
lines changed

‎contracts/implementation.sol

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,9 +9,10 @@ contract PartialMerkleTreeImplementation {
99
constructor () public {
1010
}
1111

12-
function initialize (bytes32 initialRoot) public {
12+
function initialize(bytes32 initialRoot) public {
1313
tree.initialize(initialRoot);
1414
}
15+
1516
function insert(bytes key, bytes value) public {
1617
tree.insert(key, value);
1718
}
@@ -20,6 +21,10 @@ contract PartialMerkleTreeImplementation {
2021
return tree.commitBranch(key, value, branchMask, siblings);
2122
}
2223

24+
function commitBranchOfNonInclusion(bytes key, bytes32 potentialSiblingLabel, bytes32 potentialSiblingValue, uint branchMask, bytes32[] siblings) public {
25+
return tree.commitBranchOfNonInclusion(key, potentialSiblingLabel, potentialSiblingValue, branchMask, siblings);
26+
}
27+
2328
function get(bytes key) public view returns (bytes) {
2429
return tree.get(key);
2530
}
@@ -44,7 +49,20 @@ contract PartialMerkleTreeImplementation {
4449
return tree.getProof(key);
4550
}
4651

52+
function getNonInclusionProof(bytes key) public view returns (
53+
bytes32 leafLabel,
54+
bytes32 leafNode,
55+
uint branchMask,
56+
bytes32[] _siblings
57+
) {
58+
return tree.getNonInclusionProof(key);
59+
}
60+
4761
function verifyProof(bytes32 rootHash, bytes key, bytes value, uint branchMask, bytes32[] siblings) public pure {
4862
PartialMerkleTree.verifyProof(rootHash, key, value, branchMask, siblings);
4963
}
64+
65+
function verifyNonInclusionProof(bytes32 rootHash, bytes key, bytes32 leafLabel, bytes32 leafNode, uint branchMask, bytes32[] siblings) public pure {
66+
PartialMerkleTree.verifyNonInclusionProof(rootHash, key, leafLabel, leafNode, branchMask, siblings);
67+
}
5068
}

‎contracts/tree.sol

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,56 @@ library PartialMerkleTree {
6464
tree.rootEdge = e;
6565
}
6666

67+
function commitBranchOfNonInclusion(
68+
Tree storage tree,
69+
bytes key,
70+
bytes32 potentialSiblingLabel,
71+
bytes32 potentialSiblingValue,
72+
uint branchMask,
73+
bytes32[] siblings
74+
) internal {
75+
D.Label memory k = D.Label(keccak256(key), 256);
76+
D.Edge memory e;
77+
// e.node(0x083d)
78+
for (uint i = 0; branchMask != 0; i++) {
79+
// retrieve edge data with branch mask
80+
uint bitSet = Utils.lowestBitSet(branchMask);
81+
branchMask &= ~(uint(1) << bitSet);
82+
(k, e.label) = Utils.splitAt(k, 255 - bitSet);
83+
uint bit;
84+
(bit, e.label) = Utils.chopFirstBit(e.label);
85+
86+
if (i == 0) {
87+
e.label.length = bitSet;
88+
e.label.data = potentialSiblingLabel;
89+
e.node = potentialSiblingValue;
90+
}
91+
92+
// find upper node with retrieved edge & sibling
93+
bytes32[2] memory edgeHashes;
94+
edgeHashes[bit] = edgeHash(e);
95+
edgeHashes[1 - bit] = siblings[siblings.length - i - 1];
96+
bytes32 upperNode = keccak256(abi.encode(edgeHashes[0], edgeHashes[1]));
97+
98+
// Update sibling information
99+
D.Node storage parentNode = tree.nodes[upperNode];
100+
101+
102+
// Put edge
103+
parentNode.children[bit] = e;
104+
// Put sibling edge if needed
105+
if (parentNode.children[1 - bit].isEmpty()) {
106+
parentNode.children[1 - bit].header = siblings[siblings.length - i - 1];
107+
}
108+
// go to upper edge
109+
e.node = keccak256(abi.encode(edgeHashes[0], edgeHashes[1]));
110+
}
111+
e.label = k;
112+
require(tree.root == edgeHash(e));
113+
tree.root = edgeHash(e);
114+
tree.rootEdge = e;
115+
}
116+
67117
function insert(Tree storage tree, bytes key, bytes value) internal {
68118
D.Label memory k = D.Label(keccak256(key), 256);
69119
bytes32 valueHash = keccak256(value);
@@ -164,6 +214,54 @@ library PartialMerkleTree {
164214
}
165215
}
166216

217+
function getNonInclusionProof(Tree storage tree, bytes key) internal view returns (
218+
bytes32 potentialSiblingLabel,
219+
bytes32 potentialSiblingValue,
220+
uint branchMask,
221+
bytes32[] _siblings
222+
){
223+
uint length;
224+
uint numSiblings;
225+
226+
// Start from root edge
227+
D.Label memory label = D.Label(keccak256(key), 256);
228+
D.Edge memory e = tree.rootEdge;
229+
bytes32[256] memory siblings;
230+
231+
while (true) {
232+
// Find at edge
233+
require(label.length >= e.label.length);
234+
D.Label memory prefix;
235+
D.Label memory suffix;
236+
(prefix, suffix) = Utils.splitCommonPrefix(label, e.label);
237+
238+
// suffix.length == 0 means that the key exists. Thus the length of the suffix should be not zero
239+
require(suffix.length != 0);
240+
241+
if (prefix.length >= e.label.length) {
242+
// Partial matched, keep finding
243+
length += prefix.length;
244+
branchMask |= uint(1) << (255 - length);
245+
length += 1;
246+
uint head;
247+
(head, label) = Utils.chopFirstBit(suffix);
248+
siblings[numSiblings++] = edgeHash(tree.nodes[e.node].children[1 - head]);
249+
e = tree.nodes[e.node].children[head];
250+
} else {
251+
// Found the potential sibling. Set data to return
252+
potentialSiblingLabel = e.label.data;
253+
potentialSiblingValue = e.node;
254+
break;
255+
}
256+
}
257+
if (numSiblings > 0)
258+
{
259+
_siblings = new bytes32[](numSiblings);
260+
for (uint i = 0; i < numSiblings; i++)
261+
_siblings[i] = siblings[i];
262+
}
263+
}
264+
167265
function verifyProof(bytes32 rootHash, bytes key, bytes value, uint branchMask, bytes32[] siblings) public pure {
168266
D.Label memory k = D.Label(keccak256(key), 256);
169267
D.Edge memory e;
@@ -183,6 +281,29 @@ library PartialMerkleTree {
183281
require(rootHash == edgeHash(e));
184282
}
185283

284+
function verifyNonInclusionProof(bytes32 rootHash, bytes key, bytes32 potentialSiblingLabel, bytes32 potentialSiblingValue, uint branchMask, bytes32[] siblings) public pure {
285+
D.Label memory k = D.Label(keccak256(key), 256);
286+
D.Edge memory e;
287+
for (uint i = 0; branchMask != 0; i++) {
288+
uint bitSet = Utils.lowestBitSet(branchMask);
289+
branchMask &= ~(uint(1) << bitSet);
290+
(k, e.label) = Utils.splitAt(k, 255 - bitSet);
291+
uint bit;
292+
(bit, e.label) = Utils.chopFirstBit(e.label);
293+
bytes32[2] memory edgeHashes;
294+
if (i == 0) {
295+
e.label.length = bitSet;
296+
e.label.data = potentialSiblingLabel;
297+
e.node = potentialSiblingValue;
298+
}
299+
edgeHashes[bit] = edgeHash(e);
300+
edgeHashes[1 - bit] = siblings[siblings.length - i - 1];
301+
e.node = keccak256(abi.encode(edgeHashes[0], edgeHashes[1]));
302+
}
303+
e.label = k;
304+
require(rootHash == edgeHash(e));
305+
}
306+
186307
function newEdge(bytes32 node, D.Label label) internal pure returns (D.Edge memory e){
187308
e.node = node;
188309
e.label = label;

‎package-lock.json

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎package.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"name": "solidity-partial-tree",
3-
"version": "1.0.0",
3+
"version": "1.1.0",
44
"description": "Solidity implementation of partial merkle tree",
55
"directories": {
66
"test": "test"

‎test/PartialMerkleTree.Test.js

Lines changed: 68 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -166,6 +166,47 @@ contract('PartialMerkleTree', async ([_, primary, nonPrimary]) => {
166166
assert.equal(await tree.doesInclude('fuz'), false)
167167
})
168168
})
169+
170+
describe('getNonInclusionProof()', async () => {
171+
let items = { key1: 'value1', key2: 'value2', key3: 'value3' }
172+
it('should return proof data when the key does not exist', async () => {
173+
for (const key of Object.keys(items)) {
174+
await tree.insert(key, items[key], { from: primary })
175+
}
176+
await tree.getNonInclusionProof('key4')
177+
})
178+
it('should not return data when the key does exist', async () => {
179+
for (const key of Object.keys(items)) {
180+
await tree.insert(key, items[key], { from: primary })
181+
}
182+
try {
183+
await tree.getNonInclusionProof('key1')
184+
assert.fail('Did not reverted')
185+
} catch (e) {
186+
assert.ok('Reverted successfully')
187+
}
188+
})
189+
})
190+
191+
describe('verifyNonInclusionProof()', async () => {
192+
it('should be passed when we use correct proof data', async () => {
193+
let items = { key1: 'value1', key2: 'value2', key3: 'value3' }
194+
for (const key of Object.keys(items)) {
195+
await tree.insert(key, items[key], { from: primary })
196+
}
197+
let rootHash = await tree.getRootHash()
198+
let [potentialSiblingLabel, potentialSiblingValue, branchMask, siblings] = await tree.getNonInclusionProof('key4')
199+
await tree.verifyNonInclusionProof(rootHash, 'key4', potentialSiblingLabel, potentialSiblingValue, branchMask, siblings)
200+
for (const key of Object.keys(items)) {
201+
try {
202+
await tree.verifyNonInclusionProof(rootHash, key, potentialSiblingLabel, potentialSiblingValue, branchMask, siblings)
203+
assert.fail('Did not reverted')
204+
} catch (e) {
205+
assert.ok('Reverted successfully')
206+
}
207+
}
208+
})
209+
})
169210
})
170211

171212
context('We can reenact merkle tree transformation by submitting only referred siblings instead of submitting all nodes', async () => {
@@ -207,12 +248,37 @@ contract('PartialMerkleTree', async ([_, primary, nonPrimary]) => {
207248
await treeB.getProof('key1')
208249
})
209250

251+
let secondPhaseOfTreeA
252+
let secondPhaseOfTreeB
210253
it('should have same root hash when we update key1', async () => {
211254
await treeA.insert('key1', 'val4')
212255
await treeB.insert('key1', 'val4')
213-
let secondPhaseOfTreeA = await treeA.getRootHash()
214-
let secondPhaseOfTreeB = await treeB.getRootHash()
256+
secondPhaseOfTreeA = await treeA.getRootHash()
257+
secondPhaseOfTreeB = await treeB.getRootHash()
215258
assert.equal(secondPhaseOfTreeA, secondPhaseOfTreeB)
216259
})
260+
261+
it('should revert before the branch data of non inclusion is committed', async () => {
262+
try {
263+
await treeB.insert('key4', 'val4')
264+
assert.fail('Did not reverted')
265+
} catch (e) {
266+
assert.ok('Reverted successfully')
267+
}
268+
})
269+
270+
let thirdPhaseOfTreeA
271+
let thirdPhaseOfTreeB
272+
it('should be able to insert a non inclusion key-value pair after committting related branch data', async () => {
273+
let [potentialSiblingLabel, potentialSiblingValue, branchMask, siblings] = await treeA.getNonInclusionProof('key4')
274+
await treeB.commitBranchOfNonInclusion('key4', potentialSiblingLabel, potentialSiblingValue, branchMask, siblings)
275+
assert.equal(await treeB.getRootHash(), secondPhaseOfTreeB)
276+
277+
await treeA.insert('key4', 'val4')
278+
await treeB.insert('key4', 'val4')
279+
thirdPhaseOfTreeA = await treeA.getRootHash()
280+
thirdPhaseOfTreeB = await treeB.getRootHash()
281+
assert.equal(thirdPhaseOfTreeA, thirdPhaseOfTreeB)
282+
})
217283
})
218284
})

0 commit comments

Comments
 (0)