A let has the function's extent, so a defer may be written in one

lib/check.ml

@@ -90,9 +90,30 @@ type ctx = {
   mutable slot_names : string option list;
   mutable scope : (string * binding) list;   (* innermost first *)
   (* Deferred forms, most recently registered first — which is also the order
-     they run in. At milestone 4 [defer] is function-scoped (see [check_fn]),
+     they run in. [defer] is function-scoped, so this list belongs to the
-     so this list belongs to the function and not to a block. *)
+     function and not to a block — see [defer_ok] for where one may be written
+     and [check_fn] for where the list is spliced onto the exit paths. *)
   mutable defers : Tast.expr list;
+  (* Where a [defer] may be written, which is exactly: a form whose extent is
+     the whole function body. Two things have that extent and only two — a
+     top-level form of the body, and a form in the body of a [let] that itself
+     has it, to any depth. A [let] always registers and its bindings outlive
+     the block textually enclosing them, because a [let] is not a frame here:
+     its bindings are function slots like any other, and nothing is released at
+     scope exit (spec-memory.md, "When storage is released").
+
+     Everything else is refused, and the two that matter are refused for a
+     reason rather than by omission. [defer] is a *compile-time* construct —
+     the cleanup is copied into every exit path — so a branch would have to
+     express "maybe registered", which it cannot, and a loop body would fire
+     once at function exit rather than once per iteration.
+
+     The flag is set immediately before each form that may carry one, never
+     once around a body: [check] clears it on entry, so a body whose first form
+     set it would otherwise refuse the second. [defer_block] names the
+     innermost construct that cleared it, so the refusal says which. *)
+  mutable defer_ok : bool;
+  mutable defer_block : string;
   (* Only for the two things a handler clause cannot do. [outer] is the
      establishing function's scope, kept so that a reference to one of its
      locals can be refused for the reason it is really refused for rather than
@@ -203,6 +224,21 @@ let scoped ctx f =
   ctx.scope <- saved;
   r
 
+(* A scope that is also a named blocker for [defer]. An arm of an [if] or a
+   [match] runs only sometimes, and "maybe registered" is not something a
+   compile-time construct can express — the cleanup is copied into every exit
+   path or into none — so the refusal is about the branch and says so.
+
+   Outside the [check] recursion on purpose: inside it the inferred type would
+   be monomorphic, and the two callers pass functions returning different
+   things. *)
+let branch ctx f =
+  let blocker = ctx.defer_block in
+  ctx.defer_block <- "a branch";
+  let r = scoped ctx f in
+  ctx.defer_block <- blocker;
+  r
+
 (* ── Type resolution ───────────────────────────────────────────────── *)
 
 let unimplemented loc what milestone =
@@ -496,6 +532,12 @@ let expect loc ~want (got : Tast.expr) =
 
 let rec check ctx ?want (e : Ast.expr) : Tast.expr =
   let loc = e.Ast.loc in
+  (* Read the permission this form was given and withdraw it in the same
+     breath, so that nothing reached from here inherits it. The two callers
+     that may grant it — [check_fn]'s body walk and [check_let]'s, below —
+     grant it again before the *next* form rather than once around the body. *)
+  let defer_ok = ctx.defer_ok in
+  ctx.defer_ok <- false;
   match e.Ast.e with
   | Ast.Int n -> int_literal loc ~want n
   | Ast.Byte b -> int_literal loc ~want ~default:Types.U8 (Int64.of_int b)
@@ -535,7 +577,9 @@ let rec check ctx ?want (e : Ast.expr) : Tast.expr =
     unimplemented loc "a quoted symbol (restart names)" 6
   | Ast.Var name -> var ctx loc ~want name
   | Ast.Do body -> block ctx ?want loc body
-  | Ast.Let (bs, body) -> check_let ctx ?want loc bs body
+  (* [defer_ok] rides through: a [let] at the top level of a function body has
+     exactly the function's extent, and so does a [let] nested inside one. *)
+  | Ast.Let (bs, body) -> check_let ctx ?want ~defer_ok loc bs body
   | Ast.If (c, t, e') -> check_if ctx ?want loc c t e'
   | Ast.While (c, body) ->
     let c = check ctx ~want:Types.Bool c in
@@ -679,14 +723,24 @@ let rec check ctx ?want (e : Ast.expr) : Tast.expr =
       (if binds = [] then invoke
        else mk loc Types.Never (Tast.Let (binds, [ invoke ])))
 
-  | Ast.Defer _ ->
+  | Ast.Defer forms ->
-    (* Registered by [check_fn], which is the only place that sees a form's
+    (* Registering is the whole of it: the forms are checked here, where they
-       position. A defer anywhere else would run at function exit rather than
+       can see the scope they are written in, and the node left behind is
-       at the exit of the block it is written in — once for a loop body that
+       [unit]. [check_fn] splices the registered list onto both exit paths.
-       runs a thousand times — so it is rejected instead of quietly differing. *)
+
+       [defer_ok] is true for a top-level form of the body and for a form in a
+       [let] whose extent is the body — see the field's comment. Anywhere else
+       the cleanup would run at function exit rather than at the exit of the
+       construct it was written in, so it is refused, and named. *)
+    if not defer_ok then
       fail loc
-      "defer must be a top-level form in a function body — block-scoped defer \
+        "defer is not allowed inside %s — a defer is copied into every exit \
-       is not implemented yet (milestone 4)"
+         path of the function, so it always registers and always runs at \
+         function exit. Write it at the top level of the function body, or in \
+         a let that is (a let has the function's extent, because nothing is \
+         released at scope exit)"
+        ctx.defer_block;
+    register_defer ctx loc forms
 
 and int_literal loc ~want ?(default = Types.I32) n =
   match want with
@@ -799,13 +853,21 @@ and borrowed ctx (a : Ast.expr) f =
     r
   end
 
-and block ctx ?want loc body =
+(* [defer_ok] is granted again before *every* form, not once before the block:
+   [check] withdraws it as it starts, so granting it once would let the first
+   form carry a defer and refuse the second — and two resources acquired in one
+   [let] is the case the relaxation exists for. *)
+and block ctx ?want ?(defer_ok = false) loc body =
   match body with
   | [] -> expect loc ~want (unit_at loc)
   | _ ->
     let rec go = function
-      | [ last ] -> let l = check ctx ?want last in [ l ], l.Tast.ty
+      | [ last ] ->
-      | x :: rest -> let x = check ctx x in
+        ctx.defer_ok <- defer_ok;
+        let l = check ctx ?want last in [ l ], l.Tast.ty
+      | x :: rest ->
+        ctx.defer_ok <- defer_ok;
+        let x = check ctx x in
         let rest, ty = go rest in x :: rest, ty
       | [] -> assert false
     in
@@ -842,7 +904,7 @@ and check_handler_bind ctx ?want loc clauses body =
             the enclosing one. *)
          let hctx =
            { env = ctx.env; ret = Types.Unit; slots = 0; slot_tys = []; slot_names = [];
-             scope = []; defers = []; outer = ctx.scope; in_handler = true; in_frames = None; in_defer = false; dead = []; borrow = false; owner = "<none>" }
+             scope = []; defers = []; outer = ctx.scope; in_handler = true; in_frames = None; in_defer = false; defer_ok = false; defer_block = "a nested form"; dead = []; borrow = false; owner = "<none>" }
          in
          (* The condition crosses as a pointer, because the handler runs while
             the signalling frame is still alive and there is nothing to copy.
@@ -964,7 +1026,19 @@ and check_restart_case ctx ?want loc body clauses =
   let ty = match !ty with Some t -> t | None -> Types.Never in
   mk loc ty (Tast.RestartCase (clauses, tbody))
 
-and check_let ctx ?want loc bs body =
+(* The forms of a [defer], checked in place and hung on the function. It emits
+   nothing where it stands, so what is left behind is [unit]. *)
+and register_defer ctx loc forms =
+  ctx.in_defer <- true;
+  let forms = map_lr (fun d -> check ctx d) forms in
+  ctx.in_defer <- false;
+  ctx.defers <- mk loc Types.Unit (Tast.Do forms) :: ctx.defers;
+  unit_at loc
+
+(* [defer_ok] says whether *this* let has the function's extent. If it does, so
+   does every form in its body, including a nested let — which is why the flag
+   is handed to the body rather than consumed here. *)
+and check_let ctx ?want ?(defer_ok = false) loc bs body =
   scoped ctx (fun () ->
     let bs =
       map_lr
@@ -981,7 +1055,7 @@ and check_let ctx ?want loc bs body =
            (slot, v))
         bs
     in
-    let body = block ctx ?want loc body in
+    let body = block ctx ?want ~defer_ok loc body in
     mk loc body.Tast.ty (Tast.Let (bs, [ body ])))
 
 (* (dotimes [i n] body...) is a counting loop, not a new IR node: bind [i] to 0
@@ -996,7 +1070,13 @@ and check_let ctx ?want loc bs body =
 and in_loop ctx f =
   let outer_slots = List.map (fun (_, b) -> b.slot) ctx.scope in
   let before = ctx.dead in
+  (* Named so that a defer written in here is refused as "a loop body" rather
+     than as a nested form: the reason is specific — it would fire once at
+     function exit rather than once per iteration — and the message says it. *)
+  let blocker = ctx.defer_block in
+  ctx.defer_block <- "a loop body";
   let r = f () in
+  ctx.defer_block <- blocker;
   List.iter
     (fun (slot, where) ->
        if (not (List.mem_assoc slot before)) && List.mem slot outer_slots then
@@ -1035,7 +1115,7 @@ and check_if ctx ?want loc c t e =
   | None ->
     (* A one-armed if produces Unit whatever the branch evaluates to: there is
        no value on the missing side. `when` desugars to this. *)
-    let t = scoped ctx (fun () -> check ctx t) in
+    let t = branch ctx (fun () -> check ctx t) in
     expect loc ~want (mk loc Types.Unit (Tast.If (c, t, unit_at loc)))
   | Some e ->
     (* Both arms start from the same dead set and the union survives: moving in
@@ -1044,7 +1124,7 @@ and check_if ctx ?want loc c t e =
        refused [(if c (free v) (free v))] and allowed the use after a one-armed
        move, which are the two ways to be wrong here. *)
     let before = ctx.dead in
-    let t = scoped ctx (fun () -> check ctx ?want t) in
+    let t = branch ctx (fun () -> check ctx ?want t) in
     let after_then = ctx.dead in
     ctx.dead <- before;
     (* With no expectation the then-branch supplies one for the else-branch,
@@ -1054,7 +1134,7 @@ and check_if ctx ?want loc c t e =
       | Some _ -> want
       | None -> if t.Tast.ty = Types.Never then None else Some t.Tast.ty
     in
-    let e = scoped ctx (fun () -> check ctx ?want:ewant e) in
+    let e = branch ctx (fun () -> check ctx ?want:ewant e) in
     ctx.dead <-
       after_then
       @ List.filter (fun (k, _) -> not (List.mem_assoc k after_then)) ctx.dead;
@@ -1175,7 +1255,7 @@ and check_match ctx ?want loc scrutinee arms =
          in
          ctx.dead <- before;
          let arm =
-           scoped ctx (fun () ->
+           branch ctx (fun () ->
              let binds =
                List.map (fun n -> bind ctx n elem ~assignable:false) binds
              in
@@ -2728,7 +2808,7 @@ let collect env (decls : Ast.decl list) =
      run without swallowing it. *)
   let infer (_, v) =
     (check { env; ret = Types.Unit; slots = 0; slot_tys = []; slot_names = []; scope = []; defers = [];
-      outer = []; in_handler = false; in_frames = None; in_defer = false; dead = []; borrow = false; owner = "<none>" } v).Tast.ty
+      outer = []; in_handler = false; in_frames = None; in_defer = false; defer_ok = false; defer_block = "a nested form"; dead = []; borrow = false; owner = "<none>" } v).Tast.ty
   in
   let pending = ref (List.rev !untyped) in
   let rec settle () =
@@ -2781,7 +2861,7 @@ let check_finite env =
 let check_fn env (fn : Ast.fn) : Tast.fn =
   let params, ret = Hashtbl.find env.fns fn.Ast.name in
   let ctx = { env; ret; slots = 0; slot_tys = []; slot_names = []; scope = []; defers = [];
-                outer = []; in_handler = false; in_frames = None; in_defer = false; dead = []; borrow = false;
+                outer = []; in_handler = false; in_frames = None; in_defer = false; defer_ok = false; defer_block = "a nested form"; dead = []; borrow = false;
                 owner = fn.Ast.name } in
   List.iter2
     (fun (p : Ast.field) ty ->
@@ -2799,28 +2879,25 @@ let check_fn env (fn : Ast.fn) : Tast.fn =
       (* The last form is the return value, unless the function returns Unit,
          in which case whatever it evaluates to is discarded. *)
       let want = if Types.equal ret Types.Unit then None else Some ret in
-      (* [defer] is recognised here and nowhere else, because this is the only
+      (* Every form here is at the top level of the function body, so every one
-         place that knows a form is at the top level of the function body. Each
+         of them may carry a [defer] — and so may a form inside a [let] written
-         one is checked in place — so it sees the scope it is written in — and
+         here, which is what [ctx.defer_ok] carries down. [check] registers it
-         then registered on the context; it emits nothing where it stands. *)
+         and yields [unit]; the permission is granted again before each form
-      let defer_here (e : Ast.expr) =
+         because [check] withdraws it as it starts.
-        match e.Ast.e with
+
-        | Ast.Defer forms ->
+         A trailing [defer] is still a [defer] and not the return value, so the
-          ctx.in_defer <- true;
+         expectation is not put to it: it would only ever report [Unit] against
-          let forms = map_lr (fun d -> check ctx d) forms in
+         the declared return type, which names the wrong problem. *)
-          ctx.in_defer <- false;
+      let is_defer (e : Ast.expr) =
-          let d = mk e.Ast.loc Types.Unit (Tast.Do forms) in
+        match e.Ast.e with Ast.Defer _ -> true | _ -> false
-          ctx.defers <- d :: ctx.defers;
-          Some (unit_at e.Ast.loc)
-        | _ -> None
       in
       let rec go = function
         | [ last ] ->
-          (match defer_here last with
+          ctx.defer_ok <- true;
-           | Some u -> [ u ]
+          [ (if is_defer last then check ctx last else check ctx ?want last) ]
-           | None -> [ check ctx ?want last ])
         | x :: rest ->
-          let x = match defer_here x with Some u -> u | None -> check ctx x in
+          ctx.defer_ok <- true;
+          let x = check ctx x in
           x :: go rest
         | [] -> assert false
       in
@@ -2873,7 +2950,7 @@ let no_move_only_global loc n (ty : Types.t) =
 
 let check_global env (d : Ast.decl) : Tast.global option =
   let ctx () = { env; ret = Types.Unit; slots = 0; slot_tys = []; slot_names = []; scope = []; defers = [];
-      outer = []; in_handler = false; in_frames = None; in_defer = false; dead = []; borrow = false; owner = "<none>" } in
+      outer = []; in_handler = false; in_frames = None; in_defer = false; defer_ok = false; defer_block = "a nested form"; dead = []; borrow = false; owner = "<none>" } in
   match d.Ast.d with
   | Ast.Defvar (n, _, init) ->
     let ty, _ = Hashtbl.find env.globals n in
@@ -2985,7 +3062,7 @@ let expression env (e : Ast.expr) :
   Tast.expr * Types.t array * string option array =
   let ctx =
     { env; ret = Types.Unit; slots = 0; slot_tys = []; slot_names = []; scope = []; defers = [];
-      outer = []; in_handler = false; in_frames = None; in_defer = false; dead = []; borrow = false; owner = "<none>" }
+      outer = []; in_handler = false; in_frames = None; in_defer = false; defer_ok = false; defer_block = "a nested form"; dead = []; borrow = false; owner = "<none>" }
   in
   let t = check ctx e in
   (t, Array.of_list (List.rev ctx.slot_tys),

test/programs/defer-let.flan

@@ -0,0 +1,72 @@
+;;;; defer in a let whose extent is the function body.
+;;;;
+;;;; A let is not a frame here: its bindings are function slots like any other,
+;;;; and nothing is released at scope exit (spec-memory.md, "When storage is
+;;;; released"). So a let at the top level of a function body has exactly the
+;;;; function's extent, a defer written in it always registers, and it is as
+;;;; safe as one written at the top level. A let nested inside such a let has
+;;;; the same extent and the same permission.
+;;;;
+;;;; The numbers differ per failure, so a wrong answer names its own cause.
+
+(defvar order i64)
+
+(defn note [n i32] (set order (+ (* order 10) (i64 n))))
+
+;;; One defer in a let, with the acquisition it is paired with above it. This
+;;; is the shape the relaxation exists for: acquire, defer the release beside
+;;; it, then use it.
+(defn one [] i64
+  (let [a 1]
+    (defer (note a))
+    (note 9))
+  order)
+
+;;; Two resources in one let. This is the case a flag granted once per block
+;;; instead of once per form gets wrong: the first defer registers and the
+;;; second is refused.
+(defn two []
+  (let [a 1 b 2]
+    (defer (note a))
+    (defer (note b))
+    (note 9)))
+
+;;; A nested let still has the function's extent, so a defer in it registers
+;;; too — and it registers *later* than the outer one, so it runs first.
+(defn nested []
+  (let [a 1]
+    (defer (note a))
+    (let [b 2]
+      (defer (note b))
+      (note 9))))
+
+;;; Registration order is one order across the boundary: a defer at the top
+;;; level and a defer inside a let interleave by where they are written, not by
+;;; which construct they are in. Written 1, 2, 3; run 3, 2, 1.
+(defn mixed []
+  (defer (note 1))
+  (let [x 2]
+    (defer (note x))
+    (defer (note 3))
+    (note 9)))
+
+;;; An early return runs them too, and innermost-first, exactly as it does for
+;;; a defer at the top level. The let's binding is still live when the defer
+;;; reads it, because the slot is the function's.
+(defn early [] i64
+  (let [a 1]
+    (defer (note a))
+    (let [b 2]
+      (defer (note b))
+      (return 7)))
+  0)
+
+(defn main [] i32
+  (set order 0) (print (one)) (println "")    ; 9 — read before the defer runs
+  (set order 0) (one)   (print order) (println "")  ; 91
+  (set order 0) (two)   (print order) (println "")  ; 921 — reverse of writing
+  (set order 0) (nested)(print order) (println "")  ; 921
+  (set order 0) (mixed) (print order) (println "")  ; 9321
+  (set order 0) (print (early)) (println "")        ; 7
+  (print order) (println "")                        ; 21
+  0)

test/test_acceptance.ml

@@ -1497,6 +1497,40 @@ ERR@7 unexpected token: not the kind the caller was reading
     let cleanup_out = "7\n21\n42\n0\n5\n0\n0\n" in
     outputs "cleanup paths" "programs/cleanup.flan" cleanup_out;
     outputs ~opt:"-O0" "cleanup paths, -O0" "programs/cleanup.flan" cleanup_out;
+
+    (* defer-let.flan: a [let] at the top level of a function body has exactly
+       the function's extent, so a defer written in it always registers and is
+       as safe as one written at the top level. Six claims, and the numbers
+       differ per failure.
+
+       The one a plausible wrong version gets wrong is [two]: a permission
+       granted once around a block rather than once before each form lets the
+       first defer through and refuses the second, and every other case here
+       still passes. *)
+    let defer_let_out = "9\n91\n921\n921\n9321\n7\n21\n" in
+    outputs "defer in a let" "programs/defer-let.flan" defer_let_out;
+    outputs ~opt:"-O0" "defer in a let, -O0" "programs/defer-let.flan"
+      defer_let_out;
+    (* The two refusals that stay, each named by what blocks it. A loop body
+       would fire once at function exit rather than once per iteration, and a
+       branch would have to express "maybe registered", which a construct
+       copied into every exit path cannot. A [let] inside either one inherits
+       the refusal, not the permission: its extent is the loop's or the arm's. *)
+    refuses_src "defer in a loop body"
+      "(defn g [] 0)\n(defn f [] (while true (defer (g))))"
+      "not allowed inside a loop body";
+    refuses_src "defer in a dotimes body"
+      "(defn g [] 0)\n(defn f [] (dotimes [i 3] (defer (g))))"
+      "not allowed inside a loop body";
+    refuses_src "defer in a branch"
+      "(defn g [] 0)\n(defn f [] (if true (defer (g)) 0))"
+      "not allowed inside a branch";
+    refuses_src "defer in a let inside a loop"
+      "(defn g [] 0)\n(defn f [] (while true (let [x 1] (defer (g)))))"
+      "not allowed inside a loop body";
+    refuses_src "defer in a let inside a branch"
+      "(defn g [] 0)\n(defn f [] (if true (let [x 1] (defer (g))) 0))"
+      "not allowed inside a branch";
     let signed_out = "-4\n-1\nbig is not small\nbig is large\n1\n" in
     outputs "signedness" "programs/signedness.flan" signed_out;
     outputs ~opt:"-O0" "signedness, -O0" "programs/signedness.flan" signed_out;

test/test_flan.ml

@@ -691,11 +691,19 @@ let () =
     ~needle:"milestone 6";
   rejects_check "try is milestone 6"    "(defn f [] i32 (try 1))"
     ~needle:"milestone 6";
-  (* dotimes and defer are implemented; what is still rejected is a defer that
+  (* dotimes and defer are implemented, and a defer in a [let] is now one of
-     is not a top-level form, because it would run at function exit rather than
+     the places it may be written: a let at the top level of a function body
-     at the exit of the block it was written in. *)
+     has exactly the function's extent (see test/programs/defer-let.flan). What
-  rejects_check "defer must be top-level"
+     is still rejected is a loop body and a branch, because a defer is copied
-    "(defn f [] (let [x 1] (defer (g))))" ~needle:"top-level";
+     into every exit path — so a loop body's would fire once at function exit
+     rather than once per iteration, and a branch cannot say "maybe
+     registered". *)
+  rejects_check "defer is refused in a loop body"
+    "(defn g [] 0) (defn f [] (while true (defer (g))))"
+    ~needle:"a loop body";
+  rejects_check "defer is refused in a branch"
+    "(defn g [] 0) (defn f [] (if true (defer (g)) 0))"
+    ~needle:"a branch";
   (* An import is resolved by [Load] before the checker runs, so one that
      reaches [Check] means a driver skipped that step. *)
   rejects_check "an unresolved import is a driver bug"