1const bigInt = require("snarkjs").bigInt;
  2
  3const SMTMemDB = require("./smt_memdb");
  4const mimc7 = require("./mimc7");
  5
  6function smtHash(arr) {
  7    let r = bigInt(0);
  8    for (let i=0; i<arr.length; i++) {
  9        r = mimc7.hash(r, bigInt(arr[i]), 91 );
 10    }
 11    return r;
 12}
 13
 14class SMT {
 15
 16    constructor(db, root) {
 17        this.db = db;
 18        this.root = root;
 19    }
 20
 21    _splitBits(_key) {
 22
 23        let k = bigInt(_key);
 24        const res = [];
 25
 26        while (!k.isZero()) {
 27            if (k.isOdd()) {
 28                res.push(true);
 29            } else {
 30                res.push(false);
 31            }
 32            k = k.shr(1);
 33        }
 34
 35        while (res.length<256) res.push(false);
 36
 37        return res;
 38    }
 39
 40    async insert(_key, _value) {
 41        const key = bigInt(_key);
 42        const value = bigInt(_value);
 43        let addedOne = false;
 44        const res = {};
 45        res.oldRoot = this.root;
 46        const newKeyBits = this._splitBits(key);
 47        const resFind = await this.find(key);
 48
 49        if (resFind.found) throw new Error("Key already exists");
 50
 51        res.sibblings = resFind.sibblings;
 52
 53        if (!resFind.isOld0) {
 54            const oldKeyits = this._splitBits(resFind.notFoundKey);
 55            for (let i= res.sibblings.length; oldKeyits[i] == newKeyBits[i]; i++) {
 56                res.sibblings.push(bigInt.zero);
 57            }
 58            res.sibblings.push(smtHash([1, resFind.notFoundKey, resFind.notFoundValue]));
 59            addedOne = true;
 60        }
 61
 62        const inserts = [];
 63
 64        let rt = smtHash([1, key, value]);
 65        inserts.push([rt,[1, key, value]] );
 66
 67        for (let i=res.sibblings.length-1; i>=0; i--) {
 68            if (newKeyBits[i]) {
 69                rt = smtHash([res.sibblings[i], rt]);
 70                inserts.push([rt,[res.sibblings[i], rt]] );
 71            } else {
 72                rt = smtHash([rt, res.sibblings[i]]);
 73                inserts.push([rt,[rt, res.sibblings[i]]] );
 74            }
 75        }
 76
 77        if (addedOne) res.sibblings.pop();
 78        while ((res.sibblings.length>0) && (res.sibblings[res.sibblings.length-1].isZero())) {
 79            res.sibblings.pop();
 80        }
 81        res.oldKey = resFind.notFoundKey;
 82        res.oldValue = resFind.notFoundValue;
 83        res.newRoot = rt;
 84        res.isOld0 = resFind.isOld0;
 85
 86        this.root = rt;
 87
 88        await this.db.save(rt, inserts);
 89
 90        return res;
 91    }
 92
 93    async find(key) {
 94        const keyBits = this._splitBits(key);
 95        return await this._find(key, keyBits, this.root, 0);
 96    }
 97
 98    async _find(key, keyBits, root, level) {
 99        if (typeof root === "undefined") root = this.root;
100
101        let res;
102        if (root.isZero()) {
103            res = {
104                found: false,
105                sibblings: [],
106                notFoundKey: key,
107                notFoundValue: bigInt.zero,
108                isOld0: true
109            };
110            return res;
111        }
112
113        const record = await this.db.get(root);
114
115        if ((record.length==3)&&(record[0].equals(bigInt.one))) {
116            if (record[1].equals(key)) {
117                res = {
118                    found: true,
119                    sibblings: [],
120                    foundValue: record[2],
121                    isOld0: false
122                };
123            } else {
124                res = {
125                    found: false,
126                    sibblings: [],
127                    notFoundKey: record[1],
128                    notFoundValue: record[2],
129                    isOld0: false
130                };
131            }
132        } else {
133            if (keyBits[level] == 0) {
134                const res = await this.find(key, keyBits, record[1], level+1);
135                res.sibblings.unshift(record[2]);
136            } else {
137                res = await this.find(key, keyBits, record[2], level+1);
138                res.sibblings.unshift(record[1]);
139            }
140        }
141        return res;
142    }
143}
144
145async function loadFromFile(fileName) {
146
147}
148
149async function newMemEmptyTrie() {
150    const db = new SMTMemDB();
151    const rt = await db.getRoot();
152    const smt = new SMT(db, rt);
153    return smt;
154}
155
156module.exports.loadFromFile = loadFromFile;
157module.exports.newMemEmptyTrie = newMemEmptyTrie;