Spec blocks

prusti-contracts/prusti-contracts/src/lib.rs

@@ -492,6 +492,14 @@ pub fn exists<T, A: core::marker::Tuple, F: Fn<A>>(_trigger_set: T, _closure: &F
     true
 }
 
+/// Specification block.
+///
+/// This is a Prusti-internal representation of `prusti_assert!` and others.
+#[cfg(feature = "prusti")]
+pub fn spec_block<A: core::marker::Tuple, F: Fn<A>>(_closure: &F) -> bool {
+    true
+}
+
 /// Creates an owned copy of a reference. This should only be used from within
 /// ghost code, as it circumvents the borrow checker.
 pub fn snap<T>(_x: &T) -> T {

prusti-contracts/prusti-specs/src/rewriter.rs

@@ -285,12 +285,14 @@ impl AstRewriter {
         let spec_id_str = spec_id.to_string();
         Ok(quote_spanned! {expr.span()=>
             {
-                #[prusti::spec_only]
-                #[prusti::#kind]
-                #[prusti::spec_id = #spec_id_str]
-                || -> bool {
-                    #expr
-                };
+                ::prusti_contracts::spec_block(&(
+                    #[prusti::spec_only]
+                    #[prusti::#kind]
+                    #[prusti::spec_id = #spec_id_str]
+                    || -> bool {
+                        #expr
+                    }
+                ));
             }
         })
     }

prusti-encoder/src/encoders/const.rs

@@ -364,7 +364,8 @@ impl TaskEncoder for ConstEnc {
                         };
                         let expr = deps.require_dep::<MirPureEnc>(task)?.expr;
                         use vir::Reify;
-                        Ok((Vec::new(), expr.reify(vcx, (uneval.def, &[])).downcast_ty()))
+                        let args = Default::default();
+                        Ok((Vec::new(), expr.reify(vcx, (uneval.def, vcx.alloc(args))).downcast_ty()))
                     } else {
                         todo!("const too generic")
                     }

prusti-encoder/src/encoders/mir_fn/method.rs

@@ -7,7 +7,7 @@ use crate::{
     encoders::{
         Impure, ImpureEncVisitor, MirLocalDefEnc, MirLocalDefEncTask, MirSpecEnc, WandEnc,
         WandEncTask,
-        mir_fn::{CallTaskDescription, RustSignature},
+        mir_fn::{CallTaskDescription, RustSignature, SpecBlocks},
         pure::spec::MirSpecEncMode,
         ty::generics::{GArgCaster, GArgsCastEnc, GArgsTy, GArgsTyEnc, GParams, GenericParamsEnc},
     },
@@ -267,6 +267,8 @@ impl TaskEncoder for MethodEnc {
                     vcx.mk_goto_stmt(&vir::CfgBlockLabelData::BasicBlock(0)),
                 ));
 
