@aztec/simulator 0.55.1 → 0.57.0

package/src/avm/opcodes/instruction_impl.ts

@@ -1,25 +1,6 @@
-import { type AvmContext } from '../avm_context.js';
-import { type MemoryValue } from '../avm_memory_types.js';
 import { OperandType } from '../serialization/instruction_serialization.js';
-import { Addressing } from './addressing_mode.js';
 import { Instruction } from './instruction.js';
 
-/** Wire format that informs deserialization for instructions with two operands. */
-export const TwoOperandWireFormat8 = [
-  OperandType.UINT8,
-  OperandType.UINT8,
-  OperandType.UINT8,
-  OperandType.UINT8,
-  OperandType.UINT8,
-];
-export const TwoOperandWireFormat16 = [
-  OperandType.UINT8,
-  OperandType.UINT8,
-  OperandType.UINT8,
-  OperandType.UINT16,
-  OperandType.UINT16,
-];
-
 /** Wire format that informs deserialization for instructions with three operands. */
 export const ThreeOperandWireFormat8 = [
   OperandType.UINT8,
@@ -38,25 +19,6 @@ export const ThreeOperandWireFormat16 = [
   OperandType.UINT16,
 ];
 
-/**
- * Covers (de)serialization for an instruction with:
- * indirect, inTag, and two operands.
- */
-export abstract class TwoOperandInstruction extends Instruction {
-  // Informs (de)serialization. See Instruction.deserialize.
-  static readonly wireFormat8: OperandType[] = TwoOperandWireFormat8;
-  static readonly wireFormat16: OperandType[] = TwoOperandWireFormat16;
-
-  constructor(
-    protected indirect: number,
-    protected inTag: number,
-    protected aOffset: number,
-    protected dstOffset: number,
-  ) {
-    super();
-  }
-}
-
 /**
  * Covers (de)serialization for an instruction with:
  * indirect, inTag, and three operands.
@@ -75,27 +37,3 @@ export abstract class ThreeOperandInstruction extends Instruction {
     super();
   }
 }
-
-export abstract class GetterInstruction extends Instruction {
-  // Informs (de)serialization. See Instruction.deserialize.
-  static readonly wireFormat: OperandType[] = [OperandType.UINT8, OperandType.UINT8, OperandType.UINT32];
-
-  constructor(protected indirect: number, protected dstOffset: number) {
-    super();
-  }
-
-  public async execute(context: AvmContext): Promise<void> {
-    const memoryOperations = { writes: 1, indirect: this.indirect };
-    const memory = context.machineState.memory.track(this.type);
-    context.machineState.consumeGas(this.gasCost(memoryOperations));
-
-    const [dstOffset] = Addressing.fromWire(this.indirect).resolve([this.dstOffset], memory);
-
-    memory.set(dstOffset, this.getValue(context));
-
-    memory.assert(memoryOperations);
-    context.machineState.incrementPc();
-  }
-
-  protected abstract getValue(env: AvmContext): MemoryValue;
-}

package/src/avm/opcodes/memory.ts

@@ -3,7 +3,6 @@ import { Field, TaggedMemory } from '../avm_memory_types.js';
 import { Opcode, OperandType } from '../serialization/instruction_serialization.js';
 import { Addressing } from './addressing_mode.js';
 import { Instruction } from './instruction.js';
-import { TwoOperandInstruction } from './instruction_impl.js';
 
 export class Set extends Instruction {
   static readonly type: string = 'SET';
@@ -63,85 +62,57 @@ export class Set extends Instruction {
   }
 
   public async execute(context: AvmContext): Promise<void> {
-    const memoryOperations = { writes: 1, indirect: this.indirect };
     const memory = context.machineState.memory.track(this.type);
-    context.machineState.consumeGas(this.gasCost(memoryOperations));
+    context.machineState.consumeGas(this.gasCost());
 
-    const [dstOffset] = Addressing.fromWire(this.indirect).resolve([this.dstOffset], memory);
+    const operands = [this.dstOffset];
+    const addressing = Addressing.fromWire(this.indirect, operands.length);
+    const [dstOffset] = addressing.resolve(operands, memory);
     const res = TaggedMemory.buildFromTagTruncating(this.value, this.inTag);
     memory.set(dstOffset, res);
 
-    memory.assert(memoryOperations);
+    memory.assert({ writes: 1, addressing });
     context.machineState.incrementPc();
   }
 }
 
-export class CMov extends Instruction {
-  static readonly type: string = 'CMOV';
-  static readonly opcode: Opcode = Opcode.CMOV;
-  // Informs (de)serialization. See Instruction.deserialize.
-  static readonly wireFormat: OperandType[] = [
+export class Cast extends Instruction {
+  static readonly type: string = 'CAST';
+  static readonly opcode = Opcode.CAST_8;
+
+  static readonly wireFormat8 = [
+    OperandType.UINT8,
+    OperandType.UINT8,
+    OperandType.UINT8,
     OperandType.UINT8,
     OperandType.UINT8,
-    OperandType.UINT32,
-    OperandType.UINT32,
-    OperandType.UINT32,
-    OperandType.UINT32,
+  ];
+  static readonly wireFormat16 = [
+    OperandType.UINT8,
+    OperandType.UINT8,
+    OperandType.UINT8,
+    OperandType.UINT16,
+    OperandType.UINT16,
   ];
 
-  constructor(
-    private indirect: number,
-    private aOffset: number,
-    private bOffset: number,
-    private condOffset: number,
-    private dstOffset: number,
-  ) {
+  constructor(private indirect: number, private dstTag: number, private srcOffset: number, private dstOffset: number) {
     super();
   }
 
   public async execute(context: AvmContext): Promise<void> {
-    const memoryOperations = { reads: 3, writes: 1, indirect: this.indirect };
-    const memory = context.machineState.memory.track(this.type);
-    context.machineState.consumeGas(this.gasCost(memoryOperations));
-
-    const [aOffset, bOffset, condOffset, dstOffset] = Addressing.fromWire(this.indirect).resolve(
-      [this.aOffset, this.bOffset, this.condOffset, this.dstOffset],
-      memory,
-    );
-
-    const a = memory.get(aOffset);
-    const b = memory.get(bOffset);
-    const cond = memory.get(condOffset);
-
-    // TODO: reconsider toBigInt() here
-    memory.set(dstOffset, cond.toBigInt() > 0 ? a : b);
-
-    memory.assert(memoryOperations);
-    context.machineState.incrementPc();
-  }
-}
-
-export class Cast extends TwoOperandInstruction {
-  static readonly type: string = 'CAST';
-  static readonly opcode = Opcode.CAST_8;
-
-  constructor(indirect: number, dstTag: number, srcOffset: number, dstOffset: number) {
-    super(indirect, dstTag, srcOffset, dstOffset);
-  }
-
-  public async execute(context: AvmContext): Promise<void> {
-    const memoryOperations = { reads: 1, writes: 1, indirect: this.indirect };
     const memory = context.machineState.memory.track(this.type);
-    context.machineState.consumeGas(this.gasCost(memoryOperations));
+    context.machineState.consumeGas(this.gasCost());
 
-    const [srcOffset, dstOffset] = Addressing.fromWire(this.indirect).resolve([this.aOffset, this.dstOffset], memory);
+    const operands = [this.srcOffset, this.dstOffset];
+    const addressing = Addressing.fromWire(this.indirect, operands.length);
+    const [srcOffset, dstOffset] = addressing.resolve(operands, memory);
 
     const a = memory.get(srcOffset);
-    const casted = TaggedMemory.buildFromTagTruncating(a.toBigInt(), this.inTag);
+    const casted = TaggedMemory.buildFromTagTruncating(a.toBigInt(), this.dstTag);
 
     memory.set(dstOffset, casted);
 
-    memory.assert(memoryOperations);
+    memory.assert({ reads: 1, writes: 1, addressing });
     context.machineState.incrementPc();
   }
 }
@@ -169,17 +140,18 @@ export class Mov extends Instruction {
   }
 
   public async execute(context: AvmContext): Promise<void> {
-    const memoryOperations = { reads: 1, writes: 1, indirect: this.indirect };
     const memory = context.machineState.memory.track(this.type);
-    context.machineState.consumeGas(this.gasCost(memoryOperations));
+    context.machineState.consumeGas(this.gasCost());
 
-    const [srcOffset, dstOffset] = Addressing.fromWire(this.indirect).resolve([this.srcOffset, this.dstOffset], memory);
+    const operands = [this.srcOffset, this.dstOffset];
+    const addressing = Addressing.fromWire(this.indirect, operands.length);
+    const [srcOffset, dstOffset] = addressing.resolve(operands, memory);
 
     const a = memory.get(srcOffset);
 
     memory.set(dstOffset, a);
 
-    memory.assert(memoryOperations);
+    memory.assert({ reads: 1, writes: 1, addressing });
     context.machineState.incrementPc();
   }
 }
@@ -208,21 +180,19 @@ export class CalldataCopy extends Instruction {
   public async execute(context: AvmContext): Promise<void> {
     const memory = context.machineState.memory.track(this.type);
     // We don't need to check tags here because: (1) the calldata is NOT in memory, and (2) we are the ones writing to destination.
-    const [cdStartOffset, copySizeOffset, dstOffset] = Addressing.fromWire(this.indirect).resolve(
-      [this.cdStartOffset, this.copySizeOffset, this.dstOffset],
-      memory,
-    );
+    const operands = [this.cdStartOffset, this.copySizeOffset, this.dstOffset];
+    const addressing = Addressing.fromWire(this.indirect, operands.length);
+    const [cdStartOffset, copySizeOffset, dstOffset] = addressing.resolve(operands, memory);
 
     const cdStart = memory.get(cdStartOffset).toNumber();
     const copySize = memory.get(copySizeOffset).toNumber();
-    const memoryOperations = { reads: 2, writes: copySize, indirect: this.indirect };
-    context.machineState.consumeGas(this.gasCost({ ...memoryOperations, dynMultiplier: copySize }));
+    context.machineState.consumeGas(this.gasCost(copySize));
 
     const transformedData = context.environment.calldata.slice(cdStart, cdStart + copySize).map(f => new Field(f));
 
     memory.setSlice(dstOffset, transformedData);
 
-    memory.assert(memoryOperations);
+    memory.assert({ reads: 2, writes: copySize, addressing });
     context.machineState.incrementPc();
   }
 }

package/src/avm/opcodes/misc.ts

@@ -33,17 +33,15 @@ export class DebugLog extends Instruction {
 
   public async execute(context: AvmContext): Promise<void> {
     const memory = context.machineState.memory.track(this.type);
-    const [messageOffset, fieldsOffset, fieldsSizeOffset] = Addressing.fromWire(this.indirect).resolve(
-      [this.messageOffset, this.fieldsOffset, this.fieldsSizeOffset],
-      memory,
-    );
+    const operands = [this.messageOffset, this.fieldsOffset, this.fieldsSizeOffset];
+    const addressing = Addressing.fromWire(this.indirect, operands.length);
+    const [messageOffset, fieldsOffset, fieldsSizeOffset] = addressing.resolve(operands, memory);
 
     const fieldsSize = memory.get(fieldsSizeOffset).toNumber();
     memory.checkTagsRange(TypeTag.UINT8, messageOffset, this.messageSize);
     memory.checkTagsRange(TypeTag.FIELD, fieldsOffset, fieldsSize);
 
-    const memoryOperations = { reads: 1 + fieldsSize + this.messageSize, writes: 0, indirect: this.indirect };
-    context.machineState.consumeGas(this.gasCost(memoryOperations));
+    context.machineState.consumeGas(this.gasCost());
 
     const rawMessage = memory.getSlice(messageOffset, this.messageSize);
     const fields = memory.getSlice(fieldsOffset, fieldsSize);
@@ -57,7 +55,7 @@ export class DebugLog extends Instruction {
 
     DebugLog.logger.verbose(formattedStr);
 
-    memory.assert(memoryOperations);
+    memory.assert({ reads: 1 + fieldsSize + this.messageSize, addressing });
     context.machineState.incrementPc();
   }
 }

package/src/avm/opcodes/multi_scalar_mul.ts

@@ -1,8 +1,6 @@
 import { Fq, Point } from '@aztec/circuits.js';
 import { Grumpkin } from '@aztec/circuits.js/barretenberg';
 
-import { strict as assert } from 'assert';
-
 import { type AvmContext } from '../avm_context.js';
 import { Field, TypeTag } from '../avm_memory_types.js';
 import { InstructionExecutionError } from '../errors.js';
@@ -37,16 +35,17 @@ export class MultiScalarMul extends Instruction {
   public async execute(context: AvmContext): Promise<void> {
     const memory = context.machineState.memory.track(this.type);
     // Resolve indirects
-    const [pointsOffset, scalarsOffset, outputOffset] = Addressing.fromWire(this.indirect).resolve(
-      [this.pointsOffset, this.scalarsOffset, this.outputOffset],
-      memory,
-    );
+    const operands = [this.pointsOffset, this.scalarsOffset, this.outputOffset, this.pointsLengthOffset];
+    const addressing = Addressing.fromWire(this.indirect, operands.length);
+    const [pointsOffset, scalarsOffset, outputOffset, pointsLengthOffset] = addressing.resolve(operands, memory);
 
     // Length of the points vector should be U32
-    memory.checkTag(TypeTag.UINT32, this.pointsLengthOffset);
+    memory.checkTag(TypeTag.UINT32, pointsLengthOffset);
     // Get the size of the unrolled (x, y , inf) points vector
-    const pointsReadLength = memory.get(this.pointsLengthOffset).toNumber();
-    assert(pointsReadLength % 3 === 0, 'Points vector offset should be a multiple of 3');
+    const pointsReadLength = memory.get(pointsLengthOffset).toNumber();
+    if (pointsReadLength % 3 !== 0) {
+      throw new InstructionExecutionError(`Points vector offset should be a multiple of 3, was ${pointsReadLength}`);
+    }
     // Divide by 3 since each point is represented as a triplet to get the number of points
     const numPoints = pointsReadLength / 3;
     // The tag for each triplet will be (Field, Field, Uint8)
@@ -54,21 +53,15 @@ export class MultiScalarMul extends Instruction {
       const offset = pointsOffset + i * 3;
       // Check (Field, Field)
       memory.checkTagsRange(TypeTag.FIELD, offset, 2);
-      // Check Uint8 (inf flag)
-      memory.checkTag(TypeTag.UINT8, offset + 2);
+      // Check Uint1 (inf flag)
+      memory.checkTag(TypeTag.UINT1, offset + 2);
     }
     // Get the unrolled (x, y, inf) representing the points
     const pointsVector = memory.getSlice(pointsOffset, pointsReadLength);
 
     // The size of the scalars vector is twice the NUMBER of points because of the scalar limb decomposition
     const scalarReadLength = numPoints * 2;
-    // Consume gas prior to performing work
-    const memoryOperations = {
-      reads: 1 + pointsReadLength + scalarReadLength /* points and scalars */,
-      writes: 3 /* output triplet */,
-      indirect: this.indirect,
-    };
-    context.machineState.consumeGas(this.gasCost(memoryOperations));
+    context.machineState.consumeGas(this.gasCost(pointsReadLength));
     // Get the unrolled scalar (lo & hi) representing the scalars
     const scalarsVector = memory.getSlice(scalarsOffset, scalarReadLength);
     memory.checkTagsRange(TypeTag.FIELD, scalarsOffset, scalarReadLength);
