circomlib/src/smt.js
src/smt.jsRaw
1const bigInt = require("snarkjs").bigInt;
2
3const SMTMemDB = require("./smt_memdb");
4const {hash0, hash1} = require("./smt_hashes_poseidon");
5
6class SMT {
7
8 constructor(db, root) {
9 this.db = db;
10 this.root = root;
11 }
12
13 _splitBits(_key) {
14
15 let k = bigInt(_key);
16 const res = [];
17
18 while (!k.isZero()) {
19 if (k.isOdd()) {
20 res.push(true);
21 } else {
22 res.push(false);
23 }
24 k = k.shr(1);
25 }
26
27 while (res.length<256) res.push(false);
28
29 return res;
30 }
31
32 async update(_key, _newValue) {
33 const key = bigInt(_key);
34 const newValue = bigInt(_newValue);
35
36
37 const resFind = await this.find(key);
38 const res = {};
39 res.oldRoot = this.root;
40 res.oldKey = key;
41 res.oldValue = resFind.foundValue;
42 res.newKey = key;
43 res.newValue = newValue;
44 res.siblings = resFind.siblings;
45
46 const ins = [];
47 const dels = [];
48
49 let rtOld = hash1(key, resFind.foundValue);
50 let rtNew = hash1(key, newValue);
51 ins.push([rtNew, [1, key, newValue ]]);
52 dels.push(rtOld);
53
54 const keyBits = this._splitBits(key);
55 for (let level = resFind.siblings.length-1; level >=0; level--) {
56 let oldNode, newNode;
57 const sibling = resFind.siblings[level];
58 if (keyBits[level]) {
59 oldNode = [sibling, rtOld];
60 newNode = [sibling, rtNew];
61 } else {
62 oldNode = [rtOld, sibling];
63 newNode = [rtNew, sibling];
64 }
65 rtOld = hash0(oldNode[0], oldNode[1]);
66 rtNew = hash0(newNode[0], newNode[1]);
67 dels.push(rtOld);
68 ins.push([rtNew, newNode]);
69 }
70
71 res.newRoot = rtNew;
72
73 await this.db.multiIns(ins);
74 await this.db.setRoot(rtNew);
75 this.root = rtNew;
76 await this.db.multiDel(dels);
77
78 return res;
79 }
80
81 async delete(_key) {
82 const key = bigInt(_key);
83
84 const resFind = await this.find(key);
85 if (!resFind.found) throw new Error("Key does not exists");
86
87 const res = {
88 siblings: [],
89 delKey: key,
90 delValue: resFind.foundValue
91 };
92
93 const dels = [];
94 const ins = [];
95 let rtOld = hash1(key, resFind.foundValue);
96 let rtNew;
97 dels.push(rtOld);
98
99 let mixed;
100 if (resFind.siblings.length > 0) {
101 const record = await this.db.get(resFind.siblings[resFind.siblings.length - 1]);
102 if ((record.length == 3)&&(record[0].equals(bigInt.one))) {
103 mixed = false;
104 res.oldKey = record[1];
105 res.oldValue = record[2];
106 res.isOld0 = false;
107 rtNew = resFind.siblings[resFind.siblings.length - 1];
108 } else if (record.length == 2) {
109 mixed = true;
110 res.oldKey = key;
111 res.oldValue = bigInt(0);
112 res.isOld0 = true;
113 rtNew = bigInt.zero;
114 } else {
115 throw new Error("Invalid node. Database corrupted");
116 }
117 } else {
118 rtNew = bigInt.zero;
119 res.oldKey = key;
120 res.oldValue = bigInt(0);
121 res.isOld0 = true;
122 }
123
124 const keyBits = this._splitBits(key);
125
126 for (let level = resFind.siblings.length-1; level >=0; level--) {
127 let newSibling = resFind.siblings[level];
128 if ((level == resFind.siblings.length-1)&&(!res.isOld0)) {
129 newSibling = bigInt.zero;
130 }
131 const oldSibling = resFind.siblings[level];
132 if (keyBits[level]) {
133 rtOld = hash0(oldSibling, rtOld);
134 } else {
135 rtOld = hash0(rtOld, oldSibling);
136 }
137 dels.push(rtOld);
138 if (!newSibling.isZero()) {
139 mixed = true;
140 }
141
142 if (mixed) {
143 res.siblings.unshift(resFind.siblings[level]);
144 let newNode;
145 if (keyBits[level]) {
146 newNode = [newSibling, rtNew];
147 } else {
148 newNode = [rtNew, newSibling];
149 }
150 rtNew = hash0(newNode[0], newNode[1]);
151 ins.push([rtNew, newNode]);
152 }
153 }
154
155 await this.db.multiIns(ins);
156 await this.db.setRoot(rtNew);
157 this.root = rtNew;
158 await this.db.multiDel(dels);
159
160 res.newRoot = rtNew;
161 res.oldRoot = rtOld;
162
163 return res;
164 }
165
166 async insert(_key, _value) {
167 const key = bigInt(_key);
168 const value = bigInt(_value);
169 let addedOne = false;
170 const res = {};
171 res.oldRoot = this.root;
172 const newKeyBits = this._splitBits(key);
173
174 let rtOld;
175
176 const resFind = await this.find(key);
177
178 if (resFind.found) throw new Error("Key already exists");
179
180 res.siblings = resFind.siblings;
181 let mixed;
182
183 if (!resFind.isOld0) {
184 const oldKeyits = this._splitBits(resFind.notFoundKey);
185 for (let i= res.siblings.length; oldKeyits[i] == newKeyBits[i]; i++) {
186 res.siblings.push(bigInt.zero);
187 }
188 rtOld = hash1(resFind.notFoundKey, resFind.notFoundValue);
189 res.siblings.push(rtOld);
190 addedOne = true;
191 mixed = false;
192 } else if (res.siblings.length >0) {
193 mixed = true;
194 rtOld = bigInt.zero;
195 }
196
197 const inserts = [];
198 const dels = [];
199
200 let rt = hash1(key, value);
201 inserts.push([rt,[1, key, value]] );
202
203 for (let i=res.siblings.length-1; i>=0; i--) {
204 if ((i<res.siblings.length-1)&&(!res.siblings[i].isZero())) {
205 mixed = true;
206 }
207 if (mixed) {
208 const oldSibling = resFind.siblings[i];
209 if (newKeyBits[i]) {
210 rtOld = hash0(oldSibling, rtOld);
211 } else {
212 rtOld = hash0(rtOld, oldSibling);
213 }
214 dels.push(rtOld);
215 }
216
217
218 let newRt;
219 if (newKeyBits[i]) {
220 newRt = hash0(res.siblings[i], rt);
221 inserts.push([newRt,[res.siblings[i], rt]] );
222 } else {
223 newRt = hash0(rt, res.siblings[i]);
224 inserts.push([newRt,[rt, res.siblings[i]]] );
225 }
226 rt = newRt;
227 }
228
229 if (addedOne) res.siblings.pop();
230 while ((res.siblings.length>0) && (res.siblings[res.siblings.length-1].isZero())) {
231 res.siblings.pop();
232 }
233 res.oldKey = resFind.notFoundKey;
234 res.oldValue = resFind.notFoundValue;
235 res.newRoot = rt;
236 res.isOld0 = resFind.isOld0;
237
238
239 await this.db.multiIns(inserts);
240 await this.db.setRoot(rt);
241 this.root = rt;
242 await this.db.multiDel(dels);
243
244 return res;
245 }
246
247 async find(key) {
248 const keyBits = this._splitBits(key);
249 return await this._find(key, keyBits, this.root, 0);
250 }
251
252 async _find(key, keyBits, root, level) {
253 if (typeof root === "undefined") root = this.root;
254
255 let res;
256 if (root.isZero()) {
257 res = {
258 found: false,
259 siblings: [],
260 notFoundKey: key,
261 notFoundValue: bigInt.zero,
262 isOld0: true
263 };
264 return res;
265 }
266
267 const record = await this.db.get(root);
268
269 if ((record.length==3)&&(record[0].equals(bigInt.one))) {
270 if (record[1].equals(key)) {
271 res = {
272 found: true,
273 siblings: [],
274 foundValue: record[2],
275 isOld0: false
276 };
277 } else {
278 res = {
279 found: false,
280 siblings: [],
281 notFoundKey: record[1],
282 notFoundValue: record[2],
283 isOld0: false
284 };
285 }
286 } else {
287 if (keyBits[level] == 0) {
288 res = await this._find(key, keyBits, record[0], level+1);
289 res.siblings.unshift(record[1]);
290 } else {
291 res = await this._find(key, keyBits, record[1], level+1);
292 res.siblings.unshift(record[0]);
293 }
294 }
295 return res;
296 }
297}
298
299async function loadFromFile(fileName) {
300
301}
302
303async function newMemEmptyTrie() {
304 const db = new SMTMemDB();
305 const rt = await db.getRoot();
306 const smt = new SMT(db, rt);
307 return smt;
308}
309
310module.exports.loadFromFile = loadFromFile;
311module.exports.newMemEmptyTrie = newMemEmptyTrie;
312module.exports.SMT = SMT;
313module.exports.SMTMemDB = SMTMemDB;