Merge pull request #439 from hacspec/fix-question-marks

Fix question marks

engine/backends/fstar/fstar_backend.ml

@@ -451,11 +451,18 @@ struct
              ( NoLetQualifier,
                [ (None, (pat, body)) ],
                F.term @@ F.AST.Seq (assertion, array) )
-    | Let { lhs; rhs; body; monadic = Some monad } ->
+    | Let { lhs; rhs; body; monadic = Some (monad, _) } ->
         let p =
           F.pat @@ F.AST.PatAscribed (ppat lhs, (pty lhs.span lhs.typ, None))
         in
-        let op = "let" ^ match monad with _ -> "*" in
+        let op =
+          "let"
+          ^
+          match monad with
+          | MResult _ -> "|"
+          | MOption -> "?"
+          | MException _ -> "!"
+        in
         F.term @@ F.AST.LetOperator ([ (F.id op, p, pexpr rhs) ], pexpr body)
     | Let { lhs; rhs; body; monadic = None } ->
         let p =

engine/lib/ast.ml

@@ -346,12 +346,10 @@ functor
       (* ControlFlow *)
       | Break of { e : expr; label : string option; witness : F.break * F.loop }
       | Return of { e : expr; witness : F.early_exit }
-      | QuestionMark of {
-          e : expr;
-          converted_typ : ty;
-              (** [converted_typ] is the converted type: when you do [e?], a convertion might be inserted by Rust on the fly (e.g. [Something::from_residual(e)]) *)
-          witness : F.question_mark;
-        }
+      | QuestionMark of { e : expr; return_typ : ty; witness : F.question_mark }
+          (** The expression `e?`. In opposition to Rust, no implicit
+      coercion is applied on the (potential) error payload of
+      `e`. Coercion should be made explicit within `e`. *)
       | Continue of {
           e : (F.state_passing_loop * expr) option;
           label : string option;

engine/lib/ast_utils.ml

@@ -757,27 +757,20 @@ module Make (F : Features.T) = struct
         (Concrete_ident.of_name (Constructor { is_struct }) constructor_name))
       is_struct args span ret_typ
 