@@ -108,7 +101,11 @@ export class MultiScalarMul extends Instruction {
 
     memory.setSlice(outputOffset, output);
 
-    memory.assert(memoryOperations);
+    memory.assert({
+      reads: 1 + pointsReadLength + scalarReadLength /* points and scalars */,
+      writes: 3 /* output triplet */,
+      addressing,
+    });
     context.machineState.incrementPc();
   }
 }

package/src/avm/opcodes/storage.ts

@@ -32,11 +32,12 @@ export class SStore extends BaseStorageInstruction {
       throw new StaticCallAlterationError();
     }
 
-    const memoryOperations = { reads: 2, indirect: this.indirect };
     const memory = context.machineState.memory.track(this.type);
-    context.machineState.consumeGas(this.gasCost({ ...memoryOperations }));
+    context.machineState.consumeGas(this.gasCost());
 
-    const [srcOffset, slotOffset] = Addressing.fromWire(this.indirect).resolve([this.aOffset, this.bOffset], memory);
+    const operands = [this.aOffset, this.bOffset];
+    const addressing = Addressing.fromWire(this.indirect, operands.length);
+    const [srcOffset, slotOffset] = addressing.resolve(operands, memory);
     memory.checkTag(TypeTag.FIELD, slotOffset);
     memory.checkTag(TypeTag.FIELD, srcOffset);
 
