Skip to content

Commit 85b7a88

Browse files
wanseobwanseob
committed
Feature for non-included keys
1 parent 117dd2c commit 85b7a88

File tree

3 files changed

+211
-6
lines changed

3 files changed

+211
-6
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
}
@@ -36,7 +41,20 @@ contract PartialMerkleTreeImplementation {
3641
return tree.getProof(key);
3742
}
3843

44+
function getNonInclusionProof(bytes key) public view returns (
45+
bytes32 leafLabel,
46+
bytes32 leafNode,
47+
uint branchMask,
48+
bytes32[] _siblings
49+
) {
50+
return tree.getNonInclusionProof(key);
51+
}
52+
3953
function verifyProof(bytes32 rootHash, bytes key, bytes value, uint branchMask, bytes32[] siblings) public pure {
4054
PartialMerkleTree.verifyProof(rootHash, key, value, branchMask, siblings);
4155
}
56+
57+
function verifyNonInclusionProof(bytes32 rootHash, bytes key, bytes32 leafLabel, bytes32 leafNode, uint branchMask, bytes32[] siblings) public pure {
58+
PartialMerkleTree.verifyNonInclusionProof(rootHash, key, leafLabel, leafNode, branchMask, siblings);
59+
}
4260
}

‎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);
@@ -152,6 +202,54 @@ library PartialMerkleTree {
152202
}
153203
}
154204

205+
function getNonInclusionProof(Tree storage tree, bytes key) internal view returns (
206+
bytes32 potentialSiblingLabel,
207+
bytes32 potentialSiblingValue,
208+
uint branchMask,
209+
bytes32[] _siblings
210+
){
211+
uint length;
212+
uint numSiblings;
213+
214+
// Start from root edge
215+
D.Label memory label = D.Label(keccak256(key), 256);
216+
D.Edge memory e = tree.rootEdge;
217+
bytes32[256] memory siblings;
218+
219+
while (true) {
220+
// Find at edge
221+
require(label.length >= e.label.length);
222+
D.Label memory prefix;
223+
D.Label memory suffix;
224+
(prefix, suffix) = Utils.splitCommonPrefix(label, e.label);
225+
226+
// suffix.length == 0 means that the key exists. Thus the length of the suffix should be not zero
227+
require(suffix.length != 0);
228+
229+
if (prefix.length >= e.label.length) {
230+
// Partial matched, keep finding
231+
length += prefix.length;
232+
branchMask |= uint(1) << (255 - length);
233+
length += 1;
234+
uint head;
235+
(head, label) = Utils.chopFirstBit(suffix);
236+
siblings[numSiblings++] = edgeHash(tree.nodes[e.node].children[1 - head]);
237+
e = tree.nodes[e.node].children[head];
238+
} else {
239+
// Found the potential sibling. Set data to return
240+
potentialSiblingLabel = e.label.data;
241+
potentialSiblingValue = e.node;
242+
break;
243+
}
244+
}
245+
if (numSiblings > 0)
246+
{
247+
_siblings = new bytes32[](numSiblings);
248+
for (uint i = 0; i < numSiblings; i++)
249+
_siblings[i] = siblings[i];
250+
}
251+
}
252+
155253
function verifyProof(bytes32 rootHash, bytes key, bytes value, uint branchMask, bytes32[] siblings) public pure {
156254
D.Label memory k = D.Label(keccak256(key), 256);
157255
D.Edge memory e;
@@ -171,6 +269,29 @@ library PartialMerkleTree {
171269
require(rootHash == edgeHash(e));
172270
}
173271

272+
function verifyNonInclusionProof(bytes32 rootHash, bytes key, bytes32 potentialSiblingLabel, bytes32 potentialSiblingValue, uint branchMask, bytes32[] siblings) public pure {
273+
D.Label memory k = D.Label(keccak256(key), 256);
274+
D.Edge memory e;
275+
for (uint i = 0; branchMask != 0; i++) {
276+
uint bitSet = Utils.lowestBitSet(branchMask);
277+
branchMask &= ~(uint(1) << bitSet);
278+
(k, e.label) = Utils.splitAt(k, 255 - bitSet);
279+
uint bit;
280+
(bit, e.label) = Utils.chopFirstBit(e.label);
281+
bytes32[2] memory edgeHashes;
282+
if (i == 0) {
283+
e.label.length = bitSet;
284+
e.label.data = potentialSiblingLabel;
285+
e.node = potentialSiblingValue;
286+
}
287+
edgeHashes[bit] = edgeHash(e);
288+
edgeHashes[1 - bit] = siblings[siblings.length - i - 1];
289+
e.node = keccak256(abi.encode(edgeHashes[0], edgeHashes[1]));
290+
}
291+
e.label = k;
292+
require(rootHash == edgeHash(e));
293+
}
294+
174295
function newEdge(bytes32 node, D.Label label) internal pure returns (D.Edge memory e){
175296
e.node = node;
176297
e.label = label;

‎test/PartialMerkleTree.Test.js‎

Lines changed: 71 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,47 @@ contract('PartialMerkleTree', async ([_, primary, nonPrimary]) => {
142142
assert.equal(web3.toUtf8(await tree.get('foo')), 'bar')
143143
})
144144
})
145+
146+
describe('getNonInclusionProof()', async () => {
147+
let items = { key1: 'value1', key2: 'value2', key3: 'value3' }
148+
it('should return proof data when the key does not exist', async () => {
149+
for (const key of Object.keys(items)) {
150+
await tree.insert(key, items[key], { from: primary })
151+
}
152+
await tree.getNonInclusionProof('key4')
153+
})
154+
it('should not return data when the key does exist', async () => {
155+
for (const key of Object.keys(items)) {
156+
await tree.insert(key, items[key], { from: primary })
157+
}
158+
try {
159+
await tree.getNonInclusionProof('key1')
160+
assert.fail('Did not reverted')
161+
} catch (e) {
162+
assert.ok('Reverted successfully')
163+
}
164+
})
165+
})
166+
167+
describe('verifyNonInclusionProof()', async () => {
168+
it('should be passed when we use correct proof data', async () => {
169+
let items = { key1: 'value1', key2: 'value2', key3: 'value3' }
170+
for (const key of Object.keys(items)) {
171+
await tree.insert(key, items[key], { from: primary })
172+
}
173+
let rootHash = await tree.getRootHash()
174+
let [potentialSiblingLabel, potentialSiblingValue, branchMask, siblings] = await tree.getNonInclusionProof('key4')
175+
await tree.verifyNonInclusionProof(rootHash, 'key4', potentialSiblingLabel, potentialSiblingValue, branchMask, siblings)
176+
for (const key of Object.keys(items)) {
177+
try {
178+
await tree.verifyNonInclusionProof(rootHash, key, potentialSiblingLabel, potentialSiblingValue, branchMask, siblings)
179+
assert.fail('Did not reverted')
180+
} catch (e) {
181+
assert.ok('Reverted successfully')
182+
}
183+
}
184+
})
185+
})
145186
})
146187

