1use flux_common::tracked_span_bug;
2use flux_middle::{
3 big_int::BigInt,
4 rty::{self, Binder, EarlyReftParam, InternalFuncKind, List, SpecFuncKind},
5};
6use flux_rustc_bridge::lowering::Lower;
7use itertools::Itertools;
8use rustc_hir::{def::DefKind, def_id::DefId};
9use rustc_type_ir::BoundVar;
10
11use super::{ConstKey, FixpointCtxt, fixpoint};
12use crate::fixpoint_encoding::FixpointSolution;
13
14impl<'genv, 'tcx, Tag> FixpointCtxt<'genv, 'tcx, Tag>
15where
16 Tag: std::hash::Hash + Eq + Copy,
17{
18 pub(crate) fn fixpoint_to_solution(
19 &mut self,
20 sol: &FixpointSolution,
21 ) -> rty::Binder<rty::Expr> {
22 let mut vars = vec![];
23 let mut sorts = vec![];
24 for (var, sort) in &sol.0 {
25 let fixpoint::Var::Local(local_var) = var else {
26 tracked_span_bug!("encountered non-local variable in binder: {var:?}");
27 };
28 vars.push(*local_var);
29 sorts.push(
30 self.fixpoint_to_sort(sort)
31 .unwrap_or_else(|_| tracked_span_bug!("failed to parse sort: {sort:?}")),
32 );
33 }
34 self.ecx.local_var_env.push_layer(vars);
35 let expr = self
36 .fixpoint_to_expr(&sol.1)
37 .unwrap_or_else(|err| tracked_span_bug!("failed to convert expr: {err:?}"));
38 self.ecx.local_var_env.pop_layer();
39 rty::Binder::bind_with_sorts(expr, &sorts)
40 }
41
42 fn fixpoint_to_sort_ctor(
43 &self,
44 ctor: &fixpoint::SortCtor,
45 ) -> Result<rty::SortCtor, FixpointParseError> {
46 match ctor {
47 fixpoint::SortCtor::Set => Ok(rty::SortCtor::Set),
48 fixpoint::SortCtor::Map => Ok(rty::SortCtor::Map),
49 fixpoint::SortCtor::Data(fixpoint::DataSort::Tuple(_)) => {
50 panic!("oh no! tuple!") }
52 fixpoint::SortCtor::Data(fixpoint::DataSort::User(opaque_id)) => {
53 let def_id = self.scx.opaque_sorts[opaque_id.as_usize()];
54 Ok(rty::SortCtor::User(def_id))
55 }
56 fixpoint::SortCtor::Data(fixpoint::DataSort::Adt(adt_id)) => {
57 let def_id = self.scx.adt_sorts[adt_id.as_usize()];
58 let Ok(adt_sort_def) = self.genv.adt_sort_def_of(def_id) else {
59 return Err(FixpointParseError::UnknownAdt(def_id));
60 };
61 Ok(rty::SortCtor::Adt(adt_sort_def))
62 }
63 }
64 }
65
66 pub(crate) fn fixpoint_to_sort(
67 &self,
68 fsort: &fixpoint::Sort,
69 ) -> Result<rty::Sort, FixpointParseError> {
70 match fsort {
71 fixpoint::Sort::Int => Ok(rty::Sort::Int),
72 fixpoint::Sort::Real => Ok(rty::Sort::Real),
73 fixpoint::Sort::Bool => Ok(rty::Sort::Bool),
74 fixpoint::Sort::Str => Ok(rty::Sort::Str),
75 fixpoint::Sort::Func(sorts) => {
76 let sort1 = self.fixpoint_to_sort(&sorts[0])?;
77 let sort2 = self.fixpoint_to_sort(&sorts[1])?;
78 let fsort = rty::FuncSort::new(vec![sort1], sort2);
79 let poly_sort = rty::PolyFuncSort::new(List::empty(), fsort);
80 Ok(rty::Sort::Func(poly_sort))
81 }
82 fixpoint::Sort::App(ctor, args) => {
83 let ctor = self.fixpoint_to_sort_ctor(ctor)?;
84 let args = args
85 .iter()
86 .map(|fsort| self.fixpoint_to_sort(fsort))
87 .try_collect()?;
88 Ok(rty::Sort::App(ctor, args))
89 }
90 fixpoint::Sort::BitVec(fsort) if let fixpoint::Sort::BvSize(size) = **fsort => {
91 Ok(rty::Sort::BitVec(rty::BvSize::Fixed(size)))
92 }
93 _ => unimplemented!("fixpoint_to_sort: {fsort:?}"),
94 }
95 }
96
97 fn is_curried_primop_app(
98 &mut self,
99 fhead: &fixpoint::Expr,
100 fargs: &[fixpoint::Expr],
101 op_args: &mut Vec<fixpoint::Expr>,
102 ) -> Option<rty::BinOp> {
103 match fhead {
104 fixpoint::Expr::Var(fixpoint::Var::Global(global_var, _))
105 | fixpoint::Expr::Var(fixpoint::Var::Const(global_var, _)) => {
106 if let Some(ConstKey::PrimOp(bin_op)) =
107 self.ecx.const_env.const_map_rev.get(global_var)
108 {
109 op_args.reverse();
110 Some(bin_op.clone())
111 } else {
112 None
113 }
114 }
115 fixpoint::Expr::App(fhead_inner, _, fargs_inner, _) => {
116 if fargs.len() == 1 {
117 op_args.push(fargs[0].clone());
118 }
119 self.is_curried_primop_app(fhead_inner, fargs_inner, op_args)
120 }
121 _ => None,
122 }
123 }
124
125 #[allow(dead_code)]
126 pub(crate) fn fixpoint_to_expr(
127 &mut self,
128 fexpr: &fixpoint::Expr,
129 ) -> Result<rty::Expr, FixpointParseError> {
130 match fexpr {
131 fixpoint::Expr::Constant(constant) => {
132 let c = match constant {
133 fixpoint::Constant::Numeral(num) => rty::Constant::Int(BigInt::from(*num)),
134 fixpoint::Constant::Real(dec) => rty::Constant::Real(rty::Real(dec.0)),
135 fixpoint::Constant::Boolean(b) => rty::Constant::Bool(*b),
136 fixpoint::Constant::String(s) => rty::Constant::Str(s.0),
137 fixpoint::Constant::BitVec(bv, size) => rty::Constant::BitVec(*bv, *size),
138 };
139 Ok(rty::Expr::constant(c))
140 }
141 fixpoint::Expr::Var(fvar) => {
142 match fvar {
143 fixpoint::Var::Underscore => {
144 unreachable!("Underscore should not appear in exprs")
145 }
146 fixpoint::Var::Global(_, def_id) => {
147 Ok(rty::Expr::global_func(SpecFuncKind::Def(*def_id)))
148 }
149 fixpoint::Var::Const(global_var, _) => {
150 if let Some(const_key) = self.ecx.const_env.const_map_rev.get(global_var) {
151 match const_key {
152 ConstKey::RustConst(def_id) => Ok(rty::Expr::const_def_id(*def_id)),
153 ConstKey::Alias(_flux_id, _args) => {
154 unreachable!("Should be special-cased as the head of an app")
155 }
156 ConstKey::Lambda(lambda) => Ok(rty::Expr::abs(lambda.clone())),
157 ConstKey::PrimOp(bin_op) => {
158 Ok(rty::Expr::internal_func(InternalFuncKind::Rel(
159 bin_op.clone(),
160 )))
161 }
162 ConstKey::Cast(..) => {
163 unreachable!(
164 "Should be specially handled as the head of a function app."
165 )
166 }
167 ConstKey::WKVar(_, _) => {
168 unreachable!("Weak kvars are not global vars");
169 }
170 }
171 } else {
172 Err(FixpointParseError::NoConst(*global_var))
173 }
174 }
175 fixpoint::Var::Local(fname) => {
176 if let Some(expr) = self.ecx.local_var_env.reverse_map.get(fname) {
177 return Ok(expr.clone());
178 }
179
180 for (depth, layer) in self.ecx.local_var_env.layers.iter().rev().enumerate()
181 {
182 for (idx, var) in layer.iter().enumerate() {
183 if fname == var {
184 return Ok(rty::Expr::bvar(
185 rty::DebruijnIndex::from_usize(depth),
186 BoundVar::from_usize(idx),
187 rty::BoundReftKind::Anon,
188 ));
189 }
190 }
191 }
192
193 Ok(rty::Expr::fvar(rty::Name::from_u32(fname.as_u32())))
202 }
204 fixpoint::Var::DataCtor(adt_id, variant_idx) => {
205 let def_id = self.scx.adt_sorts[adt_id.as_usize()];
206 Ok(rty::Expr::ctor_enum(def_id, *variant_idx))
207 }
208 fixpoint::Var::TupleCtor { .. }
209 | fixpoint::Var::TupleProj { .. }
210 | fixpoint::Var::DataProj { .. }
211 | fixpoint::Var::UIFRel(_) => {
212 unreachable!(
213 "Trying to convert an atomic var, but reached a var that should only occur as the head of an app (and be special-cased in conversion as a result)"
214 )
215 }
216 fixpoint::Var::Param(EarlyReftParam { index, name }) => {
217 Ok(rty::Expr::early_param(*index, *name))
218 }
219 fixpoint::Var::ConstGeneric(const_generic) => {
220 Ok(rty::Expr::const_generic(*const_generic))
221 }
222 fixpoint::Var::WKVar(..) => {
223 unreachable!(
224 "Weak kvar ids should be converted as part of fixpoint::Expr::WKVar"
225 );
226 }
227 }
228 }
229 fixpoint::Expr::App(fhead, _sort_args, fargs, _out_sort) => {
230 let mut op_args = vec![];
231 if let Some(bin_op) = self.is_curried_primop_app(fhead, fargs, &mut op_args) {
232 if op_args.len() != 2 {
233 return Err(FixpointParseError::PrimOpArityMismatch(fargs.len()));
234 } else {
235 let e1 = self.fixpoint_to_expr(&op_args[0])?;
236 let e2 = self.fixpoint_to_expr(&op_args[1])?;
237 return Ok(rty::Expr::binary_op(bin_op, e1, e2));
238 }
239 }
240 match &**fhead {
241 fixpoint::Expr::Var(fixpoint::Var::TupleProj { arity, field }) => {
242 if fargs.len() == 1 {
243 let earg = self.fixpoint_to_expr(&fargs[0])?;
244 Ok(rty::Expr::field_proj(
245 earg,
246 rty::FieldProj::Tuple { arity: *arity, field: *field },
247 ))
248 } else {
249 Err(FixpointParseError::ProjArityMismatch(fargs.len()))
250 }
251 }
252 fixpoint::Expr::Var(fixpoint::Var::DataProj { adt_id, field }) => {
253 if fargs.len() == 1 {
254 let earg = self.fixpoint_to_expr(&fargs[0])?;
255 Ok(rty::Expr::field_proj(
256 earg,
257 rty::FieldProj::Adt {
258 def_id: self.scx.adt_sorts[adt_id.as_usize()],
259 field: *field,
260 },
261 ))
262 } else {
263 Err(FixpointParseError::ProjArityMismatch(fargs.len()))
264 }
265 }
266 fixpoint::Expr::Var(fixpoint::Var::TupleCtor { arity }) => {
267 if fargs.len() == *arity {
268 let eargs = fargs
269 .iter()
270 .map(|farg| self.fixpoint_to_expr(farg))
271 .try_collect()?;
272 Ok(rty::Expr::tuple(eargs))
273 } else {
274 Err(FixpointParseError::TupleCtorArityMismatch(*arity, fargs.len()))
275 }
276 }
277 fixpoint::Expr::Var(fixpoint::Var::DataCtor(adt_id, field)) => {
278 let eargs = fargs
279 .iter()
280 .map(|farg| self.fixpoint_to_expr(farg))
281 .try_collect()?;
282 let def_id = self.scx.adt_sorts[adt_id.as_usize()];
283 match self.genv.tcx().def_kind(def_id) {
284 DefKind::Struct => Ok(rty::Expr::ctor_struct(def_id, eargs)),
285 DefKind::Enum => {
286 let ctor = rty::Ctor::Enum(def_id, *field);
287 Ok(rty::Expr::ctor(ctor, eargs))
288 }
289 _ => Err(FixpointParseError::InvalidDefKindForCtor),
290 }
291 }
292 fixpoint::Expr::Var(fixpoint::Var::UIFRel(fbinrel)) => {
293 if fargs.len() == 2 {
294 let e1 = self.fixpoint_to_expr(&fargs[0])?;
295 let e2 = self.fixpoint_to_expr(&fargs[1])?;
296 let binrel = match fbinrel {
297 fixpoint::BinRel::Eq => rty::BinOp::Eq,
298 fixpoint::BinRel::Ne => rty::BinOp::Ne,
299 fixpoint::BinRel::Gt => rty::BinOp::Gt(rty::Sort::Str),
306 fixpoint::BinRel::Ge => rty::BinOp::Ge(rty::Sort::Str),
307 fixpoint::BinRel::Lt => rty::BinOp::Lt(rty::Sort::Str),
308 fixpoint::BinRel::Le => rty::BinOp::Le(rty::Sort::Str),
309 };
310 Ok(rty::Expr::binary_op(binrel, e1, e2))
311 } else {
312 Err(FixpointParseError::UIFRelArityMismatch(fargs.len()))
313 }
314 }
315 fixpoint::Expr::Var(fixpoint::Var::Const(global_var, _)) => {
316 if let Some(const_key) = self.ecx.const_env.const_map_rev.get(global_var) {
317 match const_key {
318 ConstKey::PrimOp(_) => {
322 unreachable!(
323 "Should have been handled by is_curried_primop_app"
324 )
325 }
326 ConstKey::Cast(sort1, sort2) => {
327 if fargs.len() != 1 {
328 Err(FixpointParseError::CastArityMismatch(fargs.len()))
329 } else {
330 Ok(rty::Expr::cast(
331 sort1.clone(),
332 sort2.clone(),
333 self.fixpoint_to_expr(&fargs[0])?,
334 ))
335 }
336 }
337 ConstKey::Alias(assoc_id, generic_args) => {
338 let lowered_args: flux_rustc_bridge::ty::GenericArgs =
339 generic_args.lower(self.genv.tcx()).unwrap();
340 let generic_args = rty::refining::Refiner::default_for_item(
341 self.genv,
342 assoc_id.parent(),
343 )
344 .unwrap()
345 .refine_generic_args(assoc_id.parent(), &lowered_args)
346 .unwrap();
347 let alias_reft =
348 rty::AliasReft { assoc_id: *assoc_id, args: generic_args };
349 let args = fargs
350 .iter()
351 .map(|farg| self.fixpoint_to_expr(farg))
352 .try_collect()?;
353 Ok(rty::Expr::alias(alias_reft, args))
354 }
355 ConstKey::WKVar(..) => {
356 unreachable!("WKVars should not appear in global vars");
357 }
358 ConstKey::RustConst(..) | ConstKey::Lambda(..) => {
359 self.fixpoint_app_to_expr(fhead, fargs)
361 }
362 }
363 } else {
364 Err(FixpointParseError::NoConst(*global_var))
365 }
366 }
367 fhead => self.fixpoint_app_to_expr(fhead, fargs),
368 }
369 }
370 fixpoint::Expr::Neg(fexpr) => {
371 let e = self.fixpoint_to_expr(fexpr)?;
372 Ok(rty::Expr::neg(&e))
373 }
374 fixpoint::Expr::BinaryOp(fbinop, boxed_args) => {
375 let binop = match fbinop {
376 fixpoint::BinOp::Add => rty::BinOp::Add(rty::Sort::Int),
380 fixpoint::BinOp::Sub => rty::BinOp::Sub(rty::Sort::Int),
381 fixpoint::BinOp::Mul => rty::BinOp::Mul(rty::Sort::Int),
382 fixpoint::BinOp::Div => rty::BinOp::Div(rty::Sort::Int),
383 fixpoint::BinOp::Mod => rty::BinOp::Mod(rty::Sort::Int),
384 };
385 let [fe1, fe2] = &**boxed_args;
386 let e1 = self.fixpoint_to_expr(fe1)?;
387 let e2 = self.fixpoint_to_expr(fe2)?;
388 Ok(rty::Expr::binary_op(binop, e1, e2))
389 }
390 fixpoint::Expr::IfThenElse(boxed_args) => {
391 let [fe1, fe2, fe3] = &**boxed_args;
392 let e1 = self.fixpoint_to_expr(fe1)?;
393 let e2 = self.fixpoint_to_expr(fe2)?;
394 let e3 = self.fixpoint_to_expr(fe3)?;
395 Ok(rty::Expr::ite(e1, e2, e3))
396 }
397 fixpoint::Expr::And(fexprs) => {
398 let exprs: Vec<rty::Expr> = fexprs
399 .iter()
400 .map(|fexpr| self.fixpoint_to_expr(fexpr))
401 .try_collect()?;
402 Ok(rty::Expr::and_from_iter(exprs))
403 }
404 fixpoint::Expr::Or(fexprs) => {
405 let exprs: Vec<rty::Expr> = fexprs
406 .iter()
407 .map(|fexpr| self.fixpoint_to_expr(fexpr))
408 .try_collect()?;
409 Ok(rty::Expr::or_from_iter(exprs))
410 }
411 fixpoint::Expr::Not(fexpr) => {
412 let e = self.fixpoint_to_expr(fexpr)?;
413 Ok(rty::Expr::not(&e))
414 }
415 fixpoint::Expr::Imp(boxed_args) => {
416 let [fe1, fe2] = &**boxed_args;
417 let e1 = self.fixpoint_to_expr(fe1)?;
418 let e2 = self.fixpoint_to_expr(fe2)?;
419 Ok(rty::Expr::binary_op(rty::BinOp::Imp, e1, e2))
420 }
421 fixpoint::Expr::Iff(boxed_args) => {
422 let [fe1, fe2] = &**boxed_args;
423 let e1 = self.fixpoint_to_expr(fe1)?;
424 let e2 = self.fixpoint_to_expr(fe2)?;
425 Ok(rty::Expr::binary_op(rty::BinOp::Iff, e1, e2))
426 }
427 fixpoint::Expr::Atom(fbinrel, boxed_args) => {
428 let binrel = match fbinrel {
429 fixpoint::BinRel::Eq => rty::BinOp::Eq,
430 fixpoint::BinRel::Ne => rty::BinOp::Ne,
431 fixpoint::BinRel::Gt => rty::BinOp::Gt(rty::Sort::Int),
439 fixpoint::BinRel::Ge => rty::BinOp::Ge(rty::Sort::Int),
440 fixpoint::BinRel::Lt => rty::BinOp::Lt(rty::Sort::Int),
441 fixpoint::BinRel::Le => rty::BinOp::Le(rty::Sort::Int),
442 };
443 let [fe1, fe2] = &**boxed_args;
444 let e1 = self.fixpoint_to_expr(fe1)?;
445 let e2 = self.fixpoint_to_expr(fe2)?;
446 Ok(rty::Expr::binary_op(binrel, e1, e2))
447 }
448 fixpoint::Expr::Let(_var, _boxed_args) => {
449 todo!("Convert `var` in e2 to locally nameless var, then fill in sort");
456 }
458 fixpoint::Expr::ThyFunc(itf) => Ok(rty::Expr::global_func(SpecFuncKind::Thy(*itf))),
459 fixpoint::Expr::IsCtor(var, fe) => {
460 let (def_id, variant_idx) = match var {
461 fixpoint::Var::DataCtor(adt_id, variant_idx) => {
462 let def_id = self.scx.adt_sorts[adt_id.as_usize()];
463 Ok((def_id, *variant_idx))
464 }
465 _ => Err(FixpointParseError::WrongVarInIsCtor(*var)),
466 }?;
467 let e = self.fixpoint_to_expr(fe)?;
468 Ok(rty::Expr::is_ctor(def_id, variant_idx, e))
469 }
470 fixpoint::Expr::Quantifier(q, binder, body) => {
471 let expr = self.fixpoint_to_bind_expr(binder, body)?;
472 match q {
473 fixpoint::Quantifier::Exists => Ok(rty::Expr::exists(expr)),
474 fixpoint::Quantifier::Forall => Ok(rty::Expr::forall(expr)),
475 }
476 }
477 fixpoint::Expr::WKVar(fixpoint::WKVar { wkvid, args }) => {
478 let e_args: Vec<rty::Expr> = args
479 .iter()
480 .map(|fexpr| self.fixpoint_to_expr(fexpr))
481 .try_collect()?;
482 if let Some(const_key) = self.ecx.const_env.wkvar_map_rev.get(wkvid) {
483 match const_key {
484 ConstKey::WKVar(wkvid, self_args) => {
485 Ok(rty::Expr::wkvar(rty::WKVar {
486 wkvid: wkvid.clone(),
487 self_args: *self_args,
488 args: List::from_vec(e_args),
489 }))
490 }
491 _ => {
492 unreachable!("Weak KVar has a const_key that is not a wkvid");
493 }
494 }
495 } else {
496 unreachable!("missing weak kvar {:?} in const_env", wkvid);
497 }
498 }
499 }
500 }
501
502 fn fixpoint_to_bind_expr(
503 &mut self,
504 binder: &[(fixpoint::Var, fixpoint::Sort)],
505 body: &fixpoint::Expr,
506 ) -> Result<rty::Binder<rty::Expr>, FixpointParseError> {
507 let mut vars = vec![];
508 let mut sorts = vec![];
509 for (var, sort) in binder {
510 let fixpoint::Var::Local(local_var) = var else {
511 return Err(FixpointParseError::WrongVarInBinder(*var));
512 };
513 vars.push(*local_var);
514 sorts.push(self.fixpoint_to_sort(sort)?);
515 }
516 self.ecx.local_var_env.push_layer(vars);
517 let body = self.fixpoint_to_expr(body)?;
518 self.ecx.local_var_env.pop_layer();
519 Ok(Binder::bind_with_sorts(body, &sorts))
520 }
521
522 fn fixpoint_app_to_expr(
523 &mut self,
524 fhead: &fixpoint::Expr,
525 fargs: &[fixpoint::Expr],
526 ) -> Result<rty::Expr, FixpointParseError> {
527 let head = self.fixpoint_to_expr(fhead)?;
528 let args = fargs
529 .iter()
530 .map(|farg| self.fixpoint_to_expr(farg))
531 .try_collect()?;
532 Ok(rty::Expr::app(head, List::empty(), args))
533 }
534}
535
536#[derive(Debug)]
537pub enum FixpointParseError {
538 UIFRelArityMismatch(usize),
541 TupleCtorArityMismatch(usize, usize),
543 ProjArityMismatch(usize),
545 NoConst(fixpoint::GlobalVar),
546 CastArityMismatch(usize),
548 PrimOpArityMismatch(usize),
549 WrongVarInIsCtor(fixpoint::Var),
552 WrongVarInBinder(fixpoint::Var),
554 UnknownAdt(DefId),
555 InvalidDefKindForCtor,
556}