Fix underflow bug

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 #### Upcoming Changes
 
+* fix: Using UINT256_HINT no longer panics when b is greater than 2^256 [#1430](https://github.com/lambdaclass/cairo-vm/pull/1430)
+
 * feat: Added a differential fuzzer for programs with whitelisted hints [#1358](https://github.com/lambdaclass/cairo-vm/pull/1358)
 
 * fix: Change return type of `get_execution_resources` to `RunnerError` [#1398](https://github.com/lambdaclass/cairo-vm/pull/1398)

cairo_programs/bad_programs/uint256_sub_b_gt_256.cairo

@@ -0,0 +1,37 @@
+
+struct MyStruct0 {
+	low: felt,
+	high: felt,
+}
+func main() {
+	let a = MyStruct0(low=54987052180710815841462937121160005261203577276882008045698301095581843457, high=10);
+	let b = MyStruct0(low=2, high=3618502788666131213697322783095070105623107215331596699973087835080668217346);
+	let res = MyStruct0(low=1, high=1);
+	hint_func(a, b, res);
+	return();
+}
+
+func hint_func(a: MyStruct0, b: MyStruct0, res: MyStruct0) -> (MyStruct0, MyStruct0, MyStruct0) {
+	alloc_locals;
+
+%{
+def split(num: int, num_bits_shift: int = 128, length: int = 2):
+    a = []
+    for _ in range(length):
+        a.append( num & ((1 << num_bits_shift) - 1) )
+        num = num >> num_bits_shift
+    return tuple(a)
+
+def pack(z, num_bits_shift: int = 128) -> int:
+    limbs = (z.low, z.high)
+    return sum(limb << (num_bits_shift * i) for i, limb in enumerate(limbs))
+
+a = pack(ids.a)
+b = pack(ids.b)
+res = (a - b)%2**256
+res_split = split(res)
+ids.res.low = res_split[0]
+ids.res.high = res_split[1]
+%}
+	return(a, b, res);
+}

vm/src/hint_processor/builtin_hint_processor/uint256_utils.rs

@@ -194,6 +194,7 @@ Implements hint:
     ids.res.high = res_split[1]
 %}
 */
+
 pub fn uint256_sub(
     vm: &mut VirtualMachine,
     ids_data: &HashMap<String, HintReference>,
@@ -208,7 +209,22 @@ pub fn uint256_sub(
         a - b
     } else {
         // wrapped a - b
-        ((BigUint::one() << 256) - b) + a
+        // b is limited to (CAIRO_PRIME - 1) << 128 which is 1 << (251 + 128 + 1)
+        //                                         251: most significant felt bit
+        //                                         128:     high field left shift
+        //                                           1:       extra bit for limit
+        let mod_256 = BigUint::one() << 256;
+        if mod_256 >= b {
+            mod_256 - b + a
+        } else {
+            let lowered_b = b.mod_floor(&mod_256);
+            // Repeat the logic from before
+            if a >= lowered_b {
+                a - lowered_b
+            } else {
+                mod_256 - lowered_b + a
+            }
+        }
     };
 
     let res = Uint256::split(&res);
@@ -633,6 +649,52 @@ mod tests {
         ];
     }
 
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    fn run_uint256_sub_b_high_gt_256_lte_a() {
+        let hint_code = hint_code::UINT256_SUB;
+        let mut vm = vm_with_range_check!();
+        //Initialize fp
+        vm.run_context.fp = 0;
+        //Create hint_data
+        let ids_data = non_continuous_ids_data![("a", 0), ("b", 2), ("res", 4)];
+        vm.segments = segments![
+            ((1, 0), ("340282366920938463463374607431768211456", 10)),
+            ((1, 1), 0),
+            ((1, 2), 0),
+            ((1, 3), ("340282366920938463463374607431768211457", 10)),
+        ];
+        //Execute the hint
+        assert_matches!(run_hint!(vm, ids_data, hint_code), Ok(()));
+        //Check hint memory inserts
+        check_memory![vm.segments.memory, ((1, 4), 0), ((1, 5), 0)];
+    }
+
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    fn run_uint256_sub_b_high_gt_256_gt_a() {
+        let hint_code = hint_code::UINT256_SUB;
+        let mut vm = vm_with_range_check!();
+        //Initialize fp
+        vm.run_context.fp = 0;
+        //Create hint_data
+        let ids_data = non_continuous_ids_data![("a", 0), ("b", 2), ("res", 4)];
+        vm.segments = segments![
+            ((1, 0), 1),
+            ((1, 1), 0),
+            ((1, 2), 0),
+            ((1, 3), ("340282366920938463463374607431768211457", 10)),
+        ];
+        //Execute the hint
+        assert_matches!(run_hint!(vm, ids_data, hint_code), Ok(()));
+        //Check hint memory inserts
+        check_memory![
+            vm.segments.memory,
+            ((1, 4), ("1", 10)),
+            ((1, 5), ("340282366920938463463374607431768211455", 10))
+        ];
+    }
+
     #[test]
     #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
     fn run_uint256_sub_missing_member() {

vm/src/vm/errors/vm_exception.rs

@@ -908,6 +908,43 @@ cairo_programs/bad_programs/bad_usort.cairo:64:5: (pc=0:60)
         assert_eq!(vm_excepction.to_string(), expected_error_string);
     }
 
+    #[test]
+    #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
+    fn run_bad_uint256_sub_check_error_displayed() {
+        #[cfg(feature = "std")]
+        let expected_error_string = r#"cairo_programs/bad_programs/uint256_sub_b_gt_256.cairo:17:1: Error at pc=0:17:
+Got an exception while executing a hint: Inconsistent memory assignment at address Relocatable { segment_index: 1, offset: 6 }. Int(1) != Int(41367660292349381832802403122744918015)
+%{
+^^
+Cairo traceback (most recent call last):
+cairo_programs/bad_programs/uint256_sub_b_gt_256.cairo:10:2: (pc=0:12)
+	hint_func(a, b, res);
+ ^******************^
+"#;
+        #[cfg(not(feature = "std"))]
+        let expected_error_string = r#"cairo_programs/bad_programs/uint256_sub_b_gt_256.cairo:17:1: Error at pc=0:17:
+Got an exception while executing a hint: Inconsistent memory assignment at address Relocatable { segment_index: 1, offset: 6 }. Int(1) != Int(41367660292349381832802403122744918015)
+Cairo traceback (most recent call last):
+cairo_programs/bad_programs/uint256_sub_b_gt_256.cairo:10:2: (pc=0:12)
+"#;
+        let program = Program::from_bytes(
+            include_bytes!("../../../../cairo_programs/bad_programs/uint256_sub_b_gt_256.json"),
+            Some("main"),
+        )
+        .unwrap();
+
+        let mut hint_processor = BuiltinHintProcessor::new_empty();
+        let mut cairo_runner = cairo_runner!(program, "all_cairo", false);
+        let mut vm = vm!();
+
+        let end = cairo_runner.initialize(&mut vm).unwrap();
+        let error = cairo_runner
+            .run_until_pc(end, &mut vm, &mut hint_processor)
+            .unwrap_err();
+        let vm_excepction = VmException::from_vm_error(&cairo_runner, &vm, error);
+        assert_eq!(vm_excepction.to_string(), expected_error_string);
+    }
+
     #[test]
     #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test)]
     fn get_value_from_simple_reference_ap_based() {