feat(avm): support variable size SET opcode (AztecProtocol#4465)

Closes AztecProtocol#4267.
Swoir · Feb 6, 2024 · 545b334 · 545b334
1 parent 5e16063
commit 545b334
Show file tree

Hide file tree

Showing 3 changed files with 124 additions and 11 deletions.
diff --git a/yarn-project/simulator/src/avm/avm_memory_types.ts b/yarn-project/simulator/src/avm/avm_memory_types.ts
@@ -322,8 +322,8 @@ export class TaggedMemory {
   }
 
   // Truncates the value to fit the type.
-  public static integralFromTag(v: bigint, tag: TypeTag): IntegralValue {
-    v = BigInt(v); // FIXME: not sure why this cast is needed, but this errors otherwise
+  public static integralFromTag(v: bigint | number, tag: TypeTag): IntegralValue {
+    v = v as bigint;
     switch (tag) {
       case TypeTag.UINT8:
         return new Uint8(v & ((1n << 8n) - 1n));

diff --git a/yarn-project/simulator/src/avm/opcodes/memory.test.ts b/yarn-project/simulator/src/avm/opcodes/memory.test.ts
@@ -14,17 +14,83 @@ describe('Memory instructions', () => {
   });
 
   describe('SET', () => {
-    it('Should (de)serialize correctly', () => {
+    it('Should (de)serialize correctly [tag=u8]', () => {
+      const buf = Buffer.from([
+        Set.opcode, // opcode
+        0x01, // indirect
+        TypeTag.UINT8, // inTag
+        ...Buffer.from('12', 'hex'),
+        ...Buffer.from('3456789a', 'hex'), // dstOffset
+      ]);
+      const inst = new Set(/*indirect=*/ 0x01, /*inTag=*/ TypeTag.UINT8, /*value=*/ 0x12, /*dstOffset=*/ 0x3456789a);
+
+      expect(Set.deserialize(buf)).toEqual(inst);
+      expect(inst.serialize()).toEqual(buf);
+    });
+
+    it('Should (de)serialize correctly [tag=u16]', () => {
+      const buf = Buffer.from([
+        Set.opcode, // opcode
+        0x01, // indirect
+        TypeTag.UINT16, // inTag
+        ...Buffer.from('1234', 'hex'),
+        ...Buffer.from('3456789a', 'hex'), // dstOffset
+      ]);
+      const inst = new Set(/*indirect=*/ 0x01, /*inTag=*/ TypeTag.UINT16, /*value=*/ 0x1234, /*dstOffset=*/ 0x3456789a);
+
+      expect(Set.deserialize(buf)).toEqual(inst);
+      expect(inst.serialize()).toEqual(buf);
+    });
+
+    it('Should (de)serialize correctly [tag=u32]', () => {
+      const buf = Buffer.from([
+        Set.opcode, // opcode
+        0x01, // indirect
+        TypeTag.UINT32, // inTag
+        ...Buffer.from('12345678', 'hex'),
+        ...Buffer.from('3456789a', 'hex'), // dstOffset
+      ]);
+      const inst = new Set(
+        /*indirect=*/ 0x01,
+        /*inTag=*/ TypeTag.UINT32,
+        /*value=*/ 0x12345678,
+        /*dstOffset=*/ 0x3456789a,
+      );
+
+      expect(Set.deserialize(buf)).toEqual(inst);
+      expect(inst.serialize()).toEqual(buf);
+    });
+
+    it('Should (de)serialize correctly [tag=u64]', () => {
+      const buf = Buffer.from([
+        Set.opcode, // opcode
+        0x01, // indirect
+        TypeTag.UINT64, // inTag
+        ...Buffer.from('1234567812345678', 'hex'),
+        ...Buffer.from('3456789a', 'hex'), // dstOffset
+      ]);
+      const inst = new Set(
+        /*indirect=*/ 0x01,
+        /*inTag=*/ TypeTag.UINT64,
+        /*value=*/ 0x1234567812345678n,
+        /*dstOffset=*/ 0x3456789a,
+      );
+
+      expect(Set.deserialize(buf)).toEqual(inst);
+      expect(inst.serialize()).toEqual(buf);
+    });
+
+    it('Should (de)serialize correctly [tag=u128]', () => {
       const buf = Buffer.from([
         Set.opcode, // opcode
         0x01, // indirect
-        TypeTag.FIELD, // inTag
+        TypeTag.UINT128, // inTag
         ...Buffer.from('12345678123456781234567812345678', 'hex'), // const (will be 128 bit)
         ...Buffer.from('3456789a', 'hex'), // dstOffset
       ]);
       const inst = new Set(
         /*indirect=*/ 0x01,
-        /*inTag=*/ TypeTag.FIELD,
+        /*inTag=*/ TypeTag.UINT128,
         /*value=*/ 0x12345678123456781234567812345678n,
         /*dstOffset=*/ 0x3456789a,
       );

diff --git a/yarn-project/simulator/src/avm/opcodes/memory.ts b/yarn-project/simulator/src/avm/opcodes/memory.ts
@@ -1,26 +1,73 @@
 import type { AvmContext } from '../avm_context.js';
 import { Field, TaggedMemory, TypeTag } from '../avm_memory_types.js';
 import { InstructionExecutionError } from '../errors.js';
-import { Opcode, OperandType } from '../serialization/instruction_serialization.js';
+import { BufferCursor } from '../serialization/buffer_cursor.js';
+import { Opcode, OperandType, deserialize, serialize } from '../serialization/instruction_serialization.js';
 import { Instruction } from './instruction.js';
 import { TwoOperandInstruction } from './instruction_impl.js';
 
+const TAG_TO_OPERAND_TYPE = new Map<TypeTag, OperandType>([
+  [TypeTag.UINT8, OperandType.UINT8],
+  [TypeTag.UINT16, OperandType.UINT16],
+  [TypeTag.UINT32, OperandType.UINT32],
+  [TypeTag.UINT64, OperandType.UINT64],
+  [TypeTag.UINT128, OperandType.UINT128],
+]);
+
+function getOperandTypeFromInTag(inTag: number | bigint): OperandType {
+  inTag = inTag as number;
+  const tagOperandType = TAG_TO_OPERAND_TYPE.get(inTag);
+  if (tagOperandType === undefined) {
+    throw new Error(`Invalid tag ${inTag} for SET.`);
+  }
+  return tagOperandType;
+}
+
 export class Set extends Instruction {
   static readonly type: string = 'SET';
   static readonly opcode: Opcode = Opcode.SET;
-  // Informs (de)serialization. See Instruction.deserialize.
-  static readonly wireFormat: OperandType[] = [
+
+  private static readonly wireFormatBeforeConst: OperandType[] = [
     OperandType.UINT8,
     OperandType.UINT8,
     OperandType.UINT8,
-    OperandType.UINT128,
-    OperandType.UINT32,
   ];
+  private static readonly wireFormatAfterConst: OperandType[] = [OperandType.UINT32];
 
-  constructor(private indirect: number, private inTag: number, private value: bigint, private dstOffset: number) {
+  constructor(
+    private indirect: number,
+    private inTag: number,
+    private value: bigint | number,
+    private dstOffset: number,
+  ) {
     super();
   }
 
+  public serialize(): Buffer {
+    const format: OperandType[] = [
+      ...Set.wireFormatBeforeConst,
+      getOperandTypeFromInTag(this.inTag),
+      ...Set.wireFormatAfterConst,
+    ];
+    return serialize(format, this);
+  }
+
+  public static deserialize<T extends { new (...args: any[]): InstanceType<T> }>(
+    this: T,
+    buf: BufferCursor | Buffer,
+  ): InstanceType<T> {
+    if (buf instanceof Buffer) {
+      buf = new BufferCursor(buf);
+    }
+    const beforeConst = deserialize(buf, Set.wireFormatBeforeConst);
+    const tag = beforeConst[beforeConst.length - 1];
+    const val = deserialize(buf, [getOperandTypeFromInTag(tag)]);
+    const afterConst = deserialize(buf, Set.wireFormatAfterConst);
+    const res = [...beforeConst, ...val, ...afterConst];
+    const args = res.slice(1) as ConstructorParameters<typeof Set>; // Remove opcode.
+    return new this(...args);
+  }
+
   async execute(context: AvmContext): Promise<void> {
     // Per the YP, the tag cannot be a field.
     if ([TypeTag.FIELD, TypeTag.UNINITIALIZED, TypeTag.INVALID].includes(this.inTag)) {