Support ? for Option and Result (#950)

tjhance · web-flow · commit 17528d1c94b1 · 2024-01-09T13:43:30.000-05:00
* support ? for Option

* support Result based on Chris's suggestion
diff --git a/source/pervasive/std_specs/control_flow.rs b/source/pervasive/std_specs/control_flow.rs
@@ -0,0 +1,72 @@
+use crate::prelude::*;
+use core::ops::Try;
+use core::ops::ControlFlow;
+use core::ops::FromResidual;
+use core::convert::Infallible;
+
+verus!{
+
+#[verifier(external_type_specification)]
+#[verifier::accept_recursive_types(B)]
+#[verifier::reject_recursive_types_in_ground_variants(C)]
+pub struct ExControlFlow<B, C>(ControlFlow<B, C>);
+
+#[verifier(external_type_specification)]
+#[verifier(external_body)]
+pub struct ExInfallible(Infallible);
+
+
+#[verifier::external_fn_specification]
+pub fn ex_result_branch<T, E>(result: Result<T, E>) -> (cf: ControlFlow<<Result<T, E> as Try>::Residual, <Result<T, E> as Try>::Output>)
+    ensures
+        cf === match result {
+            Ok(v) => ControlFlow::Continue(v),
+            Err(e) => ControlFlow::Break(Err(e)),
+        },
+{
+    result.branch()
+}
+
+#[verifier::external_fn_specification]
+pub fn ex_option_branch<T>(option: Option<T>) -> (cf: ControlFlow<<Option<T> as Try>::Residual, <Option<T> as Try>::Output>)
+    ensures
+        cf === match option {
+            Some(v) => ControlFlow::Continue(v),
+            None => ControlFlow::Break(None),
+        },
+{
+    option.branch()
+}
+
+#[verifier::external_fn_specification]
+pub fn ex_option_from_residual<T>(option: Option<Infallible>) -> (option2: Option<T>)
+    ensures
+        option.is_none(),
+        option2.is_none(),
+{
+    Option::from_residual(option)
+}
+
+pub spec fn spec_from<S, T>(value: T, ret: S) -> bool;
+
+#[verifier::broadcast_forall]
+#[verifier::external_body]
+pub proof fn spec_from_blanket_identity<T>(t: T, s: T)
+    ensures
+        spec_from::<T, T>(t, s) ==> t == s
+{
+}
+
+#[verifier::external_fn_specification]
+pub fn ex_result_from_residual<T, E, F: From<E>>(result: Result<Infallible, E>)
+      -> (result2: Result<T, F>)
+    ensures
+        match (result, result2) {
+            (Err(e), Err(e2)) => spec_from::<F, E>(e, e2),
+            _ => false,
+        },
+{
+    Result::from_residual(result)
+}
+
+}
diff --git a/source/pervasive/std_specs/mod.rs b/source/pervasive/std_specs/mod.rs
@@ -4,6 +4,7 @@ pub mod result;
 pub mod option;
 pub mod atomic;
 pub mod bits;
+pub mod control_flow;
 
 #[cfg(feature = "alloc")]
 pub mod vec;
diff --git a/source/rust_verify/src/config.rs b/source/rust_verify/src/config.rs
@@ -114,6 +114,7 @@ pub fn enable_default_features_and_verus_attr(
         "register_tool",
         "tuple_trait",
         "custom_inner_attributes",
+        "try_trait_v2",
     ] {
         rustc_args.push("-Z".to_string());
         rustc_args.push(format!("crate-attr=feature({})", feature));
diff --git a/source/rust_verify/src/fn_call_to_vir.rs b/source/rust_verify/src/fn_call_to_vir.rs
@@ -21,7 +21,7 @@ use air::ast_util::str_ident;
 use rustc_ast::LitKind;
 use rustc_hir::def::Res;
 use rustc_hir::{Expr, ExprKind, Node, QPath};
-use rustc_middle::ty::{GenericArgKind, TyKind};
+use rustc_middle::ty::{GenericArg, GenericArgKind, TyKind};
 use rustc_span::def_id::DefId;
 use rustc_span::source_map::Spanned;
 use rustc_span::Span;
@@ -96,9 +96,6 @@ pub(crate) fn fn_call_to_vir<'tcx>(
                 ),
             );
         }
