Define rty::Formula and replace UnboundAssumption and Refinement with it

src/analyze/basic_block.rs

@@ -13,10 +13,12 @@ use crate::analyze;
 use crate::chc;
 use crate::pretty::PrettyDisplayExt as _;
 use crate::refine::{
-    self, BasicBlockType, Env, PlaceType, PlaceTypeVar, TempVarIdx, TemplateTypeGenerator,
-    UnboundAssumption, UnrefinedTypeGenerator, Var,
+    self, Assumption, BasicBlockType, Env, PlaceType, PlaceTypeVar, TempVarIdx,
+    TemplateTypeGenerator, UnrefinedTypeGenerator, Var,
+};
+use crate::rty::{
+    self, ClauseBuilderExt as _, ClauseScope as _, ShiftExistential as _, Subtyping as _,
 };
-use crate::rty::{self, ClauseBuilderExt as _, ClauseScope as _, Subtyping as _};
 
 mod drop_point;
 mod visitor;
@@ -339,11 +341,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         self.ctx.extend_clauses(clauses);
     }
 
-    fn with_assumptions<F, T>(
-        &mut self,
-        assumptions: Vec<impl Into<UnboundAssumption>>,
-        callback: F,
-    ) -> T
+    fn with_assumptions<F, T>(&mut self, assumptions: Vec<impl Into<Assumption>>, callback: F) -> T
     where
         F: FnOnce(&mut Self) -> T,
     {
@@ -354,7 +352,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         result
     }
 
