1use 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
31pub 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 VerifyOp::And | VerifyOp::Or | VerifyOp::Implies => {
54 constrain_as_bool(left, constraints)?;
55 constrain_as_bool(right, constraints)?;
56 }
57 VerifyOp::Add | VerifyOp::Sub | VerifyOp::Mul | VerifyOp::Div | VerifyOp::FloorDiv => {
59 constrain_as_int(left, constraints)?;
60 constrain_as_int(right, constraints)?;
61 }
62 VerifyOp::Gt | VerifyOp::Lt | VerifyOp::Gte | VerifyOp::Lte => {
64 constrain_as_int(left, constraints)?;
65 constrain_as_int(right, constraints)?;
66 }
67 VerifyOp::Eq | VerifyOp::Neq => {
69 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 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 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 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; 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 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 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, }
278}