import { DEFAULT_KEY_LENGTH , EMPTY_HASH , EMPTY_VALUE , NodeSide } from './constants' ; import { Leaf } from './leaf' ; import { Database , Proof , Query } from './types' ; import { parseBranchData , parseLeafData , isLeaf , binaryExpansion , sortByBitmapAndKey , binaryStringToBuffer , bufferToBinaryString , treeSort , } from './utils' ; import { binarySearch } from '../utils' ; import { Branch } from './branch' ; import { Empty } from './empty' ; type TreeNode = Branch | Leaf | Empty ; type SingleProof = { key : Buffer ; value : Buffer ; binaryBitmap : string ; ancestorHashes : Buffer [ ] ; siblingHashes : Buffer [ ] ; } ; type QueryWithHeight = { key : Buffer ; value : Buffer ; binaryBitmap : string ; siblingHashes : Buffer [ ] ; height : number ; } ; export class SparseMerkleTree { private readonly _db : Database ; private readonly _keyLength : number ; private _rootHash : Buffer ; public constructor ( options ) { this . _db = options . db ; this . _keyLength = options . keyLength ?? DEFAULT_KEY_LENGTH ; this . _rootHash = options . rootHash ?? EMPTY_HASH ; } public get rootHash ( ) : Buffer { return this . _rootHash ; } public get keyLength ( ) : number { return this . _keyLength ; } public async getNode ( nodeHash ) { if ( nodeHash . equals ( EMPTY_HASH ) ) { return new Empty ( ) ; } const data = await this . _db . get ( nodeHash ) ; if ( ! data ) { throw new Error ( `Node with input hash: ${ nodeHash . toString ( 'hex' ) } does not exist in the tree` ) ; } if ( isLeaf ( data ) ) { const { key , value } = parseLeafData ( data , this . keyLength ) ; return new Leaf ( key , value ) ; } const { leftHash , rightHash } = parseBranchData ( data ) ; return new Branch ( leftHash , rightHash ) ; } public async update ( key , value ) { if ( value . length === 0 ) { throw new Error ( 'Value cannot be empty' ) ; } if ( key . byteLength !== this . keyLength ) { throw new Error ( `Key is not equal to defined key length of ${ this . keyLength }` ) ; } let rootNode = await this . getNode ( this . _rootHash ) ; let currentNode = rootNode ; const newLeaf = new Leaf ( key , value ) ; await this . _db . set ( newLeaf . hash , newLeaf . data ) ; const binaryKey = binaryExpansion ( key , this . keyLength ) ; if ( currentNode instanceof Empty ) { rootNode = newLeaf ; this . _rootHash = rootNode . hash ; return rootNode ; } let h = 0 ; const ancestorNodes = [ ] ; while ( currentNode instanceof Branch ) { const d = binaryKey . charAt ( h ) ; ancestorNodes . push ( currentNode ) ; if ( d === '0' ) { currentNode = await this . getNode ( currentNode . leftHash ) ; } else if ( d === '1' ) { currentNode = await this . getNode ( currentNode . rightHash ) ; } h += 1 ; } let bottomNode = new Empty ( ) ; if ( currentNode instanceof Empty ) { bottomNode = newLeaf ; } else if ( currentNode . key === key ) { bottomNode = newLeaf ; } else { const currentNodeBinaryKey = binaryExpansion ( currentNode . key , this . keyLength ) ; while ( binaryKey . charAt ( h ) === currentNodeBinaryKey . charAt ( h ) ) { const newBranch = new Branch ( EMPTY_HASH , EMPTY_HASH ) ; ancestorNodes . push ( newBranch ) ; h += 1 ; } const d = binaryKey . charAt ( h ) ; if ( d === '0' ) { bottomNode = new Branch ( newLeaf . hash , currentNode . hash ) ; await this . _db . set ( bottomNode . hash , bottomNode . data ) ; } else if ( d === '1' ) { bottomNode = new Branch ( currentNode . hash , newLeaf . hash ) ; await this . _db . set ( bottomNode . hash , bottomNode . data ) ; } } while ( h > 0 ) { const p = ancestorNodes [ h - 1 ] ; const d = binaryKey . charAt ( h - 1 ) ; if ( d === '0' ) { ( p as Branch ) . update ( bottomNode . hash , NodeSide . LEFT ) ; } else if ( d === '1' ) { ( p as Branch ) . update ( bottomNode . hash , NodeSide . RIGHT ) ; } await this . _db . set ( p . hash , ( p as Branch ) . data ) ; bottomNode = p ; h -= 1 ; } this . _rootHash = bottomNode . hash ; return bottomNode ; } public async remove ( key ) { if ( key . length !== this . keyLength ) { throw new Error ( `Key is not equal to defined key length of ${ this . keyLength }` ) ; } const ancestorNodes = [ ] ; const binaryKey = binaryExpansion ( key , this . keyLength ) ; let currentNode = await this . getNode ( this . _rootHash ) ; let h = 0 ; let currentNodeSibling = new Empty ( ) ; while ( currentNode instanceof Branch ) { ancestorNodes . push ( currentNode ) ; const d = binaryKey [ h ] ; if ( d === '0' ) { currentNodeSibling = await this . getNode ( currentNode . rightHash ) ; currentNode = await this . getNode ( currentNode . leftHash ) ; } else if ( d === '1' ) { currentNodeSibling = await this . getNode ( currentNode . leftHash ) ; currentNode = await this . getNode ( currentNode . rightHash ) ; } h += 1 ; } if ( currentNode instanceof Empty ) { return currentNode ; } if ( ! currentNode . key . equals ( key ) ) { return currentNode ; } let bottomNode = new Empty ( ) ; if ( currentNodeSibling instanceof Branch ) { await this . _db . del ( currentNode . hash ) ; } else if ( currentNodeSibling instanceof Leaf ) { await this . _db . del ( currentNode . hash ) ; bottomNode = currentNodeSibling ; h -= 1 ; const emptyHash = new Empty ( ) . hash ; while ( h > 0 ) { const p = ancestorNodes [ h - 1 ] as Branch ; if ( p instanceof Branch && ! p . leftHash . equals ( emptyHash ) && ! p . rightHash . equals ( emptyHash ) ) { break ; } await this . _db . del ( p . hash ) ; h -= 1 ; } } while ( h > 0 ) { const p = ancestorNodes [ h - 1 ] ; const d = binaryKey . charAt ( h - 1 ) ; if ( d === '0' ) { ( p as Branch ) . update ( bottomNode . hash , NodeSide . LEFT ) ; } else if ( d === '1' ) { ( p as Branch ) . update ( bottomNode . hash , NodeSide . RIGHT ) ; } await this . _db . set ( p . hash , ( p as Branch ) . data ) ; bottomNode = p ; h -= 1 ; } this . _rootHash = bottomNode . hash ; return bottomNode ; } public async generateSingleProof ( queryKey ) { const rootNode = await this . getNode ( this . _rootHash ) ; let currentNode = rootNode ; if ( currentNode instanceof Empty ) { return { key : queryKey , value : EMPTY_VALUE , binaryBitmap : bufferToBinaryString ( EMPTY_VALUE ) , siblingHashes : [ ] , ancestorHashes : [ ] , } ; } let h = 0 ; const siblingHashes = [ ] ; const ancestorHashes = [ ] ; let binaryBitmap = '' ; const binaryKey = binaryExpansion ( queryKey , this . keyLength ) ; while ( currentNode instanceof Branch ) { ancestorHashes . push ( currentNode . hash ) ; const d = binaryKey . charAt ( h ) ; let currentNodeSibling = new Empty ( ) ; if ( d === '0' ) { currentNodeSibling = await this . getNode ( currentNode . rightHash ) ; currentNode = await this . getNode ( currentNode . leftHash ) ; } else if ( d === '1' ) { currentNodeSibling = await this . getNode ( currentNode . leftHash ) ; currentNode = await this . getNode ( currentNode . rightHash ) ; } if ( currentNodeSibling instanceof Empty ) { binaryBitmap = `0${ binaryBitmap }` ; } else { binaryBitmap = `1${ binaryBitmap }` ; siblingHashes . push ( currentNodeSibling . hash ) ; } h += 1 ; } if ( currentNode instanceof Empty ) { return { siblingHashes , ancestorHashes , binaryBitmap , key : queryKey , value : EMPTY_VALUE , } ; } if ( ! currentNode . key . equals ( queryKey ) ) { ancestorHashes . push ( currentNode . hash ) ; return { siblingHashes , ancestorHashes , binaryBitmap , key : currentNode . key , value : currentNode . value , } ; } ancestorHashes . push ( currentNode . hash ) ; return { siblingHashes , ancestorHashes , binaryBitmap , key : currentNode . key , value : currentNode . value , } ; } public async generateMultiProof ( queryKeys ) { const partialQueries = [ ] ; for ( const queryKey of queryKeys ) { const query = await this . generateSingleProof ( queryKey ) ; partialQueries . push ( query ) ; } const queries = [ ... partialQueries ] . map ( sp => ( { bitmap : binaryStringToBuffer ( sp . binaryBitmap ) , key : sp . key , value : sp . value , } ) ) ; const siblingHashes = [ ] ; const ancestorHashes = [ ... partialQueries ] . map ( sp => sp . ancestorHashes ) . flat ( ) ; let sortedQueries = [ ... partialQueries ] . map ( sp => ( { binaryBitmap : sp . binaryBitmap , key : sp . key , value : sp . value , siblingHashes : sp . siblingHashes , height : sp . binaryBitmap . length , } ) ) ; sortedQueries = sortByBitmapAndKey ( sortedQueries ) ; while ( sortedQueries . length > 0 ) { const sp = sortedQueries . shift ( ) ! ; if ( sp . height === 0 ) { continue ; } const b = sp . binaryBitmap . charAt ( sp . binaryBitmap . length - sp . height ) ; if ( b === '1' ) { const nodeHash = sp . siblingHashes . pop ( ) ! ; let isPresentInSiblingHashes = false ; let isPresentInAncestorHashes = false ; for ( const i of siblingHashes ) { if ( i . equals ( nodeHash ) ) { isPresentInSiblingHashes = true ; break ; } } for ( const i of ancestorHashes ) { if ( i . equals ( nodeHash ) ) { isPresentInAncestorHashes = true ; break ; } } if ( ! isPresentInSiblingHashes && ! isPresentInAncestorHashes ) { siblingHashes . push ( nodeHash ) ; } } sp . height -= 1 ; if ( sortedQueries . length > 0 ) { const sortedQueriesWithBinaryKey = sortedQueries . map ( query => ( { binaryKey : binaryExpansion ( query . key , this . keyLength ) , binaryBitmap : query . binaryBitmap , value : query . value , siblingHashes : query . siblingHashes , height : query . height , } ) ) ; const spWithBinaryKey = { binaryKey : binaryExpansion ( sp . key , this . keyLength ) , binaryBitmap : sp . binaryBitmap , value : sp . value , siblingHashes : sp . siblingHashes , height : sp . height , } ; const insertIndex = binarySearch ( sortedQueriesWithBinaryKey , callback => treeSort ( spWithBinaryKey , callback ) < 0 ) ; if ( insertIndex === sortedQueries . length ) { sortedQueries . push ( sp ) ; } else { const keyPrefix = binaryExpansion ( sp . key , this . keyLength ) . substring ( 0 , sp . height ) ; const query = sortedQueries [ insertIndex ] ; if ( ! binaryExpansion ( query . key , this . keyLength ) . endsWith ( keyPrefix , query . height ) ) { sortedQueries . splice ( insertIndex , 0 , sp ) ; } } } else { sortedQueries . push ( sp ) ; } } return { siblingHashes , queries } ; } }