fix(avm): fix usage of Fr with tagged memory

yarn-project/acir-simulator/src/avm/avm_memory_types.ts

@@ -2,41 +2,46 @@ import { Fr } from '@aztec/foundation/fields';
 
 import { strict as assert } from 'assert';
 
-export interface MemoryValue {
-  add(rhs: MemoryValue): MemoryValue;
-  sub(rhs: MemoryValue): MemoryValue;
-  mul(rhs: MemoryValue): MemoryValue;
-  div(rhs: MemoryValue): MemoryValue;
+export abstract class MemoryValue {
+  public abstract add(rhs: MemoryValue): MemoryValue;
+  public abstract sub(rhs: MemoryValue): MemoryValue;
+  public abstract mul(rhs: MemoryValue): MemoryValue;
+  public abstract div(rhs: MemoryValue): MemoryValue;
+
+  public abstract equals(rhs: MemoryValue): boolean;
+  public abstract lt(rhs: MemoryValue): boolean;
+
+  // We need this to be able to build an instance of the subclasses.
+  public abstract build(n: bigint): MemoryValue;
 
   // Use sparingly.
-  toBigInt(): bigint;
+  public abstract toBigInt(): bigint;
 }
 
-export interface IntegralValue extends MemoryValue {
-  shl(rhs: IntegralValue): IntegralValue;
-  shr(rhs: IntegralValue): IntegralValue;
-  and(rhs: IntegralValue): IntegralValue;
-  or(rhs: IntegralValue): IntegralValue;
-  xor(rhs: IntegralValue): IntegralValue;
-  not(): IntegralValue;
+export abstract class IntegralValue extends MemoryValue {
+  public abstract shl(rhs: IntegralValue): IntegralValue;
+  public abstract shr(rhs: IntegralValue): IntegralValue;
+  public abstract and(rhs: IntegralValue): IntegralValue;
+  public abstract or(rhs: IntegralValue): IntegralValue;
+  public abstract xor(rhs: IntegralValue): IntegralValue;
+  public abstract not(): IntegralValue;
 }
 
 // TODO: Optimize calculation of mod, etc. Can only do once per class?
-abstract class UnsignedInteger implements IntegralValue {
+abstract class UnsignedInteger extends IntegralValue {
   private readonly bitmask: bigint;
   private readonly mod: bigint;
 
   protected constructor(private n: bigint, private bits: bigint) {
+    super();
     assert(bits > 0);
     // x % 2^n == x & (2^n - 1)
     this.mod = 1n << bits;
     this.bitmask = this.mod - 1n;
     assert(n < this.mod);
   }
 
-  // We need this to be able to build an instance of the subclass
-  // and not of type UnsignedInteger.
-  protected abstract build(n: bigint): UnsignedInteger;
+  public abstract build(n: bigint): UnsignedInteger;
 
   public add(rhs: UnsignedInteger): UnsignedInteger {
     assert(this.bits == rhs.bits);
@@ -90,12 +95,18 @@ abstract class UnsignedInteger implements IntegralValue {
     return this.build(~this.n & this.bitmask);
   }
 
-  public toBigInt(): bigint {
-    return this.n;
+  public equals(rhs: UnsignedInteger): boolean {
+    assert(this.bits == rhs.bits);
+    return this.n === rhs.n;
+  }
+
+  public lt(rhs: UnsignedInteger): boolean {
+    assert(this.bits == rhs.bits);
+    return this.n < rhs.n;
   }
 
-  public equals(rhs: UnsignedInteger) {
-    return this.bits == rhs.bits && this.toBigInt() == rhs.toBigInt();
+  public toBigInt(): bigint {
+    return this.n;
   }
 }
 
@@ -104,7 +115,7 @@ export class Uint8 extends UnsignedInteger {
     super(BigInt(n), 8n);
   }
 
-  protected build(n: bigint): Uint8 {
+  public build(n: bigint): Uint8 {
     return new Uint8(n);
   }
 }
@@ -114,7 +125,7 @@ export class Uint16 extends UnsignedInteger {
     super(BigInt(n), 16n);
   }
 
-  protected build(n: bigint): Uint16 {
+  public build(n: bigint): Uint16 {
     return new Uint16(n);
   }
 }
@@ -124,7 +135,7 @@ export class Uint32 extends UnsignedInteger {
     super(BigInt(n), 32n);
   }
 
-  protected build(n: bigint): Uint32 {
+  public build(n: bigint): Uint32 {
     return new Uint32(n);
   }
 }
@@ -134,7 +145,7 @@ export class Uint64 extends UnsignedInteger {
     super(BigInt(n), 64n);
   }
 
-  protected build(n: bigint): Uint64 {
+  public build(n: bigint): Uint64 {
     return new Uint64(n);
   }
 }
@@ -144,19 +155,24 @@ export class Uint128 extends UnsignedInteger {
     super(BigInt(n), 128n);
   }
 
-  protected build(n: bigint): Uint128 {
+  public build(n: bigint): Uint128 {
     return new Uint128(n);
   }
 }
 
-export class Field implements MemoryValue {
+export class Field extends MemoryValue {
   public static readonly MODULUS: bigint = Fr.MODULUS;
   private readonly rep: Fr;
 
   constructor(v: number | bigint | Fr) {
+    super();
     this.rep = new Fr(v);
   }
 
+  public build(n: bigint): Field {
+    return new Field(n);
+  }
+
   public add(rhs: Field): Field {
     return new Field(this.rep.add(rhs.rep));
   }
@@ -173,6 +189,14 @@ export class Field implements MemoryValue {
     return new Field(this.rep.div(rhs.rep));
   }
 
+  public equals(rhs: Field): boolean {
+    return this.rep.equals(rhs.rep);
+  }
+
+  public lt(rhs: Field): boolean {
+    return this.rep.lt(rhs.rep);
+  }
+
   public toBigInt(): bigint {
     return this.rep.toBigInt();
   }

yarn-project/acir-simulator/src/avm/opcodes/external_calls.test.ts

@@ -40,16 +40,16 @@ describe('External Calls', () => {
       const addr = new Fr(123456n);
 
       const argsOffset = 2;
-      const args = [new Fr(1n), new Fr(2n), new Fr(3n)];
+      const args = [new Field(1n), new Field(2n), new Field(3n)];
       const argsSize = args.length;
 
       const retOffset = 8;
       const retSize = 2;
 
       const successOffset = 7;
 
-      machineState.memory.set(0, gas);
-      machineState.memory.set(1, addr);
+      machineState.memory.set(0, new Field(gas));
+      machineState.memory.set(1, new Field(addr));
       machineState.memory.setSlice(2, args);
 
       const otherContextInstructions: [Opcode, any[]][] = [
@@ -72,10 +72,10 @@ describe('External Calls', () => {
       await instruction.execute(machineState, journal);
 
       const successValue = machineState.memory.get(successOffset);
-      expect(successValue).toEqual(new Fr(1n));
+      expect(successValue).toEqual(new Field(1n));
 
       const retValue = machineState.memory.getSlice(retOffset, retSize);
-      expect(retValue).toEqual([new Fr(1n), new Fr(2n)]);
+      expect(retValue).toEqual([new Field(1n), new Field(2n)]);
 
       // Check that the storage call has been merged into the parent journal
       const { storageWrites } = journal.flush();
@@ -126,7 +126,7 @@ describe('External Calls', () => {
 
       // No revert has occurred, but the nested execution has failed
       const successValue = machineState.memory.get(successOffset);
-      expect(successValue).toEqual(new Fr(0n));
+      expect(successValue).toEqual(new Field(0n));
     });
   });
 });

yarn-project/acir-simulator/src/avm/opcodes/external_calls.ts

@@ -39,10 +39,11 @@ export class Call extends Instruction {
 
     // We only take as much data as was specified in the return size -> TODO: should we be reverting here
     const returnData = returnObject.output.slice(0, this.retSize);
+    const convertedReturnData = returnData.map(f => new Field(f));
 
     // Write our return data into memory
-    machineState.memory.set(this.successOffset, new Fr(success));
-    machineState.memory.setSlice(this.retOffset, returnData);
+    machineState.memory.set(this.successOffset, new Field(success ? 1 : 0));
+    machineState.memory.setSlice(this.retOffset, convertedReturnData);
 
     if (success) {
       avmContext.mergeJournal();
@@ -84,10 +85,11 @@ export class StaticCall extends Instruction {
 
     // We only take as much data as was specified in the return size -> TODO: should we be reverting here
     const returnData = returnObject.output.slice(0, this.retSize);
+    const convertedReturnData = returnData.map(f => new Field(f));
 
     // Write our return data into memory
-    machineState.memory.set(this.successOffset, new Fr(success));
-    machineState.memory.setSlice(this.retOffset, returnData);
+    machineState.memory.set(this.successOffset, new Field(success ? 1 : 0));
+    machineState.memory.setSlice(this.retOffset, convertedReturnData);
 
     if (success) {
       avmContext.mergeJournal();

yarn-project/acir-simulator/src/avm/opcodes/storage.test.ts

@@ -63,6 +63,6 @@ describe('Storage Instructions', () => {
     expect(journal.readStorage).toBeCalledWith(address, new Fr(a.toBigInt()));
 
     const actual = machineState.memory.get(1);
-    expect(actual).toEqual(expectedResult);
+    expect(actual).toEqual(new Field(expectedResult));
   });
 });

yarn-project/acir-simulator/src/avm/opcodes/storage.ts

@@ -1,6 +1,7 @@
 import { Fr } from '@aztec/foundation/fields';
 
 import { AvmMachineState } from '../avm_machine_state.js';
+import { Field } from '../avm_memory_types.js';
 import { AvmInterpreterError } from '../interpreter/interpreter.js';
 import { AvmJournal } from '../journal/journal.js';
 import { Instruction } from './instruction.js';
@@ -44,9 +45,12 @@ export class SLoad extends Instruction {
   async execute(machineState: AvmMachineState, journal: AvmJournal): Promise<void> {
     const slot = machineState.memory.get(this.slotOffset);
 
-    const data = journal.readStorage(machineState.executionEnvironment.storageAddress, new Fr(slot.toBigInt()));
+    const data: Fr = await journal.readStorage(
+      machineState.executionEnvironment.storageAddress,
+      new Fr(slot.toBigInt()),
+    );
 
-    machineState.memory.set(this.destOffset, await data);
+    machineState.memory.set(this.destOffset, new Field(data));
 
     this.incrementPc(machineState);
   }