Fix unsoundness when expanding variant of polymorphic enum

src/chc.rs

@@ -997,6 +997,12 @@ impl Pred {
 /// An atom is a predicate applied to a list of terms.
 #[derive(Debug, Clone)]
 pub struct Atom<V = TermVarIdx> {
+    /// With `guard`, this represents `guard => pred(args)`.
+    ///
+    /// As long as there is no pvar in the `guard`, it forms a valid CHC. However, in that case,
+    /// it becomes odd to call this an `Atom`... It is our TODO to clean this up by either
+    /// getting rid of the `guard` or renaming `Atom`.
+    pub guard: Option<Box<Formula<V>>>,
     pub pred: Pred,
     pub args: Vec<Term<V>>,
 }
@@ -1008,7 +1014,12 @@ where
     D::Doc: Clone,
 {
     fn pretty(self, allocator: &'a D) -> pretty::DocBuilder<'a, D, termcolor::ColorSpec> {
-        if self.pred.is_infix() {
+        let guard = if let Some(guard) = &self.guard {
+            guard.pretty(allocator).append(allocator.text(" ⇒"))
+        } else {
+            allocator.nil()
+        };
+        let atom = if self.pred.is_infix() {
             self.args[0]
                 .pretty_atom(allocator)
                 .append(allocator.line())
@@ -1027,42 +1038,58 @@ where
             } else {
                 p.append(allocator.line()).append(inner.nest(2)).group()
             }
-        }
+        };
+        guard.append(allocator.line()).append(atom).group()
     }
 }
 
 impl<V> Atom<V> {
     pub fn new(pred: Pred, args: Vec<Term<V>>) -> Self {
-        Atom { pred, args }
+        Atom {
+            guard: None,
+            pred,
+            args,
+        }
     }
 
-    pub fn top() -> Self {
+    pub fn with_guard(guard: Formula<V>, pred: Pred, args: Vec<Term<V>>) -> Self {
         Atom {
-            pred: KnownPred::TOP.into(),
-            args: vec![],
+            guard: Some(Box::new(guard)),
+            pred,
+            args,
         }
     }
 
+    pub fn top() -> Self {
+        Atom::new(KnownPred::TOP.into(), vec![])
+    }
+
     pub fn bottom() -> Self {
-        Atom {
-            pred: KnownPred::BOTTOM.into(),
-            args: vec![],
-        }
+        Atom::new(KnownPred::BOTTOM.into(), vec![])
     }
 
     pub fn is_top(&self) -> bool {
-        self.pred.is_top()
+        if let Some(guard) = &self.guard {
+            guard.is_bottom() || self.pred.is_top()
+        } else {
+            self.pred.is_top()
+        }
     }
 
     pub fn is_bottom(&self) -> bool {
-        self.pred.is_bottom()
+        if let Some(guard) = &self.guard {
+            guard.is_top() && self.pred.is_bottom()
+        } else {
+            self.pred.is_bottom()
+        }
     }
 
     pub fn subst_var<F, W>(self, mut f: F) -> Atom<W>
     where
         F: FnMut(V) -> Term<W>,
     {
         Atom {
+            guard: self.guard.map(|fo| Box::new(fo.subst_var(&mut f))),
             pred: self.pred,
             args: self.args.into_iter().map(|t| t.subst_var(&mut f)).collect(),
         }
@@ -1073,13 +1100,37 @@ impl<V> Atom<V> {
         F: FnMut(V) -> W,
     {
         Atom {
+            guard: self.guard.map(|fo| Box::new(fo.map_var(&mut f))),
             pred: self.pred,
             args: self.args.into_iter().map(|t| t.map_var(&mut f)).collect(),
         }
     }
 
     pub fn fv(&self) -> impl Iterator<Item = &V> {
-        self.args.iter().flat_map(|t| t.fv())
+        let guard_fvs: Box<dyn Iterator<Item = &V>> = if let Some(guard) = &self.guard {
+            Box::new(guard.fv())
+        } else {
+            Box::new(std::iter::empty())
+        };
+        self.args.iter().flat_map(|t| t.fv()).chain(guard_fvs)
+    }
+
+    pub fn guarded(self, new_guard: Formula<V>) -> Atom<V> {
+        let Atom {
+            guard: self_guard,
+            pred,
+            args,
+        } = self;
+        let guard = if let Some(self_guard) = self_guard {
+            self_guard.and(new_guard)
+        } else {
+            new_guard
+        };
+        Atom {
+            guard: Some(Box::new(guard)),
+            pred,
+            args,
+        }
     }
 }
 
@@ -1456,6 +1507,22 @@ impl<V> Body<V> {
     }
 }
 
+impl<V> Body<V>
+where
+    V: Var,
+{
+    pub fn guarded(self, guard: Formula<V>) -> Body<V> {
+        let Body { atoms, formula } = self;
+        Body {
+            atoms: atoms
+                .into_iter()
+                .map(|a| a.guarded(guard.clone()))
+                .collect(),
+            formula: guard.not().or(formula),
+        }
+    }
+}
+
 impl<'a, 'b, D, V> Pretty<'a, D, termcolor::ColorSpec> for &'b Body<V>
 where
     V: Var,

src/chc/smtlib2.rs

@@ -223,6 +223,10 @@ pub struct Atom<'ctx, 'a> {
 
 impl<'ctx, 'a> std::fmt::Display for Atom<'ctx, 'a> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        if let Some(guard) = &self.inner.guard {
+            let guard = Formula::new(self.ctx, self.clause, guard);
+            write!(f, "(=> {} ", guard)?;
+        }
         if self.inner.pred.is_negative() {
             write!(f, "(not ")?;
         }
@@ -244,6 +248,9 @@ impl<'ctx, 'a> std::fmt::Display for Atom<'ctx, 'a> {
         if self.inner.pred.is_negative() {
             write!(f, ")")?;
         }
+        if self.inner.guard.is_some() {
+            write!(f, ")")?;
+        }
         Ok(())
     }
 }

src/chc/unbox.rs

@@ -24,9 +24,10 @@ fn unbox_term(term: Term) -> Term {
 }
 
 fn unbox_atom(atom: Atom) -> Atom {
-    let Atom { pred, args } = atom;
+    let Atom { guard, pred, args } = atom;
+    let guard = guard.map(|fo| Box::new(unbox_formula(*fo)));
     let args = args.into_iter().map(unbox_term).collect();
-    Atom { pred, args }
+    Atom { guard, pred, args }
 }
 
 fn unbox_datatype_sort(sort: DatatypeSort) -> DatatypeSort {

src/refine/env.rs

@@ -730,6 +730,12 @@ impl Env {
         let def = self.enum_defs[&ty.symbol].clone();
         let matcher_pred = chc::MatcherPred::new(ty.symbol.clone(), ty.arg_sorts());
 
+        let discr_var = self
+            .temp_vars
+            .push(TempVarBinding::Type(rty::RefinedType::unrefined(
+                rty::Type::int(),
+            )));
+
         let mut variants = IndexVec::new();
         for variant_def in &def.variants {
             let mut fields = IndexVec::new();
@@ -738,7 +744,12 @@ impl Env {
                 fields.push(x);
                 let mut field_ty = rty::RefinedType::unrefined(field_ty.clone().vacuous());
                 field_ty.instantiate_ty_params(ty.args.clone());
-                self.bind_impl(x.into(), field_ty.boxed(), depth);
+                let guarded_field_ty = field_ty.guarded(
+                    chc::Term::var(discr_var.into())
+                        .equal_to(chc::Term::int(variant_def.discr as i64))
+                        .into(),
+                );
+                self.bind_impl(x.into(), guarded_field_ty.boxed(), depth);
             }
             variants.push(FlowBindingVariant { fields });
         }
@@ -773,11 +784,6 @@ impl Env {
         assumption
             .body
             .push_conj(chc::Atom::new(matcher_pred.into(), pred_args));
-        let discr_var = self
-            .temp_vars
-            .push(TempVarBinding::Type(rty::RefinedType::unrefined(
-                rty::Type::int(),
-            )));
         assumption
             .body
             .push_conj(

src/rty.rs

@@ -1172,6 +1172,19 @@ impl<V> Formula<V> {
     }
 }
 
+impl<V> Formula<V>
+where
+    V: chc::Var,
+{
+    pub fn guarded(self, guard: chc::Formula<V>) -> Self {
+        let Formula { existentials, body } = self;
+        Formula {
+            existentials,
+            body: body.guarded(guard),
+        }
+    }
+}
+
 impl<V> Formula<V>
 where
     V: ShiftExistential,
@@ -1351,6 +1364,15 @@ impl<FV> RefinedType<FV> {
         RefinedType { ty, refinement }
     }
 
+    pub fn guarded(self, guard: chc::Formula<FV>) -> Self
+    where
+        FV: chc::Var,
+    {
+        let RefinedType { ty, refinement } = self;
+        let refinement = refinement.guarded(guard.map_var(RefinedTypeVar::Free));
+        RefinedType { ty, refinement }
+    }
+
     /// Returns a dereferenced type of the immutable reference or owned pointer.
     ///
     /// e.g. `{ v: Box<T> | φ }  -->  { v: T | φ[box v/v] }`

tests/ui/fail/adt_variant_without_params.rs

@@ -0,0 +1,13 @@
+//@error-in-other-file: Unsat
+
+enum X<T> {
+    None1,
+    None2,
+    Some(T),
+}
+
+fn main() {
+    let mut opt: X<i32> = X::None1;
+    opt = X::None2;
+    assert!(matches!(opt, X::None1));
+}

tests/ui/pass/adt_variant_without_params.rs

@@ -0,0 +1,13 @@
+//@check-pass
+
+enum X<T> {
+    None1,
+    None2,
+    Some(T),
+}
+
+fn main() {
+    let mut opt: X<i32> = X::None1;
+    opt = X::None2;
+    assert!(matches!(opt, X::None2));
+}