import * as AsyncLock from 'async-lock' ; import { DB , SparseMerkleTree , SparseMerkleTreeImpl , MerkleTreeInclusionProof , } from '@pigi/core-db' ; import { DefaultSignatureVerifier , serializeObject , SignatureVerifier , BigNumber , ONE , ZERO , getLogger , runInDomain , NULL_ADDRESS , } from '@pigi/core-utils' ; import { Address , Balances , Swap , isSwapTransaction , Transfer , isTransferTransaction , RollupTransaction , SignedTransaction , TokenType , State , StateUpdate , StateSnapshot , InclusionProof , StateMachineCapacityError , SignatureError , abiEncodeTransaction , abiEncodeState , parseStateFromABI , } from '../index' ; import { UNISWAP_ADDRESS , UNI_TOKEN_TYPE , PIGI_TOKEN_TYPE , NON_EXISTENT_SLOT_INDEX , } from './utils' ; import { InsufficientBalanceError , InvalidTransactionTypeError , NegativeAmountError , RollupStateMachine , SlippageError , } from '../types' ; const log = getLogger ( 'rollup-state-machine' ) ; /**
 * A Tree-backed Rollup State Machine, facilitating state transitions for
 * swaps and transactions for Uniswap.
 */ export class DefaultRollupStateMachine implements RollupStateMachine { public static readonly ROOT_KEY : Buffer = Buffer . from ( 'state_machine_root' ) ; public static readonly LAST_OPEN_KEY : Buffer = Buffer . from ( 'last_open_key' ) ; public static readonly ADDRESS_TO_KEYS_COUNT_KEY : Buffer = Buffer . from ( 'address_to_keys_count' ) ; private static readonly lockKey : string = 'lock' ; private lastOpenKey : BigNumber ; private usedKeys : Set < string > ; private addressesToKeys : Map < Address , BigNumber > ; private maxAddresses : BigNumber ; private tree : SparseMerkleTree ; private readonly lock : AsyncLock ; param genesisState param db param aggregatorAddress param signatureVerifier param swapFeeBasisPoints param treeHeight returns public static async create ( genesisState , db , aggregatorAddress , signatureVerifier = DefaultSignatureVerifier . instance ( ) , swapFeeBasisPoints = 30 , treeHeight = 32 ) { const stateMachine = new DefaultRollupStateMachine ( db , aggregatorAddress , signatureVerifier , swapFeeBasisPoints , treeHeight ) ; const previousStateExists = await stateMachine . init ( ) ; if ( ! previousStateExists && ! ! genesisState . length ) { for ( const state of genesisState ) { await stateMachine . setAddressState ( state . pubkey , state . balances ) ; } } return stateMachine ; } private constructor ( private readonly db , private readonly aggregatorAddress , private readonly signatureVerifier , private readonly swapFeeBasisPoints , private readonly treeHeight = 32 ) { this . maxAddresses = new BigNumber ( Math . pow ( 2 , this . treeHeight ) - 1 ) ; this . lock = new AsyncLock ( { domainReentrant : true , } ) ; } returns private async init ( ) { const storedRoot = await this . db . get ( DefaultRollupStateMachine . ROOT_KEY ) ; this . tree = await SparseMerkleTreeImpl . create ( this . db , storedRoot , this . treeHeight ) ; this . addressesToKeys = new Map < Address , BigNumber > ( ) ; this . usedKeys = new Set < string > ( ) ; this . lastOpenKey = ZERO ; if ( ! storedRoot ) { log . info ( `No existing state root found, starting RollupStateMachine fresh` ) ; return false ; } const [ lastKeyBuffer , addressToKeyCountBuffer ] = await Promise . all ( [ this . db . get ( DefaultRollupStateMachine . LAST_OPEN_KEY ) , this . db . get ( DefaultRollupStateMachine . ADDRESS_TO_KEYS_COUNT_KEY ) , ] ) ; this . lastOpenKey = new BigNumber ( lastKeyBuffer ) ; const addressCount = parseInt ( addressToKeyCountBuffer . toString ( ) , 10 ) ; log . info ( `RollupStateMachine found root ${ storedRoot . toString ( 'hex' ) }, last open key: ${ this . lastOpenKey . toString ( 10 ) }, and address key count: ${ addressCount }. Initializing with this state.` ) ; const addressPromises = [ ] ; for ( let i = 0 ; i < addressCount ; i ++ ) { addressPromises . push ( this . db . get ( DefaultRollupStateMachine . getAddressMapDBKey ( i ) ) ) ; } const addressToKeysBuffers = await Promise . all ( addressPromises ) ; for ( const addressKeyBuf of addressToKeysBuffers ) { const addressAndKey = DefaultRollupStateMachine . deserializeAddressToKeyFromDB ( addressKeyBuf ) ; this . addressesToKeys . set ( addressAndKey [ 0 ] , addressAndKey [ 1 ] ) ; } for ( const key of this . addressesToKeys . values ( ) ) { this . usedKeys . add ( key . toString ( ) ) ; if ( key . gt ( this . lastOpenKey ) ) { this . lastOpenKey = key ; } } return true ; } param address returns public async getState ( address ) { const [ accountState , proof , stateRoot ] = await this . lock . acquire ( DefaultRollupStateMachine . lockKey , async ( ) => { const key = this . getAddressKey ( address ) ; if ( ! ! key ) { const leaf = await this . tree . getLeaf ( key ) ; if ( ! ! leaf ) { const merkleProof = await this . tree . getMerkleProof ( key , leaf ) ; return [ leaf , merkleProof , merkleProof . rootHash . toString ( 'hex' ) ] ; } } return [ undefined , undefined , ( await this . tree . getRootHash ( ) ) . toString ( 'hex' ) , ] ; } ) ; let state ; let inclusionProof ; let slotIndex ; if ( ! accountState ) { state = undefined ; inclusionProof = undefined ; slotIndex = NON_EXISTENT_SLOT_INDEX ; } else { state = DefaultRollupStateMachine . deserializeState ( accountState ) ; inclusionProof = proof . siblings . map ( ( x ) => x . toString ( 'hex' ) ) ; slotIndex = this . getAddressKey ( address ) . toNumber ( ) ; } return { slotIndex , state , stateRoot , inclusionProof , } ; } param transactions returns public async applyTransactions ( transactions ) { return runInDomain ( undefined , async ( ) => { return this . lock . acquire ( DefaultRollupStateMachine . lockKey , async ( ) => { const stateUpdates = [ ] ; for ( const tx of transactions ) { stateUpdates . push ( await this . applyTransaction ( tx ) ) ; } return stateUpdates ; } ) ; } ) ; } param signedTransaction returns public async applyTransaction ( signedTransaction ) { let signer ; log . debug ( `Validating Signature: ${ JSON . stringify ( signedTransaction ) }` ) ; signer = this . signatureVerifier . verifyMessage ( abiEncodeTransaction ( signedTransaction . transaction ) , signedTransaction . signature ) ; if ( signer !== signedTransaction . transaction . sender && signer !== this . aggregatorAddress ) { log . info ( `Received transaction with invalid signature: ${ serializeObject ( signedTransaction ) }, which recovered a signer of ${ signer }` ) ; throw new SignatureError ( ) ; } return this . lock . acquire ( DefaultRollupStateMachine . lockKey , async ( ) => { const stateUpdate = { transaction : signedTransaction } ; const transaction = signedTransaction . transaction ; let updatedStates ; if ( isTransferTransaction ( transaction ) ) { stateUpdate [ 'receiverCreated' ] = ! this . getAddressKey ( transaction . recipient ) ; updatedStates = await this . applyTransfer ( transaction ) ; stateUpdate [ 'receiverSlotIndex' ] = this . getAddressKey ( transaction . recipient ) . toNumber ( ) ; } else if ( isSwapTransaction ( transaction ) ) { updatedStates = await this . applySwap ( signer , transaction ) ; stateUpdate [ 'receiverCreated' ] = false ; stateUpdate [ 'receiverSlotIndex' ] = this . getAddressKey ( UNISWAP_ADDRESS ) . toNumber ( ) ; } else { throw new InvalidTransactionTypeError ( ) ; } const root = await this . tree . getRootHash ( ) ; await this . db . put ( DefaultRollupStateMachine . ROOT_KEY , root ) ; const senderState = updatedStates [ 0 ] ; const receiverState = updatedStates [ 1 ] ; stateUpdate [ 'senderSlotIndex' ] = this . getAddressKey ( transaction . sender ) . toNumber ( ) ; stateUpdate [ 'senderState' ] = senderState ; stateUpdate [ 'receiverState' ] = receiverState ; const inclusionProof = async ( state ) : Promise < InclusionProof > => { const proof = await this . tree . getMerkleProof ( this . getAddressKey ( state . pubkey ) , DefaultRollupStateMachine . serializeBalances ( state . pubkey , state . balances ) ) ; return proof . siblings . map ( ( p ) => p . toString ( 'hex' ) ) ; } ; [ stateUpdate [ 'senderStateInclusionProof' ] , stateUpdate [ 'receiverStateInclusionProof' ] , ] = await Promise . all ( [ inclusionProof ( senderState ) , inclusionProof ( receiverState ) , ] ) ; stateUpdate [ 'stateRoot' ] = ( await this . tree . getRootHash ( ) ) . toString ( 'hex' ) ; return stateUpdate ; } ) ; } public async getStateRoot ( ) { const lockedRoot = await this . lock . acquire ( DefaultRollupStateMachine . lockKey , async ( ) => { return this . tree . getRootHash ( ) ; } ) ; return lockedRoot ; } public getNextNewAccountSlot ( ) { return this . lastOpenKey . toNumber ( ) + 1 ; } public async getSnapshotFromSlot ( key ) { const [ accountState , proof , stateRoot ] = await this . lock . acquire ( DefaultRollupStateMachine . lockKey , async ( ) => { const bigKey = new BigNumber ( key , 10 ) ; let leaf = await this . tree . getLeaf ( bigKey ) ; if ( ! leaf || leaf . equals ( SparseMerkleTreeImpl . siblingBuffer ) ) { leaf = SparseMerkleTreeImpl . emptyBuffer ; } const merkleProof = await this . tree . getMerkleProof ( bigKey , leaf ) ; if ( ! merkleProof ) { const msg = `Unable to get merkle proof for key ${ bigKey . toNumber ( ) }. Leaf data: [${ ! leaf ? 'undefined' : leaf . toString ( ) }]` ; log . error ( msg ) ; throw Error ( msg ) ; } return [ leaf , merkleProof , merkleProof . rootHash . toString ( 'hex' ) ] ; } ) ; let state ; let inclusionProof ; state = ! ! accountState && ! accountState . equals ( SparseMerkleTreeImpl . emptyBuffer ) ? DefaultRollupStateMachine . deserializeState ( accountState ) : { pubkey : NULL_ADDRESS , balances : { [ UNI_TOKEN_TYPE ] : 0 , [ PIGI_TOKEN_TYPE ] : 0 } , } ; inclusionProof = proof . siblings . map ( ( x ) => x . toString ( 'hex' ) ) ; return { slotIndex : key , state , stateRoot , inclusionProof , } ; } param address returns private async getBalances ( address ) { const key = this . getAddressKey ( address ) ; if ( ! ! key ) { const leaf = await this . tree . getLeaf ( key ) ; if ( ! ! leaf ) { return DefaultRollupStateMachine . deserializeState ( leaf ) . balances ; } } return { [ UNI_TOKEN_TYPE ] : 0 , [ PIGI_TOKEN_TYPE ] : 0 } ; } param address param balances returns private async setAddressState ( address , balances ) { const addressKey = await this . getOrCreateAddressKey ( address ) ; const serializedBalances = DefaultRollupStateMachine . serializeBalances ( address , balances ) ; const result = await this . tree . update ( addressKey , serializedBalances ) ; if ( ! result ) { log . error ( `ERROR UPDATING TREE, address: [${ address }], key: [${ addressKey }], balances: [${ serializeObject ( balances ) }]` ) ; } else { log . debug ( `${ address } with key ${ addressKey } balance updated to ${ serializeObject ( balances ) }` ) ; } return result ; } param address param tokenType param balance returns private async hasBalance ( address , tokenType , balance ) { const balances = await this . getBalances ( address ) ; return tokenType in balances && balances [ tokenType ] >= balance ; } param transfer returns private async applyTransfer ( transfer ) { if ( transfer . amount < 1 ) { throw new NegativeAmountError ( ) ; } if ( ! ( await this . hasBalance ( transfer . sender , transfer . tokenType , transfer . amount ) ) ) { throw new InsufficientBalanceError ( ) ; } const senderBalances = await this . getBalances ( transfer . sender ) ; const recipientBalances = await this . getBalances ( transfer . recipient ) ; senderBalances [ transfer . tokenType ] -= transfer . amount ; recipientBalances [ transfer . tokenType ] += transfer . amount ; await Promise . all ( [ this . setAddressState ( transfer . sender , senderBalances ) , this . setAddressState ( transfer . recipient , recipientBalances ) , ] ) ; return [ DefaultRollupStateMachine . getStateFromBalances ( transfer . sender , senderBalances ) , DefaultRollupStateMachine . getStateFromBalances ( transfer . recipient , recipientBalances ) , ] ; } param sender param swap returns private async applySwap ( sender , swap ) { if ( swap . inputAmount < 1 ) { throw new NegativeAmountError ( ) ; } if ( ! this . hasBalance ( sender , swap . tokenType , swap . inputAmount ) ) { throw new InsufficientBalanceError ( ) ; } return this . updateBalancesFromSwap ( swap , sender ) ; } param swap param sender returns private async updateBalancesFromSwap ( swap , sender ) { const uniswapBalances = await this . getBalances ( UNISWAP_ADDRESS ) ; const inputTokenType = swap . tokenType ; const outputTokenType = swap . tokenType === UNI_TOKEN_TYPE ? PIGI_TOKEN_TYPE : UNI_TOKEN_TYPE ; const invariant = uniswapBalances [ UNI_TOKEN_TYPE ] * uniswapBalances [ PIGI_TOKEN_TYPE ] ; const totalInput = this . assessSwapFee ( swap . inputAmount ) + uniswapBalances [ inputTokenType ] ; const newOutputBalance = Math . ceil ( invariant / totalInput ) ; const outputAmount = uniswapBalances [ outputTokenType ] - newOutputBalance ; if ( outputAmount < swap . minOutputAmount ) { throw new SlippageError ( ) ; } const senderBalances = await this . getBalances ( sender ) ; senderBalances [ inputTokenType ] -= swap . inputAmount ; senderBalances [ outputTokenType ] += outputAmount ; uniswapBalances [ inputTokenType ] += swap . inputAmount ; uniswapBalances [ outputTokenType ] = newOutputBalance ; await Promise . all ( [ this . setAddressState ( sender , senderBalances ) , this . setAddressState ( UNISWAP_ADDRESS , uniswapBalances ) , ] ) ; return [ DefaultRollupStateMachine . getStateFromBalances ( sender , senderBalances ) , DefaultRollupStateMachine . getStateFromBalances ( UNISWAP_ADDRESS , uniswapBalances ) , ] ; } param amountBeforeFee return private assessSwapFee ( amountBeforeFee ) { if ( this . swapFeeBasisPoints === 0 ) { return amountBeforeFee ; } return amountBeforeFee * ( ( 10000 - this . swapFeeBasisPoints ) / 10000 ) ; } private getAddressKey ( address ) { return this . addressesToKeys . get ( address ) ; } private async getOrCreateAddressKey ( address ) { const existingKey = this . getAddressKey ( address ) ; if ( ! ! existingKey ) { return existingKey ; } let newKey = this . lastOpenKey . toString ( ) ; while ( this . usedKeys . has ( newKey ) ) { this . lastOpenKey = this . lastOpenKey . add ( ONE ) ; if ( this . lastOpenKey . gt ( this . maxAddresses ) ) { throw new StateMachineCapacityError ( ) ; } newKey = this . lastOpenKey . toString ( ) ; } this . addressesToKeys . set ( address , this . lastOpenKey ) ; this . usedKeys . add ( newKey ) ; await this . db . put ( DefaultRollupStateMachine . getAddressMapDBKey ( this . addressesToKeys . size - 1 ) , DefaultRollupStateMachine . serializeAddressToKeyForDB ( address , this . lastOpenKey ) ) ; await Promise . all ( [ this . db . put ( DefaultRollupStateMachine . ADDRESS_TO_KEYS_COUNT_KEY , Buffer . from ( this . addressesToKeys . size . toString ( 10 ) ) ) , this . db . put ( DefaultRollupStateMachine . LAST_OPEN_KEY , this . lastOpenKey . toBuffer ( ) ) , ] ) ; return this . addressesToKeys . get ( address ) ; } /********************
     * STATIC UTILITIES *
     ********************/ public static serializeBalances ( address , balances ) { return Buffer . from ( abiEncodeState ( DefaultRollupStateMachine . getStateFromBalances ( address , balances ) ) ) ; } public static deserializeState ( state ) { return parseStateFromABI ( state . toString ( ) ) ; } public static getStateFromBalances ( pubKey , balances ) { return { pubkey : pubKey , balances , } ; } public static getAddressMapDBKey ( index ) { return Buffer . from ( `ADDR_IDX_${ index }` ) ; } public static serializeAddressToKeyForDB ( address , key ) { return Buffer . from ( JSON . stringify ( [ address , key . toString ( ) ] ) ) ; } public static deserializeAddressToKeyFromDB ( buf ) { const parsed = JSON . parse ( buf . toString ( ) ) ; return [ parsed [ 0 ] , new BigNumber ( parsed [ 1 ] , 'hex' ) ] ; } /***********
     * GETTERS *
     ***********/ public getLastOpenKey ( ) { return this . lastOpenKey . clone ( ) ; } public getUsedKeys ( ) { return new Set < string > ( this . usedKeys ) ; } public getAddressesToKeys ( ) { return new Map < Address , BigNumber > ( this . addressesToKeys ) ; } }