147188
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]) => {
166207
siblingsForKey1 = proof[1]
167208
})
168209

169-
it('should start with same root hash by initialization', async()=> {
210+
it('should start with same root hash by initialization', async () => {
170211
//initilaze with the first root hash
171212
await treeB.initialize(firstPhaseOfTreeA)
172213
assert.equal(await treeB.getRootHash(), firstPhaseOfTreeA)
173214
})
174215

175-
it('should not change root after committing branch data', async ()=> {
216+
it('should not change root after committing branch data', async () => {
176217
// commit branch data
177218
await treeB.commitBranch('key1', referredValueForKey1, branchMaskForKey1, siblingsForKey1)
178219
assert.equal(await treeB.getRootHash(), firstPhaseOfTreeA)
179220
})
180221

181-
it('should be able to return proof data', async ()=> {
222+
it('should be able to return proof data', async () => {
182223
// commit branch data
183224
await treeB.getProof('key1')
184225
})
185226

227+
let secondPhaseOfTreeA
228+
let secondPhaseOfTreeB
186229
it('should have same root hash when we update key1', async () => {
187230
await treeA.insert('key1', 'val4')
188231
await treeB.insert('key1', 'val4')
189-
let secondPhaseOfTreeA = await treeA.getRootHash()
190-
let secondPhaseOfTreeB = await treeB.getRootHash()
232+
secondPhaseOfTreeA = await treeA.getRootHash()
233+
secondPhaseOfTreeB = await treeB.getRootHash()
191234
assert.equal(secondPhaseOfTreeA, secondPhaseOfTreeB)
192235
})
236+
237+
it('should revert before the branch data of non inclusion is committed', async () => {
238+
try {
239+
await treeB.insert('key4', 'val4')
240+
assert.fail('Did not reverted')
241+
} catch (e) {
242+
assert.ok('Reverted successfully')
243+
}
244+
})
245+
246+
let thirdPhaseOfTreeA
247+
let thirdPhaseOfTreeB
248+
it('should be able to insert a non inclusion key-value pair after committting related branch data', async () => {
249+
let [potentialSiblingLabel, potentialSiblingValue, branchMask, siblings] = await treeA.getNonInclusionProof('key4')
250+
await treeB.commitBranchOfNonInclusion('key4', potentialSiblingLabel, potentialSiblingValue, branchMask, siblings)
251+
assert.equal(await treeB.getRootHash(), secondPhaseOfTreeB)
252+
253+
await treeA.insert('key4', 'val4')
254+
await treeB.insert('key4', 'val4')
255+
thirdPhaseOfTreeA = await treeA.getRootHash()
256+
thirdPhaseOfTreeB = await treeB.getRootHash()
257+
assert.equal(thirdPhaseOfTreeA, thirdPhaseOfTreeB)
258+
})
193259
})
194260
})

0 commit comments

Comments
 (0)