@@ -44,7 +45,7 @@ export class SStore extends BaseStorageInstruction {
     const value = memory.get(srcOffset).toFr();
     context.persistableState.writeStorage(context.environment.storageAddress, slot, value);
 
-    memory.assert(memoryOperations);
+    memory.assert({ reads: 2, addressing });
     context.machineState.incrementPc();
   }
 }
@@ -58,11 +59,12 @@ export class SLoad extends BaseStorageInstruction {
   }
 
   public async execute(context: AvmContext): Promise<void> {
-    const memoryOperations = { writes: 1, reads: 1, indirect: this.indirect };
     const memory = context.machineState.memory.track(this.type);
-    context.machineState.consumeGas(this.gasCost({ ...memoryOperations }));
+    context.machineState.consumeGas(this.gasCost());
 
-    const [slotOffset, dstOffset] = Addressing.fromWire(this.indirect).resolve([this.aOffset, this.bOffset], memory);
+    const operands = [this.aOffset, this.bOffset];
+    const addressing = Addressing.fromWire(this.indirect, operands.length);
+    const [slotOffset, dstOffset] = addressing.resolve(operands, memory);
     memory.checkTag(TypeTag.FIELD, slotOffset);
 
     const slot = memory.get(slotOffset).toFr();
@@ -70,6 +72,6 @@ export class SLoad extends BaseStorageInstruction {
     memory.set(dstOffset, new Field(value));
 
     context.machineState.incrementPc();
-    memory.assert(memoryOperations);
+    memory.assert({ writes: 1, reads: 1, addressing });
   }
 }

