Add builtins for some checked ops such as checked_add

2 files changed, 60 insertions, 0 deletions

diff --git a/backends/lean/Base/Primitives/Scalar.lean b/backends/lean/Base/Primitives/Scalar.lean
index 3d90f1a5..f46aded9 100644
--- a/backends/lean/Base/Primitives/Scalar.lean
+++ b/backends/lean/Base/Primitives/Scalar.lean
@@ -568,6 +568,35 @@ instance {ty} : HOr (Scalar ty) (Scalar ty) (Scalar ty) where
 instance {ty} : HAnd (Scalar ty) (Scalar ty) (Scalar ty) where
   hAnd x y := Scalar.and x y
 
+-- core checked arithmetic operations
+
+/- A helper function that converts failure to none and success to some
+   TODO: move up to Base module? -/
+@[simp] abbrev Option.ofResult : Result (Scalar ty) → Result (Option (Scalar ty))
+  | Result.ret z => Result.ret (some z)
+  | Result.fail _ => Result.ret none
+  | Result.div => Result.div
+
+/- [core::num::{T}::checked_add] -/
+def core.num.checked_add (x y : Scalar ty) : Result (Option (Scalar ty)) :=
+  Option.ofResult (x + y)
+
+/- [core::num::{T}::checked_sub] -/
+def core.num.checked_sub (x y : Scalar ty) : Result (Option (Scalar ty)) :=
+  Option.ofResult (x - y)
+
+/- [core::num::{T}::checked_mul] -/
+def core.num.checked_mul (x y : Scalar ty) : Result (Option (Scalar ty)) :=
+  Option.ofResult (x * y)
+
+/- [core::num::{T}::checked_div] -/
+def core.num.checked_div (x y : Scalar ty) : Result (Option (Scalar ty)) :=
+  Option.ofResult (x / y)
+
+/- [core::num::{T}::checked_mod] -/
+def core.num.checked_mod (x y : Scalar ty) : Result (Option (Scalar ty)) :=
+  Option.ofResult (x % y)
+
 -- Generic theorem - shouldn't be used much
 @[pspec]
 theorem Scalar.add_spec {ty} {x y : Scalar ty}
diff --git a/compiler/ExtractBuiltin.ml b/compiler/ExtractBuiltin.ml
index 401d0137..a2983573 100644
--- a/compiler/ExtractBuiltin.ml
+++ b/compiler/ExtractBuiltin.ml
@@ -240,6 +240,27 @@ let builtin_funs () : (pattern * bool list option * builtin_fun_info) list =
     let f = { extract_name = basename } in
     (rust_name, filter, f)
   in
+  let mk_scalar_fun (rust_name_prefix : string) (rust_name_suffix : string)
+      (extract_name : string option) (filter : bool list option) :
+      (pattern * bool list option * builtin_fun_info) list =
+    List.map
+      (fun ty ->
+        mk_fun (rust_name_prefix ^ ty ^ rust_name_suffix) extract_name filter)
+      [
+        "usize";
+        "u8";
+        "u16";
+        "u32";
+        "u64";
+        "u128";
+        "isize";
+        "i8";
+        "i16";
+        "i32";
+        "i64";
+        "i128";
+      ]
+  in
   [
     mk_fun "core::mem::replace" None None;
     mk_fun "core::slice::{[@T]}::len"
@@ -325,6 +346,16 @@ let builtin_funs () : (pattern * bool list option * builtin_fun_info) list =
        [@T]>}::index_mut"
       (Some "core_slice_index_Slice_index_mut") None;
   ]
+  @ mk_scalar_fun "core::num::{" "}::checked_add" (Some "core.num.checked_add")
+      None
+  @ mk_scalar_fun "core::num::{" "}::checked_sub" (Some "core.num.checked_sub")
+      None
+  @ mk_scalar_fun "core::num::{" "}::checked_mul" (Some "core.num.checked_mul")
+      None
+  @ mk_scalar_fun "core::num::{" "}::checked_div" (Some "core.num.checked_div")
+      None
+  @ mk_scalar_fun "core::num::{" "}::checked_rem" (Some "core.num.checked_rem")
+      None
 
 let mk_builtin_funs_map () =
   let m =