snarky-smt
Version:
Sparse Merkle Tree for SnarkyJS
193 lines (192 loc) • 6.78 kB
JavaScript
import { Field, Poseidon } from 'snarkyjs';
import { EMPTY_VALUE } from '../constant';
import { MerkleTreeUtils } from './proofs';
import { ProvableMerkleTreeUtils } from './verify_circuit';
export { DeepMerkleSubTree };
/**
* DeepMerkleSubTree is a deep merkle subtree for working on only a few leafs.
*
* @class DeepMerkleSubTree
* @template V
*/
class DeepMerkleSubTree {
/**
* Creates an instance of DeepMerkleSubTree.
* @param {Field} root
* @param {number} height
* @param {Provable<V>} valueType
* @param {{ hasher: Hasher; hashValue: boolean }} [options={
* hasher: Poseidon.hash,
* hashValue: true,
* }] hasher: The hash function to use, defaults to Poseidon.hash;
* hashValue: whether to hash the value, the default is true.
* @memberof DeepMerkleSubTree
*/
constructor(root, height, valueType, options = {
hasher: Poseidon.hash,
hashValue: true,
}) {
this.root = root;
this.nodeStore = new Map();
this.valueStore = new Map();
this.height = height;
this.hasher = options.hasher;
this.hashValue = options.hashValue;
this.valueType = valueType;
}
/**
* Get current root.
*
* @return {*} {Field}
* @memberof DeepMerkleSubTree
*/
getRoot() {
return this.root;
}
/**
* Get height of the tree.
*
* @return {*} {number}
* @memberof DeepMerkleSubTree
*/
getHeight() {
return this.height;
}
getValueField(value) {
let valueHashOrValueField = EMPTY_VALUE;
if (value !== undefined) {
let valueFields = this.valueType.toFields(value);
valueHashOrValueField = valueFields[0];
if (this.hashValue) {
valueHashOrValueField = this.hasher(valueFields);
}
}
return valueHashOrValueField;
}
/**
* Check whether there is a corresponding key and value in the tree
*
* @param {bigint} index
* @param {V} value
* @return {*} {boolean}
* @memberof DeepMerkleSubTree
*/
has(index, value) {
const path = Field(index);
const valueField = this.getValueField(value);
let v = this.valueStore.get(path.toString());
if (v === undefined || !v.equals(valueField).toBoolean()) {
return false;
}
return true;
}
/**
* Add a branch to the tree, a branch is generated by smt.prove.
*
* @param {BaseMerkleProof} proof
* @param {bigint} index
* @param {V} [value]
* @param {boolean} [ignoreInvalidProof=false]
* @return {*}
* @memberof DeepMerkleSubTree
*/
addBranch(proof, index, value, ignoreInvalidProof = false) {
const path = Field(index);
const valueField = this.getValueField(value);
let { ok, updates } = MerkleTreeUtils.verifyProofByFieldWithUpdates(proof, this.root, index, valueField, this.hasher);
if (!ok) {
if (!ignoreInvalidProof) {
throw new Error(`invalid proof, proof path: ${path.toString()}, valueField: ${valueField.toString()}`);
}
else {
return;
}
}
for (let i = 0, len = updates.length; i < len; i++) {
let v = updates[i];
this.nodeStore.set(v[0].toString(), v[1]);
}
this.valueStore.set(path.toString(), valueField);
}
/**
* Create a merkle proof for a key against the current root.
*
* @param {bigint} index
* @return {*} {BaseMerkleProof}
* @memberof DeepMerkleSubTree
*/
prove(index) {
const path = Field(index);
let pathStr = path.toString();
let valueHash = this.valueStore.get(pathStr);
if (valueHash === undefined) {
throw new Error(`The DeepSubTree does not contain a branch of the path: ${pathStr}`);
}
const pathBits = path.toBits(this.height);
let sideNodes = [];
let nodeHash = this.root;
for (let i = 0; i < this.height; i++) {
const currentValue = this.nodeStore.get(nodeHash.toString());
if (currentValue === undefined) {
throw new Error('Make sure you have added the correct proof, key and value using the addBranch method');
}
if (pathBits[i].toBoolean()) {
sideNodes.push(currentValue[0]);
nodeHash = currentValue[1];
}
else {
sideNodes.push(currentValue[1]);
nodeHash = currentValue[0];
}
}
class MerkleProof_ extends ProvableMerkleTreeUtils.MerkleProof(this.height) {
}
return new MerkleProof_(this.root, sideNodes);
}
/**
* Update a new value for a key in the tree and return the new root of the tree.
*
* @param {bigint} index
* @param {V} [value]
* @return {*} {Field}
* @memberof DeepMerkleSubTree
*/
update(index, value) {
const path = Field(index);
const valueField = this.getValueField(value);
const pathBits = path.toBits(this.height);
let sideNodes = [];
let nodeHash = this.root;
for (let i = 0; i < this.height; i++) {
const currentValue = this.nodeStore.get(nodeHash.toString());
if (currentValue === undefined) {
throw new Error('Make sure you have added the correct proof, key and value using the addBranch method');
}
if (pathBits[i].toBoolean()) {
sideNodes.push(currentValue[0]);
nodeHash = currentValue[1];
}
else {
sideNodes.push(currentValue[1]);
nodeHash = currentValue[0];
}
}
let currentHash = valueField;
this.nodeStore.set(currentHash.toString(), [currentHash]);
for (let i = this.height - 1; i >= 0; i--) {
let sideNode = sideNodes[i];
let currentValue = [];
if (pathBits[i].toBoolean()) {
currentValue = [sideNode, currentHash];
}
else {
currentValue = [currentHash, sideNode];
}
currentHash = this.hasher(currentValue);
this.nodeStore.set(currentHash.toString(), currentValue);
}
this.valueStore.set(path.toString(), valueField);
this.root = currentHash;
return this.root;
}
}