Fixing const folded NonZero types.

corelib/src/zeroable.cairo

@@ -76,6 +76,14 @@ pub(crate) impl Felt252Zeroable = zero_based::ZeroableImpl<felt252>;
 /// This type guarantees that the wrapped value is never zero.
 #[derive(Copy, Drop)]
 pub extern type NonZero<T>;
+impl NonZeroNeg<T, +Neg<T>, +TryInto<T, NonZero<T>>> of Neg<NonZero<T>> {
+    fn neg(a: NonZero<T>) -> NonZero<T> {
+        // TODO(orizi): Optimize using bounded integers.
+        let value: T = a.into();
+        let negated: T = -value;
+        negated.try_into().unwrap()
+    }
+}
 
 /// Represents the result of checking whether a value is zero.
 pub(crate) enum IsZeroResult<T> {

crates/cairo-lang-lowering/src/optimizations/const_folding.rs

@@ -230,7 +230,7 @@
             || id == self.bounded_int_add
             || id == self.bounded_int_sub
         {
-            let lhs = self.as_int(stmt.inputs[0].var_id)?;
+            let (lhs, nz_ty) = self.as_int_ex(stmt.inputs[0].var_id)?;
             let rhs = self.as_int(stmt.inputs[1].var_id)?;
             let value = if id == self.bounded_int_add {
                 lhs + rhs
@@ -240,8 +240,10 @@
                 lhs * rhs
             };
             let output = stmt.outputs[0];
-            let ty = self.variables[output].ty;
-            let value = ConstValue::Int(value, ty);
+            let mut value = ConstValue::Int(value, self.variables[output].ty);
+            if nz_ty {
+                value = ConstValue::NonZero(Box::new(value));
+            }
             self.var_info.insert(output, VarInfo::Const(value.clone()));
             Some(Statement::Const(StatementConst { value, output }))
         } else if id == self.storage_base_address_from_felt252 {
@@ -355,7 +357,7 @@
             })
         } else if id == self.bounded_int_constrain {
             let input_var = info.inputs[0].var_id;
-            let value = self.as_int(input_var)?;
+            let (value, nz_ty) = self.as_int_ex(input_var)?;
             let semantic_id =
                 extract_matches!(info.function.lookup_intern(self.db), FunctionLongId::Semantic);
             let generic_arg = semantic_id.get_concrete(self.db.upcast()).generic_args[1];
@@ -365,7 +367,10 @@
                 .unwrap();
             let arm_idx = if value < &constrain_value { 0 } else { 1 };
             let output = info.arms[arm_idx].var_ids[0];
-            let value = ConstValue::Int(value.clone(), self.variables[output].ty);
+            let mut value = ConstValue::Int(value.clone(), self.variables[output].ty);
+            if nz_ty {
+                value = ConstValue::NonZero(Box::new(value));
+            }
             self.var_info.insert(output, VarInfo::Const(value.clone()));
             Some((
                 Some(Statement::Const(StatementConst { value, output })),
@@ -381,9 +386,25 @@
         try_extract_matches!(self.var_info.get(&var_id)?, VarInfo::Const)
     }
 
+    /// Return the const value as a int if it exists and is an integer, aditionally, if it is of a
+    /// non-zero type.
+    fn as_int_ex(&self, var_id: VariableId) -> Option<(&BigInt, bool)> {
+        match self.as_const(var_id)? {
+            ConstValue::Int(value, _) => Some((value, false)),
+            ConstValue::NonZero(const_value) => {
+                if let ConstValue::Int(value, _) = const_value.as_ref() {
+                    Some((value, true))
+                } else {
+                    None
+                }
+            }
+            _ => None,
+        }
+    }
+
     /// Return the const value as a int if it exists and is an integer.
     fn as_int(&self, var_id: VariableId) -> Option<&BigInt> {
-        if let ConstValue::Int(value, _) = self.as_const(var_id)? { Some(value) } else { None }
+        Some(self.as_int_ex(var_id)?.0)
     }
 
     /// Replaces the inputs in place if they are in the var_info map.

crates/cairo-lang-lowering/src/optimizations/test_data/const_folding

@@ -3348,7 +3348,7 @@ End:
 
 //! > ==========================================================================
 
-//! > `bounded_int_constrain` and `bounded_int_mul`.
+//! > `AbsAndSign` on positive const.
 
 //! > test_runner_name
 test_match_optimizer
@@ -3439,6 +3439,97 @@ End:
 
 //! > ==========================================================================
 
+//! > `AbsAndSign` on negative const.
+
+//! > test_runner_name
+test_match_optimizer
+
+//! > function
+fn foo() -> (u8, bool) {
+    core::integer::AbsAndSign::<i8, u8>::abs_and_sign(-6_i8)
+}
+
+//! > function_name
+foo
+
+//! > module_code
+
+//! > semantic_diagnostics
+
+//! > before
+Parameters:
+blk0 (root):
+Statements:
+  (v0: core::integer::i8) <- -6
+End:
+  Match(match core::internal::bounded_int::bounded_int_constrain::<core::integer::i8, 0, core::internal::bounded_int::constrain0::Impl::<core::integer::i8, -128, 127>>(v0) {
+    Result::Ok(v1) => blk1,
+    Result::Err(v2) => blk2,
+  })
+
+blk1:
+Statements:
+  (v3: core::internal::bounded_int::BoundedInt::<-1, -1>) <- core::internal::bounded_int::minus_1::const_as_immediate::<core::internal::bounded_int::minus_1::Const::<core::internal::bounded_int::BoundedInt::<-1, -1>, -1>>()
+  (v4: core::internal::bounded_int::BoundedInt::<1, 128>) <- core::internal::bounded_int::bounded_int_mul::<core::internal::bounded_int::BoundedInt::<-128, -1>, core::internal::bounded_int::BoundedInt::<-1, -1>, core::internal::bounded_int::MulMinus1::<-128, -1, core::internal::bounded_int::neg_felt252::Impl::<-128, 128>, core::internal::bounded_int::neg_felt252::Impl::<-1, 1>>>(v1, v3)
+  (v5: core::integer::u8) <- core::integer::upcast::<core::internal::bounded_int::BoundedInt::<1, 128>, core::integer::u8>(v4)
+  (v6: ()) <- struct_construct()
+  (v7: core::bool) <- bool::True(v6)
+  (v8: (core::integer::u8, core::bool)) <- struct_construct(v5, v7)
+End:
+  Goto(blk3, {v8 -> v9})
+
+blk2:
+Statements:
+  (v10: core::integer::u8) <- core::integer::upcast::<core::internal::bounded_int::BoundedInt::<0, 127>, core::integer::u8>(v2)
+  (v11: ()) <- struct_construct()
+  (v12: core::bool) <- bool::False(v11)
+  (v13: (core::integer::u8, core::bool)) <- struct_construct(v10, v12)
+End:
+  Goto(blk3, {v13 -> v9})
+
+blk3:
+Statements:
+End:
+  Return(v9)
+
+//! > after
+Parameters:
+blk0 (root):
+Statements:
+  (v0: core::integer::i8) <- -6
+  (v1: core::internal::bounded_int::BoundedInt::<-128, -1>) <- -6
+End:
+  Goto(blk1, {})
+
+blk1:
+Statements:
+  (v3: core::internal::bounded_int::BoundedInt::<-1, -1>) <- core::internal::bounded_int::minus_1::const_as_immediate::<core::internal::bounded_int::minus_1::Const::<core::internal::bounded_int::BoundedInt::<-1, -1>, -1>>()
+  (v4: core::internal::bounded_int::BoundedInt::<1, 128>) <- core::internal::bounded_int::bounded_int_mul::<core::internal::bounded_int::BoundedInt::<-128, -1>, core::internal::bounded_int::BoundedInt::<-1, -1>, core::internal::bounded_int::MulMinus1::<-128, -1, core::internal::bounded_int::neg_felt252::Impl::<-128, 128>, core::internal::bounded_int::neg_felt252::Impl::<-1, 1>>>(v1, v3)
+  (v5: core::integer::u8) <- core::integer::upcast::<core::internal::bounded_int::BoundedInt::<1, 128>, core::integer::u8>(v4)
+  (v6: ()) <- struct_construct()
+  (v7: core::bool) <- bool::True(v6)
+  (v8: (core::integer::u8, core::bool)) <- struct_construct(v5, v7)
+End:
+  Goto(blk3, {v8 -> v9})
+
+blk2:
+Statements:
+  (v10: core::integer::u8) <- core::integer::upcast::<core::internal::bounded_int::BoundedInt::<0, 127>, core::integer::u8>(v2)
+  (v11: ()) <- struct_construct()
+  (v12: core::bool) <- bool::False(v11)
+  (v13: (core::integer::u8, core::bool)) <- struct_construct(v10, v12)
+End:
+  Goto(blk3, {v13 -> v9})
+
+blk3:
+Statements:
+End:
+  Return(v9)
+
+//! > lowering_diagnostics
+
+//! > ==========================================================================
+
 //! > Bitwise not const folding.
 
 //! > test_runner_name
@@ -3479,3 +3570,169 @@ End:
   Return(v3)
 
 //! > lowering_diagnostics
+
+//! > ==========================================================================
+
+//! > bounded_int_constrain on NonZero below.
+
+//! > test_runner_name
+test_match_optimizer
+
+//! > function
+fn foo() -> Result<Box<NonZero<BoundedInt<-0x80, -1>>>, Box<NonZero<BoundedInt<0, 0x7f>>>> {
+    match core::internal::bounded_int::constrain::<NonZero<i8>, 0>(-5) {
+        Result::Ok(v) => Result::Ok(BoxTrait::new(v)),
+        Result::Err(v) => Result::Err(BoxTrait::new(v)),
+    }
+}
+
+//! > function_name
+foo
+
+//! > module_code
+use core::internal::bounded_int::BoundedInt;
+
+//! > semantic_diagnostics
+
+//! > before
+Parameters:
+blk0 (root):
+Statements:
+  (v0: core::zeroable::NonZero::<core::integer::i8>) <- NonZero(-5)
+End:
+  Match(match core::internal::bounded_int::bounded_int_constrain::<core::zeroable::NonZero::<core::integer::i8>, 0, core::internal::bounded_int::NonZeroConstrainHelper::<core::integer::i8, 0, core::internal::bounded_int::constrain0::Impl::<core::integer::i8, -128, 127>>>(v0) {
+    Result::Ok(v1) => blk1,
+    Result::Err(v2) => blk2,
+  })
+
+blk1:
+Statements:
+  (v3: core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<-128, -1>>>) <- core::box::into_box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<-128, -1>>>(v1)
+  (v4: core::result::Result::<core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<-128, -1>>>, core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<0, 127>>>>) <- Result::Ok(v3)
+End:
+  Goto(blk3, {v4 -> v5})
+
+blk2:
+Statements:
+  (v6: core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<0, 127>>>) <- core::box::into_box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<0, 127>>>(v2)
+  (v7: core::result::Result::<core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<-128, -1>>>, core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<0, 127>>>>) <- Result::Err(v6)
+End:
+  Goto(blk3, {v7 -> v5})
+
+blk3:
+Statements:
+End:
+  Return(v5)
+
+//! > after
+Parameters:
+blk0 (root):
+Statements:
+  (v0: core::zeroable::NonZero::<core::integer::i8>) <- NonZero(-5)
+  (v1: core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<-128, -1>>) <- NonZero(-5)
+End:
+  Goto(blk1, {})
+
+blk1:
+Statements:
+  (v3: core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<-128, -1>>>) <- NonZero(-5).into_box()
+  (v4: core::result::Result::<core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<-128, -1>>>, core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<0, 127>>>>) <- Result::Ok(v3)
+End:
+  Goto(blk3, {v4 -> v5})
+
+blk2:
+Statements:
+  (v6: core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<0, 127>>>) <- core::box::into_box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<0, 127>>>(v2)
+  (v7: core::result::Result::<core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<-128, -1>>>, core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<0, 127>>>>) <- Result::Err(v6)
+End:
+  Goto(blk3, {v7 -> v5})
+
+blk3:
+Statements:
+End:
+  Return(v5)
+
+//! > lowering_diagnostics
+
+//! > ==========================================================================
+
+//! > bounded_int_constrain on NonZero above.
+
+//! > test_runner_name
+test_match_optimizer
+
+//! > function
+fn foo() -> Result<Box<NonZero<BoundedInt<-0x80, -1>>>, Box<NonZero<BoundedInt<0, 0x7f>>>> {
+    match core::internal::bounded_int::constrain::<NonZero<i8>, 0>(5) {
+        Result::Ok(v) => Result::Ok(BoxTrait::new(v)),
+        Result::Err(v) => Result::Err(BoxTrait::new(v)),
+    }
+}
+
+//! > function_name
+foo
+
+//! > module_code
+use core::internal::bounded_int::BoundedInt;
+
+//! > semantic_diagnostics
+
+//! > before
+Parameters:
+blk0 (root):
+Statements:
+  (v0: core::zeroable::NonZero::<core::integer::i8>) <- NonZero(5)
+End:
+  Match(match core::internal::bounded_int::bounded_int_constrain::<core::zeroable::NonZero::<core::integer::i8>, 0, core::internal::bounded_int::NonZeroConstrainHelper::<core::integer::i8, 0, core::internal::bounded_int::constrain0::Impl::<core::integer::i8, -128, 127>>>(v0) {
+    Result::Ok(v1) => blk1,
+    Result::Err(v2) => blk2,
+  })
+
+blk1:
+Statements:
+  (v3: core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<-128, -1>>>) <- core::box::into_box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<-128, -1>>>(v1)
+  (v4: core::result::Result::<core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<-128, -1>>>, core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<0, 127>>>>) <- Result::Ok(v3)
+End:
+  Goto(blk3, {v4 -> v5})
+
+blk2:
+Statements:
+  (v6: core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<0, 127>>>) <- core::box::into_box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<0, 127>>>(v2)
+  (v7: core::result::Result::<core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<-128, -1>>>, core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<0, 127>>>>) <- Result::Err(v6)
+End:
+  Goto(blk3, {v7 -> v5})
+
+blk3:
+Statements:
+End:
+  Return(v5)
+
+//! > after
+Parameters:
+blk0 (root):
+Statements:
+  (v0: core::zeroable::NonZero::<core::integer::i8>) <- NonZero(5)
+  (v2: core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<0, 127>>) <- NonZero(5)
+End:
+  Goto(blk2, {})
+
+blk1:
+Statements:
+  (v3: core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<-128, -1>>>) <- core::box::into_box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<-128, -1>>>(v1)
+  (v4: core::result::Result::<core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<-128, -1>>>, core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<0, 127>>>>) <- Result::Ok(v3)
+End:
+  Goto(blk3, {v4 -> v5})
+
+blk2:
+Statements:
+  (v6: core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<0, 127>>>) <- NonZero(5).into_box()
+  (v7: core::result::Result::<core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<-128, -1>>>, core::box::Box::<core::zeroable::NonZero::<core::internal::bounded_int::BoundedInt::<0, 127>>>>) <- Result::Err(v6)
+End:
+  Goto(blk3, {v7 -> v5})
+
+blk3:
+Statements:
+End:
+  Return(v5)
+
+//! > lowering_diagnostics