Track zero-alloc across aliases (I think this slows down compilation …

…dramatically...)

ocaml/lambda/translcore.ml

@@ -369,7 +369,7 @@ let zero_alloc_of_application
       | Check_default | No_check -> false
       | Check_all | Check_opt_only -> true
     in
-    begin match Zero_alloc.get val_zero_alloc with
+    begin match Zero_alloc.get_default val_zero_alloc with
     | Check c when c.arity = num_args && (use_opt || not c.opt) ->
       let assume : Zero_alloc.assume =
         { strict = c.strict;

ocaml/testsuite/tests/typing-zero-alloc/signatures.ml

@@ -1443,3 +1443,115 @@ module type S_outer_subst =
   sig module M_inner : sig val f : unit -> unit end end
 module M : S_outer_subst
 |}]
+
+(* Test 18: rebinding *)
+module type S_rebinding_base = sig
+  val[@zero_alloc] f : 'a -> 'a
+end
+
+module type S_rebinding_opt = sig
+  val[@zero_alloc opt] f : 'a -> 'a
+end
+
+module M1 = struct
+  let g x = x
+  let f = g
+end
+
+module _ : S_rebinding_base = M1
+module _ : S_rebinding_opt = M1
+
+module M2 = struct
+  let g x = x
+  let f = g
+end
+
+module _ : S_rebinding_opt = M2
+module _ : S_rebinding_base = M2
+
+[%%expect{|
+module type S_rebinding_base = sig val f : 'a -> 'a [@@zero_alloc] end
+module type S_rebinding_opt = sig val f : 'a -> 'a [@@zero_alloc opt] end
+module M1 : sig val g : 'a -> 'a val f : 'a -> 'a end
+module M2 : sig val g : 'a -> 'a val f : 'a -> 'a end
+|}]
+
+(* Test 19: rebinding errors *)
+module type S_rebinding_errors_base = sig
+  val[@zero_alloc] f : 'a -> 'a
+end
+
+module type S_rebinding_errors_opt = sig
+  val[@zero_alloc opt] f : 'a -> 'a
+end
+
+[%%expect {|
+module type S_rebinding_errors_base = sig val f : 'a -> 'a [@@zero_alloc] end
+module type S_rebinding_errors_opt =
+  sig val f : 'a -> 'a [@@zero_alloc opt] end
+|}]
+
+module M1 : S_rebinding_errors_base = struct
+  let[@zero_alloc opt] g x = x
+  let f = g
+end
+
+[%%expect{|
+Lines 1-4, characters 38-3:
+1 | ......................................struct
+2 |   let[@zero_alloc opt] g x = x
+3 |   let f = g
+4 | end
+Error: Signature mismatch:
+       Modules do not match:
+         sig
+           val g : 'a -> 'a [@@zero_alloc opt]
+           val f : 'a -> 'a [@@zero_alloc opt]
+         end
+       is not included in
+         S_rebinding_errors_base
+       Values do not match:
+         val f : 'a -> 'a [@@zero_alloc opt]
+       is not included in
+         val f : 'a -> 'a [@@zero_alloc]
+       The former provides a weaker "zero_alloc" guarantee than the latter.
+|}]
+
+module M1 : S_rebinding_errors_base = struct
+  let g x = x
+  let[@zero_alloc opt] f = g
+end
+
+(* CR nroberts: this output looks wrong... *)
+[%%expect{|
+module M1 : S_rebinding_errors_base
+|}]
+
+(* Test 20: complex rebinding *)
+module type S_rebinding_complex = sig
+  val[@zero_alloc] f : 'a -> 'a
+  val[@zero_alloc opt] g : 'a -> 'a
+end
+
+[%%expect{|
+module type S_rebinding_complex =
+  sig val f : 'a -> 'a [@@zero_alloc] val g : 'a -> 'a [@@zero_alloc opt] end
+|}]
+
+module M1 : S_rebinding_complex = struct
+  let g x = x
+  let f = g
+end
+
+[%%expect{|
+module M1 : S_rebinding_complex
+|}]
+
+module M2 : S_rebinding_complex = struct
+  let f x = x
+  let g = f
+end
+
+[%%expect{|
+module M2 : S_rebinding_complex
+|}]

ocaml/typing/printtyp.ml

@@ -2132,7 +2132,7 @@ let tree_of_value_description id decl =
     count 0 decl.val_type
   in
   let attrs =
-    match Zero_alloc.get decl.val_zero_alloc with
+    match Zero_alloc.get_default decl.val_zero_alloc with
     | Default_zero_alloc | Ignore_assert_all -> []
     | Check { strict; opt; arity; _ } ->
       [{ oattr_name =

ocaml/typing/subst.ml

@@ -825,7 +825,7 @@ let rec subst_lazy_value_description s descr =
          more strictly than the signature indicates, which is sound). *)
      (match s.additional_action with
       | Prepare_for_saving _ ->
-        Zero_alloc.create_const (Zero_alloc.get descr.val_zero_alloc)
+        Zero_alloc.create_const (Zero_alloc.get_default descr.val_zero_alloc)
       | _ -> descr.val_zero_alloc);
     val_attributes = attrs s descr.val_attributes;
     val_uid = descr.val_uid;

