feat: add gestResOperandRelocatable method (#131)

* feat: add gestResOperandRelocatable method

* refactor: add a common method to get a SegmentValue and two methods to enforce Felt or Relocatable

src/vm/virtualMachine.test.ts

@@ -721,7 +721,7 @@ describe('VirtualMachine', () => {
           new Felt(21n),
         ],
       ])(
-        'should properly read ResOperand',
+        'should properly read ResOperand as Felt',
         (resOperand: ResOperand, expected: Felt) => {
           const vm = new VirtualMachine();
           vm.memory.addSegment();
@@ -735,6 +735,119 @@ describe('VirtualMachine', () => {
           expect(vm.getResOperandValue(resOperand)).toEqual(expected);
         }
       );
+
+      test.each([
+        [
+          {
+            type: OpType.Deref,
+            cell: { register: Register.Ap, offset: 0 },
+          },
+          new Relocatable(2, 4),
+        ],
+        [
+          {
+            type: OpType.DoubleDeref,
+            cell: { register: Register.Ap, offset: 1 },
+            offset: -1,
+          },
+          new Relocatable(2, 4),
+        ],
+        [
+          {
+            type: OpType.BinOp,
+            op: Operation.Add,
+            a: { register: Register.Fp, offset: 0 },
+            b: { type: OpType.Immediate, value: new Felt(5n) },
+          },
+          new Relocatable(2, 9),
+        ],
+        [
+          {
+            type: OpType.BinOp,
+            op: Operation.Add,
+            a: { register: Register.Ap, offset: 0 },
+            b: {
+              type: OpType.Deref,
+              cell: { register: Register.Ap, offset: 2 },
+            },
+          },
+          new Relocatable(2, 16),
+        ],
+      ])(
+        'should properly read ResOperand as Relocatable',
+        (resOperand: ResOperand, expected: Relocatable) => {
+          const vm = new VirtualMachine();
+          vm.memory.addSegment();
+          vm.memory.addSegment();
+          const address0 = new Relocatable(2, 4);
+          const address1 = new Relocatable(1, 1);
+          const value = new Felt(12n);
+          vm.memory.assertEq(vm.ap, address0);
+          vm.memory.assertEq(vm.ap.add(1), address1);
+          vm.memory.assertEq(vm.ap.add(2), value);
+          expect(vm.getResOperandRelocatable(resOperand)).toEqual(expected);
+        }
+      );
+
+      test('should throw ExpectedRelocatable when extracting from an Immediate ResOperand', () => {
+        const vm = new VirtualMachine();
+        vm.memory.addSegment();
+        vm.memory.addSegment();
+        const address0 = new Relocatable(2, 4);
+        const address1 = new Relocatable(1, 1);
+        const value = new Felt(12n);
+        vm.memory.assertEq(vm.ap, address0);
+        vm.memory.assertEq(vm.ap.add(1), address1);
+        vm.memory.assertEq(vm.ap.add(2), value);
+
+        const resOperand = {
+          type: OpType.Immediate,
+          value: new Felt(5n),
+        };
+        expect(() => vm.getResOperandRelocatable(resOperand)).toThrow(
+          new ExpectedRelocatable(resOperand.value)
+        );
+      });
+
+      test.each([
+        [
+          {
+            type: OpType.BinOp,
+            op: Operation.Mul,
+            a: { register: Register.Fp, offset: 0 },
+            b: { type: OpType.Immediate, value: new Felt(5n) },
+          },
+          new Relocatable(2, 4),
+        ],
+        [
+          {
+            type: OpType.BinOp,
+            op: Operation.Mul,
+            a: { register: Register.Ap, offset: 0 },
+            b: {
+              type: OpType.Deref,
+              cell: { register: Register.Ap, offset: 2 },
+            },
+          },
+          new Relocatable(2, 4),
+        ],
+      ])(
+        'should throw ExpectedFelt with BinOp + Operation.Mul ResOperand',
+        (resOperand: ResOperand, receivedValue: Relocatable) => {
+          const vm = new VirtualMachine();
+          vm.memory.addSegment();
+          vm.memory.addSegment();
+          const address0 = new Relocatable(2, 4);
+          const address1 = new Relocatable(1, 1);
+          const value = new Felt(12n);
+          vm.memory.assertEq(vm.ap, address0);
+          vm.memory.assertEq(vm.ap.add(1), address1);
+          vm.memory.assertEq(vm.ap.add(2), value);
+          expect(() => vm.getResOperandRelocatable(resOperand)).toThrow(
+            new ExpectedFelt(receivedValue)
+          );
+        }
+      );
     });
   });
 });

src/vm/virtualMachine.ts

@@ -550,7 +550,31 @@ export class VirtualMachine {
   }
 
   /**
-   * Return the value defined by `resOperand`
+   * Get the Felt defined by `resOperand`
+   *
+   * @param resOperand - The ResOperand to extract a Felt from.
+   * @returns {Felt} The value expressed by the given ResOperand.
+   */
+  getResOperandValue(resOperand: ResOperand): Felt {
+    const value = this.getResOperandSegmentValue(resOperand);
+    if (!isFelt(value)) throw new ExpectedFelt(value);
+    return value;
+  }
+
+  /**
+   * Get the Relocatable defined by `resOperand`
+   *
+   * @param resOperand - The ResOperand to extract a Relocatable from.
+   * @returns {Relocatable} The value expressed by the given ResOperand.
+   */
+  getResOperandRelocatable(resOperand: ResOperand): Relocatable {
+    const value = this.getResOperandSegmentValue(resOperand);
+    if (!isRelocatable(value)) throw new ExpectedRelocatable(value);
+    return value;
+  }
+
+  /**
+   * Return the SegmentValue value defined by `resOperand`
    *
    * Generic patterns:
    * - Deref: `[register + offset]`
@@ -561,26 +585,30 @@ export class VirtualMachine {
    * - BinOp (Mul): `[register1 + offset1] * [register2 + offset2]`
    * or `[register1 + offset1] * immediate`
    *
+   * @param {ResOperand} resOperand - The ResOperand to extract a Felt from.
+   * @returns {Felt} The value expressed by the given ResOperand.
+   * @throws {ExpectedFelt} If ResOperand is a BinOp (Mul) with `a` being a Relocatable.
+   *
    * NOTE: used in Cairo hints
    */
-  getResOperandValue(resOperand: ResOperand): Felt {
+  getResOperandSegmentValue(resOperand: ResOperand): SegmentValue {
     switch (resOperand.type) {
       case OpType.Deref:
-        return this.getFelt((resOperand as Deref).cell);
+        return this.getSegmentValue((resOperand as Deref).cell);
 
       case OpType.DoubleDeref:
         const dDeref = resOperand as DoubleDeref;
         const deref = this.getRelocatable(dDeref.cell);
         const value = this.memory.get(deref.add(dDeref.offset));
-        if (!value || !isFelt(value)) throw new ExpectedFelt(value);
+        if (!value) throw new UndefinedSegmentValue();
         return value;
 
       case OpType.Immediate:
         return (resOperand as Immediate).value;
 
       case OpType.BinOp:
         const binOp = resOperand as BinOp;
-        const a = this.getFelt(binOp.a);
+        const a = this.getSegmentValue(binOp.a);
 
         let b: Felt | undefined = undefined;
         switch (binOp.b.type) {
@@ -601,6 +629,7 @@ export class VirtualMachine {
             return a.add(b);
 
           case Operation.Mul:
+            if (!isFelt(a)) throw new ExpectedFelt(a);
             return a.mul(b);
         }
     }