test: add tests for i32 pow2/ipow intrinsics

codegen/masm/intrinsics/i32.masm

@@ -288,53 +288,55 @@ export.pow2 # [n]
     push.31
     u32checked_lt  # [n < 31, pow]
     assert         # [n]
-    push.1 swap.1 # [n, 1]
+    push.1 swap.1  # [n, 1]
     u32checked_shl # [1 << n]
 end
 
 # Compute a^b, where `b` must be a positive i32 value < 31
 export.ipow # [b, a]
-    dup.0 eq.0      # [b == 0, b, a]
-    dup.2 eq.0      # [a == 0, b == 0, b, a]
+    dup.0 push.31 u32checked_lt assert # assert that `b` is < 31
+    dup.0 eq.0       # [b == 0, b, a]
+    dup.2 eq.0       # [a == 0, b == 0, b, a]
     or               # [a == 0 || b == 0, b, a]
     # if a == 0, the result is always 0; otherwise if b == 0, then the result is always 1
     if.true
         eq.0                # [b == 0, a]
-        push.1 push.0      # [0, 1, b == 0, a]
+        push.1 push.0       # [0, 1, b == 0, a]
         swap.2              # [b == 0, 1, 0, a]
         cdrop               # [1 or 0, a]
-        swap.1 drop        # [1 or 0]
+        swap.1 drop         # [1 or 0]
     else  # [b, a]
-      # for all other values, we do exponentiation by squaring
-      push.1         # [acc, b, a]
-      dup.1          # [b, acc, b, a]
-      push.1         # [1, b, acc, b, a]
-      u32checked_gt  # [b > 1, acc, b, a]
-      while.true     # [acc, b, a => base]
-          dup.2 dup.1      # [acc, base, acc, b, base]
-          mul              # [base * acc, acc, b, base]
-          dup.2            # [b, base * acc, acc, b, base]
-          push.1           # [1, b, base * acc, acc, b, base]
-          u32checked_and   # [b & 1, base * acc, acc, b, base]
-          eq.1             # [b & 1 == 1, base * acc, acc, b, base]
-          cdrop            # [acc, b, base]
-          u32assert
-          swap.1           # [b, acc, base]
-          u32checked_div.2 # [b /= 2, acc, base]
-          movup.2 dup.0   # [base, base, b, acc]
-          u32checked_mul   # [base * base, b, acc]
-          swap.1           # [b, base, acc]
-          movup.2          # [acc, b, base]
-          dup.1 push.1    # [1, b, acc, b, base]
-          u32checked_gt    # [b > 1, acc, b, base]
-      end
-      swap.1 drop   # [acc, base]
-      u32checked_mul # [acc * base]
+        # for all other values, we do exponentiation by squaring
+        push.1         # [acc, b, a]
+        dup.1          # [b, acc, b, a]
+        push.1         # [1, b, acc, b, a]
+        u32checked_gt  # [b > 1, acc, b, a]
+        while.true     # [acc, b, a => base]
+            dup.2 dup.1      # [acc, base, acc, b, base]
+            u32wrapping_mul  # [base * acc, acc, b, base]
+            dup.2            # [b, base * acc, acc, b, base]
+            push.1           # [1, b, base * acc, acc, b, base]
+            u32checked_and   # [b & 1, base * acc, acc, b, base]
+            eq.1             # [b & 1 == 1, base * acc, acc, b, base]
+            cdrop            # [acc, b, base]
+            swap.1           # [b, acc, base]
+            u32checked_div.2 # [b /= 2, acc, base]
+            movup.2 dup.0    # [base, base, b, acc]
+            u32wrapping_mul  # [base * base, b, acc]
+            swap.1           # [b, base, acc]
+            movup.2          # [acc, b, base]
+            dup.1 push.1     # [1, b, acc, b, base]
+            u32checked_gt    # [b > 1, acc, b, base]
+        end
+        swap.1 drop     # [acc, base]
+        u32wrapping_mul # [acc * base]
     end
 end
 
 # Arithmetic shift-right, i.e. `a >> b` preserves the signedness of the value
 #
+# This function will assert if `b` is > 31.
+#
 # This implementation is checked, so it will assert if the inputs are invalid
 export.checked_shr # [b, a]
     # validate the shift is valid

codegen/masm/src/tests.rs

@@ -703,4 +703,57 @@ proptest! {
         let result = u32::try_from(raw_result).map_err(|_| raw_result).map(|res| res as i32);
         prop_assert_eq!(result, Ok(a / b.get()));
     }
+
+    #[test]
+    fn i32_pow2(a in 0u32..30u32) {
+        let mut harness = TestByEmulationHarness::default();
+
+        harness
+            .emulator
+            .load_module(
+                Box::new(
+                    Module::load_intrinsic("intrinsics::i32", &harness.context.session.codemap)
+                        .expect("parsing failed"),
+                )
+                .freeze(),
+            )
+            .expect("failed to load intrinsics::i32");
+
+        let a_felt = Felt::new(a as u64);
+        let pow2 = "intrinsics::i32::pow2".parse().unwrap();
+        let mut stack = harness.invoke(pow2, &[a_felt]).expect("execution failed");
+        harness.emulator.stop();
+
+        prop_assert_eq!(stack.len(), 1);
+        let raw_result = stack.pop().unwrap().as_int();
+        let result = u32::try_from(raw_result).map_err(|_| raw_result).map(|res| res as i32);
+        prop_assert_eq!(result, Ok(2i32.pow(a)));
+    }
+
+    #[test]
+    fn i32_ipow(a: i32, b in 0u32..30u32) {
+        let mut harness = TestByEmulationHarness::default();
+
+        harness
+            .emulator
+            .load_module(
+                Box::new(
+                    Module::load_intrinsic("intrinsics::i32", &harness.context.session.codemap)
+                        .expect("parsing failed"),
+                )
+                .freeze(),
+            )
+            .expect("failed to load intrinsics::i32");
+
+        let a_felt = Felt::new(a as u32 as u64);
+        let b_felt = Felt::new(b as u64);
+        let ipow = "intrinsics::i32::ipow".parse().unwrap();
+        let mut stack = harness.invoke(ipow, &[b_felt, a_felt]).expect("execution failed");
+        harness.emulator.stop();
+
+        prop_assert_eq!(stack.len(), 1);
+        let raw_result = stack.pop().unwrap().as_int();
+        let result = u32::try_from(raw_result).map_err(|_| raw_result).map(|res| res as i32);
+        prop_assert_eq!(result, Ok(a.wrapping_pow(b)));
+    }
 }