ocaml/typing/typecore.ml

@@ -1137,7 +1137,7 @@ let add_pattern_variables ?check ?check_as env pv =
        Env.add_value ?check ~mode:pv_mode pv_id
          {val_type = pv_type; val_kind = Val_reg; Types.val_loc = pv_loc;
           val_attributes = pv_attributes; val_modalities = Modality.Value.id;
-          val_zero_alloc = Zero_alloc.default;
+          val_zero_alloc = Zero_alloc.create_rvar ();
           val_uid = pv_uid
          } env
     )
@@ -9045,7 +9045,7 @@ and type_n_ary_function
     in
     let zero_alloc =
       match zero_alloc with
-      | Default_zero_alloc -> Zero_alloc.create_var loc syntactic_arity
+      | Default_zero_alloc -> Zero_alloc.create_lvar loc syntactic_arity
       | (Check _ | Assume _ | Ignore_assert_all) ->
         Zero_alloc.create_const zero_alloc
     in

ocaml/typing/typedtree.ml

@@ -1084,6 +1084,8 @@ let let_bound_idents_with_modes_sorts_and_checks bindings =
     List.fold_left (fun checks vb ->
       iter_pattern_full ~both_sides_of_or:true f vb.vb_sort vb.vb_pat;
        match vb.vb_pat.pat_desc, vb.vb_expr.exp_desc with
+       | Tpat_var (id, _, _, _), Texp_ident (_, _, vd, _, _) ->
+           Ident.Map.add id vd.val_zero_alloc checks
        | Tpat_var (id, _, _, _), Texp_function fn ->
          let zero_alloc =
            match Zero_alloc.get fn.zero_alloc with

ocaml/typing/typemod.ml

@@ -563,7 +563,7 @@ let rec remove_modality_and_zero_alloc_variables_sg env ~zap_modality sg =
           |> zap_modality |> Mode.Modality.Value.of_const
         in
         let val_zero_alloc =
-          Zero_alloc.create_const (Zero_alloc.get desc.val_zero_alloc)
+          Zero_alloc.create_const (Zero_alloc.get_default desc.val_zero_alloc)
         in
         let desc = {desc with val_modalities; val_zero_alloc} in
         Sig_value (id, desc, vis)
@@ -2966,7 +2966,7 @@ and type_structure ?(toplevel = None) funct_body anchor env sstr =
                    convert "Assume"s in structures to the equivalent "Check" for
                    the signature. *)
                 let open Builtin_attributes in