-        Some(RustItem::TryTraitBranch) => {
-            return err_span(expr.span, "Verus does not yet support the ? operator");
-        }
         Some(RustItem::IntoIterFn) => {
             return err_span(
                 expr.span,
@@ -173,6 +170,8 @@ pub(crate) fn fn_call_to_vir<'tcx>(
     // If the resolution is statically known, we record the resolved function for the
     // to be used by lifetime_generate.
 
+    let node_substs = fix_node_substs(tcx, bctx.types, node_substs, rust_item, &args, expr);
+
     let target_kind = if tcx.trait_of_item(f).is_none() {
         vir::ast::CallTargetKind::Static
     } else {
@@ -1652,6 +1651,33 @@ fn mk_is_smaller_than<'tcx>(
     return Ok(dec_exp);
 }
 
+pub(crate) fn fix_node_substs<'tcx, 'a>(
+    tcx: rustc_middle::ty::TyCtxt<'tcx>,
+    types: &'tcx rustc_middle::ty::TypeckResults<'tcx>,
+    node_substs: &'tcx rustc_middle::ty::List<rustc_middle::ty::GenericArg<'tcx>>,
+    rust_item: Option<RustItem>,
+    args: &'a [&'tcx Expr<'tcx>],
+    expr: &'a Expr<'tcx>,
+) -> &'tcx rustc_middle::ty::List<rustc_middle::ty::GenericArg<'tcx>> {
+    match rust_item {
+        Some(RustItem::TryTraitBranch) => {
+            // I don't understand why, but in this case, node_substs is empty instead
+            // of having the type argument. Let's fix it here.
+            // `branch` has type `fn branch(self) -> ...`
+            // so we can get the Self argument from the first argument.
+            let generic_arg = GenericArg::from(types.expr_ty_adjusted(&args[0]));
+            tcx.mk_args(&[generic_arg])
+        }
+        Some(RustItem::ResidualTraitFromResidual) => {
+            // `fn from_residual(residual: R) -> Self;`
+            let generic_arg0 = GenericArg::from(types.expr_ty(expr));
+            let generic_arg1 = GenericArg::from(types.expr_ty_adjusted(&args[0]));
+            tcx.mk_args(&[generic_arg0, generic_arg1])
+        }
+        _ => node_substs,
+    }
+}
+
 fn mk_typ_args<'tcx>(
     bctx: &BodyCtxt<'tcx>,
     substs: &rustc_middle::ty::List<rustc_middle::ty::GenericArg<'tcx>>,
diff --git a/source/rust_verify/src/lifetime_generate.rs b/source/rust_verify/src/lifetime_generate.rs
@@ -409,6 +409,27 @@ fn erase_ty<'tcx>(ctxt: &Context<'tcx>, state: &mut State, ty: &Ty<'tcx>) -> Typ
         TyKind::Alias(rustc_middle::ty::AliasKind::Projection, t) => {
             // Note: even if rust_to_vir_base decides to normalize t,
             // we don't have to normalize t here, since we're generating Rust code, not VIR.
+            // However, normalizing means we might reach less stuff so it's
+            // still useful.
+
+            // Try normalization:
+            use crate::rustc_trait_selection::infer::TyCtxtInferExt;
+            use crate::rustc_trait_selection::traits::NormalizeExt;
+            if let Some(fun_id) = state.enclosing_fun_id {
+                let param_env = ctxt.tcx.param_env(fun_id);
+                let infcx = ctxt.tcx.infer_ctxt().ignoring_regions().build();
+                let cause = rustc_infer::traits::ObligationCause::dummy();
+                let at = infcx.at(&cause, param_env);
+                let resolved_ty = infcx.resolve_vars_if_possible(*ty);
+                if !rustc_middle::ty::TypeVisitableExt::has_escaping_bound_vars(&resolved_ty) {
+                    let norm = at.normalize(*ty);
+                    if norm.value != *ty {
+                        return erase_ty(ctxt, state, &norm.value);
+                    }
+                }
+            }
+
+            // If normalization isn't possible:
             let assoc_item = ctxt.tcx.associated_item(t.def_id);
             let name = state.typ_param(assoc_item.name.to_string(), None);
             let trait_def = ctxt.tcx.generics_of(t.def_id).parent;
@@ -722,10 +743,21 @@ fn erase_call<'tcx>(
 
             // Maybe resolve from trait function to a specific implementation
 
-            let mut node_substs = node_substs;
+            let node_substs = node_substs;
             let mut fn_def_id = fn_def_id.expect("call id");
 
             let param_env = ctxt.tcx.param_env(state.enclosing_fun_id.expect("enclosing_fun_id"));
+
+            let rust_item = crate::verus_items::get_rust_item(ctxt.tcx, fn_def_id);
+            let mut node_substs = crate::fn_call_to_vir::fix_node_substs(
+                ctxt.tcx,
+                ctxt.types(),
+                node_substs,
+                rust_item,
+                &args_slice.iter().collect::<Vec<_>>(),
+                expr,
+            );
+
             let normalized_substs = ctxt.tcx.normalize_erasing_regions(param_env, node_substs);
             let inst = rustc_middle::ty::Instance::resolve(
                 ctxt.tcx,
@@ -1738,6 +1770,8 @@ fn erase_fn_common<'tcx>(
             &mut typ_params,
             &mut generic_bounds,
         );
+
+        state.enclosing_fun_id = Some(id);
         let mut params: Vec<Param> = Vec::new();
         for ((input, param), param_info) in
             inputs.iter().zip(f_vir.x.params.iter()).zip(params_info.iter())
@@ -1772,6 +1806,7 @@ fn erase_fn_common<'tcx>(
         } else {
             Some((None, erase_ty(ctxt, state, &fn_sig.output().skip_binder())))
         };
+        state.enclosing_fun_id = None;
         let decl = FunDecl {
             sig_span: sig_span,
             name_span,
diff --git a/source/rust_verify/src/verus_items.rs b/source/rust_verify/src/verus_items.rs
@@ -546,6 +546,7 @@ pub(crate) enum RustItem {
     IntIntrinsic(RustIntIntrinsicItem),
     AllocGlobal,
     TryTraitBranch,
+    ResidualTraitFromResidual,
     IntoIterFn,
     Destruct,
 }
@@ -582,6 +583,9 @@ pub(crate) fn get_rust_item<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> Option<Ru
     if tcx.lang_items().branch_fn() == Some(def_id) {
         return Some(RustItem::TryTraitBranch);
     }
+    if tcx.lang_items().from_residual_fn() == Some(def_id) {
+        return Some(RustItem::ResidualTraitFromResidual);
+    }
     if tcx.lang_items().into_iter_fn() == Some(def_id) {
         return Some(RustItem::IntoIterFn);
     }
diff --git a/source/rust_verify_test/tests/std.rs b/source/rust_verify_test/tests/std.rs
@@ -195,3 +195,93 @@ test_verify_one_file! {
         }
     } => Err(err) => assert_one_fails(err)
 }
+
+test_verify_one_file! {
+    #[test] question_mark_option verus_code! {
+        use vstd::*;
+
+        fn test() -> (res: Option<u32>)
+            ensures res.is_none()
+        {
+            let x: Option<u8> = None;
+            let y = x?;
+
+            assert(false);
+            return None;
+        }
+
+        fn test2() -> (res: Option<u32>)
+        {
+            let x: Option<u8> = Some(5);
+            let y = x?;
+
+            assert(false); // FAILS
+            return None;
+        }
+
+        fn test3() -> (res: Option<u32>)
+            ensures res.is_some(),
+        {
+            let x: Option<u8> = None;
+            let y = x?; // FAILS
+
+            return Some(13);
+        }
+
+        fn test4() -> (res: Option<u32>)
+            ensures false,
+        {
+            let x: Option<u8> = Some(12);
+            let y = x?;
+
+            assert(y == 12);
+
+            loop { }
+        }
+    } => Err(err) => assert_fails(err, 2)
+}
+
+test_verify_one_file! {
+    #[test] question_mark_result verus_code! {
+        use vstd::*;
+
+        fn test() -> (res: Result<u32, bool>)
+            ensures res === Err(false),
+        {
+            let x: Result<u8, bool> = Err(false);
+            let y = x?;
+
+            assert(false);
+            return Err(true);
+        }
+
+        fn test2() -> (res: Result<u32, bool>)
+        {
+            let x: Result<u8, bool> = Ok(5);
+            let y = x?;
+
+            assert(false); // FAILS
+            return Err(false);
+        }
+
+        fn test3() -> (res: Result<u32, bool>)
+            ensures res.is_ok(),
+        {
+            let x: Result<u8, bool> = Err(false);
+            let y = x?; // FAILS
+
+            return Ok(13);
+        }
+
+        fn test4() -> (res: Result<u32, bool>)
+            ensures false,
+        {
+            let x: Result<u8, bool> = Ok(12);
+            let y = x?;
+
+            assert(y == 12);
+
+            loop { }
+        }
+    } => Err(err) => assert_fails(err, 2)
+}

Original file line number	Diff line number	Diff line change
`@@ -546,6 +546,7 @@ pub(crate) enum RustItem {`
`546`	`546`	`IntIntrinsic(RustIntIntrinsicItem),`
`547`	`547`	`AllocGlobal,`
`548`	`548`	`TryTraitBranch,`
	`549`	`+ ResidualTraitFromResidual,`
`549`	`550`	`IntoIterFn,`
`550`	`551`	`Destruct,`
`551`	`552`	`}`
`@@ -582,6 +583,9 @@ pub(crate) fn get_rust_item<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> Option<Ru`
`582`	`583`	`if tcx.lang_items().branch_fn() == Some(def_id) {`
`583`	`584`	`return Some(RustItem::TryTraitBranch);`
`584`	`585`	`}`
	`586`	`+ if tcx.lang_items().from_residual_fn() == Some(def_id) {`
	`587`	`+ return Some(RustItem::ResidualTraitFromResidual);`
	`588`	`+ }`
`585`	`589`	`if tcx.lang_items().into_iter_fn() == Some(def_id) {`
`586`	`590`	`return Some(RustItem::IntoIterFn);`
`587`	`591`	`}`