package/src/avm/serialization/bytecode_serialization.ts

@@ -1,33 +1,26 @@
-import { PedersenCommitment } from '../opcodes/commitment.js';
-import { DAGasLeft, L2GasLeft } from '../opcodes/context_getters.js';
-import { EcAdd } from '../opcodes/ec_add.js';
-import { Keccak, KeccakF1600, Pedersen, Poseidon2, Sha256 } from '../opcodes/hashing.js';
 import {
   Add,
-  Address,
   And,
-  BlockNumber,
-  CMov,
   Call,
   CalldataCopy,
   Cast,
-  ChainId,
   DebugLog,
   Div,
+  EcAdd,
   EmitNoteHash,
   EmitNullifier,
   EmitUnencryptedLog,
   Eq,
-  FeePerDAGas,
-  FeePerL2Gas,
   FieldDiv,
-  FunctionSelector,
   GetContractInstance,
+  GetEnvVar,
   Instruction,
   InternalCall,
   InternalReturn,
   Jump,
   JumpI,
+  Keccak,
+  KeccakF1600,
   L1ToL2MessageExists,
   Lt,
   Lte,
@@ -37,22 +30,21 @@ import {
   NoteHashExists,
   NullifierExists,
   Or,
+  Pedersen,
+  PedersenCommitment,
+  Poseidon2,
   Return,
   Revert,
   SLoad,
   SStore,
   SendL2ToL1Message,
-  Sender,
   Set,
+  Sha256Compression,
   Shl,
   Shr,
   StaticCall,
-  StorageAddress,
   Sub,
-  Timestamp,
   ToRadixLE,
-  TransactionFee,
-  Version,
   Xor,
 } from '../opcodes/index.js';
 import { MultiScalarMul } from '../opcodes/multi_scalar_mul.js';
@@ -104,25 +96,10 @@ const INSTRUCTION_SET = () =>
     [Opcode.SHR_16, Shr.as(Shr.wireFormat16).deserialize],
     [Opcode.CAST_8, Cast.as(Cast.wireFormat8).deserialize],
     [Opcode.CAST_16, Cast.as(Cast.wireFormat16).deserialize],
-    [Address.opcode, Instruction.deserialize.bind(Address)],
-    [StorageAddress.opcode, Instruction.deserialize.bind(StorageAddress)],
-    [Sender.opcode, Instruction.deserialize.bind(Sender)],
-    [FunctionSelector.opcode, Instruction.deserialize.bind(FunctionSelector)],
-    [TransactionFee.opcode, Instruction.deserialize.bind(TransactionFee)],
-    // Execution Environment - Globals
-    [ChainId.opcode, Instruction.deserialize.bind(ChainId)],
-    [Version.opcode, Instruction.deserialize.bind(Version)],
-    [BlockNumber.opcode, Instruction.deserialize.bind(BlockNumber)],
-    [Timestamp.opcode, Instruction.deserialize.bind(Timestamp)],
-    [FeePerL2Gas.opcode, Instruction.deserialize.bind(FeePerL2Gas)],
-    [FeePerDAGas.opcode, Instruction.deserialize.bind(FeePerDAGas)],
-    // Execution Environment - Calldata
+    // Execution Environment
+    [Opcode.GETENVVAR_16, GetEnvVar.as(GetEnvVar.wireFormat16).deserialize],
     [CalldataCopy.opcode, Instruction.deserialize.bind(CalldataCopy)],
 
-    // Machine State
-    // Machine State - Gas
-    [L2GasLeft.opcode, Instruction.deserialize.bind(L2GasLeft)],
-    [DAGasLeft.opcode, Instruction.deserialize.bind(DAGasLeft)],
     // Machine State - Internal Control Flow
     [Jump.opcode, Instruction.deserialize.bind(Jump)],
     [JumpI.opcode, Instruction.deserialize.bind(JumpI)],
@@ -136,7 +113,6 @@ const INSTRUCTION_SET = () =>
     [Opcode.SET_FF, Set.as(Set.wireFormatFF).deserialize],
     [Opcode.MOV_8, Mov.as(Mov.wireFormat8).deserialize],
     [Opcode.MOV_16, Mov.as(Mov.wireFormat16).deserialize],
-    [CMov.opcode, Instruction.deserialize.bind(CMov)],
 
     // World State
     [SLoad.opcode, Instruction.deserialize.bind(SLoad)], // Public Storage
@@ -167,7 +143,7 @@ const INSTRUCTION_SET = () =>
     [EcAdd.opcode, Instruction.deserialize.bind(EcAdd)],
     [Keccak.opcode, Instruction.deserialize.bind(Keccak)],
     [Poseidon2.opcode, Instruction.deserialize.bind(Poseidon2)],
-    [Sha256.opcode, Instruction.deserialize.bind(Sha256)],
+    [Sha256Compression.opcode, Instruction.deserialize.bind(Sha256Compression)],
     [Pedersen.opcode, Instruction.deserialize.bind(Pedersen)],
     [MultiScalarMul.opcode, Instruction.deserialize.bind(MultiScalarMul)],
     [PedersenCommitment.opcode, Instruction.deserialize.bind(PedersenCommitment)],

package/src/avm/serialization/instruction_serialization.ts

@@ -39,21 +39,8 @@ export enum Opcode {
   CAST_8,
   CAST_16,
   // Execution environment
-  ADDRESS,
-  STORAGEADDRESS,
-  SENDER,
-  FUNCTIONSELECTOR,
-  TRANSACTIONFEE,
-  CHAINID,
-  VERSION,
-  BLOCKNUMBER,
-  TIMESTAMP,
-  FEEPERL2GAS,
-  FEEPERDAGAS,
+  GETENVVAR_16,
   CALLDATACOPY,
-  // Gas
-  L2GASLEFT,
-  DAGASLEFT,
   // Control flow
   JUMP_16,
   JUMPI_16,
@@ -68,7 +55,6 @@ export enum Opcode {
   SET_FF,
   MOV_8,
   MOV_16,
-  CMOV,
   // World state
   SLOAD,
   SSTORE,
@@ -92,16 +78,14 @@ export enum Opcode {
   // Gadgets
   KECCAK,
   POSEIDON2,
-  SHA256, // temp - may be removed, but alot of contracts rely on it
+  SHA256COMPRESSION,
+  KECCAKF1600,
   PEDERSEN, // temp - may be removed, but alot of contracts rely on it
   ECADD,
   MSM,
   PEDERSENCOMMITMENT,
   // Conversion
   TORADIXLE,
-  // Future Gadgets -- pending changes in noir
-  SHA256COMPRESSION,
-  KECCAKF1600, // Here for when we eventually support this
 }
 
 // Possible types for an instruction's operand in its wire format. (Keep in sync with CPP code.

package/src/avm/test_utils.ts

@@ -4,12 +4,11 @@ import { type ContractInstanceWithAddress } from '@aztec/types/contracts';
 import { type jest } from '@jest/globals';
 import { mock } from 'jest-mock-extended';
 
-import { type CommitmentsDB, type PublicContractsDB, type PublicStateDB } from '../public/db_interfaces.js';
+import { type WorldStateDB } from '../public/public_db_sources.js';
 import { type PublicSideEffectTraceInterface } from '../public/side_effect_trace_interface.js';
-import { type HostStorage } from './journal/host_storage.js';
 
-export function mockGetBytecode(hs: HostStorage, bytecode: Buffer) {
-  (hs as jest.Mocked<HostStorage>).contractsDb.getBytecode.mockResolvedValue(bytecode);
+export function mockGetBytecode(worldStateDB: WorldStateDB, bytecode: Buffer) {
+  (worldStateDB as jest.Mocked<WorldStateDB>).getBytecode.mockResolvedValue(bytecode);
 }
 
 export function mockTraceFork(trace: PublicSideEffectTraceInterface, nestedTrace?: PublicSideEffectTraceInterface) {
@@ -18,18 +17,18 @@ export function mockTraceFork(trace: PublicSideEffectTraceInterface, nestedTrace
   );
 }
 
-export function mockStorageRead(hs: HostStorage, value: Fr) {
-  (hs.publicStateDb as jest.Mocked<PublicStateDB>).storageRead.mockResolvedValue(value);
+export function mockStorageRead(worldStateDB: WorldStateDB, value: Fr) {
+  (worldStateDB as jest.Mocked<WorldStateDB>).storageRead.mockResolvedValue(value);
 }
 
-export function mockStorageReadWithMap(hs: HostStorage, mockedStorage: Map<bigint, Fr>) {
-  (hs.publicStateDb as jest.Mocked<PublicStateDB>).storageRead.mockImplementation((_address, slot) =>
+export function mockStorageReadWithMap(worldStateDB: WorldStateDB, mockedStorage: Map<bigint, Fr>) {
+  (worldStateDB as jest.Mocked<WorldStateDB>).storageRead.mockImplementation((_address, slot) =>
     Promise.resolve(mockedStorage.get(slot.toBigInt()) ?? Fr.ZERO),
   );
 }
 
-export function mockNoteHashExists(hs: HostStorage, _leafIndex: Fr, value?: Fr) {
-  (hs.commitmentsDb as jest.Mocked<CommitmentsDB>).getCommitmentValue.mockImplementation((index: bigint) => {
+export function mockNoteHashExists(worldStateDB: WorldStateDB, _leafIndex: Fr, value?: Fr) {
+  (worldStateDB as jest.Mocked<WorldStateDB>).getCommitmentValue.mockImplementation((index: bigint) => {
     if (index == _leafIndex.toBigInt()) {
       return Promise.resolve(value);
     } else {
@@ -39,12 +38,17 @@ export function mockNoteHashExists(hs: HostStorage, _leafIndex: Fr, value?: Fr)
   });
 }
 
-export function mockNullifierExists(hs: HostStorage, leafIndex: Fr, _value?: Fr) {
-  (hs.commitmentsDb as jest.Mocked<CommitmentsDB>).getNullifierIndex.mockResolvedValue(leafIndex.toBigInt());
+export function mockNullifierExists(worldStateDB: WorldStateDB, leafIndex: Fr, _value?: Fr) {
+  (worldStateDB as jest.Mocked<WorldStateDB>).getNullifierIndex.mockResolvedValue(leafIndex.toBigInt());
 }
 
-export function mockL1ToL2MessageExists(hs: HostStorage, leafIndex: Fr, value: Fr, valueAtOtherIndices?: Fr) {
-  (hs.commitmentsDb as jest.Mocked<CommitmentsDB>).getL1ToL2LeafValue.mockImplementation((index: bigint) => {
+export function mockL1ToL2MessageExists(
+  worldStateDB: WorldStateDB,
+  leafIndex: Fr,
+  value: Fr,
+  valueAtOtherIndices?: Fr,
+) {
+  (worldStateDB as jest.Mocked<WorldStateDB>).getL1ToL2LeafValue.mockImplementation((index: bigint) => {
     if (index == leafIndex.toBigInt()) {
       return Promise.resolve(value);
     } else {
@@ -55,6 +59,6 @@ export function mockL1ToL2MessageExists(hs: HostStorage, leafIndex: Fr, value: F
   });
 }
 
-export function mockGetContractInstance(hs: HostStorage, contractInstance: ContractInstanceWithAddress) {
-  (hs.contractsDb as jest.Mocked<PublicContractsDB>).getContractInstance.mockResolvedValue(contractInstance);
+export function mockGetContractInstance(worldStateDB: WorldStateDB, contractInstance: ContractInstanceWithAddress) {
+  (worldStateDB as jest.Mocked<WorldStateDB>).getContractInstance.mockResolvedValue(contractInstance);
 }