diff --git a/prusti-contracts/prusti-contracts-proc-macros/src/lib.rs b/prusti-contracts/prusti-contracts-proc-macros/src/lib.rs index 9a50a2f3f8c..2ceb1ed9b2a 100644 --- a/prusti-contracts/prusti-contracts-proc-macros/src/lib.rs +++ b/prusti-contracts/prusti-contracts-proc-macros/src/lib.rs @@ -58,6 +58,12 @@ pub fn body_invariant(_tokens: TokenStream) -> TokenStream { TokenStream::new() } +#[cfg(not(feature = "prusti"))] +#[proc_macro] +pub fn loop_invariant(_tokens: TokenStream) -> TokenStream { + TokenStream::new() +} + #[cfg(not(feature = "prusti"))] #[proc_macro] pub fn prusti_assert(_tokens: TokenStream) -> TokenStream { @@ -189,6 +195,12 @@ pub fn body_invariant(tokens: TokenStream) -> TokenStream { prusti_specs::body_invariant(tokens.into()).into() } +#[cfg(feature = "prusti")] +#[proc_macro] +pub fn loop_invariant(tokens: TokenStream) -> TokenStream { + prusti_specs::loop_invariant_expr(tokens.into()).into() +} + #[cfg(feature = "prusti")] #[proc_macro] pub fn prusti_assert(tokens: TokenStream) -> TokenStream { diff --git a/prusti-contracts/prusti-contracts/src/lib.rs b/prusti-contracts/prusti-contracts/src/lib.rs index 2bd7ad3edc8..78f81fc42d0 100644 --- a/prusti-contracts/prusti-contracts/src/lib.rs +++ b/prusti-contracts/prusti-contracts/src/lib.rs @@ -28,6 +28,9 @@ pub use prusti_contracts_proc_macros::invariant; /// A macro for writing a loop body invariant. pub use prusti_contracts_proc_macros::body_invariant; +/// A macro to annotate loop invariants (Holds at start and end of the loop) +pub use prusti_contracts_proc_macros::loop_invariant; + /// A macro for writing assertions using the full prusti specifications pub use prusti_contracts_proc_macros::prusti_assert; diff --git a/prusti-contracts/prusti-specs/src/lib.rs b/prusti-contracts/prusti-specs/src/lib.rs index a16b21cd4dc..3bf905f125c 100644 --- a/prusti-contracts/prusti-specs/src/lib.rs +++ b/prusti-contracts/prusti-specs/src/lib.rs @@ -420,6 +420,10 @@ pub fn body_variant(tokens: TokenStream) -> TokenStream { } pub fn body_invariant(tokens: TokenStream) -> TokenStream { + generate_expression_closure(&AstRewriter::process_body_invariant, tokens) +} + +pub fn loop_invariant_expr(tokens: TokenStream) -> TokenStream { generate_expression_closure(&AstRewriter::process_loop_invariant, tokens) } diff --git a/prusti-contracts/prusti-specs/src/rewriter.rs b/prusti-contracts/prusti-specs/src/rewriter.rs index df115c5e6e6..ede3c481b59 100644 --- a/prusti-contracts/prusti-specs/src/rewriter.rs +++ b/prusti-contracts/prusti-specs/src/rewriter.rs @@ -240,7 +240,7 @@ impl AstRewriter { } /// Parse a loop invariant into a Rust expression - pub fn process_loop_invariant( + pub fn process_body_invariant( &mut self, spec_id: SpecificationId, tokens: TokenStream, @@ -248,6 +248,15 @@ impl AstRewriter { self.process_prusti_expression(quote! {loop_body_invariant_spec}, spec_id, tokens) } + /// Parse a proper loop invariant into a Rust expression + pub fn process_loop_invariant( + &mut self, + spec_id: SpecificationId, + tokens: TokenStream, + ) -> syn::Result { + self.process_prusti_expression(quote! {loop_invariant_spec}, spec_id, tokens) + } + /// Parse a prusti assertion into a Rust expression pub fn process_prusti_assertion( &mut self, diff --git a/prusti-encoder/src/encoders/impure/loop.rs b/prusti-encoder/src/encoders/impure/loop.rs index 70ca55b3961..73a657b9392 100644 --- a/prusti-encoder/src/encoders/impure/loop.rs +++ b/prusti-encoder/src/encoders/impure/loop.rs @@ -7,32 +7,89 @@ use pcg::{ r#loop::LoopId, utils::{maybe_old::MaybeOldPlace, maybe_remote::MaybeRemotePlace, Place, SnapshotLocation}, }; -use prusti_rustc_interface::middle::mir; +use prusti_rustc_interface::{ + middle::{ + mir::{self, visit::Visitor}, + ty::{self, TyKind}, + }, + span::def_id::DefId, +}; +use std::collections::HashSet; use task_encoder::TaskEncoder; use vir::{CastType, Reify}; -use crate::encoders::{ - indirect::{IndirectKey, IndirectPredicatesEnc}, - rust_ty_predicates::{RustTyPredicatesEnc, RustTyPredicatesEncOutputRef}, - ImpureEncVisitor, +use crate::{ + encoders::{ + indirect::{IndirectKey, IndirectPredicatesEnc}, + lifted::rust_ty_cast::RustTyCastersEnc, + mir_pure::{MirPureEnc, MirPureEncTask, PureKind}, + rust_ty_predicates::{RustTyPredicatesEnc, RustTyPredicatesEncOutputRef}, + rust_ty_snapshots::RustTySnapshotsEnc, + spec, ImpureEncVisitor, + }, + CastTypePure, }; +type ExprInput<'vir> = (DefId, &'vir [vir::ExprSnap<'vir>]); +type ExprRet<'vir> = vir::ExprGenBool<'vir, ExprInput<'vir>, vir::ExprKind<'vir>>; + pub(super) enum WandOldOuter<'vir> { LetBind(Vec<(&'vir str, vir::ExprSnap<'vir>)>), Label(Option<&'vir str>), } +struct CollectedLocals { + locals: HashSet, +} + +impl CollectedLocals { + fn new() -> Self { + Self { + locals: HashSet::new(), + } + } +} + +impl<'tcx> Visitor<'tcx> for CollectedLocals { + fn visit_local( + &mut self, + local: mir::Local, + _context: mir::visit::PlaceContext, + _location: mir::Location, + ) { + self.locals.insert(local); + } +} + impl<'vir, 'enc, E: TaskEncoder> ImpureEncVisitor<'vir, 'enc, E> { + fn collect_used_locals_in_loop( + &self, + loop_id: LoopId, + ) -> std::collections::HashSet { + let mut visitor = CollectedLocals::new(); + + for (block_idx, block_data) in self.body.basic_blocks.iter_enumerated() { + if !self.loop_analysis.in_loop(block_idx, loop_id) { + continue; + } + + visitor.visit_basic_block_data(block_idx, block_data); + } + + visitor.locals + } + /// Calculate invariant at loop head pub(crate) fn get_loop_inv( &mut self, - _lh: LoopId, + lh: LoopId, cfpcs: &PcgBasicBlock<'vir>, ) -> &'vir [vir::ExprBool<'vir>] { let mut inv = Vec::new(); let start = &cfpcs.statements[0]; let state = &start.states[EvalStmtPhase::PreOperands]; + let used_locals = self.collect_used_locals_in_loop(lh); // let borrows = &*start.borrows[EvalStmtPhase::PreOperands]; // self.stmt(self.vcx.mk_comment_stmt( // vir::vir_format!(self.vcx, "_borrows: {:#?}", borrows), @@ -46,6 +103,9 @@ impl<'vir, 'enc, E: TaskEncoder> ImpureEncVisitor<'vir, 'enc, E> { if !state.capabilities().is_exclusive(*place) { continue; } + if !used_locals.contains(&place.local) { + continue; + } let (place_res, snap, _, _) = self.encode_place_snap(*place); let ty = (*place).ty(self.pcg_ctxt()); let ty_out = self.deps.require_ref::(ty.ty).unwrap(); @@ -91,6 +151,12 @@ impl<'vir, 'enc, E: TaskEncoder> ImpureEncVisitor<'vir, 'enc, E> { } inv.push(wand); } + + let loop_invariants_map = self.build_loop_invariants_map(); + if let Some(loop_invariants) = loop_invariants_map.get(&lh) { + inv.extend(loop_invariants.iter().cloned()); + } + self.vcx.alloc_slice(&inv) } @@ -231,4 +297,225 @@ impl<'vir, 'enc, E: TaskEncoder> ImpureEncVisitor<'vir, 'enc, E> { } } } + + fn build_loop_invariants_map( + &mut self, + ) -> std::collections::HashMap>> { + let mut loop_invariants_map = std::collections::HashMap::new(); + + for (block_idx, block_data) in self.body.basic_blocks.iter_enumerated() { + for stmt in &block_data.statements { + if let mir::StatementKind::Assign(box (_, rvalue)) = &stmt.kind { + if let mir::Rvalue::Aggregate( + box mir::AggregateKind::Closure(cl_def_id, cl_args), + ref upvar_operands, + ) = rvalue + { + let is_loop_invariant = spec::with_type_spec(|def_spec| { + if let Some(loop_spec) = def_spec.get_loop_spec(cl_def_id) { + assert!(!matches!(loop_spec, prusti_interface::specs::typed::LoopSpecification::BodyInvariant(_)), "body_invariant! currently not supported"); + matches!(loop_spec, prusti_interface::specs::typed::LoopSpecification::LoopInvariant(_)) + } else { + false + } + }); + + if is_loop_invariant { + if let Some(innermost_loop_id) = + self.loop_analysis.innermost_loop(block_idx) + { + let invariant_expr = self.encode_loop_invariant_closure( + *cl_def_id, + *cl_args, + &upvar_operands.raw, + ); + let concrete_expr = invariant_expr; + + loop_invariants_map + .entry(innermost_loop_id) + .or_insert_with(Vec::new) + .push(concrete_expr); + } + } + } + } + } + } + + loop_invariants_map + } + + fn encode_loop_invariant_closure( + &mut self, + cl_def_id: DefId, + _cl_args: ty::GenericArgsRef<'vir>, + upvar_operands: &[mir::Operand<'vir>], + ) -> vir::ExprBool<'vir> { + let tcx = self.vcx.tcx(); + let closure_ty = tcx.type_of(cl_def_id).instantiate_identity(); + + let (qvar_tys, upvar_rust_tys_from_closure_sig) = match closure_ty.kind() { + TyKind::Closure(_, gen_args) => ( + match gen_args.as_closure().sig().skip_binder().inputs()[0].kind() { + TyKind::Tuple(list) => list, + _ => unreachable!("Invariant closure signature malformed: qvars not a tuple"), + }, + gen_args.as_closure().upvar_tys().iter().collect::>(), + ), + _ => panic!("Illegal loop invariant closure type: {:?}", closure_ty), + }; + + let qvars = self.vcx.alloc_slice( + &qvar_tys + .iter() + .enumerate() + .map(|(idx, qvar_ty)| { + let ty_out = self + .deps + .require_ref::(qvar_ty) + .unwrap(); + self.vcx.mk_local_decl( + vir::vir_format!(self.vcx, "qvar_{idx}"), + ty_out.generic_snapshot.snapshot, + ) + }) + .collect::>(), + ); + let mut ref_to_original_place_map: std::collections::HashMap< + mir::Place<'vir>, + mir::Place<'vir>, + > = std::collections::HashMap::new(); + for (_block_idx, block_data) in self.body.basic_blocks.iter_enumerated() { + for stmt in &block_data.statements { + if let mir::StatementKind::Assign(box (place, rvalue)) = &stmt.kind { + if let mir::Rvalue::Ref(_, _, original_place) = rvalue { + if let Some(existing_place) = + ref_to_original_place_map.insert(*place, *original_place) + { + panic!("Collision in ref_to_original_place_map: place {:?} already mapped to {:?}, trying to map to {:?}", + place, existing_place, original_place); + } + } + } + } + } + + let mut fields_for_closure_struct = Vec::new(); + + for (idx, upvar_operand) in upvar_operands.iter().enumerate() { + let upvar_mir_place = match upvar_operand { + mir::Operand::Move(p) | mir::Operand::Copy(p) => *p, + mir::Operand::Constant(_) => { + panic!("Constant upvars in loop invariant closure not yet handled") + } + }; + + let original_mir_place = + ref_to_original_place_map + .get(&upvar_mir_place) + .expect(&format!( + "Could not find original place for upvar {:?}", + upvar_mir_place + )); + let original_place_viper_ref = self.encode_place((*original_mir_place).into()).expr; + let original_place_rust_ty = original_mir_place.ty(self.body, self.vcx.tcx()).ty; + let (_, direct_snap_of_original_place, _, _) = + self.encode_place_snap((*original_mir_place).into()); + let val_caster = self + .deps + .require_local::>(original_place_rust_ty) + .unwrap(); + let param_for_snap_original = + val_caster.cast_to_generic_if_necessary(self.vcx, direct_snap_of_original_place); + let upvar_ref_rust_ty = upvar_rust_tys_from_closure_sig[idx]; + let s_ref_imm_enc_for_field = self + .deps + .require_local::(upvar_ref_rust_ty) + .unwrap(); + let field_s_ref_immutable_expr = (s_ref_imm_enc_for_field + .generic_snapshot + .specifics + .expect_immref() + .prim_to_snap)( + original_place_viper_ref, param_for_snap_original + ); + + fields_for_closure_struct.push(field_s_ref_immutable_expr); + } + + let closure_struct_snapshots_enc = self + .deps + .require_local::(closure_ty) + .unwrap(); + let closure_struct_val_expr = vir::with_vcx(|vcx| { + (closure_struct_snapshots_enc + .generic_snapshot + .specifics + .expect_structlike() + .field_snaps_to_snap)(vcx.alloc_slice(&fields_for_closure_struct.iter().map(|f| f.upcast_ty()).collect::>())) + }); + + let closure_caster = self + .deps + .require_local::>(closure_ty) + .unwrap(); + let closure_struct_as_param_expr = + closure_caster.cast_to_generic_if_necessary(self.vcx, closure_struct_val_expr.upcast_ty()); + let outer_ref_to_closure_rust_ty = tcx.mk_ty_from_kind(ty::TyKind::Ref( + tcx.lifetimes.re_erased, + closure_ty, + ty::Mutability::Not, + )); + let outer_s_ref_imm_enc = self + .deps + .require_local::(outer_ref_to_closure_rust_ty) + .unwrap(); + + let final_reify_arg0 = (outer_s_ref_imm_enc + .generic_snapshot + .specifics + .expect_immref() + .prim_to_snap)(self.vcx.mk_null(), closure_struct_as_param_expr); + + let final_reify_arg0_generic = final_reify_arg0.upcast_ty(); + let mut reify_args = vec![final_reify_arg0_generic]; + reify_args.extend( + qvars + .iter() + .map(|qvar| self.vcx.mk_local_ex(qvar.name, qvar.ty)), + ); + + let body = self + .deps + .require_local::(MirPureEncTask { + encoding_depth: 1, + kind: PureKind::Closure, + parent_def_id: cl_def_id, + param_env: tcx.param_env(cl_def_id), + substs: ty::List::identity_for_item(self.vcx.tcx(), cl_def_id), + caller_def_id: Some(self.def_id), + }) + .unwrap() + .expr; + + let reified_body = body + .reify(self.vcx, (cl_def_id, self.vcx.alloc_slice(&reify_args))) + .downcast_ty(); + + let bool_snapshots_enc = self + .deps + .require_local::(tcx.types.bool) + .unwrap(); + let bool_primitive_enc = bool_snapshots_enc + .generic_snapshot + .specifics + .expect_primitive(); + + self.vcx.mk_forall_expr( + qvars, + &[], + (bool_primitive_enc + .snap_to_prim)(reified_body).downcast_ty(), + ) + } } diff --git a/prusti-encoder/src/encoders/mir_impure.rs b/prusti-encoder/src/encoders/mir_impure.rs index d7689fee2de..67f7c356721 100644 --- a/prusti-encoder/src/encoders/mir_impure.rs +++ b/prusti-encoder/src/encoders/mir_impure.rs @@ -1223,6 +1223,14 @@ impl<'vir, 'enc, E: TaskEncoder> mir::visit::Visitor<'vir> for ImpureEncVisitor< //mir::Rvalue::ShallowInitBox(Operand<'vir>, Ty<'vir>) => {} //mir::Rvalue::CopyForDeref(Place<'vir>) => {} other => { + if let ty::TyKind::Closure(def_id, _) = rvalue_ty.kind() { + let has_loop_spec = crate::encoders::spec::with_type_spec(|def_spec| { + def_spec.get_loop_spec(def_id).is_some() + }); + if has_loop_spec { + return; + } + } let e_rvalue_ty = self.deps.require_ref::(rvalue_ty).unwrap(); tracing::error!("unsupported rvalue {other:?}"); self.vcx.mk_todo_expr(vir::vir_format!(self.vcx, "rvalue {rvalue:?}"), e_rvalue_ty.snapshot()) diff --git a/prusti-interface/src/specs/mod.rs b/prusti-interface/src/specs/mod.rs index 4e9baf0b948..96fcbb04914 100644 --- a/prusti-interface/src/specs/mod.rs +++ b/prusti-interface/src/specs/mod.rs @@ -81,6 +81,7 @@ pub struct SpecCollector<'a, 'tcx> { /// Map from functions/loops/types to their specifications. procedure_specs: FxHashMap, loop_specs: Vec, + loop_invariants: Vec, loop_variants: Vec, type_specs: FxHashMap, prusti_assertions: Vec, @@ -98,6 +99,7 @@ impl<'a, 'tcx> SpecCollector<'a, 'tcx> { spec_functions: FxHashMap::default(), procedure_specs: FxHashMap::default(), loop_specs: vec![], + loop_invariants: vec![], loop_variants: vec![], type_specs: FxHashMap::default(), prusti_assertions: vec![], @@ -220,7 +222,13 @@ impl<'a, 'tcx> SpecCollector<'a, 'tcx> { for local_id in self.loop_specs.iter() { def_spec.loop_specs.insert( local_id.to_def_id(), - typed::LoopSpecification::Invariant(*local_id), + typed::LoopSpecification::BodyInvariant(*local_id), + ); + } + for local_id in self.loop_invariants.iter() { + def_spec.loop_specs.insert( + local_id.to_def_id(), + typed::LoopSpecification::LoopInvariant(*local_id), ); } for local_id in self.loop_variants.iter() { @@ -487,6 +495,9 @@ impl<'a, 'tcx> intravisit::Visitor<'tcx> for SpecCollector<'a, 'tcx> { if has_prusti_attr(attrs, "loop_body_invariant_spec") { self.loop_specs.push(local_id); } + if has_prusti_attr(attrs, "loop_invariant_spec") { + self.loop_invariants.push(local_id); + } if has_prusti_attr(attrs, "loop_body_variant_spec") { self.loop_variants.push(local_id); } diff --git a/prusti-interface/src/specs/typed.rs b/prusti-interface/src/specs/typed.rs index b2cca2e9025..16ba9d643a4 100644 --- a/prusti-interface/src/specs/typed.rs +++ b/prusti-interface/src/specs/typed.rs @@ -264,7 +264,8 @@ impl ProcedureSpecificationKind { #[derive(Debug, Clone)] pub enum LoopSpecification { - Invariant(LocalDefId), + BodyInvariant(LocalDefId), // body_invariant!() - true at call site in loop body + LoopInvariant(LocalDefId), // invariant!() - true at loop head/start and end of iteration Variant(LocalDefId), } diff --git a/prusti-tests/tests/v2/pass/loop-invariants/nested.rs b/prusti-tests/tests/v2/pass/loop-invariants/nested.rs new file mode 100644 index 00000000000..d4e6b3aa606 --- /dev/null +++ b/prusti-tests/tests/v2/pass/loop-invariants/nested.rs @@ -0,0 +1,19 @@ +fn test() { + let mut idx: usize = 0; + let length: usize = 10;; + while idx < length { + loop_invariant!(idx <= length); + loop_invariant!(0 <= length); + idx += 1; + let mut jdx: usize = 0; + let inner_length: usize = 5; + while jdx < inner_length { + loop_invariant!(jdx <= inner_length); + jdx += 1; + } + } +} + +fn main() { + test(); +} \ No newline at end of file diff --git a/prusti-tests/tests/v2/pass/loop-invariants/simple.rs b/prusti-tests/tests/v2/pass/loop-invariants/simple.rs new file mode 100644 index 00000000000..19dd34b31b5 --- /dev/null +++ b/prusti-tests/tests/v2/pass/loop-invariants/simple.rs @@ -0,0 +1,12 @@ +fn test() { + let mut idx: usize = 0; + let length: usize = 10;; + while idx < length { + loop_invariant!(idx <= length); + idx += 1; + } +} + +fn main() { + test(); +} \ No newline at end of file