import { EVMResult } from '@nomiclabs/ethereumjs-vm/dist/evm/evm' ; import { VmError } from '@nomiclabs/ethereumjs-vm/dist/exceptions' ; import BN from 'bn.js' ; import { ethers } from 'ethers' ; import { findLast } from 'lodash' ; import { Observable , withLatestFrom } from 'rxjs' ; import { ContractCall , ProgrammedReturnValue , WhenCalledWithChain } from '../index' ; import { WatchableFunctionLogic } from '../logic/watchable-function-logic' ; import { fromHexString } from '../utils' ; const EMPTY_ANSWER = fromHexString ( '0x' + '00' . repeat ( 2048 ) ) ; class ProgrammedAnswer { value ? : any ; shouldRevert : boolean ; constructor ( value ? , shouldRevert = false ) { this . value = value ; this . shouldRevert = shouldRevert ; } } export class ProgrammableFunctionLogic extends WatchableFunctionLogic { protected encoder : ( value : any ) => string ; protected defaultAnswer : ProgrammedAnswer | undefined ; protected answerByIndex : { [ index ] : ProgrammedAnswer ; } = { } ; protected answerByArgs : { answer : ProgrammedAnswer ; args : unknown [ ] ; } [ ] = [ ] ; constructor ( name , calls$ , results$ , encoder ) { super ( name , calls$ ) ; this . encoder = encoder ; results$ . pipe ( withLatestFrom ( calls$ ) ) . subscribe ( async ( [ result , call ] ) => { await this . modifyAnswer ( result , call ) ; } ) ; } returns ( value ? ) { this . defaultAnswer = new ProgrammedAnswer ( value , false ) ; } returnsAtCall ( callIndex , value ? ) { this . answerByIndex [ callIndex ] = new ProgrammedAnswer ( value , false ) ; } whenCalledWith ( ... args ) { return { returns : ( value ? ) => { this . answerByArgs . push ( { args , answer : new ProgrammedAnswer ( value , false ) , } ) ; } , reverts : ( reason ? ) => { this . answerByArgs . push ( { args , answer : new ProgrammedAnswer ( reason , true ) , } ) ; } , } ; } reverts ( reason ? ) { this . defaultAnswer = new ProgrammedAnswer ( reason , true ) ; } revertsAtCall ( callIndex , reason ? ) { this . answerByIndex [ callIndex ] = new ProgrammedAnswer ( reason , true ) ; } reset ( ) { super . reset ( ) ; this . defaultAnswer = undefined ; this . answerByIndex = { } ; this . answerByArgs = [ ] ; } private async modifyAnswer ( result , call ) { const answer = this . getCallAnswer ( call ) ; if ( answer ) { result . gasUsed = new BN ( 0 ) ; result . execResult . gasUsed = new BN ( 0 ) ; if ( answer . shouldRevert ) { result . execResult . exceptionError = new VmError ( 'smock revert' as any ) ; result . execResult . returnValue = this . encodeRevertReason ( answer . value ) ; } else { result . execResult . exceptionError = undefined ; result . execResult . returnValue = await this . encodeValue ( answer . value , call ) ; } } } private getCallAnswer ( call ) { let answer ; answer = this . answerByIndex [ this . getCallCount ( ) - 1 ] ; if ( answer ) return answer ; answer = findLast ( this . answerByArgs , ( option ) => this . isDeepEqual ( option . args , call . args ) ) ?. answer ; if ( answer ) return answer ; return this . defaultAnswer ; } private async encodeValue ( value , call ) { if ( value === undefined ) return EMPTY_ANSWER ; let toEncode = typeof value === 'function' ? await value ( call . args ) : value ; let encodedReturnValue = '0x' ; try { encodedReturnValue = this . encoder ( toEncode ) ; } catch ( err ) { if ( err . code === 'INVALID_ARGUMENT' ) { if ( typeof toEncode !== 'string' ) { throw new Error ( `Failed to encode return value for ${ this . name }` ) ; } encodedReturnValue = toEncode ; } else { throw err ; } } return fromHexString ( encodedReturnValue ) ; } private encodeRevertReason ( reason ) { if ( reason === undefined ) return EMPTY_ANSWER ; const errorInterface = new ethers . utils . Interface ( [ { inputs : [ { name : '_reason' , type : 'string' , } , ] , name : 'Error' , outputs : [ ] , stateMutability : 'nonpayable' , type : 'function' , } , ] ) ; return fromHexString ( errorInterface . encodeFunctionData ( 'Error' , [ reason ] ) ) ; } } export class SafeProgrammableContract extends ProgrammableFunctionLogic { protected defaultAnswer = new ProgrammedAnswer ( ) ; reset ( ) { super . reset ( ) ; this . defaultAnswer = new ProgrammedAnswer ( ) ; } }