A grown parameter the function returns, or whose field it reassigns, is not warned at

lib/check.ml

@@ -13317,6 +13317,55 @@ let check_union_members env =
 
 (* ── Declarations: pass 2, check bodies ────────────────────────────── *)
 
+(* The names a body hands back or stores into: every name mentioned in a value
+   it answers — its last form's tails, a [return]'s value — and the name at
+   the root of every [set] place. A parameter among them is not warned at for
+   growing: the grown copy goes back to the caller, or the copy is the
+   function's own business. *)
+let escaping_names ~returns (body : Ast.expr list) : string list =
+  let names = ref [] in
+  let rec mentions (e : Ast.expr) =
+    (match e.Ast.e with Ast.Var n -> names := n :: !names | _ -> ());
+    ignore (Ast.map_children (fun x -> mentions x; x) e)
+  in
+  let rec tails (e : Ast.expr) =
+    match e.Ast.e with
+    | Ast.Do es | Ast.Let (_, es) ->
+      (match List.rev es with x :: _ -> tails x | [] -> ())
+    | Ast.If (_, a, b) -> tails a; Option.iter tails b
+    | Ast.Match (_, arms) ->
+      List.iter
+        (fun (a : Ast.arm) ->
+           match List.rev a.Ast.body with x :: _ -> tails x | [] -> ())
+        arms
+    (* A value that is the parameter, a field of it, or a literal built with
+       it. A call's result is its callee's business, and a unit form — the
+       push itself, last in a function that returns nothing — answers
+       nothing. *)
+    | Ast.Var _ | Ast.Field _ | Ast.Struct _ | Ast.Bare _ | Ast.Arr _
+    | Ast.MapLit _ -> mentions e
+    | _ -> ()
+  in
+  let rec root (e : Ast.expr) =
+    match e.Ast.e with
+    | Ast.Var n -> names := n :: !names
+    | Ast.Field (x, _) -> root x
+    | Ast.Call ({ Ast.e = Ast.Var ("at" | "deref"); _ }, x :: _) -> root x
+    | _ -> ()
+  in
+  let rec walk (e : Ast.expr) =
+    (match e.Ast.e with
+     | Ast.Return (Some x) -> tails x
+     | Ast.Set (Ast.Pvar n, _) -> names := n :: !names
+     | Ast.Set ((Ast.Pfield (x, _) | Ast.Pindex (x, _) | Ast.Pderef x
+                | Ast.Pslot (x, _)), _) -> root x
+     | _ -> ());
+    ignore (Ast.map_children (fun x -> walk x; x) e)
+  in
+  List.iter walk body;
+  if returns then (match List.rev body with x :: _ -> tails x | [] -> ());
+  !names
+
 let rec check_fn env (fn : Ast.fn) : Tast.fn =
   let params, ret = Hashtbl.find env.fns fn.Ast.name in
   let ctx = { (invented_ctx env ret) with owner = fn.Ast.name } in
@@ -13339,6 +13388,7 @@ let rec check_fn env (fn : Ast.fn) : Tast.fn =
        end;
        ignore (bind ctx p.Ast.fname ty ~assignable:false))
     fn.Ast.params params;
+  let grow_before = !grow_warnings in
   let grow_saved = !grow_params in
   grow_params :=
     ( ctx,
@@ -13347,7 +13397,32 @@ let rec check_fn env (fn : Ast.fn) : Tast.fn =
            Option.map (fun b -> (b.slot, p)) (List.assoc_opt p.Ast.fname ctx.scope))
         fn.Ast.params )
     :: grow_saved;
-  Fun.protect ~finally:(fun () -> grow_params := grow_saved) @@ fun () ->
+  let escaping =
+    lazy
+      (let names =
+         escaping_names ~returns:(not (Types.equal ret Types.Unit)) fn.Ast.fbody
+       in
+       List.filter_map
+         (fun (p : Ast.field) ->
+            if List.mem p.Ast.fname names then Some p.Ast.floc else None)
+         fn.Ast.params)
+  in
+  Fun.protect
+    ~finally:(fun () ->
+        grow_params := grow_saved;
+        let added =
+          List.filteri
+            (fun i _ -> i < List.length !grow_warnings - List.length grow_before)
+            !grow_warnings
+        in
+        if added <> [] then
+          grow_warnings :=
+            List.filter
+              (fun (d : Loc.diag) ->
+                 not (List.mem d.Loc.dloc (Lazy.force escaping)))
+              added
+            @ grow_before)
+  @@ fun () ->
   let body =
     match fn.Ast.fbody with
     | [] ->

test/test_flan.ml

@@ -5915,6 +5915,22 @@ let () =
    | Some ds ->
      check (Printf.sprintf "one field grow warning, not %d" (List.length ds)) false
    | None -> check "the grown-field program checks" false);
+  (* Not when the grown copy goes back to the caller — the parameter, or the
+     struct holding the field, is what the function answers — nor when the
+     field is given a container of the function's own before it grows. *)
+  check "a grown parameter the function returns is not warned at"
+    (grown "(defstruct Bag [items (Vec i32)])\n\
+            (defn add [v (Vec i32) x i32] (Vec i32) (push v x) v)\n\
+            (defn early [v (Vec i32) c bool] (Vec i32) (push v 1) \
+               (when c (return v)) v)\n\
+            (defn bag [b Bag] Bag (push (.items b) 1) b)\n\
+            (defn items [b Bag] (Vec i32) (push (.items b) 1) (.items b))"
+     = Some []);
+  check "a field reassigned before it grows is not warned at"
+    (grown "(defstruct Bag [items (Vec i32)])\n\
+            (defn f [b Bag] ()\n\
+            \  (set (.items b) (vec-new i32)) (push (.items b) 1) (free (.items b)))"
+     = Some []);
   check "a pointer parameter and a local are not warned at"
     (grown "(defn f [v (Ptr (Vec i32))] ()\n\
             \  (push (deref v) 1) (let [w (vec-new i32)] (push w 1) (free w)))"