flux_infer/
refine_tree.rs

1use std::{
2    cell::RefCell,
3    ops::ControlFlow,
4    rc::{Rc, Weak},
5};
6
7use flux_common::{index::IndexVec, iter::IterExt, tracked_span_bug};
8use flux_config::OverflowMode;
9use flux_macros::DebugAsJson;
10use flux_middle::{
11    global_env::GlobalEnv,
12    pretty::{PrettyCx, PrettyNested, format_cx},
13    queries::QueryResult,
14    rty::{
15        BaseTy, EVid, Expr, ExprKind, KVid, Name, NameProvenance, PrettyVar, Sort, Ty, TyKind, Var,
16        fold::{TypeFoldable, TypeSuperVisitable, TypeVisitable, TypeVisitor},
17    },
18};
19use itertools::Itertools;
20use rustc_data_structures::{
21    snapshot_map::SnapshotMap,
22    unord::{UnordMap, UnordSet},
23};
24use rustc_hash::FxHashSet;
25use rustc_index::newtype_index;
26use serde::Serialize;
27
28use crate::{
29    evars::EVarStore,
30    fixpoint_encoding::{FixpointCtxt, fixpoint},
31    infer::{Tag, TypeTrace},
32};
33
34/// A *refine*ment *tree* tracks the "tree-like structure" of refinement variables and predicates
35/// generated during refinement type-checking. This tree can be encoded as a fixpoint constraint
36/// whose satisfiability implies the safety of a function.
37///
38/// We try to hide the representation of the tree as much as possible and only a couple of operations
39/// can be used to manipulate the structure of the tree explicitly. Instead, the tree is mostly
40/// constructed implicitly via a restricted api provided by [`Cursor`]. Some methods operate on *nodes*
41/// of the tree which we try to keep abstract, but it is important to remember that there's an
42/// underlying tree.
43///
44/// The current implementation uses [`Rc`] and [`RefCell`] to represent the tree, but we ensure
45/// statically that the [`RefineTree`] is the single owner of the data and require a mutable reference
46/// to it for all mutations, i.e., we could in theory replace the [`RefCell`] with an [`UnsafeCell`]
47/// (or a [`GhostCell`]).
48///
49/// [`UnsafeCell`]: std::cell::UnsafeCell
50/// [`GhostCell`]: https://docs.rs/ghost-cell/0.2.3/ghost_cell/ghost_cell/struct.GhostCell.html
51pub struct RefineTree {
52    root: NodePtr,
53}
54
55impl RefineTree {
56    pub(crate) fn new(params: Vec<(Var, Sort)>, assumptions: Vec<Expr>) -> RefineTree {
57        let root = Node {
58            kind: NodeKind::Root(params, assumptions),
59            nbindings: 0,
60            parent: None,
61            children: vec![],
62        };
63        let root = NodePtr(Rc::new(RefCell::new(root)));
64        RefineTree { root }
65    }
66
67    pub(crate) fn simplify(&mut self, genv: GlobalEnv) {
68        self.root
69            .borrow_mut()
70            .simplify(SimplifyPhase::Full(genv), &mut SnapshotMap::default());
71        self.root.borrow_mut().simplify_bot();
72        self.root.borrow_mut().simplify_top();
73    }
74
75    pub(crate) fn to_fixpoint(
76        &self,
77        cx: &mut FixpointCtxt<Tag>,
78    ) -> QueryResult<fixpoint::Constraint> {
79        Ok(self
80            .root
81            .borrow()
82            .to_fixpoint(cx)?
83            .unwrap_or(fixpoint::Constraint::TRUE))
84    }
85
86    pub(crate) fn cursor_at_root(&mut self) -> Cursor<'_> {
87        Cursor { ptr: NodePtr(Rc::clone(&self.root)), tree: self }
88    }
89
90    pub(crate) fn replace_evars(&mut self, evars: &EVarStore) -> Result<(), EVid> {
91        self.root.borrow_mut().replace_evars(evars)
92    }
93}
94
95/// A cursor into the [refinement tree]. More specifically, a [`Cursor`] represents a path from the
96/// root to some internal node in a [refinement tree].
97///
98/// [refinement tree]: RefineTree
99pub struct Cursor<'a> {
100    tree: &'a mut RefineTree,
101    ptr: NodePtr,
102}
103
104impl Cursor<'_> {
105    /// Moves the cursor to the specified [marker]. If `clear_children` is `true`, all children of
106    /// the node are removed after moving the cursor, invalidating any markers pointing to a node
107    /// within those children.
108    ///
109    /// [marker]: Marker
110    pub(crate) fn move_to(&mut self, marker: &Marker, clear_children: bool) -> Option<Cursor<'_>> {
111        let ptr = marker.ptr.upgrade()?;
112        if clear_children {
113            ptr.borrow_mut().children.clear();
114        }
115        Some(Cursor { ptr, tree: self.tree })
116    }
117
118    /// Returns a marker to the current node
119    #[must_use]
120    pub(crate) fn marker(&self) -> Marker {
121        Marker { ptr: NodePtr::downgrade(&self.ptr) }
122    }
123
124    #[must_use]
125    pub(crate) fn branch(&mut self) -> Cursor<'_> {
126        Cursor { tree: self.tree, ptr: NodePtr::clone(&self.ptr) }
127    }
128
129    pub(crate) fn vars(&self) -> impl Iterator<Item = (Var, Sort)> {
130        // TODO(nilehmann) we could incrementally cache the scope
131        self.ptr.scope().into_iter()
132    }
133
134    #[expect(dead_code, reason = "used for debugging")]
135    pub(crate) fn push_trace(&mut self, trace: TypeTrace) {
136        self.ptr = self.ptr.push_node(NodeKind::Trace(trace));
137    }
138
139    /// Defines a fresh refinement variable with the given `sort` and advance the cursor to the new
140    /// node. It returns the freshly generated name for the variable.
141    pub(crate) fn define_var(&mut self, sort: &Sort, provenance: NameProvenance) -> Name {
142        let fresh = Name::from_usize(self.ptr.next_name_idx());
143        self.ptr = self
144            .ptr
145            .push_node(NodeKind::ForAll(fresh, sort.clone(), provenance));
146        fresh
147    }
148
149    /// Pushes an [assumption] and moves the cursor into the new node.
150    ///
151    /// [assumption]: NodeKind::Assumption
152    pub(crate) fn assume_pred(&mut self, pred: impl Into<Expr>) {
153        let pred = pred.into();
154        if !pred.is_trivially_true() {
155            self.ptr = self.ptr.push_node(NodeKind::Assumption(pred));
156        }
157    }
158
159    /// Pushes a predicate that must be true assuming variables and predicates in the current branch
160    /// of the tree (i.e., it pushes a [`NodeKind::Head`]). This methods does not advance the cursor.
161    pub(crate) fn check_pred(&mut self, pred: impl Into<Expr>, tag: Tag) {
162        let pred = pred.into();
163        if !pred.is_trivially_true() {
164            self.ptr.push_node(NodeKind::Head(pred, tag));
165        }
166    }
167
168    /// Convenience method to push an assumption followed by a predicate that needs to be checked.
169    /// This method does not advance the cursor.
170    pub(crate) fn check_impl(&mut self, pred1: impl Into<Expr>, pred2: impl Into<Expr>, tag: Tag) {
171        self.ptr
172            .push_node(NodeKind::Assumption(pred1.into()))
173            .push_node(NodeKind::Head(pred2.into(), tag));
174    }
175
176    pub(crate) fn assume_invariants(
177        &mut self,
178        genv: GlobalEnv,
179        ty: &Ty,
180        overflow_checking: OverflowMode,
181    ) {
182        struct Visitor<'cursor, 'tree, 'genv, 'tcx> {
183            genv: GlobalEnv<'genv, 'tcx>,
184            cursor: &'cursor mut Cursor<'tree>,
185            overflow_mode: OverflowMode,
186        }
187        impl TypeVisitor for Visitor<'_, '_, '_, '_> {
188            fn visit_bty(&mut self, bty: &BaseTy) -> ControlFlow<!> {
189                match bty {
190                    BaseTy::Adt(adt_def, substs) if adt_def.is_box() => substs.visit_with(self),
191                    BaseTy::Ref(_, ty, _) => ty.visit_with(self),
192                    BaseTy::Tuple(tys) => tys.visit_with(self),
193                    _ => ControlFlow::Continue(()),
194                }
195            }
196
197            fn visit_ty(&mut self, ty: &Ty) -> ControlFlow<!> {
198                if let TyKind::Indexed(bty, idx) = ty.kind()
199                    && !idx.has_escaping_bvars()
200                {
201                    for invariant in bty.invariants(self.genv, self.overflow_mode) {
202                        let invariant = invariant.apply(idx);
203                        self.cursor.assume_pred(&invariant);
204                    }
205                }
206                ty.super_visit_with(self)
207            }
208        }
209        let _ =
210            ty.visit_with(&mut Visitor { genv, cursor: self, overflow_mode: overflow_checking });
211    }
212}
213
214/// A marker is a pointer to a node in the [refinement tree] that can be used to query information
215/// about that node or to move the cursor. A marker may become invalid if the underlying node is
216/// [cleared].
217///
218/// [cleared]: Cursor::move_to
219/// [refinement tree]: RefineTree
220pub struct Marker {
221    ptr: WeakNodePtr,
222}
223
224impl Marker {
225    /// Returns the [`scope`] at the marker if it is still valid or [`None`] otherwise.
226    ///
227    /// [`scope`]: Scope
228    pub fn scope(&self) -> Option<Scope> {
229        Some(self.ptr.upgrade()?.scope())
230    }
231
232    pub fn has_free_vars<T: TypeVisitable>(&self, t: &T) -> bool {
233        let ptr = self
234            .ptr
235            .upgrade()
236            .unwrap_or_else(|| tracked_span_bug!("`has_free_vars` called on invalid `Marker`"));
237
238        let nbindings = ptr.next_name_idx();
239
240        !t.fvars().into_iter().all(|name| name.index() < nbindings)
241    }
242}
243
244/// A list of refinement variables and their sorts.
245#[derive(PartialEq, Eq)]
246pub struct Scope {
247    params: Vec<(Var, Sort)>,
248    bindings: IndexVec<Name, Sort>,
249}
250
251impl Scope {
252    pub(crate) fn iter(&self) -> impl Iterator<Item = (Var, Sort)> + '_ {
253        itertools::chain(
254            self.params.iter().cloned(),
255            self.bindings
256                .iter_enumerated()
257                .map(|(name, sort)| (Var::Free(name), sort.clone())),
258        )
259    }
260
261    fn into_iter(self) -> impl Iterator<Item = (Var, Sort)> {
262        itertools::chain(
263            self.params,
264            self.bindings
265                .into_iter_enumerated()
266                .map(|(name, sort)| (Var::Free(name), sort.clone())),
267        )
268    }
269
270    /// Whether `t` has any free variables not in this scope
271    pub fn has_free_vars<T: TypeFoldable>(&self, t: &T) -> bool {
272        !self.contains_all(t.fvars())
273    }
274
275    fn contains_all(&self, iter: impl IntoIterator<Item = Name>) -> bool {
276        iter.into_iter().all(|name| self.contains(name))
277    }
278
279    fn contains(&self, name: Name) -> bool {
280        name.index() < self.bindings.len()
281    }
282}
283
284struct Node {
285    kind: NodeKind,
286    /// Number of bindings between the root and this node's parent, i.e., we have
287    /// as an invariant that `nbindings` equals the number of [`NodeKind::ForAll`]
288    /// nodes from the parent of this node to the root.
289    nbindings: usize,
290    parent: Option<WeakNodePtr>,
291    children: Vec<NodePtr>,
292}
293
294#[derive(Clone)]
295struct NodePtr(Rc<RefCell<Node>>);
296
297impl NodePtr {
298    fn downgrade(this: &Self) -> WeakNodePtr {
299        WeakNodePtr(Rc::downgrade(&this.0))
300    }
301
302    fn push_node(&mut self, kind: NodeKind) -> NodePtr {
303        debug_assert!(!matches!(self.borrow().kind, NodeKind::Head(..)));
304        let node = Node {
305            kind,
306            nbindings: self.next_name_idx(),
307            parent: Some(NodePtr::downgrade(self)),
308            children: vec![],
309        };
310        let node = NodePtr(Rc::new(RefCell::new(node)));
311        self.borrow_mut().children.push(NodePtr::clone(&node));
312        node
313    }
314
315    fn next_name_idx(&self) -> usize {
316        self.borrow().nbindings + usize::from(self.borrow().is_forall())
317    }
318
319    fn scope(&self) -> Scope {
320        let mut params = None;
321        let parents = ParentsIter::new(self.clone());
322        let bindings = parents
323            .filter_map(|node| {
324                let node = node.borrow();
325                match &node.kind {
326                    NodeKind::Root(p, _) => {
327                        params = Some(p.clone());
328                        None
329                    }
330                    NodeKind::ForAll(_, sort, _) => Some(sort.clone()),
331                    _ => None,
332                }
333            })
334            .collect_vec()
335            .into_iter()
336            .rev()
337            .collect();
338        Scope { bindings, params: params.unwrap_or_default() }
339    }
340}
341
342struct WeakNodePtr(Weak<RefCell<Node>>);
343
344impl WeakNodePtr {
345    fn upgrade(&self) -> Option<NodePtr> {
346        Some(NodePtr(self.0.upgrade()?))
347    }
348}
349
350enum NodeKind {
351    /// List of const and refinement generics, and facts assumed about them (e.g., `N >= 0` for
352    /// `const N: usize`)
353    Root(Vec<(Var, Sort)>, Vec<Expr>),
354    ForAll(Name, Sort, NameProvenance),
355    Assumption(Expr),
356    Head(Expr, Tag),
357    True,
358    /// Used for debugging. See [`TypeTrace`]
359    Trace(TypeTrace),
360}
361
362impl std::ops::Index<Name> for Scope {
363    type Output = Sort;
364
365    fn index(&self, name: Name) -> &Self::Output {
366        &self.bindings[name]
367    }
368}
369
370impl std::ops::Deref for NodePtr {
371    type Target = Rc<RefCell<Node>>;
372
373    fn deref(&self) -> &Self::Target {
374        &self.0
375    }
376}
377
378#[derive(Clone, Copy)]
379enum SimplifyPhase<'genv, 'tcx> {
380    /// Normalize and simplify inner `Expr`
381    Full(GlobalEnv<'genv, 'tcx>),
382    /// Only propagate `true` (TOP) and `false` (BOT)
383    Partial,
384}
385
386impl Node {
387    fn simplify(&mut self, phase: SimplifyPhase, assumed_preds: &mut SnapshotMap<Expr, ()>) {
388        // First, simplify the node itself
389        match &mut self.kind {
390            NodeKind::Head(pred, tag) => {
391                let pred = match phase {
392                    SimplifyPhase::Full(genv) => pred.reduce(genv).simplify(assumed_preds),
393                    SimplifyPhase::Partial => pred.clone(),
394                };
395                if pred.is_trivially_true() {
396                    self.kind = NodeKind::True;
397                } else {
398                    self.kind = NodeKind::Head(pred, *tag);
399                }
400            }
401            NodeKind::Assumption(pred) => {
402                if let SimplifyPhase::Full(genv) = phase {
403                    *pred = pred.reduce(genv).simplify(assumed_preds);
404                }
405                pred.visit_conj(|conjunct| {
406                    assumed_preds.insert(conjunct.erase_spans(), ());
407                });
408            }
409            NodeKind::Root(_, assumptions) => {
410                for pred in assumptions {
411                    pred.visit_conj(|conjunct| {
412                        assumed_preds.insert(conjunct.erase_spans(), ());
413                    });
414                }
415            }
416            _ => {}
417        }
418        let is_false_asm =
419            matches!(&self.kind, NodeKind::Assumption(pred) if pred.is_trivially_false());
420
421        // Then simplify the children
422        // (the order matters here because we need to collect assumed preds first)
423        for child in &self.children {
424            let current_version = assumed_preds.snapshot();
425            child.borrow_mut().simplify(phase, assumed_preds);
426            assumed_preds.rollback_to(current_version);
427        }
428
429        // Then remove any unnecessary children
430        match &mut self.kind {
431            NodeKind::Head(..) | NodeKind::True => {}
432            NodeKind::Assumption(_)
433            | NodeKind::Trace(_)
434            | NodeKind::Root(..)
435            | NodeKind::ForAll(..) => {
436                self.children
437                    .extract_if(.., |child| {
438                        is_false_asm || matches!(&child.borrow().kind, NodeKind::True)
439                    })
440                    .for_each(drop);
441            }
442        }
443        if !self.is_leaf() && self.children.is_empty() {
444            self.kind = NodeKind::True;
445        }
446    }
447
448    fn is_leaf(&self) -> bool {
449        matches!(self.kind, NodeKind::Head(..) | NodeKind::True)
450    }
451
452    fn replace_evars(&mut self, evars: &EVarStore) -> Result<(), EVid> {
453        for child in &self.children {
454            child.borrow_mut().replace_evars(evars)?;
455        }
456        match &mut self.kind {
457            NodeKind::Assumption(pred) => *pred = evars.replace_evars(pred)?,
458            NodeKind::Head(pred, _) => {
459                *pred = evars.replace_evars(pred)?;
460            }
461            NodeKind::Trace(trace) => {
462                evars.replace_evars(trace)?;
463            }
464            NodeKind::Root(_, assumptions) => {
465                for pred in assumptions {
466                    *pred = evars.replace_evars(pred)?;
467                }
468            }
469            NodeKind::ForAll(..) | NodeKind::True => {}
470        }
471        Ok(())
472    }
473
474    fn to_fixpoint(&self, cx: &mut FixpointCtxt<Tag>) -> QueryResult<Option<fixpoint::Constraint>> {
475        let cstr = match &self.kind {
476            NodeKind::Trace(_) | NodeKind::ForAll(_, Sort::Loc, _) => {
477                children_to_fixpoint(cx, &self.children)?
478            }
479
480            NodeKind::Root(params, assumptions) => {
481                // declare pretty-vars for params
482                for (var, sort) in params {
483                    if let Var::EarlyParam(param) = var
484                        && !sort.is_loc()
485                    {
486                        cx.with_early_param(param);
487                    }
488                }
489
490                let Some(children) = children_to_fixpoint(cx, &self.children)? else {
491                    return Ok(None);
492                };
493                let mut constr = children;
494                if !assumptions.is_empty() {
495                    let preds = cx.assumptions_to_fixpoint(assumptions)?;
496                    constr = fixpoint::Constraint::ForAll(
497                        fixpoint::Bind {
498                            name: fixpoint::Var::Underscore,
499                            sort: fixpoint::Sort::Int,
500                            preds,
501                        },
502                        Box::new(constr),
503                    );
504                }
505                for (var, sort) in params.iter().rev() {
506                    if sort.is_loc() {
507                        continue;
508                    }
509                    constr = fixpoint::Constraint::ForAll(
510                        fixpoint::Bind {
511                            name: cx.var_to_fixpoint(var),
512                            sort: cx.sort_to_fixpoint(sort),
513                            preds: vec![],
514                        },
515                        Box::new(constr),
516                    );
517                }
518                Some(constr)
519            }
520            NodeKind::ForAll(name, sort, provenance) => {
521                cx.with_name_map(*name, *provenance, |cx, fresh| -> QueryResult<_> {
522                    let Some(children) = children_to_fixpoint(cx, &self.children)? else {
523                        return Ok(None);
524                    };
525                    Ok(Some(fixpoint::Constraint::ForAll(
526                        fixpoint::Bind {
527                            name: fixpoint::Var::Local(fresh),
528                            sort: cx.sort_to_fixpoint(sort),
529                            preds: vec![],
530                        },
531                        Box::new(children),
532                    )))
533                })?
534            }
535            NodeKind::Assumption(pred) => {
536                let preds = cx.assumptions_to_fixpoint([pred])?;
537                let Some(cstr) = children_to_fixpoint(cx, &self.children)? else {
538                    return Ok(None);
539                };
540                Some(fixpoint::Constraint::ForAll(
541                    fixpoint::Bind {
542                        name: fixpoint::Var::Underscore,
543                        sort: fixpoint::Sort::Int,
544                        preds,
545                    },
546                    Box::new(cstr),
547                ))
548            }
549            NodeKind::Head(pred, tag) => {
550                Some(cx.head_to_fixpoint(pred, |span| tag.with_dst(span))?)
551            }
552            NodeKind::True => None,
553        };
554        Ok(cstr)
555    }
556
557    /// Returns `true` if the node kind is [`ForAll`].
558    ///
559    /// [`ForAll`]: NodeKind::ForAll
560    fn is_forall(&self) -> bool {
561        matches!(self.kind, NodeKind::ForAll(..))
562    }
563
564    /// Returns `true` if the node kind is [`Head`].
565    ///
566    /// [`Head`]: NodeKind::Head
567    fn is_head(&self) -> bool {
568        matches!(self.kind, NodeKind::Head(..))
569    }
570}
571
572fn children_to_fixpoint(
573    cx: &mut FixpointCtxt<Tag>,
574    children: &[NodePtr],
575) -> QueryResult<Option<fixpoint::Constraint>> {
576    let mut children = children
577        .iter()
578        .filter_map(|node| node.borrow().to_fixpoint(cx).transpose())
579        .try_collect_vec()?;
580    let cstr = match children.len() {
581        0 => None,
582        1 => children.pop(),
583        _ => Some(fixpoint::Constraint::conj(children)),
584    };
585    Ok(cstr)
586}
587
588struct ParentsIter {
589    ptr: Option<NodePtr>,
590}
591
592impl ParentsIter {
593    fn new(ptr: NodePtr) -> Self {
594        Self { ptr: Some(ptr) }
595    }
596}
597
598impl Iterator for ParentsIter {
599    type Item = NodePtr;
600
601    fn next(&mut self) -> Option<Self::Item> {
602        if let Some(ptr) = self.ptr.take() {
603            self.ptr = ptr.borrow().parent.as_ref().and_then(WeakNodePtr::upgrade);
604            Some(ptr)
605        } else {
606            None
607        }
608    }
609}
610
611mod pretty {
612    use std::fmt::{self, Write};
613
614    use flux_middle::pretty::*;
615    use pad_adapter::PadAdapter;
616
617    use super::*;
618
619    fn bindings_chain(ptr: &NodePtr) -> (Vec<(Name, Sort)>, Vec<NodePtr>) {
620        fn go(ptr: &NodePtr, mut bindings: Vec<(Name, Sort)>) -> (Vec<(Name, Sort)>, Vec<NodePtr>) {
621            let node = ptr.borrow();
622            if let NodeKind::ForAll(name, sort, _) = &node.kind {
623                bindings.push((*name, sort.clone()));
624                if let [child] = &node.children[..] {
625                    go(child, bindings)
626                } else {
627                    (bindings, node.children.clone())
628                }
629            } else {
630                (bindings, vec![NodePtr::clone(ptr)])
631            }
632        }
633        go(ptr, vec![])
634    }
635
636    fn preds_chain(ptr: &NodePtr) -> (Vec<Expr>, Vec<NodePtr>) {
637        fn go(ptr: &NodePtr, mut preds: Vec<Expr>) -> (Vec<Expr>, Vec<NodePtr>) {
638            let node = ptr.borrow();
639            if let NodeKind::Assumption(pred) = &node.kind {
640                preds.push(pred.clone());
641                if let [child] = &node.children[..] {
642                    go(child, preds)
643                } else {
644                    (preds, node.children.clone())
645                }
646            } else {
647                (preds, vec![NodePtr::clone(ptr)])
648            }
649        }
650        go(ptr, vec![])
651    }
652
653    impl Pretty for RefineTree {
654        fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
655            w!(cx, f, "{:?}", &self.root)
656        }
657    }
658
659    impl Pretty for NodePtr {
660        fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
661            let node = self.borrow();
662            match &node.kind {
663                NodeKind::Trace(trace) => {
664                    w!(cx, f, "@ {:?}", ^trace)?;
665                    w!(cx, with_padding(f), "\n{:?}", join!("\n", &node.children))
666                }
667                NodeKind::Root(bindings, assumptions) => {
668                    w!(cx, f,
669                        "∀ {}.",
670                        ^bindings
671                            .iter()
672                            .format_with(", ", |(name, sort), f| {
673                                f(&format_args_cx!(cx, "{:?}: {:?}", ^name, sort))
674                            })
675                    )?;
676                    if !assumptions.is_empty() {
677                        w!(cx, f, " {:?} =>", join!(" ∧ ", assumptions))?;
678                    }
679                    fmt_children(&node.children, cx, f)
680                }
681                NodeKind::ForAll(name, sort, _) => {
682                    let (bindings, children) = if cx.bindings_chain {
683                        bindings_chain(self)
684                    } else {
685                        (vec![(*name, sort.clone())], node.children.clone())
686                    };
687
688                    w!(cx, f,
689                        "∀ {}.",
690                        ^bindings
691                            .into_iter()
692                            .format_with(", ", |(name, sort), f| {
693                                f(&format_args_cx!(cx, "{:?}: {:?}", ^name, sort))
694                            })
695                    )?;
696                    fmt_children(&children, cx, f)
697                }
698                NodeKind::Assumption(pred) => {
699                    let (preds, children) = if cx.preds_chain {
700                        preds_chain(self)
701                    } else {
702                        (vec![pred.clone()], node.children.clone())
703                    };
704                    let guard = Expr::and_from_iter(preds).simplify(&SnapshotMap::default());
705                    w!(cx, f, "{:?} =>", parens!(guard, !guard.is_atom()))?;
706                    fmt_children(&children, cx, f)
707                }
708                NodeKind::Head(pred, tag) => {
709                    let pred = if cx.simplify_exprs {
710                        pred.simplify(&SnapshotMap::default())
711                    } else {
712                        pred.clone()
713                    };
714                    w!(cx, f, "{:?}", parens!(pred, !pred.is_atom()))?;
715                    if cx.tags {
716                        w!(cx, f, " ~ {:?}", tag)?;
717                    }
718                    Ok(())
719                }
720                NodeKind::True => {
721                    w!(cx, f, "true")
722                }
723            }
724        }
725    }
726
727    fn fmt_children(
728        children: &[NodePtr],
729        cx: &PrettyCx,
730        f: &mut fmt::Formatter<'_>,
731    ) -> fmt::Result {
732        match children {
733            [] => w!(cx, f, " true"),
734            [n] => {
735                if n.borrow().is_head() {
736                    w!(cx, f, " {:?}", n)
737                } else {
738                    w!(cx, with_padding(f), "\n{:?}", n)
739                }
740            }
741            _ => w!(cx, with_padding(f), "\n{:?}", join!("\n", children)),
742        }
743    }
744
745    impl Pretty for Cursor<'_> {
746        fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
747            let mut elements = vec![];
748            for node in ParentsIter::new(NodePtr::clone(&self.ptr)) {
749                let n = node.borrow();
750                match &n.kind {
751                    NodeKind::Root(bindings, assumptions) => {
752                        // We reverse here because is reversed again at the end
753                        for pred in assumptions.iter().rev() {
754                            elements.push(format_cx!(cx, "{:?}", pred));
755                        }
756                        for (name, sort) in bindings.iter().rev() {
757                            elements.push(format_cx!(cx, "{:?}: {:?}", ^name, sort));
758                        }
759                    }
760                    NodeKind::ForAll(name, sort, _) => {
761                        elements.push(format_cx!(cx, "{:?}: {:?}", ^name, sort));
762                    }
763                    NodeKind::Assumption(pred) => {
764                        elements.push(format_cx!(cx, "{:?}", pred));
765                    }
766                    _ => {}
767                }
768            }
769            write!(f, "{{{}}}", elements.into_iter().rev().format(", "))
770        }
771    }
772
773    impl Pretty for Scope {
774        fn fmt(&self, cx: &PrettyCx, f: &mut fmt::Formatter<'_>) -> fmt::Result {
775            write!(
776                f,
777                "[bindings = {}, reftgenerics = {}]",
778                self.bindings
779                    .iter_enumerated()
780                    .format_with(", ", |(name, sort), f| {
781                        f(&format_args_cx!(cx, "{:?}: {:?}", ^name, sort))
782                    }),
783                self.params
784                    .iter()
785                    .format_with(", ", |(param_const, sort), f| {
786                        f(&format_args_cx!(cx, "{:?}: {:?}", ^param_const, sort))
787                    }),
788            )
789        }
790    }
791
792    fn with_padding<'a, 'b>(f: &'a mut fmt::Formatter<'b>) -> PadAdapter<'a, 'b, 'static> {
793        PadAdapter::with_padding(f, "  ")
794    }
795
796    impl_debug_with_default_cx!(
797        RefineTree => "refine_tree",
798        Cursor<'_> => "cursor",
799        Scope,
800    );
801}
802
803/// An explicit representation of a path in the [`RefineTree`] used for debugging/tracing/serialization ONLY.
804#[derive(Serialize, DebugAsJson)]
805pub struct RefineCtxtTrace {
806    bindings: Vec<RcxBind>,
807    exprs: Vec<String>,
808}
809
810#[derive(Serialize)]
811struct RcxBind {
812    name: String,
813    sort: String,
814}
815
816impl RefineCtxtTrace {
817    pub fn new(cx: &mut PrettyCx, cursor: &Cursor) -> Self {
818        let parents = ParentsIter::new(NodePtr::clone(&cursor.ptr)).collect_vec();
819        let mut bindings = vec![];
820        let mut exprs = vec![];
821
822        parents.into_iter().rev().for_each(|ptr| {
823            let node = ptr.borrow();
824            match &node.kind {
825                NodeKind::ForAll(name, sort, provenance) => {
826                    let name = cx
827                        .pretty_var_env
828                        .set(PrettyVar::Local(*name), provenance.opt_symbol());
829                    let sort = format_cx!(cx, "{:?}", sort);
830                    let bind = RcxBind { name, sort };
831                    bindings.push(bind);
832                }
833                NodeKind::Assumption(e)
834                    if !e.simplify(&SnapshotMap::default()).is_trivially_true() =>
835                {
836                    e.visit_conj(|e| {
837                        exprs.push(e.nested_string(cx));
838                    });
839                }
840                NodeKind::Root(binds, assumptions) => {
841                    for (name, sort) in binds {
842                        let name = if let Var::EarlyParam(param) = name {
843                            cx.pretty_var_env
844                                .set(PrettyVar::Param(*param), Some(param.name))
845                        } else {
846                            format_cx!(cx, "{:?}", name)
847                        };
848                        let sort = format_cx!(cx, "{:?}", sort);
849                        let bind = RcxBind { name, sort };
850                        bindings.push(bind);
851                    }
852                    for e in assumptions {
853                        e.visit_conj(|e| {
854                            exprs.push(e.nested_string(cx));
855                        });
856                    }
857                }
858                _ => (),
859            }
860        });
861        Self { bindings, exprs }
862    }
863}
864
865impl Node {
866    /// replace bot-kvars with false
867    fn simplify_bot(&mut self) {
868        let graph = ConstraintDeps::new(self);
869        let bots = graph.bot_kvars();
870        self.simplify_with_assignment(&bots);
871        self.simplify(SimplifyPhase::Partial, &mut SnapshotMap::default());
872    }
873
874    /// replace top-kvars with true
875    fn simplify_top(&mut self) {
876        let graph = ConstraintDeps::new(self);
877        let tops = graph.top_kvars();
878        self.simplify_with_assignment(&tops);
879        self.simplify(SimplifyPhase::Partial, &mut SnapshotMap::default());
880    }
881
882    /// simplifies assumptions and heads using the TOP/BOT kvar assignment; follow
883    /// with a call to `simplify` to delete constraints with FALSE assm.
884    fn simplify_with_assignment(&mut self, assignment: &Assignment) {
885        match &mut self.kind {
886            NodeKind::Head(pred, tag) => {
887                let pred = assignment.simplify(pred);
888                self.kind = NodeKind::Head(pred, *tag);
889            }
890            NodeKind::Assumption(pred) => {
891                let pred = assignment.simplify(pred);
892                self.kind = NodeKind::Assumption(pred);
893            }
894            _ => {}
895        }
896        for child in &self.children {
897            child.borrow_mut().simplify_with_assignment(assignment);
898        }
899    }
900}
901
902#[derive(Debug)]
903struct ConstraintDeps {
904    /// assumptions for each clause
905    assumptions: IndexVec<ClauseId, FxHashSet<KVid>>,
906    /// head of each clause
907    heads: IndexVec<ClauseId, Head>,
908}
909
910impl ConstraintDeps {
911    fn new(node: &Node) -> Self {
912        let mut graph = Self { assumptions: IndexVec::default(), heads: IndexVec::default() };
913        graph.build(node, &mut vec![]);
914        graph
915    }
916
917    fn insert_clause(&mut self, assumptions: &[KVid], head: Head) {
918        self.assumptions.push(assumptions.iter().copied().collect());
919        self.heads.push(head);
920    }
921
922    fn build(&mut self, node: &Node, assumptions: &mut Vec<KVid>) {
923        let n = assumptions.len();
924        match &node.kind {
925            NodeKind::Head(expr, _) => {
926                expr.visit_conj(|e| {
927                    if let ExprKind::KVar(kvar) = e.kind() {
928                        self.insert_clause(assumptions, Head::KVar(kvar.kvid));
929                    } else {
930                        self.insert_clause(assumptions, Head::Conc);
931                    }
932                });
933            }
934            NodeKind::Assumption(expr) => {
935                expr.visit_conj(|e| {
936                    if let ExprKind::KVar(kvar) = e.kind() {
937                        assumptions.push(kvar.kvid);
938                    }
939                });
940            }
941            _ => {}
942        };
943
944        for child in &node.children {
945            self.build(&child.borrow(), assumptions);
946        }
947
948        assumptions.truncate(n); // restore ctx
949    }
950
951    /// set of edges where kvid appears as ASSM
952    fn kv_lhs(&self) -> UnordMap<KVid, Vec<ClauseId>> {
953        let mut res: UnordMap<KVid, Vec<ClauseId>> = UnordMap::default();
954        for (clause_id, kvids) in self.assumptions.iter_enumerated() {
955            for kvid in kvids {
956                res.entry(*kvid).or_default().push(clause_id);
957            }
958        }
959        res
960    }
961
962    /// set of edges where kvid appears as HEAD
963    fn kv_rhs(&self) -> UnordMap<KVid, Vec<ClauseId>> {
964        let mut res: UnordMap<KVid, Vec<ClauseId>> = UnordMap::default();
965        for (clause_id, head) in self.heads.iter_enumerated() {
966            if let Head::KVar(kvid) = head {
967                res.entry(*kvid).or_default().push(clause_id);
968            }
969        }
970        res
971    }
972
973    /// Computes the set of all kvars that can be assigned to Bot (False),
974    /// because they are not (transitively) reachable from any concrete ASSUMPTION.
975    fn bot_kvars(self) -> Assignment {
976        // set of BOT kvars (initially, all)
977        let mut assignment = Assignment::new(Label::Bot);
978
979        let kv_lhs = self.kv_lhs();
980
981        // set of BOT kvars in LHS of each constraint with KVar HEAD
982        let mut bot_assms: IndexVec<ClauseId, FxHashSet<KVid>> = self.assumptions;
983
984        // set of constraints `cid` whose bot-assms is empty
985        let mut candidates: Vec<ClauseId> = bot_assms
986            .iter_enumerated()
987            .filter_map(|(cid, lhs)| if lhs.is_empty() { Some(cid) } else { None })
988            .collect();
989
990        // while there is a candidate constraint, that has NO BOT kvars in lhs
991        while let Some(cid) = candidates.pop() {
992            if let Head::KVar(kvid) = self.heads[cid] {
993                // un-BOT the head kvar
994                assignment.remove(kvid);
995                // remove the head kvar from all (bot) assumptions where it currently occurs
996                for cid in kv_lhs.get(&kvid).unwrap_or(&vec![]) {
997                    // if cid HEAD is a kvar
998                    if let Head::KVar(rhs_kvid) = self.heads[*cid] {
999                        let assms = &mut bot_assms[*cid];
1000                        assms.remove(&kvid);
1001                        if assignment.has_label(rhs_kvid) && assms.is_empty() {
1002                            candidates.push(*cid);
1003                        }
1004                    };
1005                }
1006            }
1007        }
1008
1009        assignment
1010    }
1011
1012    /// Computes the set of all kvars that can be assigned to Top (True),
1013    /// because they do not (transitively) reach any concrete HEAD.
1014    fn top_kvars(self) -> Assignment {
1015        // initialize
1016        let mut assignment = Assignment::new(Label::Top);
1017
1018        let kv_rhs = self.kv_rhs();
1019
1020        // set of kvar {k | cid in graph.edges, c.rhs is concrete, k in c.lhs }
1021        let mut candidates = vec![];
1022        for (cid, head) in self.heads.iter_enumerated() {
1023            if matches!(head, Head::Conc) {
1024                for kvid in &self.assumptions[cid] {
1025                    candidates.push(*kvid);
1026                }
1027            }
1028        }
1029
1030        // set each kvar that transitively reaches a concrete HEAD to NON-BOT
1031        while let Some(kvid) = candidates.pop() {
1032            // set that kvar to non-top
1033            assignment.remove(kvid);
1034
1035            // for each constraint where kvid appears as head
1036            for cid in kv_rhs.get(&kvid).unwrap_or(&vec![]) {
1037                // add kvars in lhs to candidates (if they have not already been solved to non-BOT)
1038                for lhs_kvid in &self.assumptions[*cid] {
1039                    if assignment.has_label(*lhs_kvid) {
1040                        candidates.push(*lhs_kvid);
1041                    }
1042                }
1043            }
1044        }
1045
1046        assignment
1047    }
1048}
1049
1050newtype_index! {
1051    struct ClauseId {}
1052}
1053
1054#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
1055enum Head {
1056    /// KVar
1057    KVar(KVid),
1058    /// A *conc*rete predicate, i.e., an [`Expr`] that's not a kvar. We don't need to know
1059    /// the exact expression, only that it's concrete.
1060    Conc,
1061}
1062
1063#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1064enum Label {
1065    /// Kvar can be solved to false
1066    Bot,
1067    /// Kvar can be solved to true
1068    Top,
1069}
1070
1071struct Assignment {
1072    /// These vars are NOT assigned `label`,
1073    /// all other `KVid` implicitly have assignment `label`
1074    vars: UnordSet<KVid>,
1075    label: Label,
1076}
1077
1078impl Assignment {
1079    fn new(label: Label) -> Self {
1080        let vars = UnordSet::new();
1081        Self { vars, label }
1082    }
1083
1084    fn has_label(&self, kvid: KVid) -> bool {
1085        !self.vars.contains(&kvid)
1086    }
1087
1088    fn remove(&mut self, kvid: KVid) {
1089        self.vars.insert(kvid);
1090    }
1091
1092    /// simplifies the given predicate expression by replacing
1093    /// kvid assigned to TOP with True, BOT with false.
1094    fn simplify(&self, pred: &Expr) -> Expr {
1095        let mut preds = vec![];
1096        for p in pred.flatten_conjs() {
1097            if let ExprKind::KVar(kvar) = p.kind()
1098                && self.has_label(kvar.kvid)
1099            {
1100                if self.label == Label::Bot {
1101                    return Expr::ff();
1102                } // else, skip pushing `p` into `preds`
1103            } else {
1104                preds.push(p.clone());
1105            }
1106        }
1107        Expr::and_from_iter(preds)
1108    }
1109}
flux_infer/refine_tree.rs

flux_infer/
refine_tree.rs