-                match[@warning "+9"] Zero_alloc.get zero_alloc with
+                match[@warning "+9"] Zero_alloc.get_default zero_alloc with
                 | Default_zero_alloc | Check _ -> zero_alloc
                 | Assume { strict; arity; loc;
                            never_returns_normally = _;
@@ -2981,9 +2981,15 @@ and type_structure ?(toplevel = None) funct_body anchor env sstr =
               let modalities =
                 maybe_infer_modalities ~loc:first_loc ~env ~md_mode ~mode
               in
+              begin
+                match Zero_alloc.sub zero_alloc vd.val_zero_alloc with
+                | Ok () -> ()
+                | Error _ ->
+                    fatal_error
+                      "Failed to constrain an unconstrained zero alloc var"
+              end;
               let vd =
                 { vd with
-                  val_zero_alloc = zero_alloc;
                   val_modalities = modalities }
               in
               Sig_value(id, vd, Exported) :: acc,

ocaml/typing/zero_alloc.ml

@@ -24,7 +24,7 @@ type const = Builtin_attributes.zero_alloc_attribute =
 
 type desc = { strict : bool; opt : bool }
 
-type var =
+type lvar =
   { loc : Location.t;
     arity : int;
     mutable desc : desc option;
@@ -35,9 +35,10 @@ type var =
 
 type t =
   | Const of const
-  | Var of var
+  | Lvar of lvar
+  | Rvar of { mutable t : t option (* not Rvar *) }
 
-let debug_printer ppf t =
+let rec debug_printer ppf t =
   let head c = match c with
     | Default_zero_alloc -> "Default"
     | Ignore_assert_all -> "Ignore"
@@ -46,34 +47,53 @@ let debug_printer ppf t =
   in
   match t with
   | Const c -> Format.fprintf ppf "Const %s" (head c)
-  | Var v ->
+  | Rvar { t = None } -> Format.fprintf ppf "Rvar"
+  | Rvar { t = Some t } -> Format.fprintf ppf "Rvar (%a)" debug_printer t
+  | Lvar v ->
     let print_desc ppf desc =
       match desc with
       | None -> Format.fprintf ppf "None"
       | Some desc ->
         Format.fprintf ppf "{ strict = %b; opt = %b }" desc.strict desc.opt
     in
-    Format.fprintf ppf "Var { arity = %d; desc = %a }" v.arity print_desc v.desc
+    Format.fprintf ppf "Lvar { arity = %d; desc = %a }" v.arity print_desc v.desc
 
 (* For backtracking *)
-type change = desc option * var
-let undo_change (d, v) = v.desc <- d
+type change =
+  | Change_desc of desc option * lvar
+  | Constrain_rvar of t
+
+let undo_change = function
+  | Change_desc (d, v) -> v.desc <- d
+  | Constrain_rvar (Rvar rv) -> rv.t <- None
+  | Constrain_rvar (Const _ | Lvar _) -> assert false
+
 let log_change = ref (fun _ -> ())
 let set_change_log f = log_change := f
 
 let create_const x = Const x
-let create_var loc arity = Var { loc; arity; desc = None }
+let create_lvar loc arity = Lvar { loc; arity; desc = None }
+let create_rvar () = Rvar { t = None }
 let default = Const Default_zero_alloc
 
-let get (t : t) =
+let rec get (t : t) =
   match t with
   | Const c -> c
-  | Var { loc; arity; desc } ->
+  | Rvar { t = None } -> Misc.fatal_error "Zero_alloc.get of unconstrained Rvar"
+  | Rvar { t = Some t } -> get t
+  | Lvar { loc; arity; desc } ->
     match desc with
     | None -> Default_zero_alloc
     | Some { strict; opt } ->
       Check { loc; arity; strict; opt }
 
+let get_default (t : t) =
+  match t with
+  | Rvar ({ t = None } as rv) ->
+      rv.t <- Some (Const Default_zero_alloc);
+      get t
+  | _ -> get t
+
 type error =
   | Less_general of { missing_entirely : bool }
   | Arity_mismatch of int * int
@@ -169,20 +189,29 @@ let sub_var_const_exn v c =
     when arity1 <> arity2 ->
     raise (Error (Arity_mismatch (arity1, arity2)))
   | { desc = None; _ }, Check { strict; opt; _ } ->
-    !log_change (None, v);
+    !log_change (Change_desc (None, v));
     v.desc <- Some { strict; opt }
   | { desc = (Some { strict = strict1; opt = opt1 } as desc); _ },
     Check { strict = strict2; opt = opt2 } ->
     let strict = strict1 || strict2 in
     let opt = opt1 && opt2 in
     if strict <> strict1 || opt <> opt1 then begin
-      !log_change (desc, v);
+      !log_change (Change_desc (desc, v));
       v.desc <- Some { strict; opt }
     end
 
-let sub_exn za1 za2 =
+let rec sub_exn za1 za2 =
   match za1, za2 with
-  | _, Var _ ->
+  | Rvar { t = Some za1 }, x -> sub_exn za1 x
+  | Rvar { t = None }, (Const _ | Lvar _) ->
+      Misc.fatal_error "Rvar appeared unconstrained on lhs of <="
+  | za1, Rvar ({ t = None } as rv2) ->
+      !log_change (Constrain_rvar za2);
+      rv2.t <- Some za1
+  | _, Rvar { t = Some za2 } ->
+      if not (za1 == za2) then
+        Misc.fatal_error "Rvar appeared on rhs of 2 different <='s"
+  | _, Lvar _ ->
     (* A fully inferred signature will never have a variable in it, so we almost
        never have to constrain by a variable, but there is one special case:
 
@@ -196,11 +225,13 @@ let sub_exn za1 za2 =
        event).
     *)
     if not (za1 == za2) then
-      Misc.fatal_error "zero_alloc: variable constraint"
+      Misc.fatal_errorf "zero_alloc: variable constraint: %a, %a"
+        debug_printer za1
+        debug_printer za2
   | _, Const (Ignore_assert_all | Assume _) ->
     Misc.fatal_error "zero_alloc: invalid constraint"
   | _, (Const Default_zero_alloc) -> ()
-  | Var v, Const c -> sub_var_const_exn v c
+  | Lvar v, Const c -> sub_var_const_exn v c
   | Const c1, Const c2 -> sub_const_const_exn c1 c2
 
 let sub za1 za2 =

ocaml/typing/zero_alloc.mli

@@ -34,12 +34,15 @@ val create_const : const -> t
 (* [create_var loc n] creates a variable. [loc] is the location of the function
    you are creating a variable for, and [n] is its syntactic arity of the
    function the variable is being created for. *)
-val create_var : Location.t -> int -> t
+val create_lvar : Location.t -> int -> t
+val create_rvar : unit -> t
 
 (* In the case [t] is a variable, [get t] returns its current contents as a
    [const] and has no effect. *)
 val get : t -> const
 
+val get_default : t -> const
+
 (* For types.ml's backtracking mechanism. *)
 type change
 val set_change_log : (change -> unit) -> unit