Support for monomorphization over regular values like trait methods

lib/Ast.ml

@@ -197,7 +197,13 @@ type expr' =
   | EIgnore of expr
 
   | EApp of expr * expr list
-  | ETApp of expr * expr list * typ_wo list
+  | ETApp of expr * expr list * expr list * typ_wo list
+    (** The arguments are:
+      - the head of the application
+      - the const generic args (TODO: switch those to cg from the get-go!
+        that'll save the need for those silly "diff" computations)
+      - additional arguments to monomorphize over NOT IN SCOPE in types
+      - type arguments *)
   | EPolyComp of poly_comp * typ_wo
   | ELet of binder * expr * expr
   | EFun of binder list * expr * typ_wo
@@ -522,12 +528,13 @@ class ['self] map = object (self: 'self)
     let e = self#visit_expr_w env e in
     DFunction (cc, flags, n_cg, n, t, lid, bs, e)
 
-  method! visit_ETApp env e es ts =
+  method! visit_ETApp env e es es' ts =
     let ts = List.map (self#visit_typ_wo env) ts in
     let es = List.map (self#visit_expr env) es in
+    let es' = List.map (self#visit_expr env) es' in
     let env = self#extend_tmany (fst env) (List.length ts) in
     let e = self#visit_expr_w env e in
-    ETApp (e, es, ts)
+    ETApp (e, es, es', ts)
 end
 
 class ['self] iter = object (self: 'self)

lib/AstToCFlat.ml

@@ -723,14 +723,14 @@ and mk_expr (env: env) (locals: locals) (e: expr): locals * CF.expr =
   | EOpen _ ->
       invalid_arg "mk_expr (EOpen)"
 
-  | EApp ({ node = ETApp ({ node = EQualified (["LowStar"; "Ignore"],"ignore"); _ }, cgs, _); _ }, [ e ]) ->
+  | EApp ({ node = ETApp ({ node = EQualified (["LowStar"; "Ignore"],"ignore"); _ }, cgs, _, _); _ }, [ e ]) ->
       assert (cgs = []);
       let locals, e = mk_expr env locals e in
       (* This is slightly ill-typed since everywhere else the result of
          intermediary sequence bits is units, but that's fine *)
       locals, CF.Sequence [ e; cflat_unit ]
 
-  | EApp ({ node = ETApp ({ node = EQualified (["Lib"; "Memzero0"],"memzero"); _ }, cgs, _); _ }, [ dst; len ]) ->
+  | EApp ({ node = ETApp ({ node = EQualified (["Lib"; "Memzero0"],"memzero"); _ }, cgs, _, _); _ }, [ dst; len ]) ->
       assert (cgs = []);
       (* TODO: now that the C backend is generic for type applications, do the
          same here and have generic support for ETApp. Idea: reuse the JSON

lib/AstToCStar.ml

@@ -358,15 +358,15 @@ and mk_expr env in_stmt e =
   | EConstant c ->
       CStar.Constant c
 
-  | EApp ({ node = ETApp (e0, cgs, ts); _ }, es) when !Options.allow_tapps || whitelisted_tapp e0 ->
+  | EApp ({ node = ETApp (e0, cgs, cgs', ts); _ }, es) when !Options.allow_tapps || whitelisted_tapp e0 ->
       (* Return type is oftentimes very useful when having to build a return value using e.g. a
          compound literal *)
       let ret_t = CStar.Type (mk_type env (MonomorphizationState.resolve e.typ)) in
-      unit_to_void env e0 (cgs @ es) (List.map (fun t -> CStar.Type (mk_type env t)) ts @ [ ret_t ])
+      unit_to_void env e0 (cgs @ cgs' @ es) (List.map (fun t -> CStar.Type (mk_type env t)) ts @ [ ret_t ])
 
-  | ETApp (e0, cgs, ts) when !Options.allow_tapps || whitelisted_tapp e0 ->
+  | ETApp (e0, cgs, cgs', ts) when !Options.allow_tapps || whitelisted_tapp e0 ->
       let ret_t = CStar.Type (mk_type env (MonomorphizationState.resolve e.typ)) in
-      unit_to_void env e0 cgs (List.map (fun t -> CStar.Type (mk_type env t)) ts @ [ ret_t ])
+      unit_to_void env e0 (cgs @ cgs') (List.map (fun t -> CStar.Type (mk_type env t)) ts @ [ ret_t ])
 
   | EApp ({ node = EOp (op, w); _ }, [ _; _ ]) when is_arith op w ->
       fst (mk_arith env e)

lib/AstToMiniRust.ml

@@ -718,8 +718,8 @@ and translate_expr_with_type (env: env) (e: Ast.expr) (t_ret: MiniRust.typ): env
       let env, e2 = translate_expr env e2 in
       env, Assign (Index (e1, MiniRust.zero_usize), e2)
 
-  | EApp ({ node = ETApp (e, cgs, ts); _ }, es) ->
-      assert (cgs = []);
+  | EApp ({ node = ETApp (e, cgs, cgs', ts); _ }, es) ->
+      assert (cgs @ cgs' = []);
       let es =
         match es with
         | [ { typ = TUnit; node } ] -> assert (node = EUnit); []
@@ -747,7 +747,7 @@ and translate_expr_with_type (env: env) (e: Ast.expr) (t_ret: MiniRust.typ): env
       let env, es = translate_expr_list env es in
       env, possibly_convert (Call (e0, [], es)) (translate_type env e.typ)
 
-  | ETApp (_, _, _) ->
+  | ETApp (_, _, _, _) ->
       failwith "TODO: ETApp"
 
   | EPolyComp (op, _t) ->

lib/Checker.ml

@@ -180,6 +180,7 @@ and check_program env r (name, decls) =
           -dast for further debugging.\n\n"
           plid (lid_of_decl d) plid (lid_of_decl d);
         Warn.maybe_fatal_error e;
+        KPrint.beprintf "--------------------------------------------------------------------------------\n\n";
         flush stderr;
         None
   ) decls in
@@ -279,7 +280,7 @@ and check env t e =
 and check' env t e =
   let c t' = check_subtype env t' t in
   match e.node with
-  | ETApp (e0, _, _) ->
+  | ETApp (e0, _, _, _) ->
       begin match e0.node with
       | EOp ((K.Eq | K.Neq), _) ->
           (* This is probably old code that predates proper treatment of polymorphic equalities.
@@ -566,7 +567,7 @@ and infer' env e =
   in
 
   match e.node with
-  | ETApp (e0, cs, ts) ->
+  | ETApp (e0, cs, cs', ts) ->
       begin match e0.node with
       | EOp ((K.Eq | K.Neq), _) ->
           (* Special incorrect encoding of polymorphic equalities *)
@@ -577,14 +578,14 @@ and infer' env e =
           let t = infer env e0 in
           if Options.debug "checker-cg" then
             KPrint.bprintf "infer-cg: t=%a, cs=%a, ts=%a, diff=%d\n" ptyp t pexprs cs ptyps ts diff;
-          let t = DeBruijn.subst_tn ts t in
-          if Options.debug "checker-cg" then
-            KPrint.bprintf "infer-cg: subst_tn --> %a\n" ptyp t;
           let t = DeBruijn.subst_ctn diff cs t in
           if Options.debug "checker-cg" then
             KPrint.bprintf "infer-cg: subst_ctn (diff=%d)--> %a\n" diff ptyp t;
+          let t = DeBruijn.subst_tn ts t in
+          if Options.debug "checker-cg" then
+            KPrint.bprintf "infer-cg: subst_tn --> %a\n" ptyp t;
           (* Now type-check the application itself, after substitution *)
-          let t = infer_app t cs in
+          let t = infer_app t (cs @ cs') in
           if Options.debug "checker-cg" then
             KPrint.bprintf "infer-cg: infer_app --> %a\n" ptyp t;
           t
@@ -1180,7 +1181,7 @@ and check_eqtype env t1 t2 =
 
 and check_subtype env t1 t2 =
   if not (subtype env t1 t2) then
-    checker_error env "subtype mismatch, %a (a.k.a. %a) vs %a (a.k.a. %a)"
+    checker_error env "subtype mismatch:\n  %a (a.k.a. %a) vs:\n  %a (a.k.a. %a)"
       ptyp t1 ptyp (expand_abbrev env t1) ptyp t2 ptyp (expand_abbrev env t2)
 
 and expand_abbrev env t =

lib/DataTypes.ml

@@ -67,7 +67,7 @@ let remove_unused_type_arguments files =
       method! visit_TApp env lid args =
         self#visit_app env lid args
 
-      method! visit_ETApp env e _ args =
+      method! visit_ETApp env e _ _ args =
         (* TODO: for now, we ignore unused const generics *)
         let lid = assert_elid e.node in
         self#visit_app (fst env) lid args
@@ -144,8 +144,8 @@ let remove_unused_type_arguments files =
       let n, (def, binders, ret) = chop 0 0 (def, binders, ret) in
       DFunction (cc, flags, n_cgs, n, ret, lid, binders, def)
 
-    method! visit_ETApp env e cgs args =
-      assert (cgs = []);
+    method! visit_ETApp env e cgs cgs' args =
+      assert (cgs @ cgs' = []);
       let lid = assert_elid e.node in
       let args = List.map (self#visit_typ_wo env) args in
       let args = KList.filter_mapi (fun i arg ->
@@ -155,7 +155,7 @@ let remove_unused_type_arguments files =
           None
       ) args in
       if List.length args > 0 then
-        ETApp (e, [], args)
+        ETApp (e, [], [], args)
       else
         e.node
   end in

lib/DeBruijn.ml

@@ -6,6 +6,7 @@
 (* ---------------------------------------------------------------------------- *)
 
 open Ast
+open PrintAst.Ops
 
 (* Counting binders. *)
 
@@ -118,6 +119,13 @@ let subst_n e es =
     subst arg k body
   ) e es
 
+let subst_n' ofs e es =
+  let l = List.length es in
+  KList.fold_lefti (fun i body arg ->
+    let k = l - i - 1 in
+    subst arg (k + ofs) body
+  ) e es
+
 (* Substitute [t2] for [i] in [t1]. *)
 
 class subst_t (t2: typ) = object
@@ -278,7 +286,7 @@ let cg_of_expr diff e =
   | EConstant (w, s) ->
       CgConst (w, s)
   | _ ->
-      failwith "Unsuitable const generic"
+      failwith (KPrint.bsprintf "Unsuitable const generic: %a" pexpr e)
 
 (* Substitute const generics *)
 class subst_ct (c: cg Lazy.t) = object (self)

lib/Helpers.ml

@@ -340,7 +340,7 @@ class ['self] readonly_visitor = object (self: 'self)
         List.for_all (self#visit_expr_w ()) es
     | EQualified lid when is_readonly_builtin_lid lid ->
         List.for_all (self#visit_expr_w ()) es
-    | ETApp ({ node = EQualified lid; _ }, _, _) when is_readonly_builtin_lid lid ->
+    | ETApp ({ node = EQualified lid; _ }, _, _, _) when is_readonly_builtin_lid lid ->
         List.for_all (self#visit_expr_w ()) es
     | _ ->
         false
@@ -741,7 +741,7 @@ let push_ignore = nest_in_return_pos TUnit (fun _ e ->
       ETApp (
         with_type (TArrow (TBound 0, TUnit))
           (EQualified (["LowStar"; "Ignore"], "ignore")),
-        [],
+        [], [],
         [ e.typ ]
       )),
     [ e ])))

lib/Inlining.ml

@@ -36,7 +36,7 @@ let reparenthesize_applications = object (self)
     let es = List.map (self#visit_expr env) es in
     let e = self#visit_expr env e in
     match e.node with
-    | EQualified lid | ETApp ({ node = EQualified lid; _ }, _, _) ->
+    | EQualified lid | ETApp ({ node = EQualified lid; _ }, _, _, _) ->
         begin try
           let n = Hashtbl.find natural_arity lid in
           let first, last = KList.split n es in
@@ -395,7 +395,7 @@ let cross_call_analysis files =
           | None ->
               super#visit_TCgApp () name ts
 
-        method! visit_ETApp ((), t) e _ ts =
+        method! visit_ETApp ((), t) e _ _ ts =
           (* Monomorphization happened long ago. If we hit this, this means we are the call-site for
              a type- or cg-polymorphic external function. The callee retains its original
              polymorphic signature (e.g. \0. () -> uint8[0] * uint8[0]), and the call-site (here)

lib/InputAstToAst.ml

@@ -160,7 +160,7 @@ and mk_expr = function
       let c = match op with K.Eq -> K.PEq | K.Neq -> K.PNeq | _ -> assert false in
       mk (EPolyComp (c, mk_typ t))
   | I.ETApp (e, es) ->
-      mk (ETApp (mk_expr e, [], List.map mk_typ es))
+      mk (ETApp (mk_expr e, [], [], List.map mk_typ es))
   | I.ELet (b, e1, e2) ->
       mk (ELet (mk_binder b, mk_expr e1, mk_expr e2))
   | I.EIfThenElse (e1, e2, e3) ->

lib/Monomorphization.ml

@@ -486,14 +486,15 @@ let functions files =
             [ d ]
       ) decls
 
-    method! visit_ETApp ((diff, _) as env) e cgs ts =
+    method! visit_ETApp ((diff, _) as env) e cgs cgs' ts =
+      let cgs' = List.map (self#visit_expr env) cgs' in
       let fail_if () =
-        if cgs <> [] then
+        if cgs @ cgs' <> [] then
           Warn.fatal_error "TODO: e=%a\ncgs=%a\nts=%a\n%a\n"
             pexpr e
             pexprs cgs
             ptyps ts
-            pexpr (with_type TUnit (ETApp (e, cgs, ts)));
+            pexpr (with_type TUnit (ETApp (e, cgs, cgs', ts)));
       in
       match e.node with
       | EQualified lid ->
@@ -506,7 +507,7 @@ let functions files =
                 (* External function. Bail. Leave cgs -- treated as normal
                    arguments when going to C. C'est la vie. *)
                 if !Options.allow_tapps || AstToCStar.whitelisted_tapp e then
-                  super#visit_ETApp env e cgs ts
+                  super#visit_ETApp env e cgs cgs' ts
                 else
                   (self#visit_expr env e).node
             | `Function (cc, flags, n_cgs, n, ret, name, binders, body) ->
@@ -524,18 +525,24 @@ let functions files =
                   let def () =
                     let ret = DeBruijn.(subst_ctn diff cgs (subst_tn ts ret)) in
                     assert (List.length cgs = n_cgs);
-                    let _, binders = KList.split (List.length cgs) binders in
+                    (* binders are the remaining binders after the cg-binders have been eliminated *)
+                    let diff = List.length binders - List.length cgs in
+                    let _, binders = KList.split (List.length cgs + List.length cgs') binders in
                     let binders = List.map (fun { node; typ } ->
                       { node; typ = DeBruijn.(subst_ctn diff cgs (subst_tn ts typ)) }
                     ) binders in
-                    (* KPrint.bprintf "about to substitute: e=%a\ncgs=%a\nts=%a\n%a\n" *)
-                    (*   pexpr e *)
-                    (*   pexprs cgs *)
-                    (*   ptyps ts *)
-                    (*   pexpr (with_type TUnit (ETApp (e, cgs, ts))); *)
-                    (* KPrint.bprintf "after to substitute: body:%a\n\n" pexpr body; *)
-                    let body = DeBruijn.(subst_cen (List.length binders) cgs (subst_ten ts body)) in
-                    (* KPrint.bprintf "after substitution: body:%a\n\n" pexpr body; *)
+                    KPrint.bprintf "about to substitute:\n  e=%a\n  cgs=%a\n cgs'=%a\n  ts=%a\n  head type=%a\n%a\n"
+                      pexpr e
+                      pexprs cgs
+                      pexprs cgs'
+                      ptyps ts
+                      ptyp e.typ
+                      pexpr (with_type TUnit (ETApp (e, cgs, cgs', ts)));
+                    KPrint.bprintf "body: %a\n\n" pexpr body;
+                    KPrint.bprintf "subst_ten ts body: %a\n\n" pexpr DeBruijn.(subst_ten ts body);
+                    KPrint.bprintf "subst_cen diff cgs (subst_ten ts body): %a\n\n" pexpr DeBruijn.(subst_cen diff cgs (subst_ten ts body));
+                    let body = DeBruijn.(subst_n' (List.length binders) (subst_cen diff cgs (subst_ten ts body)) cgs') in
+                    KPrint.bprintf "after substitution: body :%a\n\n" pexpr body;
                     let body = self#visit_expr env body in
                     DFunction (cc, flags, 0, 0, ret, name, binders, body)
                   in

lib/PrintAst.ml

@@ -173,14 +173,14 @@ and print_typ_paren = function
       print_typ t
 
 and print_cg = function
-  | CgVar i -> int i
-  | CgConst c -> print_constant c
+  | CgVar i -> dollar ^^ int i
+  | CgConst c -> dollar ^^ print_constant c
 
 and print_typ = function
   | TInt w -> print_width w
   | TBuf (t, bool) -> (if bool then string "const" else empty) ^/^ print_typ t ^^ star
   | TArray (t, k) -> print_typ t ^^ lbracket ^^ print_constant k ^^ rbracket
-  | TCgArray (t, v) -> print_typ t ^^ lbracket ^^ int v ^^ rbracket
+  | TCgArray (t, v) -> print_typ t ^^ lbracket ^^ dollar ^^ int v ^^ rbracket
   | TCgApp (t, cg) -> print_typ t ^^ brackets (brackets (print_cg cg))
   | TUnit -> string "()"
   | TQualified name -> string (string_of_lident name)
@@ -205,7 +205,6 @@ and print_let_binding (binder, e1) =
   jump (print_expr e1))
 
 and print_expr { node; typ } =
-  (* print_typ typ ^^ colon ^^ space ^^ *)
   match node with
   | EComment (s, e, s') ->
       surround 2 1 (string s) (print_expr e) (string s')
@@ -233,10 +232,10 @@ and print_expr { node; typ } =
       dquote ^^ string s ^^ dquote
   | EApp (e, es) ->
       print_app print_expr e print_expr es
-  | ETApp (e, es, ts) ->
+  | ETApp (e, es, es', ts) ->
       print_cg_app (fun (e, ts) ->
         print_app print_expr e (fun t -> group (langle ^/^ print_typ t ^/^ rangle)) ts
-      ) (e, ts) (fun e -> brackets (brackets (print_expr e))) es
+      ) (e, ts) (fun e -> brackets (brackets (print_expr e))) (es @ es')
   | ELet (binder, e1, e2) ->
       group (print_let_binding (binder, e1) ^/^ string "in") ^^ hardline ^^
       group (print_expr e2)
@@ -389,6 +388,7 @@ and print_pat p =
 let print_files = print_files print_decl
 
 module Ops = struct
+  let print_cgs = separate_map (comma ^^ space) print_cg
   let print_typs = separate_map (comma ^^ space) print_typ
   let print_exprs = separate_map (comma ^^ space) print_expr
   let print_lidents = separate_map (comma  ^^ space) print_lident
@@ -398,6 +398,8 @@ module Ops = struct
   let pcase = printf_of_pprint print_case
   let ptyp = printf_of_pprint print_typ
   let ptyps = printf_of_pprint print_typs
+  let pcg = printf_of_pprint print_cg
+  let pcgs = printf_of_pprint print_cgs
   let pptyp = printf_of_pprint_pretty print_typ
   let pexpr = printf_of_pprint print_expr
   let pbind = printf_of_pprint print_binder

lib/Simplify.ml

@@ -76,7 +76,7 @@ class ['self] safe_use = object (self: 'self)
     match e.node with
     | EOp _ -> super#visit_EApp env e es
     | EQualified lid when Helpers.is_readonly_builtin_lid lid -> super#visit_EApp env e es
-    | ETApp ({ node = EQualified lid; _ }, _, _) when Helpers.is_readonly_builtin_lid lid -> super#visit_EApp env e es
+    | ETApp ({ node = EQualified lid; _ }, _, _, _) when Helpers.is_readonly_builtin_lid lid -> super#visit_EApp env e es
     | _ -> self#unordered env (e :: es)
 
   method! visit_ELet env _ e1 e2 = self#sequential env e1 (Some e2)
@@ -398,7 +398,7 @@ let remove_ignore_unit = object
 
   method! visit_EApp env hd args =
     match hd.node, args with
-    | ETApp ({ node = EQualified (["LowStar"; "Ignore"], "ignore"); _}, _, [ TUnit ]), [ { node = EUnit; _ } ] ->
+    | ETApp ({ node = EQualified (["LowStar"; "Ignore"], "ignore"); _}, _, _, [ TUnit ]), [ { node = EUnit; _ } ] ->
         EUnit
     | _ ->
         super#visit_EApp env hd args
@@ -1127,9 +1127,9 @@ and hoist_stmt loc e =
 and hoist_expr loc pos e =
   let mk node = { node; typ = e.typ } in
   match e.node with
-  | ETApp (e, cgs, ts) ->
+  | ETApp (e, cgs, cgs', ts) ->
       let lhs, e = hoist_expr loc Unspecified e in
-      lhs, mk (ETApp (e, cgs, ts))
+      lhs, mk (ETApp (e, cgs, cgs', ts))
 
   | EBufNull
   | EAbort _