Model Bool < U32 < Felt during type inference

Author: fegge

src/types/domain.rs

@@ -23,24 +23,31 @@ impl InferredType {
     /// Combine two values from different control-flow paths.
     ///
     /// `Unknown` means at least one path is opaque, so the result is unknown.
-    /// Distinct known subtypes join to `Felt`.
+    /// Distinct known subtypes join to their least upper bound in the lattice
+    /// `Bool <: U32 <: Felt`, with `Address` as a sibling of `U32`.
     pub fn combine_paths(self, other: Self) -> Self {
         match (self, other) {
             (a, b) if a == b => a,
             (Self::Unknown, _) | (_, Self::Unknown) => Self::Unknown,
+            // Bool <: U32
+            (Self::Bool, Self::U32) | (Self::U32, Self::Bool) => Self::U32,
             _ => Self::Felt,
         }
     }
 
     /// Refine a stored type with newly inferred information.
     ///
     /// Existing `Unknown` values are replaced by `new_type`. Existing known
-    /// values are widened conservatively if they disagree.
+    /// values are widened to their least upper bound in the lattice
+    /// `Bool <: U32 <: Felt` if they disagree, with `Address` as a sibling
+    /// of `U32`.
     pub fn refine(self, new_type: Self) -> Self {
         match (self, new_type) {
             (Self::Unknown, t) => t,
             (t, Self::Unknown) => t,
             (a, b) if a == b => a,
+            // Bool <: U32
+            (Self::Bool, Self::U32) | (Self::U32, Self::Bool) => Self::U32,
             _ => Self::Felt,
         }
     }
@@ -140,7 +147,7 @@ pub fn check_compatibility(actual: InferredType, expected: TypeRequirement) -> C
         },
         TypeRequirement::U32 => match actual {
             InferredType::Unknown => Compatibility::Indeterminate,
-            InferredType::U32 => Compatibility::Compatible,
+            InferredType::U32 | InferredType::Bool => Compatibility::Compatible,
             _ => Compatibility::Incompatible,
         },
         TypeRequirement::Address => match actual {
@@ -187,6 +194,8 @@ impl VarKey {
 mod tests {
     use super::*;
 
+    // -- check_compatibility --------------------------------------------------
+
     #[test]
     fn felt_accepts_promotable_types() {
         assert_eq!(
@@ -203,16 +212,36 @@ mod tests {
         );
     }
 
+    #[test]
+    fn bool_is_compatible_with_u32() {
+        assert_eq!(
+            check_compatibility(InferredType::Bool, TypeRequirement::U32),
+            Compatibility::Compatible
+        );
+    }
+
     #[test]
     fn mismatched_exact_types_are_incompatible() {
         assert_eq!(
             check_compatibility(InferredType::Felt, TypeRequirement::Bool),
             Compatibility::Incompatible
         );
+        assert_eq!(
+            check_compatibility(InferredType::Felt, TypeRequirement::U32),
+            Compatibility::Incompatible
+        );
+        assert_eq!(
+            check_compatibility(InferredType::U32, TypeRequirement::Bool),
+            Compatibility::Incompatible
+        );
         assert_eq!(
             check_compatibility(InferredType::Bool, TypeRequirement::Address),
             Compatibility::Incompatible
         );
+        assert_eq!(
+            check_compatibility(InferredType::Address, TypeRequirement::Bool),
+            Compatibility::Incompatible
+        );
     }
 
     #[test]
@@ -226,4 +255,124 @@ mod tests {
             Compatibility::Indeterminate
         );
     }
+
+    // -- combine_paths --------------------------------------------------------
+
+    #[test]
+    fn combine_paths_same_type() {
+        assert_eq!(
+            InferredType::Bool.combine_paths(InferredType::Bool),
+            InferredType::Bool
+        );
+        assert_eq!(
+            InferredType::U32.combine_paths(InferredType::U32),
+            InferredType::U32
+        );
+        assert_eq!(
+            InferredType::Felt.combine_paths(InferredType::Felt),
+            InferredType::Felt
+        );
+    }
+
+    #[test]
+    fn combine_paths_bool_u32_yields_u32() {
+        assert_eq!(
+            InferredType::Bool.combine_paths(InferredType::U32),
+            InferredType::U32
+        );
+        assert_eq!(
+            InferredType::U32.combine_paths(InferredType::Bool),
+            InferredType::U32
+        );
+    }
+
+    #[test]
+    fn combine_paths_with_felt_yields_felt() {
+        assert_eq!(
+            InferredType::Bool.combine_paths(InferredType::Felt),
+            InferredType::Felt
+        );
+        assert_eq!(
+            InferredType::U32.combine_paths(InferredType::Felt),
+            InferredType::Felt
+        );
+        assert_eq!(
+            InferredType::Address.combine_paths(InferredType::Felt),
+            InferredType::Felt
+        );
+    }
+
+    #[test]
+    fn combine_paths_address_siblings_yield_felt() {
+        assert_eq!(
+            InferredType::Bool.combine_paths(InferredType::Address),
+            InferredType::Felt
+        );
+        assert_eq!(
+            InferredType::U32.combine_paths(InferredType::Address),
+            InferredType::Felt
+        );
+    }
+
+    #[test]
+    fn combine_paths_unknown_yields_unknown() {
+        assert_eq!(
+            InferredType::Bool.combine_paths(InferredType::Unknown),
+            InferredType::Unknown
+        );
+        assert_eq!(
+            InferredType::Unknown.combine_paths(InferredType::U32),
+            InferredType::Unknown
+        );
+    }
+
+    // -- refine ---------------------------------------------------------------
+
+    #[test]
+    fn refine_unknown_takes_new() {
+        assert_eq!(
+            InferredType::Unknown.refine(InferredType::Bool),
+            InferredType::Bool
+        );
+        assert_eq!(
+            InferredType::Unknown.refine(InferredType::U32),
+            InferredType::U32
+        );
+    }
+
+    #[test]
+    fn refine_bool_u32_yields_u32() {
+        assert_eq!(
+            InferredType::Bool.refine(InferredType::U32),
+            InferredType::U32
+        );
+        assert_eq!(
+            InferredType::U32.refine(InferredType::Bool),
+            InferredType::U32
+        );
+    }
+
+    #[test]
+    fn refine_with_felt_yields_felt() {
+        assert_eq!(
+            InferredType::Bool.refine(InferredType::Felt),
+            InferredType::Felt
+        );
+        assert_eq!(
+            InferredType::U32.refine(InferredType::Felt),
+            InferredType::Felt
+        );
+    }
+
+    #[test]
+    fn refine_address_siblings_yield_felt() {
+        assert_eq!(
+            InferredType::Bool.refine(InferredType::Address),
+            InferredType::Felt
+        );
+        assert_eq!(
+            InferredType::U32.refine(InferredType::Address),
+            InferredType::Felt
+        );
+    }
 }

src/types/intra.rs

@@ -314,8 +314,13 @@ impl<'a> ProcTypeAnalyzer<'a> {
     }
 
     /// Infer type for a constant expression.
+    ///
+    /// `Felt(0)` and `Felt(1)` are inferred as `Bool` since they are the
+    /// most precise type for these values. The lattice `Bool <: U32 <: Felt`
+    /// ensures they widen correctly in arithmetic or u32 contexts.
     fn infer_constant_type(&self, constant: &Constant) -> InferredType {
         match constant {
+            Constant::Felt(0 | 1) => InferredType::Bool,
             Constant::Felt(_) | Constant::Defined(_) => InferredType::Felt,
             Constant::Word(_) => InferredType::Unknown,
         }

tests/type_inference.rs

@@ -64,7 +64,7 @@ fn setup_decompiler() -> Decompiler<'static> {
         end
 
         pub proc caller_bad_and_bool
-            push.1
+            push.3
             push.2
             exec.needs_and_bool
         end
@@ -116,13 +116,12 @@ fn setup_decompiler() -> Decompiler<'static> {
             push.2
             push.3
             u32wrapping_add
-            push.1.1
-            eq
+            push.2
             exec.needs_bool_u32
         end
 
         pub proc out_mixed
-            push.1
+            push.2
             push.1.1
             eq
         end