-  let call' f (args : expr list) span ret_typ =
+  let call' ?impl f (args : expr list) span ret_typ =
     let typ = TArrow (List.map ~f:(fun arg -> arg.typ) args, ret_typ) in
     let e = GlobalVar f in
     {
-      e =
-        App
-          {
-            f = { e; typ; span };
-            args;
-            generic_args = [];
-            impl =
-              None
-              (* TODO: see issue #328, and check that for evrey call to `call'` *);
-          };
+      e = App { f = { e; typ; span }; args; generic_args = []; impl };
       typ = ret_typ;
       span;
     }
 
-  let call ?(kind : Concrete_ident.Kind.t = Value)
+  let call ?(kind : Concrete_ident.Kind.t = Value) ?impl
       (f_name : Concrete_ident.name) (args : expr list) span ret_typ =
-    call' (`Concrete (Concrete_ident.of_name kind f_name)) args span ret_typ
+    call' ?impl
+      (`Concrete (Concrete_ident.of_name kind f_name))
+      args span ret_typ
 
   let string_lit span (s : string) : expr =
     { span; typ = TStr; e = Literal (String s) }

engine/lib/phases/phase_cf_into_monads.ml

@@ -213,13 +213,7 @@ struct
       | Break _ ->
           Error.unimplemented ~issue_id:96
             ~details:"TODO: Monad for loop-related control flow" span
-      | QuestionMark { e; converted_typ; _ } ->
-          let e = dexpr e in
-          let converted_typ = dty span converted_typ in
-          if [%equal: B.ty] converted_typ e.typ then e
-          else
-            UB.call Core__ops__try_trait__FromResidual__from_residual [ e ] span
-              converted_typ
+      | QuestionMark { e; _ } -> dexpr e
       | Return { e; _ } ->
           let open KnownMonads in
           let e = dexpr e in

engine/lib/phases/phase_reconstruct_question_marks.ml

@@ -30,19 +30,97 @@ module%inlined_contents Make (FA : Features.T) = struct
 
       open A
 
-      type t = {
-        start : expr;
-        end_ : expr;
-        var : local_ident;
-        var_typ : ty;
-        body : expr;
-        state : loop_state option;
-        label : string option;
-        witness : FA.loop;
-      }
-      [@@deriving show]
+      (** The types supported for [e] in a [e?] expression *)
+      type qm_kind = QMResult of { success : ty; error : ty } | QMOption of ty
 
+      (** Interpret a type [t] as a [qm_kind] *)
+      let qm_kind_of_typ span (t : ty) : qm_kind =
+        let is_result = Global_ident.eq_name Core__result__Result in
+        let is_option = Global_ident.eq_name Core__option__Option in
+        match t with
+        | TApp { ident; args = [ GType s; GType e ] } when is_result ident ->
+            QMResult { success = s; error = e }
+        | TApp { ident; args = [ GType s ] } when is_option ident -> QMOption s
+        | _ ->
+            Error.assertion_failure span
+              ("expected something of type Option<_> or Result<_, _>, instead, \
+                got: "
+              ^ [%show: ty] t)
+
+      (** Expects [impl] to be an impl. expr. for the trait
+      `std::ops::FromResidual` for the type [Result<_, _>], and
+      extract its parent [From] impl expr *)
+      let expect_residual_impl_result (impl : impl_expr) : impl_expr option =
+        match impl with
+        | ImplApp
+            { impl = Concrete { trait; _ }; args = [ _; _; _; from_impl ] }
+          when Concrete_ident.eq_name Core__result__Impl_27 trait ->
+            Some from_impl
+        | _ -> None
+
+      (** Expects [t] to be [Result<S, E>], and returns [(S, E)] *)
+      let expect_result_type (t : ty) : (ty * ty) option =
+        match t with
+        | TApp { ident; args = [ GType s; GType e ] }
+          when Global_ident.eq_name Core__result__Result ident ->
+            Some (s, e)
+        | _ -> None
+
+      (** Construct [Result<S,E>] *)
+      let make_result_type (success : ty) (error : ty) : ty =
+        let ident = Global_ident.of_name Type Core__result__Result in
+        TApp { ident; args = [ GType success; GType error ] }
+
+      (** Retype a [Err::<_, E>(x)] literal, as [Err::<success, E>(x)] *)
+      let retype_err_literal (e : expr) (success : ty) (error : ty) =
+        match e.e with
+        | Construct { constructor; _ }
+          when Global_ident.eq_name Core__result__Result__Err constructor ->
+            { e with typ = make_result_type success error }
+        | _ -> e
+
+      (** [map_err e error_dest impl] creates the expression
+      [e.map_err(from)] with the proper types and impl
+      informations. *)
+      let map_err (e : expr) (error_dest : ty) impl : expr option =
+        let* success, error_src = expect_result_type e.typ in
+        let* impl = expect_residual_impl_result impl in
+        let from_typ = TArrow ([ error_src ], error_dest) in
+        let from =
+          UA.call ~kind:(AssociatedItem Value) ~impl Core__convert__From__from
+            [] e.span from_typ
+        in
+        let call =
+          UA.call Core__result__Impl__map_err [ e; from ] e.span
+            (make_result_type success error_dest)
+        in
+        Some call
+
+      (** [extract e] returns [Some (x, ty)] if [e] was a `y?`
+      desugared by rustc. `y` is `x` plus possibly a coercion. [ty] is
+      the return type of the function. *)
       let extract (e : expr) : (expr * ty) option =
+        let extract_return (e : expr) =
+          match e.e with
+          | Return
+              {
+                e =
+                  {
+                    e =
+                      App
+                        {
+                          f = { e = GlobalVar f };
+                          args = [ { e = LocalVar residual_var; _ } ];
+                          impl = Some impl;
+                        };
+                    typ = return_typ;
+                    _;
+                  };
+                _;
+              } ->
+              Some (f, residual_var, return_typ, impl)
+          | _ -> None
+        in
         let extract_pat_app_bd (p : pat) : (global_ident * local_ident) option =
           match p.p with
           | PConstruct
@@ -72,38 +150,7 @@ module%inlined_contents Make (FA : Features.T) = struct
                 { e = App { f = { e = GlobalVar n; _ }; args = [ expr ] }; _ };
               arms =
                 [
-                  {
-                    arm =
-                      {
-                        arm_pat = pat_break;
-                        body =
-                          {
-                            e =
-                              Return
-                                {
-                                  e =
-                                    {
-                                      e =
-                                        App
-                                          {
-                                            f = { e = GlobalVar f };
-                                            args =
-                                              [
-                                                {
-                                                  e = LocalVar residual_arg_var;
-                                                  _;
-                                                };
-                                              ];
-                                          };
-                                      typ = converted_typ;
-                                      _;
-                                    };
-                                  _;
-                                };
-                          };
-                      };
-                    _;
-                  };
+                  { arm = { arm_pat = pat_break; body }; _ };
                   {
                     arm =
                       {
@@ -115,22 +162,40 @@ module%inlined_contents Make (FA : Features.T) = struct
                 ];
             }
         (*[@ocamlformat "disable"]*)
-          when Global_ident.eq_name Core__ops__try_trait__Try__branch n -> (
-            match
-              (extract_pat_app_bd pat_break, extract_pat_app_bd pat_continue)
-            with
-            | Some (f_break, residual_arg_var'), Some (f_continue, continue_var')
-              when Global_ident.eq_name
-                     Core__ops__control_flow__ControlFlow__Break f_break
-                   && Global_ident.eq_name
-                        Core__ops__control_flow__ControlFlow__Continue
-                        f_continue
-                   && Global_ident.eq_name
-                        Core__ops__try_trait__FromResidual__from_residual f
-                   && [%equal: local_ident] residual_arg_var residual_arg_var'
-                   && [%equal: local_ident] continue_var continue_var' ->
-                Some (expr, converted_typ)
-            | _ -> None)
+          when Global_ident.eq_name Core__ops__try_trait__Try__branch n ->
+            let* body =
+              UA.Expect.concrete_app1 Rust_primitives__hax__never_to_any body
+            in
+            let* f, residual_var, fun_return_typ, residual_impl =
+              extract_return body
+            in
+            let* f_break, residual_var' = extract_pat_app_bd pat_break in
+            let* f_continue, continue_var' = extract_pat_app_bd pat_continue in
+            let*? _ = [%equal: local_ident] residual_var residual_var' in
+            let*? _ = [%equal: local_ident] continue_var continue_var' in
+            let*? _ =
+              Global_ident.eq_name Core__ops__control_flow__ControlFlow__Break
+                f_break
+              && Global_ident.eq_name
+                   Core__ops__control_flow__ControlFlow__Continue f_continue
+              && Global_ident.eq_name
+                   Core__ops__try_trait__FromResidual__from_residual f
+            in
+            let expr =
+              let kind = qm_kind_of_typ e.span in
+              match (kind expr.typ, kind fun_return_typ) with
+              | ( QMResult { error = local_err; _ },
+                  QMResult { error = return_err; _ } ) ->
+                  let expr = retype_err_literal expr e.typ local_err in
+                  map_err expr return_err residual_impl
+                  |> Option.value ~default:expr
+              | QMOption _, QMOption _ -> expr
+              | _ ->
+                  Error.assertion_failure e.span
+                    "expected expr.typ and fun_return_typ to be both Options \
+                     or both Results"
+            in
+            Some (expr, fun_return_typ)
         | _ -> None
     end
 
@@ -139,13 +204,13 @@ module%inlined_contents Make (FA : Features.T) = struct
     let rec dexpr_unwrapped (expr : A.expr) : B.expr =
       let h = [%inline_body dexpr_unwrapped] in
       match QuestionMarks.extract expr with
-      | Some (e, converted_typ) ->
+      | Some (e, return_typ) ->
           {
             e =
               QuestionMark
                 {
                   e = dexpr e;
-                  converted_typ = dty e.span converted_typ;
+                  return_typ = dty e.span return_typ;
                   witness = Features.On.question_mark;
                 };
             span = expr.span;

engine/lib/phases/phase_reconstruct_question_marks.mli

@@ -1,9 +1,33 @@
+(** In THIR, there are no construct for question marks. Instead, Rustc
+desugars `e?` into the following:
+
+{@rust[
+match core::ops::try_trait::branch(y) {
+    core::ops::control_flow::Break(residual) => {
+        never_to_any(
+            {return core::ops::try_trait::from_residual(residual)},
+        )
+    }
+    core::ops::control_flow::Continue(val) => val,
+})
+]}
+
+This phase does the opposite rewrite.
+
+While `e?` in Rust might implies an implicit coercion, in our AST, a
+question mark is expected to already be of the right type. This phase
+inlines a coercion (of the shape `x.map_err(from)`, in the case of a
+`Result`).
+*)
+
 open! Prelude
 
+(** This phase can be applied to any feature set. *)
 module Make (F : Features.T) : sig
   include module type of struct
     module FA = F
 
+    (** This phase outputs an AST with question marks. *)
     module FB = struct
       include F
       include Features.On.Question_mark

engine/lib/side_effect_utils.ml

@@ -243,17 +243,14 @@ struct
         method! visit_expr (env : CollectContext.t) e =
           match e.e with
           | LocalVar v -> (e, no_lbs (SideEffects.reads v e.typ))
-          | QuestionMark { e = e'; converted_typ; witness } ->
+          | QuestionMark { e = e'; return_typ; witness } ->
               HoistSeq.one env (super#visit_expr env e') (fun e' effects ->
                   let effects =
                     m#plus effects
                       (no_lbs
-                         { SideEffects.zero with return = Some converted_typ })
+                         { SideEffects.zero with return = Some return_typ })
                   in
-                  ( {
-                      e with
-                      e = QuestionMark { e = e'; converted_typ; witness };
-                    },
+                  ( { e with e = QuestionMark { e = e'; return_typ; witness } },
                     effects ))
           | Return { e = e'; witness } ->
               HoistSeq.one env (super#visit_expr env e') (fun e' effects ->

engine/lib/subtype.ml

@@ -217,11 +217,11 @@ struct
           }
     | Return { e; witness } ->
         Return { e = dexpr e; witness = S.early_exit span witness }
-    | QuestionMark { e; converted_typ; witness } ->
+    | QuestionMark { e; return_typ; witness } ->
         QuestionMark
           {
             e = dexpr e;
-            converted_typ = dty span converted_typ;
+            return_typ = dty span return_typ;
             witness = S.question_mark span witness;
           }
     | Continue { e; label; witness = w1, w2 } ->

engine/lib/utils.ml

@@ -23,6 +23,12 @@ let tup2 a b = (a, b)
 let ( let* ) x f = Option.bind ~f x
 let some_if_true = function true -> Some () | _ -> None
 
+(** [let*? () = guard in body] acts as a guard: if [guard] holds, then
+[body] is executed, otherwise [None] is returned. *)
+let ( let*? ) (type a) (x : bool) (f : unit -> a option) =
+  let* () = some_if_true x in
+  f ()
+
 let map_first_letter (f : string -> string) (s : string) =
   let first, rest = String.(prefix s 1, drop_prefix s 1) in
   f first ^ rest