From bcc780eff0363bd4d01bce51c1ec9eb1aabd1825 Mon Sep 17 00:00:00 2001
From: fcarreiro <facundo@aztecprotocol.com>
Date: Fri, 2 Feb 2024 10:47:08 +0000
Subject: [PATCH] feat(avm): complete SET instruction and handle invalid tags

---
 .../src/avm/opcodes/memory.test.ts            | 22 +++++++++++++++++++
 .../acir-simulator/src/avm/opcodes/memory.ts  |  8 +++++--
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/yarn-project/acir-simulator/src/avm/opcodes/memory.test.ts b/yarn-project/acir-simulator/src/avm/opcodes/memory.test.ts
index ce4e6165db4e..438e7d9eadb5 100644
--- a/yarn-project/acir-simulator/src/avm/opcodes/memory.test.ts
+++ b/yarn-project/acir-simulator/src/avm/opcodes/memory.test.ts
@@ -6,6 +6,7 @@ import { AvmMachineState } from '../avm_machine_state.js';
 import { Field, TypeTag, Uint8, Uint16, Uint32, Uint64, Uint128 } from '../avm_memory_types.js';
 import { initExecutionEnvironment } from '../fixtures/index.js';
 import { AvmJournal } from '../journal/journal.js';
+import { InstructionExecutionError } from './instruction.js';
 import { CMov, CalldataCopy, Cast, Mov, Set } from './memory.js';
 
 describe('Memory instructions', () => {
@@ -64,6 +65,27 @@ describe('Memory instructions', () => {
       expect(actual).toEqual(new Uint32(1234n));
       expect(tag).toEqual(TypeTag.UINT32);
     });
+
+    it('should correctly set value and tag (truncating)', async () => {
+      await new Set(/*indirect=*/ 0, /*inTag=*/ TypeTag.UINT16, /*value=*/ 0x12345678n, /*offset=*/ 1).execute(
+        machineState,
+        journal,
+      );
+
+      const actual = machineState.memory.get(1);
+      const tag = machineState.memory.getTag(1);
+
+      expect(actual).toEqual(new Uint16(0x5678));
+      expect(tag).toEqual(TypeTag.UINT16);
+    });
+
+    it('should throw if tag is FIELD, UNINITIALIZED, INVALID', async () => {
+      for (const tag of [TypeTag.FIELD, TypeTag.UNINITIALIZED, TypeTag.INVALID]) {
+        await expect(
+          new Set(/*indirect=*/ 0, /*inTag=*/ tag, /*value=*/ 1234n, /*offset=*/ 1).execute(machineState, journal),
+        ).rejects.toThrow(InstructionExecutionError);
+      }
+    });
   });
 
   describe('CAST', () => {
diff --git a/yarn-project/acir-simulator/src/avm/opcodes/memory.ts b/yarn-project/acir-simulator/src/avm/opcodes/memory.ts
index 31b05c959690..fd26f92060c2 100644
--- a/yarn-project/acir-simulator/src/avm/opcodes/memory.ts
+++ b/yarn-project/acir-simulator/src/avm/opcodes/memory.ts
@@ -2,7 +2,7 @@ import { AvmMachineState } from '../avm_machine_state.js';
 import { Field, TaggedMemory, TypeTag } from '../avm_memory_types.js';
 import { AvmJournal } from '../journal/index.js';
 import { Opcode, OperandType } from '../serialization/instruction_serialization.js';
-import { Instruction } from './instruction.js';
+import { Instruction, InstructionExecutionError } from './instruction.js';
 import { TwoOperandInstruction } from './instruction_impl.js';
 
 export class Set extends Instruction {
@@ -22,8 +22,12 @@ export class Set extends Instruction {
   }
 
   async execute(machineState: AvmMachineState, _journal: AvmJournal): Promise<void> {
-    const res = TaggedMemory.integralFromTag(this.value, this.inTag);
+    // Per the YP, the tag cannot be a field.
+    if ([TypeTag.FIELD, TypeTag.UNINITIALIZED, TypeTag.INVALID].includes(this.inTag)) {
+      throw new InstructionExecutionError(`Invalid tag ${TypeTag[this.inTag]} for SET.`);
+    }
 
+    const res = TaggedMemory.integralFromTag(this.value, this.inTag);
     machineState.memory.set(this.dstOffset, res);
 
     this.incrementPc(machineState);