-    fn with_assumption<F, T>(&mut self, assumption: impl Into<UnboundAssumption>, callback: F) -> T
+    fn with_assumption<F, T>(&mut self, assumption: impl Into<Assumption>, callback: F) -> T
     where
         F: FnOnce(&mut Self) -> T,
     {
@@ -767,23 +765,23 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
 #[derive(Debug, Clone)]
 pub struct UnbindAtoms<T> {
     existentials: IndexVec<rty::ExistentialVarIdx, chc::Sort>,
-    formula: rty::FormulaWithAtoms<rty::RefinedTypeVar<Var>>,
+    body: chc::Body<rty::RefinedTypeVar<Var>>,
     target_equations: Vec<(rty::RefinedTypeVar<T>, chc::Term<rty::RefinedTypeVar<Var>>)>,
 }
 
 impl<T> Default for UnbindAtoms<T> {
     fn default() -> Self {
         UnbindAtoms {
             existentials: Default::default(),
-            formula: Default::default(),
+            body: Default::default(),
             target_equations: Default::default(),
         }
     }
 }
 
 impl<T> UnbindAtoms<T> {
     pub fn push(&mut self, target: rty::RefinedTypeVar<T>, var_ty: PlaceType) {
-        self.formula.push_conj(
+        self.body.push_conj(
             var_ty
                 .formula
                 .map_var(|v| v.shift_existential(self.existentials.len()).into()),
@@ -802,13 +800,10 @@ impl<T> UnbindAtoms<T> {
             ty: src_ty,
             refinement,
         } = ty;
-        let rty::Refinement {
-            existentials,
-            formula,
-        } = refinement;
+        let rty::Refinement { existentials, body } = refinement;
 
-        self.formula
-            .push_conj(formula.map_var(|v| v.shift_existential(self.existentials.len())));
+        self.body
+            .push_conj(body.map_var(|v| v.shift_existential(self.existentials.len())));
         self.existentials.extend(existentials);
 
         let mut substs = HashMap::new();
@@ -817,12 +812,12 @@ impl<T> UnbindAtoms<T> {
             substs.insert(v, ev);
         }
 
-        let mut formula = self.formula.map_var(|v| match v {
+        let mut body = self.body.map_var(|v| match v {
             rty::RefinedTypeVar::Value => rty::RefinedTypeVar::Value,
             rty::RefinedTypeVar::Free(v) => rty::RefinedTypeVar::Existential(substs[&v]),
             rty::RefinedTypeVar::Existential(ev) => rty::RefinedTypeVar::Existential(ev),
         });
-        formula.push_conj(
+        body.push_conj(
             self.target_equations
                 .into_iter()
                 .map(|(t, term)| {
@@ -839,7 +834,7 @@ impl<T> UnbindAtoms<T> {
                 .collect::<Vec<_>>(),
         );
 
-        let refinement = rty::Refinement::with_formula(self.existentials, formula);
+        let refinement = rty::Refinement::new(self.existentials, body);
         // TODO: polymorphic datatypes: template needed?
         rty::RefinedType::new(src_ty.assert_closed().vacuous(), refinement)
     }
@@ -852,7 +847,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         expected_params: &IndexVec<rty::FunctionParamIdx, rty::RefinedType<rty::FunctionParamIdx>>,
     ) {
         let mut param_terms = HashMap::<rty::FunctionParamIdx, chc::Term<PlaceTypeVar>>::new();
-        let mut assumption = UnboundAssumption::default();
+        let mut assumption = Assumption::default();
 
         let bb_ty = self.basic_block_ty(self.basic_block).clone();
         let params = &bb_ty.as_ref().params;
@@ -878,7 +873,7 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
                 }
 
                 let local_ty = self.env.local_type(local);
-                assumption.formula.push_conj(
+                assumption.body.push_conj(
                     local_ty
                         .formula
                         .map_var(|v| v.shift_existential(assumption.existentials.len())),
@@ -892,11 +887,8 @@ impl<'tcx, 'ctx> Analyzer<'tcx, 'ctx> {
         }
 
         for (idx, param) in expected_params.iter_enumerated() {
-            let rty::Refinement {
-                existentials,
-                formula,
-            } = param.refinement.clone();
-            assumption.formula.push_conj(formula.subst_var(|v| match v {
+            let rty::Refinement { existentials, body } = param.refinement.clone();
+            assumption.body.push_conj(body.subst_var(|v| match v {
                 rty::RefinedTypeVar::Value => param_terms[&idx].clone(),
                 rty::RefinedTypeVar::Free(v) => param_terms[&v].clone(),
                 rty::RefinedTypeVar::Existential(ev) => chc::Term::var(PlaceTypeVar::Existential(

src/chc.rs

@@ -1193,16 +1193,16 @@ impl<V> Formula<V> {
         }
     }
 
-    pub fn atoms(&self) -> impl Iterator<Item = &Atom<V>> {
-        self.atoms_impl()
+    pub fn iter_atoms(&self) -> impl Iterator<Item = &Atom<V>> {
+        self.iter_atoms_impl()
     }
 
-    fn atoms_impl(&self) -> Box<dyn Iterator<Item = &Atom<V>> + '_> {
+    fn iter_atoms_impl(&self) -> Box<dyn Iterator<Item = &Atom<V>> + '_> {
         match self {
             Formula::Atom(atom) => Box::new(std::iter::once(atom)),
-            Formula::Not(fo) => Box::new(fo.atoms()),
-            Formula::And(fs) => Box::new(fs.iter().flat_map(Formula::atoms)),
-            Formula::Or(fs) => Box::new(fs.iter().flat_map(Formula::atoms)),
+            Formula::Not(fo) => Box::new(fo.iter_atoms()),
+            Formula::And(fs) => Box::new(fs.iter().flat_map(Formula::iter_atoms)),
+            Formula::Or(fs) => Box::new(fs.iter().flat_map(Formula::iter_atoms)),
         }
     }
 
@@ -1248,13 +1248,154 @@ impl<V> Formula<V> {
     }
 }
 
+#[derive(Debug, Clone)]
+pub struct Body<V = TermVarIdx> {
+    pub atoms: Vec<Atom<V>>,
+    // `formula` doesn't contain PredVar
+    pub formula: Formula<V>,
+}
+
+impl<V> Default for Body<V> {
+    fn default() -> Self {
+        Body {
+            atoms: Default::default(),
+            formula: Default::default(),
+        }
+    }
+}
+
+impl<V> From<Atom<V>> for Body<V> {
+    fn from(atom: Atom<V>) -> Self {
+        Body {
+            atoms: vec![atom],
+            formula: Formula::top(),
+        }
+    }
+}
+
+impl<V> From<Vec<Atom<V>>> for Body<V> {
+    fn from(atoms: Vec<Atom<V>>) -> Self {
+        Body {
+            atoms,
+            formula: Formula::top(),
+        }
+    }
+}
+
+impl<V> From<Formula<V>> for Body<V> {
+    fn from(formula: Formula<V>) -> Self {
+        Body {
+            atoms: vec![],
+            formula,
+        }
+    }
+}
+
+impl<V> Body<V> {
+    pub fn new(atoms: Vec<Atom<V>>, formula: Formula<V>) -> Self {
+        Body { atoms, formula }
+    }
+
+    pub fn top() -> Self {
+        Body {
+            atoms: vec![],
+            formula: Formula::top(),
+        }
+    }
+
+    pub fn bottom() -> Self {
+        Body {
+            atoms: vec![],
+            formula: Formula::bottom(),
+        }
+    }
+
+    pub fn is_top(&self) -> bool {
+        self.formula.is_top() && self.atoms.iter().all(|a| a.is_top())
+    }
+
+    pub fn is_bottom(&self) -> bool {
+        self.formula.is_bottom() || self.atoms.iter().any(|a| a.is_bottom())
+    }
+
+    pub fn push_conj(&mut self, other: impl Into<Body<V>>) {
+        let Body { atoms, formula } = other.into();
+        self.atoms.extend(atoms);
+        self.formula.push_conj(formula);
+    }
+
+    pub fn map_var<F, W>(self, mut f: F) -> Body<W>
+    where
+        F: FnMut(V) -> W,
+    {
+        Body {
+            atoms: self.atoms.into_iter().map(|a| a.map_var(&mut f)).collect(),
+            formula: self.formula.map_var(f),
+        }
+    }
+
+    pub fn subst_var<F, W>(self, mut f: F) -> Body<W>
+    where
+        F: FnMut(V) -> Term<W>,
+    {
+        Body {
+            atoms: self
+                .atoms
+                .into_iter()
+                .map(|a| a.subst_var(&mut f))
+                .collect(),
+            formula: self.formula.subst_var(f),
+        }
+    }
+
+    pub fn simplify(&mut self) {
+        self.formula.simplify();
+        self.atoms.retain(|a| !a.is_top());
+        if self.is_bottom() {
+            self.atoms = vec![Atom::bottom()];
+            self.formula = Formula::top();
+        }
+    }
+
+    pub fn iter_atoms(&self) -> impl Iterator<Item = &Atom<V>> {
+        self.formula.iter_atoms().chain(&self.atoms)
+    }
+}
+
+impl<'a, 'b, D, V> Pretty<'a, D, termcolor::ColorSpec> for &'b Body<V>
+where
+    V: Var,
+    D: pretty::DocAllocator<'a, termcolor::ColorSpec>,
+    D::Doc: Clone,
+{
+    fn pretty(self, allocator: &'a D) -> pretty::DocBuilder<'a, D, termcolor::ColorSpec> {
+        let atoms = allocator.intersperse(
+            &self.atoms,
+            allocator
+                .text("∧")
+                .enclose(allocator.line(), allocator.space()),
+        );
+        let formula = self.formula.pretty(allocator);
+        if self.atoms.is_empty() {
+            formula
+        } else if self.formula.is_top() {
+            atoms.group()
+        } else {
+            atoms
+                .append(allocator.line())
+                .append(allocator.text("∧"))
+                .append(allocator.line())
+                .append(formula)
+                .group()
+        }
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct Clause {
     pub vars: IndexVec<TermVarIdx, Sort>,
     pub head: Atom<TermVarIdx>,
-    pub body_atoms: Vec<Atom<TermVarIdx>>,
-    // body_formula doesn't contain PredVar
-    pub body_formula: Formula<TermVarIdx>,
+    pub body: Body<TermVarIdx>,
     pub debug_info: DebugInfo,
 }
 
@@ -1272,7 +1413,7 @@ where
                 allocator.text(",").append(allocator.line()),
             )
             .group();
-        let body = self.body().pretty(allocator);
+        let body = self.body.pretty(allocator);
         let imp = self
             .head
             .pretty(allocator)
@@ -1291,21 +1432,8 @@ where
 }
 
 impl Clause {
-    pub fn body(&self) -> Formula {
-        Formula::And(
-            self.body_atoms
-                .clone()
-                .into_iter()
-                .map(Formula::from)
-                .collect(),
-        )
-        .and(self.body_formula.clone())
-    }
-
     pub fn is_nop(&self) -> bool {
-        self.head.is_top()
-            || self.body_atoms.iter().any(Atom::is_bottom)
-            || self.body_formula.is_bottom()
+        self.head.is_top() || self.body.is_bottom()
     }
 
     fn term_sort(&self, term: &Term<TermVarIdx>) -> Sort {