Implement more expressions in formula_fn

src/analyze/annot.rs

@@ -125,6 +125,14 @@ pub fn box_model_new_path() -> [Symbol; 3] {
     ]
 }
 
+pub fn array_model_store_path() -> [Symbol; 3] {
+    [
+        Symbol::intern("thrust"),
+        Symbol::intern("def"),
+        Symbol::intern("array_store"),
+    ]
+}
+
 pub fn exists_path() -> [Symbol; 3] {
     [
         Symbol::intern("thrust"),

src/analyze/annot_fn.rs

@@ -333,10 +333,16 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
                 }
                 rustc_hir::UnOp::Deref => {
                     let operand_ty = self.expr_ty(operand);
+                    let term = self.to_term(operand);
+                    if matches!(
+                        operand_ty.kind(),
+                        mir_ty::TyKind::Ref(_, _, mir_ty::Mutability::Not)
+                    ) {
+                        return FormulaOrTerm::Term(term.box_current());
+                    }
                     let adt = operand_ty
                         .ty_adt_def()
                         .expect("deref operand must be a model type");
-                    let term = self.to_term(operand);
                     if Some(adt.did()) == self.def_ids.mut_model() {
                         FormulaOrTerm::Term(term.mut_current())
                     } else if Some(adt.did()) == self.def_ids.box_model() {
@@ -349,6 +355,10 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
                     }
                 }
             },
+            ExprKind::AddrOf(rustc_hir::BorrowKind::Ref, rustc_hir::Mutability::Not, operand) => {
+                let operand = self.to_term(operand);
+                FormulaOrTerm::Term(operand.boxed())
+            }
             ExprKind::Lit(lit) => match lit.node {
                 rustc_ast::LitKind::Int(i, _) => {
                     let n = i64::try_from(i.get())
@@ -386,6 +396,23 @@ impl<'tcx> AnnotFnTranslator<'tcx> {
                 let term = self.to_term(expr);
                 FormulaOrTerm::Term(term.tuple_proj(index))
             }
+            ExprKind::Index(array, index, _) => {
+                let array_term = self.to_term(array);
+                let index_term = self.to_term(index);
+                FormulaOrTerm::Term(array_term.select(index_term))
+            }
+            ExprKind::MethodCall(method, receiver, args, _) => {
+                if let Some(def_id) = self.typeck.type_dependent_def_id(hir.hir_id) {
+                    if Some(def_id) == self.def_ids.array_model_store() {
+                        assert_eq!(args.len(), 2, "array_store takes exactly 2 arguments");
+                        let array_term = self.to_term(receiver);
+                        let index_term = self.to_term(&args[0]);
+                        let value_term = self.to_term(&args[1]);
+                        return FormulaOrTerm::Term(array_term.store(index_term, value_term));
+                    }
+                }
+                unimplemented!("unsupported method call in formula: {:?}", method)
+            }
             ExprKind::Call(func_expr, args) => {
                 if let ExprKind::Path(qpath) = &func_expr.kind {
                     let res = self.typeck.qpath_res(qpath, func_expr.hir_id);
@@ -434,16 +461,20 @@ 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,
-                            ));
+                        if let rustc_hir::def::DefKind::Ctor(ctor_of, _) = def_kind {
+                            let terms = args.iter().map(|e| self.to_term(e)).collect();
+                            match ctor_of {
+                                rustc_hir::def::CtorOf::Variant => {
+                                    return FormulaOrTerm::Term(self.variant_ctor_term(
+                                        def_id,
+                                        self.expr_ty(hir),
+                                        terms,
+                                    ));
+                                }
+                                rustc_hir::def::CtorOf::Struct => {
+                                    return FormulaOrTerm::Term(chc::Term::tuple(terms));
+                                }
+                            }
                         }
                     }
                 }

src/analyze/did_cache.rs

@@ -22,6 +22,7 @@ struct DefIds {
 
     mut_model_new: OnceCell<Option<DefId>>,
     box_model_new: OnceCell<Option<DefId>>,
+    array_model_store: OnceCell<Option<DefId>>,
 
     exists: OnceCell<Option<DefId>>,
 }
@@ -163,6 +164,13 @@ impl<'tcx> DefIdCache<'tcx> {
             .get_or_init(|| self.annotated_def(&crate::analyze::annot::box_model_new_path()))
     }
 
+    pub fn array_model_store(&self) -> Option<DefId> {
+        *self
+            .def_ids
+            .array_model_store
+            .get_or_init(|| self.annotated_def(&crate::analyze::annot::array_model_store_path()))
+    }
+
     pub fn exists(&self) -> Option<DefId> {
         *self
             .def_ids

std.rs

@@ -137,10 +137,35 @@ mod thrust_models {
             }
         }
 
+        impl<I, T> std::ops::Index<I> for Array<I, T> {
+            type Output = T;
+
+            #[thrust::ignored]
+            fn index(&self, _index: I) -> &Self::Output {
+                unimplemented!()
+            }
+        }
+
+        impl<I, T> Array<I, T> {
+            #[allow(dead_code)]
+            #[thrust::def::array_store]
+            #[thrust::ignored]
+            pub fn store(&self, _index: I, _value: T) -> Self {
+                unimplemented!()
+            }
+        }
+
         #[thrust::def::closure_model]
         pub struct Closure<T>(PhantomData<T>);
 
-        pub struct Vec<T>(pub Array<Int, T>, pub usize);
+        pub struct Vec<T>(pub Array<Int, T>, pub Int);
+
+        impl<T, U> PartialEq<U> for Vec<T> where U: super::Model<Ty = Self> {
+            #[thrust::ignored]
+            fn eq(&self, _other: &U) -> bool {
+                unimplemented!()
+            }
+        }
     }
 
     impl Model for model::Int {
@@ -219,6 +244,10 @@ mod thrust_models {
         type Ty = model::Box<T>;
     }
 
+    impl<I, T> Model for model::Array<I, T> {
+        type Ty = model::Array<I, T>;
+    }
+
     impl<T> Model for Vec<T> where T: Model {
         type Ty = model::Vec<<T as Model>::Ty>;
     }