feat(avm): add internal jump and return, adjust control flow (#4140)

fixes: #4132
AztecProtocol · Jan 18, 2024 · b77afb1 · b77afb1
1 parent de6e2ed
commit b77afb1
Show file tree

Hide file tree

Showing 11 changed files with 409 additions and 46 deletions.
diff --git a/yarn-project/acir-simulator/src/avm/avm_machine_state.ts b/yarn-project/acir-simulator/src/avm/avm_machine_state.ts
@@ -12,11 +12,22 @@ export class AvmMachineState {
   /** - */
   public memory: Fr[];
 
+  /**
+   * When an internal_call is invoked, the internal call stack is added to with the current pc + 1
+   * When internal_return is invoked, the latest value is popped from the internal call stack and set to the pc.
+   */
+  public internalCallStack: number[];
+
   /** - */
   public pc: number;
   /** - */
   public callStack: number[];
 
+  /**
+   * If an instruction triggers a halt, then it ends execution of the VM
+   */
+  public halted: boolean;
+
   /**
    * Create a new avm context
    * @param calldata -
@@ -25,9 +36,12 @@ export class AvmMachineState {
     this.calldata = calldata;
     this.returnData = [];
     this.memory = [];
+    this.internalCallStack = [];
 
     this.pc = 0;
     this.callStack = [];
+
+    this.halted = false;
   }
 
   /**

diff --git a/yarn-project/acir-simulator/src/avm/interpreter/interpreter.test.ts b/yarn-project/acir-simulator/src/avm/interpreter/interpreter.test.ts
@@ -5,7 +5,7 @@ import { mock } from 'jest-mock-extended';
 import { AvmMachineState } from '../avm_machine_state.js';
 import { AvmStateManager } from '../avm_state_manager.js';
 import { Add } from '../opcodes/arithmetic.js';
-import { Return } from '../opcodes/control_flow.js';
+import { Jump, Return } from '../opcodes/control_flow.js';
 import { Instruction } from '../opcodes/instruction.js';
 import { CalldataCopy } from '../opcodes/memory.js';
 import { AvmInterpreter } from './interpreter.js';
@@ -34,4 +34,20 @@ describe('interpreter', () => {
     expect(returnData.length).toBe(1);
     expect(returnData).toEqual([new Fr(3)]);
   });
+
+  it('Should revert with an invalid jump', () => {
+    const calldata: Fr[] = [];
+    const stateManager = mock<AvmStateManager>();
+
+    const invalidJumpDestination = 22;
+
+    const instructions: Instruction[] = [new Jump(invalidJumpDestination)];
+
+    const context = new AvmMachineState(calldata);
+    const interpreter = new AvmInterpreter(context, stateManager, instructions);
+
+    const avmReturnData = interpreter.run();
+
+    expect(avmReturnData.reverted).toBe(true);
+  });
 });
diff --git a/yarn-project/acir-simulator/src/avm/interpreter/interpreter.ts b/yarn-project/acir-simulator/src/avm/interpreter/interpreter.ts
@@ -30,8 +30,18 @@ export class AvmInterpreter {
    */
   run(): AvmMessageCallResult {
     try {
-      for (const instruction of this.instructions) {
+      while (!this.machineState.halted && this.machineState.pc < this.instructions.length) {
+        const instruction = this.instructions[this.machineState.pc];
+
+        if (!instruction) {
+          throw new InvalidInstructionError(this.machineState.pc);
+        }
+
         instruction.execute(this.machineState, this.stateManager);
+
+        if (this.machineState.pc >= this.instructions.length) {
+          throw new InvalidProgramCounterError(this.machineState.pc, this.instructions.length);
+        }
       }
 
       const returnData = this.machineState.getReturnData();
@@ -52,3 +62,22 @@ export class AvmInterpreter {
     return this.machineState.getReturnData();
   }
 }
+
+/**
+ * Error is thrown when the program counter goes to an invalid location.
+ * There is no instruction at the provided pc
+ */
+class InvalidProgramCounterError extends Error {
+  constructor(pc: number, max: number) {
+    super(`Invalid program counter ${pc}, max is ${max}`);
+  }
+}
+
+/**
+ * This assertion should never be hit - there should always be a valid instruction
+ */
+class InvalidInstructionError extends Error {
+  constructor(pc: number) {
+    super(`Invalid instruction at ${pc}`);
+  }
+}
diff --git a/yarn-project/acir-simulator/src/avm/opcodes/arithmetic.ts b/yarn-project/acir-simulator/src/avm/opcodes/arithmetic.ts
@@ -5,59 +5,73 @@ import { AvmStateManager } from '../avm_state_manager.js';
 import { Instruction } from './instruction.js';
 
 /** -*/
-export class Add implements Instruction {
+export class Add extends Instruction {
   static type: string = 'ADD';
   static numberOfOperands = 3;
 
-  constructor(private aOffset: number, private bOffset: number, private destOffset: number) {}
+  constructor(private aOffset: number, private bOffset: number, private destOffset: number) {
+    super();
+  }
 
   execute(machineState: AvmMachineState, _stateManager: AvmStateManager): void {
     const a = machineState.readMemory(this.aOffset);
     const b = machineState.readMemory(this.bOffset);
 
-    const dest = new Fr(a.toBigInt() + (b.toBigInt() % Fr.MODULUS));
+    const dest = new Fr((a.toBigInt() + b.toBigInt()) % Fr.MODULUS);
     machineState.writeMemory(this.destOffset, dest);
+
+    this.incrementPc(machineState);
   }
 }
 
 /** -*/
-export class Sub implements Instruction {
+export class Sub extends Instruction {
   static type: string = 'SUB';
   static numberOfOperands = 3;
 
-  constructor(private aOffset: number, private bOffset: number, private destOffset: number) {}
+  constructor(private aOffset: number, private bOffset: number, private destOffset: number) {
+    super();
+  }
 
   execute(machineState: AvmMachineState, _stateManager: AvmStateManager): void {
     const a = machineState.readMemory(this.aOffset);
     const b = machineState.readMemory(this.bOffset);
 
-    const dest = new Fr(a.toBigInt() - (b.toBigInt() % Fr.MODULUS));
+    const dest = new Fr((a.toBigInt() - b.toBigInt()) % Fr.MODULUS);
     machineState.writeMemory(this.destOffset, dest);
+
+    this.incrementPc(machineState);
   }
 }
 
 /** -*/
-export class Mul implements Instruction {
+export class Mul extends Instruction {
   static type: string = 'MUL';
   static numberOfOperands = 3;
 
-  constructor(private aOffset: number, private bOffset: number, private destOffset: number) {}
+  constructor(private aOffset: number, private bOffset: number, private destOffset: number) {
+    super();
+  }
 
   execute(machineState: AvmMachineState, _stateManager: AvmStateManager): void {
     const a: Fr = machineState.readMemory(this.aOffset);
     const b: Fr = machineState.readMemory(this.bOffset);
 
     const dest = new Fr((a.toBigInt() * b.toBigInt()) % Fr.MODULUS);
     machineState.writeMemory(this.destOffset, dest);
+
+    this.incrementPc(machineState);
   }
 }
 
 /** -*/
-export class Div implements Instruction {
+export class Div extends Instruction {
   static type: string = 'DIV';
   static numberOfOperands = 3;
 
-  constructor(private aOffset: number, private bOffset: number, private destOffset: number) {}
+  constructor(private aOffset: number, private bOffset: number, private destOffset: number) {
+    super();
+  }
 
   execute(machineState: AvmMachineState, _stateManager: AvmStateManager): void {
     const a: Fr = machineState.readMemory(this.aOffset);
@@ -66,5 +80,7 @@ export class Div implements Instruction {
     // TODO(https://github.com/AztecProtocol/aztec-packages/issues/3993): proper field division
     const dest = new Fr(a.toBigInt() / b.toBigInt());
     machineState.writeMemory(this.destOffset, dest);
+
+    this.incrementPc(machineState);
   }
 }
diff --git a/yarn-project/acir-simulator/src/avm/opcodes/bitwise.ts b/yarn-project/acir-simulator/src/avm/opcodes/bitwise.ts
@@ -5,96 +5,122 @@ import { AvmStateManager } from '../avm_state_manager.js';
 import { Instruction } from './instruction.js';
 
 /** - */
-export class And implements Instruction {
+export class And extends Instruction {
   static type: string = 'AND';
   static numberOfOperands = 3;
 
-  constructor(private aOffset: number, private bOffset: number, private destOffset: number) {}
+  constructor(private aOffset: number, private bOffset: number, private destOffset: number) {
+    super();
+  }
 
   execute(machineState: AvmMachineState, _stateManager: AvmStateManager): void {
     const a: Fr = machineState.readMemory(this.aOffset);
     const b: Fr = machineState.readMemory(this.bOffset);
 
     const dest = new Fr(a.toBigInt() & b.toBigInt());
     machineState.writeMemory(this.destOffset, dest);
+
+    this.incrementPc(machineState);
   }
 }
 
 /** - */
-export class Or implements Instruction {
+export class Or extends Instruction {
   static type: string = 'OR';
   static numberOfOperands = 3;
 
-  constructor(private aOffset: number, private bOffset: number, private destOffset: number) {}
+  constructor(private aOffset: number, private bOffset: number, private destOffset: number) {
+    super();
+  }
 
   execute(machineState: AvmMachineState, _stateManager: AvmStateManager): void {
     const a: Fr = machineState.readMemory(this.aOffset);
     const b: Fr = machineState.readMemory(this.bOffset);
 
     const dest = new Fr(a.toBigInt() | b.toBigInt());
     machineState.writeMemory(this.destOffset, dest);
+
+    this.incrementPc(machineState);
   }
 }
 
 /** - */
-export class Xor implements Instruction {
+export class Xor extends Instruction {
   static type: string = 'XOR';
   static numberOfOperands = 3;
 
-  constructor(private aOffset: number, private bOffset: number, private destOffset: number) {}
+  constructor(private aOffset: number, private bOffset: number, private destOffset: number) {
+    super();
+  }
 
   execute(machineState: AvmMachineState, _stateManager: AvmStateManager): void {
     const a: Fr = machineState.readMemory(this.aOffset);
     const b: Fr = machineState.readMemory(this.bOffset);
 
     const dest = new Fr(a.toBigInt() ^ b.toBigInt());
     machineState.writeMemory(this.destOffset, dest);
+
+    this.incrementPc(machineState);
   }
 }
 
 /** - */
-export class Not implements Instruction {
+export class Not extends Instruction {
   static type: string = 'NOT';
   static numberOfOperands = 2;
 
-  constructor(private aOffset: number, private bOffset: number, private destOffset: number) {}
+  constructor(private aOffset: number, private destOffset: number) {
+    super();
+  }
 
   execute(machineState: AvmMachineState, _stateManager: AvmStateManager): void {
     const a: Fr = machineState.readMemory(this.aOffset);
 
-    const dest = new Fr(~a.toBigInt());
+    // TODO: hack -> until proper field arithmetic is implemented
+    const result = ~a.toBigInt();
+    const dest = new Fr(result < 0 ? Fr.MODULUS + /* using a + as result is -ve*/ result : result);
     machineState.writeMemory(this.destOffset, dest);
+
+    this.incrementPc(machineState);
   }
 }
 
 /** -*/
-export class Shl implements Instruction {
+export class Shl extends Instruction {
   static type: string = 'SHL';
   static numberOfOperands = 3;
 
-  constructor(private aOffset: number, private bOffset: number, private destOffset: number) {}
+  constructor(private aOffset: number, private bOffset: number, private destOffset: number) {
+    super();
+  }
 
   execute(machineState: AvmMachineState, _stateManager: AvmStateManager): void {
     const a: Fr = machineState.readMemory(this.aOffset);
     const b: Fr = machineState.readMemory(this.bOffset);
 
     const dest = new Fr(a.toBigInt() << b.toBigInt());
     machineState.writeMemory(this.destOffset, dest);
+
+    this.incrementPc(machineState);
   }
 }
 
 /** -*/
-export class Shr implements Instruction {
+export class Shr extends Instruction {
   static type: string = 'SHR';
   static numberOfOperands = 3;
 
-  constructor(private aOffset: number, private bOffset: number, private destOffset: number) {}
+  constructor(private aOffset: number, private bOffset: number, private destOffset: number) {
+    super();
+  }
 
   execute(machineState: AvmMachineState, _stateManager: AvmStateManager): void {
     const a: Fr = machineState.readMemory(this.aOffset);
     const b: Fr = machineState.readMemory(this.bOffset);
 
     const dest = new Fr(a.toBigInt() >> b.toBigInt());
     machineState.writeMemory(this.destOffset, dest);
+
+    this.incrementPc(machineState);
   }
 }