feat(avm): tagged memory

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

@@ -1,5 +1,7 @@
 import { Fr } from '@aztec/foundation/fields';
 
+import { TaggedMemory } from './avm_memory_types.js';
+
 /**
  * Store's data for an Avm execution frame
  */
@@ -8,9 +10,8 @@ export class AvmMachineState {
   public readonly calldata: Fr[];
   private returnData: Fr[];
 
-  // TODO: implement tagged memory
   /** - */
-  public memory: Fr[];
+  public readonly memory: TaggedMemory;
 
   /**
    * When an internal_call is invoked, the internal call stack is added to with the current pc + 1
@@ -35,7 +36,7 @@ export class AvmMachineState {
   constructor(calldata: Fr[]) {
     this.calldata = calldata;
     this.returnData = [];
-    this.memory = [];
+    this.memory = new TaggedMemory();
     this.internalCallStack = [];
 
     this.pc = 0;
@@ -57,37 +58,4 @@ export class AvmMachineState {
   public getReturnData(): Fr[] {
     return this.returnData;
   }
-
-  /** -
-   * @param offset -
-   */
-  public readMemory(offset: number): Fr {
-    // TODO: check offset is within bounds
-    return this.memory[offset] ?? Fr.ZERO;
-  }
-
-  /** -
-   * @param offset -
-   * @param size -
-   */
-  public readMemoryChunk(offset: number, size: number): Fr[] {
-    // TODO: bounds -> initialise to 0
-    return this.memory.slice(offset, offset + size);
-  }
-
-  /** -
-   * @param offset -
-   * @param value -
-   */
-  public writeMemory(offset: number, value: Fr): void {
-    this.memory[offset] = value;
-  }
-
-  /** -
-   * @param offset -
-   * @param values -
-   */
-  public writeMemoryChunk(offset: number, values: Fr[]): void {
-    this.memory.splice(offset, values.length, ...values);
-  }
 }

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

@@ -0,0 +1,22 @@
+import { Field, Uint8 } from './avm_memory_types.js';
+
+// TODO: complete
+describe('Uint8', () => {
+  it('Unsigned 8 max value', () => {
+    expect(new Uint8(255).toBigInt()).toEqual(255n);
+  });
+
+  it('Unsigned 8 bit add', () => {
+    expect(new Uint8(50).add(new Uint8(20))).toEqual(new Uint8(70));
+  });
+
+  it('Unsigned 8 bit add wraps', () => {
+    expect(new Uint8(200).add(new Uint8(100))).toEqual(new Uint8(44));
+  });
+});
+
+describe('Field', () => {
+  it('Add correctly without wrapping', () => {
+    expect(new Field(27).add(new Field(48))).toEqual(new Field(75));
+  });
+});

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

