Support enum construction in formula_fn

Author: coord-e

src/analyze/annot_fn.rs

@@ -8,6 +8,7 @@ use rustc_middle::ty::{self as mir_ty, TyCtxt};
 use crate::analyze::did_cache::DefIdCache;
 use crate::annot::AnnotFormula;
 use crate::chc;
+use crate::refine::TypeBuilder;
 use crate::rty;
 
 #[derive(Debug, Clone)]
@@ -127,6 +128,7 @@ pub struct AnnotFnTranslator<'tcx> {
     generic_args: mir_ty::GenericArgsRef<'tcx>,
 
     def_ids: DefIdCache<'tcx>,
+    type_builder: TypeBuilder<'tcx>,
     env: HashMap<HirId, chc::Term<rty::FunctionParamIdx>>,
 }
 
@@ -138,13 +140,15 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
         let generic_args = tcx.mk_args(&[]);
         let typeck = tcx.typeck(local_def_id);
         let def_ids = DefIdCache::new(tcx);
+        let type_builder = TypeBuilder::new(tcx, def_ids.clone(), local_def_id.to_def_id());
         let mut translator = Self {
             tcx,
             local_def_id,
             typeck,
             body,
             generic_args,
             def_ids,
+            type_builder,
             env: HashMap::default(),
         };
         translator.build_env_from_params();
@@ -224,6 +228,28 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
             .expect("expected a term")
     }
 
+    fn variant_ctor_term(
+        &self,
+        ctor_did: rustc_span::def_id::DefId,
+        result_ty: mir_ty::Ty<'tcx>,
+        field_terms: Vec<chc::Term<rty::FunctionParamIdx>>,
+    ) -> chc::Term<rty::FunctionParamIdx> {
+        let variant_did = self.tcx.parent(ctor_did);
+        let adt_did = self.tcx.parent(variant_did);
+        let d_sym = crate::refine::datatype_symbol(self.tcx, adt_did);
+        let variant_name = self.tcx.item_name(variant_did);
+        let v_sym = chc::DatatypeSymbol::new(format!("{}.{}", d_sym, variant_name));
+        let sort_args = if let mir_ty::TyKind::Adt(_, generic_args) = result_ty.kind() {
+            generic_args
+                .types()
+                .map(|ty| self.type_builder.build(ty).to_sort())
+                .collect()
+        } else {
+            panic!("expected an ADT type for variant constructor")
+        };
+        chc::Term::datatype_ctor(d_sym, sort_args, v_sym, field_terms)
+    }
+
     fn to_formula_or_term(
         &self,
         hir: &'tcx rustc_hir::Expr<'tcx>,
@@ -319,20 +345,21 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
                 rustc_ast::LitKind::Bool(b) => FormulaOrTerm::Literal(b),
                 _ => unimplemented!("unsupported literal in formula: {:?}", lit),
             },
-            ExprKind::Path(qpath) => {
-                if let rustc_hir::def::Res::Local(hir_id) =
-                    self.typeck.qpath_res(&qpath, hir.hir_id)
-                {
-                    FormulaOrTerm::Term(
-                        self.env
-                            .get(&hir_id)
-                            .expect("unbound variable in formula")
-                            .clone(),
-                    )
-                } else {
-                    unimplemented!("unsupported path in formula: {:?}", qpath);
+            ExprKind::Path(qpath) => match self.typeck.qpath_res(&qpath, hir.hir_id) {
+                rustc_hir::def::Res::Local(hir_id) => FormulaOrTerm::Term(
+                    self.env
+                        .get(&hir_id)
+                        .expect("unbound variable in formula")
+                        .clone(),
+                ),
+                rustc_hir::def::Res::Def(
+                    rustc_hir::def::DefKind::Ctor(rustc_hir::def::CtorOf::Variant, _),
+                    ctor_did,
+                ) => {
+                    FormulaOrTerm::Term(self.variant_ctor_term(ctor_did, self.expr_ty(hir), vec![]))
                 }
-            }
+                _ => unimplemented!("unsupported path in formula: {:?}", qpath),
+            },
             ExprKind::Tup(exprs) => {
                 let terms = exprs.iter().map(|e| self.to_term(e)).collect();
                 FormulaOrTerm::Term(chc::Term::tuple(terms))
@@ -349,7 +376,7 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
             ExprKind::Call(func_expr, args) => {
                 if let ExprKind::Path(qpath) = &func_expr.kind {
                     let res = self.typeck.qpath_res(qpath, func_expr.hir_id);
-                    if let rustc_hir::def::Res::Def(_, def_id) = res {
+                    if let rustc_hir::def::Res::Def(def_kind, def_id) = res {
                         if Some(def_id) == self.def_ids.mut_model_new() {
                             assert_eq!(args.len(), 2, "Mut::new takes exactly 2 arguments");
                             let t1 = self.to_term(&args[0]);
@@ -361,6 +388,17 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
                             let t = self.to_term(&args[0]);
                             return FormulaOrTerm::Term(chc::Term::box_(t));
                         }
+                        if matches!(
+                            def_kind,
+                            rustc_hir::def::DefKind::Ctor(rustc_hir::def::CtorOf::Variant, _)
+                        ) {
+                            let field_terms = args.iter().map(|arg| self.to_term(arg)).collect();
+                            return FormulaOrTerm::Term(self.variant_ctor_term(
+                                def_id,
+                                self.expr_ty(hir),
+                                field_terms,
+                            ));
+                        }
                     }
                 }
                 unimplemented!("unsupported call in formula: {:?}", func_expr)

tests/ui/fail/annot_enum_simple_formula_fn.rs

@@ -0,0 +1,40 @@
+//@error-in-other-file: Unsat
+
+#[derive(PartialEq)]
+pub enum X {
+    A(i64),
+    B(bool),
+}
+
+impl thrust_models::Model for X {
+    type Ty = X;
+}
+
+#[thrust::formula_fn]
+fn _thrust_requires_test(x: X) -> bool {
+    x == X::A(1)
+}
+
+#[thrust::formula_fn]
+fn _thrust_ensures_test(_result: (), _x: X) -> bool {
+    true
+}
+
+#[allow(path_statements)]
+fn test(x: X) {
+    #[thrust::requires_path]
+    _thrust_requires_test;
+
+    #[thrust::ensures_path]
+    _thrust_ensures_test;
+
+    if let X::A(i) = x {
+        assert!(i == 2);
+    } else {
+        loop {}
+    }
+}
+
+fn main() {
+    test(X::A(1));
+}

tests/ui/pass/annot_enum_simple_formula_fn.rs

@@ -0,0 +1,40 @@
+//@check-pass
+
+#[derive(PartialEq)]
+pub enum X {
+    A(i64),
+    B(bool),
+}
+
+impl thrust_models::Model for X {
+    type Ty = X;
+}
+
+#[thrust::formula_fn]
+fn _thrust_requires_test(x: X) -> bool {
+    x == X::A(1)
+}
+
+#[thrust::formula_fn]
+fn _thrust_ensures_test(_result: (), _x: X) -> bool {
+    true
+}
+
+#[allow(path_statements)]
+fn test(x: X) {
+    #[thrust::requires_path]
+    _thrust_requires_test;
+
+    #[thrust::ensures_path]
+    _thrust_ensures_test;
+
+    if let X::A(i) = x {
+        assert!(i == 1);
+    } else {
+        loop {}
+    }
+}
+
+fn main() {
+    test(X::A(1));
+}