Skip to main content

logicaffeine_verify/
type_infer.rs

1//! Constraint-based type inference for VerifyExpr.
2//!
3//! Walks the expression tree, collects type constraints from operators,
4//! unifies constraints per variable, detects conflicts, and propagates
5//! bitvector widths.
6
7use crate::ir::{VerifyExpr, VerifyOp, VerifyType};
8use std::collections::HashMap;
9
10#[derive(Debug, Clone)]
11pub enum TypeError {
12    Conflict {
13        var: String,
14        expected: VerifyType,
15        found: VerifyType,
16    },
17}
18
19impl std::fmt::Display for TypeError {
20    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
21        match self {
22            TypeError::Conflict { var, expected, found } => {
23                write!(f, "Type conflict for '{}': expected {:?}, found {:?}", var, expected, found)
24            }
25        }
26    }
27}
28
29impl std::error::Error for TypeError {}
30
31/// Infer types for all variables in a VerifyExpr.
32///
33/// Returns a map from variable name to inferred type.
34/// Returns TypeError if a variable is used with conflicting types.
35pub fn infer_types(expr: &VerifyExpr) -> Result<HashMap<String, VerifyType>, TypeError> {
36    let mut constraints: HashMap<String, VerifyType> = HashMap::new();
37    collect_constraints(expr, &mut constraints)?;
38    Ok(constraints)
39}
40
41fn collect_constraints(
42    expr: &VerifyExpr,
43    constraints: &mut HashMap<String, VerifyType>,
44) -> Result<(), TypeError> {
45    match expr {
46        VerifyExpr::Bool(_) | VerifyExpr::Int(_) => Ok(()),
47
48        VerifyExpr::Var(_) => Ok(()),
49
50        VerifyExpr::Binary { op, left, right } => {
51            match op {
52                // Boolean operators: operands must be Bool
53                VerifyOp::And | VerifyOp::Or | VerifyOp::Implies => {
54                    constrain_as_bool(left, constraints)?;
55                    constrain_as_bool(right, constraints)?;
56                }
57                // Arithmetic operators: operands must be Int
58                VerifyOp::Add | VerifyOp::Sub | VerifyOp::Mul | VerifyOp::Div | VerifyOp::FloorDiv => {
59                    constrain_as_int(left, constraints)?;
60                    constrain_as_int(right, constraints)?;
61                }
62                // Comparison operators: operands must be Int
63                VerifyOp::Gt | VerifyOp::Lt | VerifyOp::Gte | VerifyOp::Lte => {
64                    constrain_as_int(left, constraints)?;
65                    constrain_as_int(right, constraints)?;
66                }
67                // Equality: operands must match but could be any sort
68                VerifyOp::Eq | VerifyOp::Neq => {
69                    // If either side is clearly Int, constrain both
70                    if expr_suggests_int(left) || expr_suggests_int(right) {
71                        constrain_as_int(left, constraints)?;
72                        constrain_as_int(right, constraints)?;
73                    }
74                }
75            }
76            collect_constraints(left, constraints)?;
77            collect_constraints(right, constraints)?;
78            Ok(())
79        }
80
81        VerifyExpr::Not(inner) => {
82            constrain_as_bool(inner, constraints)?;
83            collect_constraints(inner, constraints)
84        }
85
86        VerifyExpr::Iff(l, r) => {
87            collect_constraints(l, constraints)?;
88            collect_constraints(r, constraints)
89        }
90
91        VerifyExpr::ForAll { body, vars } | VerifyExpr::Exists { body, vars } => {
92            for (name, ty) in vars {
93                add_constraint(name, ty.clone(), constraints)?;
94            }
95            collect_constraints(body, constraints)
96        }
97
98        VerifyExpr::Apply { args, .. } | VerifyExpr::ApplyInt { args, .. } => {
99            for arg in args {
100                collect_constraints(arg, constraints)?;
101            }
102            Ok(())
103        }
104
105        // Bitvector operations
106        VerifyExpr::BitVecConst { .. } => Ok(()),
107
108        VerifyExpr::BitVecBinary { op: _, left, right } => {
109            let width = bv_width_hint(left).or_else(|| bv_width_hint(right));
110            if let Some(w) = width {
111                constrain_as_bv(left, w, constraints)?;
112                constrain_as_bv(right, w, constraints)?;
113            }
114            collect_constraints(left, constraints)?;
115            collect_constraints(right, constraints)
116        }
117
118        VerifyExpr::BitVecExtract { high, low: _, operand } => {
119            // Operand must be a BV with width > high
120            let min_width = high + 1;
121            constrain_as_bv(operand, min_width, constraints)?;
122            collect_constraints(operand, constraints)
123        }
124
125        VerifyExpr::BitVecConcat(l, r) => {
126            collect_constraints(l, constraints)?;
127            collect_constraints(r, constraints)
128        }
129
130        // Array operations
131        VerifyExpr::Select { array, index } => {
132            if let VerifyExpr::Var(name) = array.as_ref() {
133                let idx_ty = infer_sort_of(index);
134                let elem_ty = VerifyType::Int; // default
135                add_constraint(name, VerifyType::Array(Box::new(idx_ty), Box::new(elem_ty)), constraints)?;
136            }
137            collect_constraints(array, constraints)?;
138            collect_constraints(index, constraints)
139        }
140
141        VerifyExpr::Store { array, index, value } => {
142            if let VerifyExpr::Var(name) = array.as_ref() {
143                let idx_ty = infer_sort_of(index);
144                let val_ty = infer_sort_of(value);
145                add_constraint(name, VerifyType::Array(Box::new(idx_ty), Box::new(val_ty)), constraints)?;
146            }
147            collect_constraints(array, constraints)?;
148            collect_constraints(index, constraints)?;
149            collect_constraints(value, constraints)
150        }
151
152        VerifyExpr::AtState { state, expr } => {
153            collect_constraints(state, constraints)?;
154            collect_constraints(expr, constraints)
155        }
156
157        VerifyExpr::Transition { from, to } => {
158            collect_constraints(from, constraints)?;
159            collect_constraints(to, constraints)
160        }
161    }
162}
163
164fn constrain_as_bool(
165    expr: &VerifyExpr,
166    constraints: &mut HashMap<String, VerifyType>,
167) -> Result<(), TypeError> {
168    if let VerifyExpr::Var(name) = expr {
169        add_constraint(name, VerifyType::Bool, constraints)?;
170    }
171    Ok(())
172}
173
174fn constrain_as_int(
175    expr: &VerifyExpr,
176    constraints: &mut HashMap<String, VerifyType>,
177) -> Result<(), TypeError> {
178    if let VerifyExpr::Var(name) = expr {
179        add_constraint(name, VerifyType::Int, constraints)?;
180    }
181    // Also handle nested binary expressions with Var leaves
182    if let VerifyExpr::Binary { left, right, .. } = expr {
183        constrain_as_int(left, constraints)?;
184        constrain_as_int(right, constraints)?;
185    }
186    Ok(())
187}
188
189fn constrain_as_bv(
190    expr: &VerifyExpr,
191    width: u32,
192    constraints: &mut HashMap<String, VerifyType>,
193) -> Result<(), TypeError> {
194    if let VerifyExpr::Var(name) = expr {
195        add_constraint(name, VerifyType::BitVector(width), constraints)?;
196    }
197    Ok(())
198}
199
200fn add_constraint(
201    name: &str,
202    ty: VerifyType,
203    constraints: &mut HashMap<String, VerifyType>,
204) -> Result<(), TypeError> {
205    if let Some(existing) = constraints.get(name) {
206        if !types_compatible(existing, &ty) {
207            return Err(TypeError::Conflict {
208                var: name.to_string(),
209                expected: existing.clone(),
210                found: ty,
211            });
212        }
213        // If new type is more specific, update
214        if type_specificity(&ty) > type_specificity(existing) {
215            constraints.insert(name.to_string(), ty);
216        }
217    } else {
218        constraints.insert(name.to_string(), ty);
219    }
220    Ok(())
221}
222
223fn types_compatible(a: &VerifyType, b: &VerifyType) -> bool {
224    match (a, b) {
225        (VerifyType::Int, VerifyType::Int) => true,
226        (VerifyType::Bool, VerifyType::Bool) => true,
227        (VerifyType::Real, VerifyType::Real) => true,
228        (VerifyType::Object, _) | (_, VerifyType::Object) => true,
229        (VerifyType::BitVector(w1), VerifyType::BitVector(w2)) => w1 == w2,
230        (VerifyType::Array(i1, e1), VerifyType::Array(i2, e2)) => {
231            types_compatible(i1, i2) && types_compatible(e1, e2)
232        }
233        _ => false,
234    }
235}
236
237fn type_specificity(ty: &VerifyType) -> u8 {
238    match ty {
239        VerifyType::Object => 0,
240        VerifyType::Bool => 1,
241        VerifyType::Int => 1,
242        VerifyType::Real => 1,
243        VerifyType::BitVector(_) => 2,
244        VerifyType::Array(_, _) => 2,
245    }
246}
247
248fn bv_width_hint(expr: &VerifyExpr) -> Option<u32> {
249    match expr {
250        VerifyExpr::BitVecConst { width, .. } => Some(*width),
251        VerifyExpr::BitVecBinary { left, right, .. } => {
252            bv_width_hint(left).or_else(|| bv_width_hint(right))
253        }
254        VerifyExpr::BitVecExtract { high, low, .. } => Some(high - low + 1),
255        _ => None,
256    }
257}
258
259fn expr_suggests_int(expr: &VerifyExpr) -> bool {
260    matches!(
261        expr,
262        VerifyExpr::Int(_)
263            | VerifyExpr::Binary {
264                op: VerifyOp::Add | VerifyOp::Sub | VerifyOp::Mul | VerifyOp::Div | VerifyOp::FloorDiv,
265                ..
266            }
267    )
268}
269
270fn infer_sort_of(expr: &VerifyExpr) -> VerifyType {
271    match expr {
272        VerifyExpr::Int(_) => VerifyType::Int,
273        VerifyExpr::Bool(_) => VerifyType::Bool,
274        VerifyExpr::BitVecConst { width, .. } => VerifyType::BitVector(*width),
275        VerifyExpr::Binary { op: VerifyOp::Add | VerifyOp::Sub | VerifyOp::Mul | VerifyOp::Div | VerifyOp::FloorDiv, .. } => VerifyType::Int,
276        _ => VerifyType::Int, // default
277    }
278}