From a94a3f86d6145771c5f6be277a9aaf0d9dd428d7 Mon Sep 17 00:00:00 2001
From: Elias Tazartes <66871571+Eikix@users.noreply.github.com>
Date: Thu, 16 Nov 2023 14:20:55 +0100
Subject: [PATCH] fix: signextend to fit specification (#808)

<!--- Please provide a general summary of your changes in the title
above -->

<!-- Give an estimate of the time you spent on this PR in terms of work
days.
Did you spend 0.5 days on this PR or rather 2 days?  -->

Time spent on this PR: 0.5d

## Pull request type

<!-- Please try to limit your pull request to one type,
submit multiple pull requests if needed. -->

Please check the type of change your PR introduces:

- [x] Bugfix
- [ ] Feature
- [ ] Code style update (formatting, renaming)
- [ ] Refactoring (no functional changes, no api changes)
- [ ] Build related changes
- [ ] Documentation content changes
- [ ] Other (please describe):

## What is the current behavior?

<!-- Please describe the current behavior that you are modifying,
or link to a relevant issue. -->

Resolves #677
---
 .../stop_and_math_operations.cairo            |  7 ++--
 src/utils/uint256.cairo                       | 39 ++++++++++++++++++-
 .../test_stop_and_math_operations.py          | 20 +++++++++-
 3 files changed, 60 insertions(+), 6 deletions(-)

diff --git a/src/kakarot/instructions/stop_and_math_operations.cairo b/src/kakarot/instructions/stop_and_math_operations.cairo
index 620f22e5f..9f9e513b9 100644
--- a/src/kakarot/instructions/stop_and_math_operations.cairo
+++ b/src/kakarot/instructions/stop_and_math_operations.cairo
@@ -8,7 +8,6 @@ from starkware.cairo.common.uint256 import (
     uint256_add,
     uint256_and,
     uint256_eq,
-    uint256_le,
     uint256_lt,
     uint256_mul_div_mod,
     uint256_mul,
@@ -32,7 +31,7 @@ from kakarot.model import model
 from kakarot.execution_context import ExecutionContext
 from kakarot.stack import Stack
 from kakarot.errors import Errors
-from utils.uint256 import uint256_exp
+from utils.uint256 import uint256_exp, uint256_signextend
 
 // @title Stop and Math operations opcodes.
 // @notice Math operations gathers Arithmetic and Comparison operations
@@ -253,11 +252,11 @@ namespace StopAndMathOperations {
         let range_check_ptr = [ap - 2];
         let popped = cast([ap - 1], Uint256*);
 
-        // TODO: see https://github.com/kkrt-labs/kakarot/issues/677
+        let result = uint256_signextend(popped[1], popped[0]);
 
         tempvar bitwise_ptr = cast([fp - 4], BitwiseBuiltin*);
         tempvar range_check_ptr = range_check_ptr;
-        tempvar result = Uint256(popped[1].low, popped[1].high);
+        tempvar result = result;
         jmp end;
 
         INVALID:
diff --git a/src/utils/uint256.cairo b/src/utils/uint256.cairo
index ba514ea14..a4b4b184c 100644
--- a/src/utils/uint256.cairo
+++ b/src/utils/uint256.cairo
@@ -1,4 +1,13 @@
-from starkware.cairo.common.uint256 import Uint256, uint256_eq, uint256_le, uint256_sub, uint256_mul
+from starkware.cairo.common.uint256 import (
+    Uint256,
+    uint256_eq,
+    uint256_le,
+    uint256_sub,
+    uint256_mul,
+    uint256_add,
+    uint256_pow2,
+    uint256_unsigned_div_rem,
+)
 from starkware.cairo.common.bool import FALSE
 
 // @notice Internal exponentiation of two 256-bit integers.
@@ -19,3 +28,31 @@ func uint256_exp{range_check_ptr}(a: Uint256, b: Uint256) -> Uint256 {
     let (res, _) = uint256_mul(a, pow);
     return res;
 }
+
+// @notice Extend a signed number which fits in N bytes to 32 bytes.
+// @param x The number to be sign extended.
+// @param byte_num The size in bytes minus one of x to consider.
+// @returns x if byteNum > 31, or x interpreted as a signed number with sign-bit at (byte_num*8+7), extended to the full 256 bits
+func uint256_signextend{range_check_ptr}(x: Uint256, byte_num: Uint256) -> Uint256 {
+    alloc_locals;
+    let (byte_num_gt_word_size) = uint256_le(Uint256(32, 0), byte_num);
+    if (byte_num_gt_word_size != 0) {
+        return x;
+    }
+
+    let sign_bit_position = byte_num.low * 8 + 7;
+
+    let (s) = uint256_pow2(Uint256(sign_bit_position, 0));
+    let (sign_bit, value) = uint256_unsigned_div_rem(x, s);
+    let (_, x_is_negative) = uint256_unsigned_div_rem(sign_bit, Uint256(2, 0));
+
+    if (x_is_negative.low == 0) {
+        return value;
+    }
+
+    let (mask) = uint256_sub(s, Uint256(1, 0));
+    let max_uint256 = Uint256(2 ** 128 - 1, 2 ** 128 - 1);
+    let (padding) = uint256_sub(max_uint256, mask);
+    let (value, _) = uint256_add(value, padding);
+    return value;
+}
diff --git a/tests/src/kakarot/instructions/test_stop_and_math_operations.py b/tests/src/kakarot/instructions/test_stop_and_math_operations.py
index e27b1bd40..33f8b85ec 100644
--- a/tests/src/kakarot/instructions/test_stop_and_math_operations.py
+++ b/tests/src/kakarot/instructions/test_stop_and_math_operations.py
@@ -37,7 +37,25 @@ class TestMathOperations:
                 (Opcodes.MULMOD, [3, 2, 2], (3 * 2) % 2),
                 (Opcodes.EXP, [3, 2], (3**2)),
                 (Opcodes.EXP, [3, 1], (3**1)),
-                (Opcodes.SIGNEXTEND, [3, 2, 1], 2),
+                (
+                    Opcodes.SIGNEXTEND,
+                    [
+                        0x01,
+                        0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF0001,
+                        3,
+                    ],
+                    0x01,
+                ),
+                (
+                    Opcodes.SIGNEXTEND,
+                    [0x00, 0xFF, 3],
+                    0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF,
+                ),
+                (
+                    Opcodes.SIGNEXTEND,
+                    [0x20, 0xFF, 3],
+                    0xFF,
+                ),
                 (Opcodes.LT, [2, 1], 0),
                 (Opcodes.LT, [1, 2], 1),
                 (Opcodes.GT, [2, 1], 1),