+                let spec_blocks = SpecBlocks::new(&body, fpcs_analysis.analysis().loop_analysis());
+
                 deps.check_cycle()?;
                 let mut visitor = ImpureEncVisitor {
                     vcx,
@@ -275,6 +277,7 @@ impl TaskEncoder for MethodEnc {
                     local_decls: &body.local_decls,
                     fpcs_analysis,
                     local_defs,
+                    spec_blocks,
                     body,
 
                     wands,

prusti-encoder/src/encoders/mir_fn/mod.rs

@@ -1,10 +1,12 @@
 mod method;
 mod function;
 mod signature;
+mod spec_blocks;
 
 pub use function::*;
 pub use method::*;
 pub use signature::*;
+pub use spec_blocks::*;
 
 use crate::encoders::ty::generics::{
     GArgs, GParams,

prusti-encoder/src/encoders/mir_fn/signature.rs

@@ -29,6 +29,7 @@ impl<'tcx> RustSignature<'tcx> {
     pub fn get_def_id_and_caller_substs(ty: ty::Ty<'tcx>) -> (DefId, ty::GenericArgsRef<'tcx>) {
         match ty.kind() {
             ty::TyKind::FnDef(def_id, substs) => (*def_id, substs),
+            ty::TyKind::Closure(def_id, substs) => (*def_id, substs),
             _ => todo!(),
         }
     }

prusti-encoder/src/encoders/mir_fn/spec_blocks.rs

@@ -0,0 +1,183 @@
+use pcg::r#loop::{LoopAnalysis, LoopId};
+use prusti_interface::{environment::EnvQuery, utils::{has_prusti_attr,}};
+use prusti_rustc_interface::{
+    data_structures::fx::{FxHashMap, FxHashSet},
+    middle::mir::{self, BasicBlock},
+    span::Span,
+};
+
+use crate::encoders::mir_fn::RustSignature;
+
+#[derive(Clone, Debug)]
+pub enum SpecBlockKind {
+    LoopInvariant,
+    GhostStart,
+    GhostEnd,
+    Assert,
+    Assume,
+    Refute,
+}
+
+#[derive(Debug)]
+pub struct LoopSpec {
+    has_body_invariant: bool,
+    pub loop_id: LoopId,
+    pub original_head_block: BasicBlock,
+    pub head_block: BasicBlock,
+    pub invariants: Vec<BasicBlock>,
+}
+
+#[derive(Clone, Debug)]
+pub struct SpecBlock {
+    pub attached_to: BasicBlock,
+    pub block: BasicBlock,
+    pub kind: SpecBlockKind,
+    pub span: Span,
+}
+
+#[derive(Default)]
+pub struct SpecBlocks {
+    pub specs_for: FxHashMap<BasicBlock, Vec<SpecBlock>>,
+    pub spec_blocks: FxHashSet<BasicBlock>,
+    pub loop_specs: FxHashMap<BasicBlock, LoopSpec>,
+}
+
+impl SpecBlocks {
+    /// Determine the spec-only blocks for the given MIR body. Spec-only blocks
+    /// are ones which consists of *only* a closure assignment of a closure
+    /// marked with the Prusti spec-only attribute. For each spec-only block we
+    /// determine which non-spec block it is attached to.
+    pub fn new<'enc, 'vir: 'enc>(
+        body: &'enc mir::Body<'vir>,
+        loop_analysis: &'enc LoopAnalysis,
+    ) -> Self {
+        use mir::visit::Visitor;
+        let mut visitor = SpecVisitor {
+            body,
+            specs_for: Default::default(),
+            spec_blocks: Default::default(),
+        };
+        visitor.visit_body(body);
+
+        // Associate specs and determine loop heads (at body invariants) for loops
+        let mut loop_specs: FxHashMap<LoopId, LoopSpec> = Default::default();
+
+        // For any loop that is not specified with a body invariant (determined
+        // above), we default to the loop head being at the loop head identified
+        // by the PCG, with no specs.
+        for (block, _) in body.basic_blocks.iter_enumerated() {
+            let Some(loop_id) = loop_analysis.loop_head_of(block) else { continue; };
+
+            loop_specs.insert(loop_id, LoopSpec {
+                has_body_invariant: false,
+                loop_id,
+                head_block: block,
+                original_head_block: block,
+                invariants: Vec::new(),
+            });
+        }
+
+        for (_, specified_blocks) in &visitor.specs_for {
+            for spec_block in specified_blocks {
+                // If this assertion ever fails, then consecutive spec blocks
+                // are actually consecutive blocks in the CFG. If this happens,
+                // we need to keep walking up the predecessors for each spec
+                // block until we find a non-spec block.
+                assert!(!visitor.spec_blocks.contains(&spec_block.attached_to));
+
+                let SpecBlockKind::LoopInvariant = spec_block.kind else {
+                    continue;
+                };
+                let loop_id = loop_analysis.innermost_loop(spec_block.block)
+                    .expect("malformed spec-only block: body invariant not in a loop");
+                let loop_spec = loop_specs.get_mut(&loop_id).unwrap();
+                loop_spec.has_body_invariant = true;
+                // TODO: is the iteration order of blocks well defined here?
+                //   do we always consider the first or last body invariant's
+                //   predecessor to be the loop head?
+                // The loop head (for our encoding and for querying the PCG) of
+                // the loop is the non-spec block preceding the body invariant.
+                // It's not the invariant block itself since that block is
+                // spec-only and guarded in `if false`.
+                loop_spec.head_block = spec_block.attached_to;
+                loop_spec.invariants.push(spec_block.block);
+            }
+        }
+
+        let loop_specs = loop_specs.into_iter()
+            .map(|(_loop_id, spec)| (spec.head_block, spec))
+            .collect();
+        Self {
+            specs_for: visitor.specs_for,
+            spec_blocks: visitor.spec_blocks,
+            loop_specs,
+        }
+    }
+}
+
+struct SpecVisitor<'enc, 'vir: 'enc> {
+    body: &'enc mir::Body<'vir>,
+    specs_for: FxHashMap<BasicBlock, Vec<SpecBlock>>,
+    spec_blocks: FxHashSet<BasicBlock>,
+}
+
+impl<'enc, 'vir: 'enc> mir::visit::Visitor<'vir> for SpecVisitor<'enc, 'vir> {
+    fn visit_terminator(&mut self, terminator: &mir::Terminator<'vir>, location: mir::Location) {
+        let mut spec_kind = None;
+        vir::with_vcx(|vcx| {
+            let env_query = EnvQuery::new(vcx.tcx());
+            match &terminator.kind {
+                mir::TerminatorKind::Call { func, .. } => {
+                    let func_ty = func.ty(self.body, vcx.tcx());
+                    let (def_id, arg_tys) = RustSignature::get_def_id_and_caller_substs(func_ty);
+                    if !env_query.is_function_in_crate(def_id, arg_tys, "prusti_contracts") {
+                        return;
+                    }
+
+                    let item_name = vcx.tcx().item_name(def_id);
+                    if item_name.as_str() != "spec_block" {
+                        return;
+                    }
+
+                    let (cl_def_id, _) = RustSignature::get_def_id_and_caller_substs(arg_tys[1].expect_ty());
+                    let cl_attrs = EnvQuery::new(vcx.tcx()).get_attributes(cl_def_id);
+
+                    spec_kind = Some(if has_prusti_attr(cl_attrs, "loop_body_invariant_spec") {
+                        SpecBlockKind::LoopInvariant
+                    } else if has_prusti_attr(cl_attrs, "ghost_begin") {
+                        SpecBlockKind::GhostStart
+                    } else if has_prusti_attr(cl_attrs, "ghost_end") {
+                        SpecBlockKind::GhostEnd
+                    } else if has_prusti_attr(cl_attrs, "prusti_assertion") {
+                        SpecBlockKind::Assert
+                    } else if has_prusti_attr(cl_attrs, "prusti_assumption") {
+                        SpecBlockKind::Assume
+                    } else if has_prusti_attr(cl_attrs, "prusti_refutation") {
+                        SpecBlockKind::Refute
+                    } else {
+                        unreachable!("malformed spec-only block: unknown spec kind");
+                    });
+                }
+                _ => (),
+            }
+        });
+        if let Some(kind) = spec_kind {
+            let nonspec_predecessor = get_single_predecessor(&self.body.basic_blocks.predecessors()[location.block]);
+            self.specs_for
+                .entry(nonspec_predecessor)
+                .or_default()
+                .push(SpecBlock {
+                    attached_to: nonspec_predecessor,
+                    block: location.block,
+                    kind,
+                    span: terminator.source_info.span,
+                });
+            self.spec_blocks.insert(location.block);
+        }
+    }
+}
+
+fn get_single_predecessor(predecessors: &[BasicBlock]) -> BasicBlock {
+    assert_eq!(predecessors.len(), 1, "malformed spec-only block: expected a single predecessor");
+    predecessors[0]
+}

prusti-encoder/src/encoders/mir_impure.rs

@@ -27,28 +27,25 @@ use pcg::{
 use prusti_interface::{PrustiError, specs::specifications::SpecQuery};
 use prusti_rustc_interface::{
     abi,
-    data_structures::fx::FxHashMap,
     index::Idx,
     middle::{
         mir,
         ty::{self, TyKind},
     },
     span::{Span, def_id::DefId},
 };
+use rustc_hash::FxHashMap;
 use prusti_utils::config;
 use task_encoder::{EncodeFullError, TaskEncoder, TaskEncoderDependencies};
 use vir::{CastType, CompType, LocalDeclData};
 
 use crate::encoders::{
-    self, FunctionCallEnc, MirBuiltinEnc, MirBuiltinEncTask, TyUseImpureEnc, WandEnc, WandEncTask,
-    mir_fn::{CallTaskDescription, RustSignature},
-    mir_shared::PureRvalueEnc,
-    ty::{
+    self, FunctionCallEnc, MirBuiltinEnc, MirBuiltinEncTask, MirPureEnc, MirPureEncTask, PureKind, TyUseImpureEnc, WandEnc, WandEncTask, mir_fn::{CallTaskDescription, RustSignature, SpecBlockKind, SpecBlocks}, mir_shared::PureRvalueEnc, ty::{
         RustTyDecomposition,
         generics::{GParams, GenericParamsEnc},
         use_impure::TyUseImpure,
         use_pure::{TyUsePure, TyUsePureEnc},
-    },
+    }
 };
 
 use super::WandEncOutput;
@@ -156,6 +153,7 @@ where
     pub local_decls: &'enc mir::LocalDecls<'vir>,
     pub fpcs_analysis: PcgOutput<'enc, 'vir>,
     pub local_defs: crate::encoders::MirLocalDefEncOutput<'vir>,
+    pub spec_blocks: SpecBlocks,
     pub body: &'enc mir::Body<'vir>,
 
     pub wands: WandEncOutput<'vir>,
@@ -1297,6 +1295,50 @@ impl<'vir, 'enc, E: TaskEncoder> ImpureEncVisitor<'vir, 'enc, E> {
         }
         */
 
+        if let Some(specs) = self.spec_blocks.specs_for.get(&block).cloned() {
+            for spec in specs {
+                let enc_output = self.deps.require_dep::<MirPureEnc>(MirPureEncTask {
+                    encoding_depth: 0,
+                    parent_def_id: self.def_id,
+                    param_env: self.vcx.tcx().param_env(self.def_id),
+                    substs: ty::List::identity_for_item(self.vcx.tcx(), self.def_id),
+                    kind: PureKind::SpecBlock(spec.block),
+                    caller_def_id: Some(self.def_id),
+                }).unwrap();
+                use vir::Reify;
+                let locals: FxHashMap<mir::Local, _> = enc_output.inputs.iter()
+                    .map(|local| (*local, self.local_defs.locals[*local].impure_snap))
+                    .collect();
+                let expr = enc_output.expr.reify(self.vcx, (self.def_id, self.vcx.alloc(locals))).downcast_ty();
+                let to_bool = self.deps
+                    .require_dep::<TyUsePureEnc>(RustTyDecomposition::from_prim_ty(
+                        self.vcx.tcx().types.bool,
+                    ))
+                    .unwrap()
+                    .expect_native()
+                    .snap_to_prim;
+                let span = spec.span;
+                match spec.kind {
+                    SpecBlockKind::Assert => {
+                        self.vcx.with_span(span, |vcx| {
+                            let error_msg = "assertion might not hold";
+                            vcx.handle_error("exhale.failed:assertion.false", move |_| {
+                                Some(vec![PrustiError::verification(error_msg, span.into())])
+                            });
+                            self.stmt(self.vcx.mk_exhale_stmt(to_bool(expr).downcast_ty()));
+                        });
+                    }
+                    SpecBlockKind::Assume => {
+                        self.stmt(self.vcx.mk_inhale_stmt(to_bool(expr).downcast_ty()));
+                    }
+                    SpecBlockKind::Refute => {
+                        self.stmt(self.vcx.mk_refute_stmt(to_bool(expr).downcast_ty()));
+                    }
+                    _ => unreachable!(),
+                }
+            }
+        }
+
         assert!(self.current_terminator.is_none());
         for (index, statement) in data.statements.iter().enumerate() {
             let location = mir::Location {