From 348163af6f2cbae358b6c6215ee5520ec89bb813 Mon Sep 17 00:00:00 2001
From: JK <jkktom@gmail.com>
Date: Mon, 2 Oct 2023 16:00:43 +0900
Subject: [PATCH] Fix as per comments of prev. PR. Fixes issue #56 adding test
 cases and code change

---
 pkg/vm/memory/memory_value.go | 56 +++++++++++++----------------------
 pkg/vm/vm.go                  |  2 +-
 2 files changed, 21 insertions(+), 37 deletions(-)

diff --git a/pkg/vm/memory/memory_value.go b/pkg/vm/memory/memory_value.go
index 2312c0aff..0d8e5277c 100644
--- a/pkg/vm/memory/memory_value.go
+++ b/pkg/vm/memory/memory_value.go
@@ -26,19 +26,16 @@ func (address *MemoryAddress) Equal(other *MemoryAddress) bool {
 }
 
 // Adds a memory address and a field element
-func (ma *MemoryAddress) Add(address *MemoryAddress, offset *f.Element) (*MemoryAddress, error) {
-	addressOffset := new(f.Element).SetUint64(address.Offset)
-	resOffset := new(f.Element).Add(addressOffset, offset)
 
-	if !resOffset.IsUint64() {
-		return nil, fmt.Errorf("new offset bigger than uint64: %s", resOffset.Text(10))
+func (address *MemoryAddress) Add(lhs *MemoryAddress, rhs *f.Element) (*MemoryAddress, error) {
+	lhsOffset := new(f.Element).SetUint64(lhs.Offset)
+	newOffset := new(f.Element).Add(lhsOffset, rhs)
+	if !newOffset.IsUint64() {
+		return nil, fmt.Errorf("new offset bigger than uint64: %s", rhs.Text(10))
 	}
-
-	resAddress := &MemoryAddress{
-		SegmentIndex: address.SegmentIndex,
-		Offset:       resOffset.Uint64(),
-	}
-	return resAddress, nil
+	address.SegmentIndex = lhs.SegmentIndex
+	address.Offset = newOffset.Uint64()
+	return address, nil
 }
 
 // Subs from a memory address a felt or another memory address in the same segment
@@ -214,40 +211,27 @@ func (mv *MemoryValue) Equal(other *MemoryValue) bool {
 	return false
 }
 
-// Adds two memory values is the second one is a Felt
 func (mv *MemoryValue) Add(lhs, rhs *MemoryValue) (*MemoryValue, error) {
 	var err error
 
 	// If both lhs and rhs are felts, perform a simple addition
-	if !lhs.IsAddress() && !rhs.IsAddress() {
-		lhsUint := lhs.felt.Uint64()
-		rhsUint := rhs.felt.Uint64()
-		sumRes := lhsUint + rhsUint
-		sumMemoryValue := MemoryValueFromUint(sumRes)
-		mv.felt = sumMemoryValue.felt
-		return mv, nil
-	}
-
-	// If lhs is an address and rhs is a felt
-	if lhs.IsAddress() && !rhs.IsAddress() {
-		mv.address, err = mv.address.Add(lhs.address, rhs.felt)
-		if err != nil {
-			return nil, err
+	if lhs.IsFelt() {
+		if rhs.IsFelt() { // Felt + Felt
+			mv.felt = MemoryValueFromUint(lhs.felt.Uint64() + rhs.felt.Uint64()).felt
+		} else { // Felt + Address
+			mv.address, err = mv.address.Add(rhs.address, lhs.felt)
 		}
-		return mv, nil
+	} else if rhs.IsFelt() { // Address + Felt
+		mv.address, err = mv.address.Add(lhs.address, rhs.felt)
+	} else { // Address + Address
+		return nil, errors.New("addition of two addresses is not supported")
 	}
 
-	// If rhs is an address and lhs is a felt
-	if !lhs.IsAddress() && rhs.IsAddress() {
-		mv.address, err = mv.address.Add(rhs.address, lhs.felt)
-		if err != nil {
-			return nil, err
-		}
-		return mv, nil
+	if err != nil {
+		return nil, err
 	}
 
-	// If both lhs and rhs are addresses, this is an unsupported operation
-	return nil, errors.New("addition of two addresses is not supported")
+	return mv, nil
 }
 
 // Subs two memory values if they're in the same segment or the rhs is a Felt.
diff --git a/pkg/vm/vm.go b/pkg/vm/vm.go
index fec0eb389..2826a036d 100644
--- a/pkg/vm/vm.go
+++ b/pkg/vm/vm.go
@@ -345,7 +345,7 @@ func (vm *VirtualMachine) computeRes(
 	case AddOperands:
 		op0 := op0Cell.Read()
 		op1 := op1Cell.Read()
-		return mem.EmptyMemoryValueAsAddress().Add(op0, op1)
+		return mem.EmptyMemoryValueAs(op0.IsAddress() || op1.IsAddress()).Add(op0, op1)
 	case MulOperands:
 		op0 := op0Cell.Read()
 		op1 := op1Cell.Read()