@@ -0,0 +1,276 @@
+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;
+
+  // Use sparingly.
+  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;
+}
+
+// TODO: Optimize calculation of mod, etc. Can only do once per class?
+abstract class UnsignedInteger implements IntegralValue {
+  private readonly bitmask: bigint;
+  private readonly mod: bigint;
+
+  protected constructor(private n: bigint, private bits: bigint) {
+    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 add(rhs: UnsignedInteger): UnsignedInteger {
+    assert(this.bits == rhs.bits);
+    return this.build((this.n + rhs.n) & this.bitmask);
+  }
+
+  public sub(rhs: UnsignedInteger): UnsignedInteger {
+    assert(this.bits == rhs.bits);
+    const res: bigint = this.n - rhs.n;
+    return this.build(res >= 0 ? res : res + this.mod);
+  }
+
+  public mul(rhs: UnsignedInteger): UnsignedInteger {
+    assert(this.bits == rhs.bits);
+    return this.build((this.n * rhs.n) & this.bitmask);
+  }
+
+  public div(rhs: UnsignedInteger): UnsignedInteger {
+    assert(this.bits == rhs.bits);
+    return this.build(this.n / rhs.n);
+  }
+
+  // No sign extension.
+  public shr(rhs: UnsignedInteger): UnsignedInteger {
+    assert(this.bits == rhs.bits);
+    // Note that this.n is > 0 by class invariant.
+    return this.build(this.n >> rhs.n);
+  }
+
+  public shl(rhs: UnsignedInteger): UnsignedInteger {
+    assert(this.bits == rhs.bits);
+    return this.build((this.n << rhs.n) & this.bitmask);
+  }
+
+  public and(rhs: UnsignedInteger): UnsignedInteger {
+    assert(this.bits == rhs.bits);
+    return this.build(this.n & rhs.n);
+  }
+
+  public or(rhs: UnsignedInteger): UnsignedInteger {
+    assert(this.bits == rhs.bits);
+    return this.build(this.n | rhs.n);
+  }
+
+  public xor(rhs: UnsignedInteger): UnsignedInteger {
+    assert(this.bits == rhs.bits);
+    return this.build(this.n ^ rhs.n);
+  }
+
+  public not(): UnsignedInteger {
+    return this.build(~this.n & this.bitmask);
+  }
+
+  public toBigInt(): bigint {
+    return this.n;
+  }
+
+  public equals(rhs: UnsignedInteger) {
+    return this.bits == rhs.bits && this.toBigInt() == rhs.toBigInt();
+  }
+}
+
+export class Uint8 extends UnsignedInteger {
+  constructor(n: number | bigint) {
+    super(BigInt(n), 8n);
+  }
+
+  protected build(n: bigint): Uint8 {
+    return new Uint8(n);
+  }
+}
+
+export class Uint16 extends UnsignedInteger {
+  constructor(n: number | bigint) {
+    super(BigInt(n), 16n);
+  }
+
+  protected build(n: bigint): Uint16 {
+    return new Uint16(n);
+  }
+}
+
+export class Uint32 extends UnsignedInteger {
+  constructor(n: number | bigint) {
+    super(BigInt(n), 32n);
+  }
+
+  protected build(n: bigint): Uint32 {
+    return new Uint32(n);
+  }
+}
+
+export class Uint64 extends UnsignedInteger {
+  constructor(n: number | bigint) {
+    super(BigInt(n), 64n);
+  }
+
+  protected build(n: bigint): Uint64 {
+    return new Uint64(n);
+  }
+}
+
+export class Uint128 extends UnsignedInteger {
+  constructor(n: number | bigint) {
+    super(BigInt(n), 128n);
+  }
+
+  protected build(n: bigint): Uint128 {
+    return new Uint128(n);
+  }
+}
+
+export class Field implements MemoryValue {
+  public static readonly MODULUS: bigint = Fr.MODULUS;
+  private readonly rep: Fr;
+
+  constructor(v: number | bigint | Fr) {
+    this.rep = new Fr(v);
+  }
+
+  public add(rhs: Field): Field {
+    return new Field(this.rep.add(rhs.rep));
+  }
+
+  public sub(rhs: Field): Field {
+    return new Field(this.rep.sub(rhs.rep));
+  }
+
+  public mul(rhs: Field): Field {
+    return new Field(this.rep.mul(rhs.rep));
+  }
+
+  public div(rhs: Field): Field {
+    return new Field(this.rep.div(rhs.rep));
+  }
+
+  public toBigInt(): bigint {
+    return this.rep.toBigInt();
+  }
+}
+
+export enum TypeTag {
+  UNINITIALIZED,
+  UINT8,
+  UINT16,
+  UINT32,
+  UINT64,
+  UINT128,
+  FIELD,
+  INVALID,
+}
+
+// TODO: Consider automatic conversion when getting undefined values.
+export class TaggedMemory {
+  static readonly MAX_MEMORY_SIZE = 1n << 32n;
+  private _mem: MemoryValue[];
+
+  constructor() {
+    this._mem = [];
+  }
+
+  public get(offset: number): MemoryValue {
+    return this.getAs<MemoryValue>(offset);
+  }
+
+  public getAs<T>(offset: number): T {
+    assert(offset < TaggedMemory.MAX_MEMORY_SIZE);
+    const e = this._mem[offset];
+    return <T>e;
+  }
+
+  public getSlice(offset: number, size: number): MemoryValue[] {
+    assert(offset < TaggedMemory.MAX_MEMORY_SIZE);
+    return this._mem.slice(offset, offset + size);
+  }
+
+  public getSliceTags(offset: number, size: number): TypeTag[] {
+    assert(offset < TaggedMemory.MAX_MEMORY_SIZE);
+    return this._mem.slice(offset, offset + size).map(TaggedMemory.getTag);
+  }
+
+  public set(offset: number, v: MemoryValue) {
+    assert(offset < TaggedMemory.MAX_MEMORY_SIZE);
+    this._mem[offset] = v;
+  }
+
+  public setSlice(offset: number, vs: MemoryValue[]) {
+    assert(offset < TaggedMemory.MAX_MEMORY_SIZE);
+    this._mem.splice(offset, vs.length, ...vs);
+  }
+
+  public getTag(offset: number): TypeTag {
+    return TaggedMemory.getTag(this._mem[offset]);
+  }
+
+  // TODO: this might be slow, but I don't want to have the types know of their tags.
+  // It might be possible to have a map<Prototype, TypeTag>.
+  public static getTag(v: MemoryValue | undefined): TypeTag {
+    let tag = TypeTag.INVALID;
+
+    if (v === undefined) {
+      tag = TypeTag.UNINITIALIZED;
+    } else if (v instanceof Field) {
+      tag = TypeTag.FIELD;
+    } else if (v instanceof Uint8) {
+      tag = TypeTag.UINT8;
+    } else if (v instanceof Uint16) {
+      tag = TypeTag.UINT16;
+    } else if (v instanceof Uint32) {
+      tag = TypeTag.UINT32;
+    } else if (v instanceof Uint64) {
+      tag = TypeTag.UINT64;
+    } else if (v instanceof Uint128) {
+      tag = TypeTag.UINT128;
+    }
+
+    return tag;
+  }
+
+  // Truncates the value to fit the type.
+  public static integralFromTag(v: bigint, tag: TypeTag): IntegralValue {
+    switch (tag) {
+      case TypeTag.UINT8:
+        return new Uint8(v & ((1n << 8n) - 1n));
+      case TypeTag.UINT16:
+        return new Uint16(v & ((1n << 16n) - 1n));
+      case TypeTag.UINT32:
+        return new Uint32(v & ((1n << 32n) - 1n));
+      case TypeTag.UINT64:
+        return new Uint64(v & ((1n << 64n) - 1n));
+      case TypeTag.UINT128:
+        return new Uint128(v & ((1n << 128n) - 1n));
+      default:
+        throw new Error(`${TypeTag[tag]} is not a valid integral type.`);
+    }
+  }
+}

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

@@ -71,8 +71,8 @@ export class AvmInterpreter {
  * Avm-specific errors should derive from this
  */
 export abstract class AvmInterpreterError extends Error {
-  constructor(message: string) {
-    super(message);
+  constructor(message: string, ...rest: any[]) {
+    super(message, ...rest);
     this.name = 'AvmInterpreterError';
   }
 }