(* Where a capturing fn's environment lives: on the frame it was written in,
   or on the collector's heap. A module of its own, beneath both [Check] and
   the emitters, because the answer depends on the build: [Check] places
   precisely, and a dev build's emitters place again under the dev rule —
   see [place]. *)

(* The environment struct a capture built. [Session]'s layout guard exempts
   these; see there. *)
let is_env_struct n = String.length n >= 4 && String.sub n 0 4 = "env/"

(* ── Where a closure's environment lives ───────────────────────────────

   A capturing [fn] is checked with its copies on the frame it was written
   in: a slot holding the environment struct, and a [Closure] carrying the
   slot's address. That is right for a value that is only called, passed
   down and let-bound — the frame outlives every use — and it costs nothing
   the static side would notice. A value that may outlive the frame needs
   its copies somewhere the frame's end does not reclaim, and for those this
   pass rewrites the literal to carry the copies themselves; the backend
   then allocates the environment from the collector. Only a closure that
   escapes allocates: the static side does not pay for the dynamic side.

   **What escapes.** The analysis follows function values back to where they
   came from — a literal, a parameter, or a lifted body's copy of a captured
   value — and asks whether any of those reaches a position that outlives
   the frame: a [set], a [return] or a function's last form, an [Option], a
   fixed array, a struct or data type field, a pointer to the slot holding
   it, anything handed to the runtime (a push, a put, a box), a restart's
   arguments, and any argument of a call through a function value. A call
   to a named function passes the question to the callee's parameter, and a
   capture passes it to the lifted body's copy — or escapes outright when
   the capturing literal itself escapes. Everything only grows, so the pass
   runs to a fixed point over the whole program.

   A function value read out of storage — a field, an element, a case — has
   no source here and needs none: nothing puts a value in storage without
   going through one of the positions above, which already sent its literal
   to the collector. *)
type fsrc = Lit of string | Par of int | Env of int

(* Whether a [Closure]'s environment is a collector allocation: it carries
   its copies rather than the address of a frame slot holding them. *)
let heap_env (env : Tast.expr) =
  match env.Tast.ty with Types.Ptr _ -> false | _ -> true

(* [dev] is a dev build's rule. There a call to a named function goes through
   its cell, and a redefinition can replace the callee with a body that keeps
   the value — while the caller, which is not recompiled, still made it on its
   frame. So a closure handed to any named call escapes, whatever the callee
   does today. A release build asks the callee. *)
let place ~dev (fns : Tast.fn list) : Tast.fn list =
  let by_name = Hashtbl.create 64 in
  List.iter (fun (f : Tast.fn) -> Hashtbl.replace by_name f.Tast.name f) fns;
  let lits = Hashtbl.create 16 in              (* escaping literals *)
  let pars = Hashtbl.create 16 in              (* (fn, i) escaping params *)
  let envs = Hashtbl.create 16 in              (* (fn, i) escaping copies *)
  let changed = ref true in
  let mark tbl k =
    if not (Hashtbl.mem tbl k) then begin
      Hashtbl.replace tbl k ();
      changed := true
    end
  in
  let is_fn (t : Types.t) = match t with Types.Fn _ -> true | _ -> false in
  let pass (fn : Tast.fn) =
    let slot = Hashtbl.create 16 in
    let add s rs =
      let old = try Hashtbl.find slot s with Not_found -> [] in
      Hashtbl.replace slot s (List.sort_uniq compare (rs @ old))
    in
    List.iteri (fun i t -> if is_fn t then add i [ Par i ]) fn.Tast.params;
    (* The environment structs this function fills, by the slot they sit
       in: a literal's [Closure] and a handler frame name the slot. *)
    let makes = Hashtbl.create 8 in
    let escape rs =
      List.iter
        (function
          | Lit n -> mark lits n
          | Par i -> mark pars (fn.Tast.name, i)
          | Env i -> mark envs (fn.Tast.name, i))
        rs
    in
    let rec roots (e : Tast.expr) =
      if not (is_fn e.Tast.ty) then []
      else
        let tail body =
          match List.rev body with x :: _ -> roots x | [] -> []
        in
        match e.Tast.e with
        | Tast.Closure (Tast.Flanfn n, _) -> [ Lit n ]
        | Tast.Local s -> (try Hashtbl.find slot s with Not_found -> [])
        | Tast.If (_, a, b) -> roots a @ roots b
        | Tast.Do body | Tast.Let (_, body) | Tast.Handled (_, body)
        | Tast.WithAlloc (_, body) -> tail body
        | Tast.Match (_, arms) ->
          List.concat_map (fun (a : Tast.arm) -> tail a.Tast.abody) arms
        | Tast.RestartCase (cs, body) ->
          roots body
          @ List.concat_map (fun (c : Tast.rclause) -> tail c.Tast.rbody) cs
        | _ -> []
    in
    let deny es = List.iter (fun e -> escape (roots e)) es in
    (* A capture: each copy escapes when the literal does, or when the
       lifted body lets its copy escape. *)
    let captured (fields : Tast.expr list) outright lifted =
      List.iteri
        (fun j (v : Tast.expr) ->
           if outright || Hashtbl.mem envs (lifted, j) then escape (roots v))
        fields
    in
    let go (e : Tast.expr) =
      match e.Tast.e with
      | Tast.Let (bs, _) ->
        List.iter
          (fun (s, (v : Tast.expr)) ->
             (match v.Tast.e with
              | Tast.Make (n, es) when is_env_struct n -> Hashtbl.replace makes s es
              (* A lifted body's copy of what it captured. *)
              | Tast.Field
                  ({ Tast.e = Tast.Deref { Tast.e = Tast.Local es; _ }; _ }, i)
                when fn.Tast.fenv = Some es -> add s [ Env i ]
              | _ -> ());
             add s (roots v))
          bs
      | Tast.Set (_, v) | Tast.Return (Some v) | Tast.Some_ v -> deny [ v ]
      | Tast.Arr es | Tast.MakeCase (_, _, es)
      | Tast.InvokeRestart (_, _, es, _, _, _) -> deny es
      | Tast.Make (n, es) -> if not (is_env_struct n) then deny es
      | Tast.Addr (Tast.Plocal s) ->
        escape (try Hashtbl.find slot s with Not_found -> [])
      | Tast.Prim (Tast.Rt _, es) | Tast.Prim (Tast.AddrOf, es) -> deny es
      | Tast.CallPtr (_, es) -> deny es
      | Tast.Call (name, es) ->
        if Hashtbl.mem by_name name && not dev then
          List.iteri (fun i a -> if Hashtbl.mem pars (name, i) then deny [ a ]) es
        else deny es
      | Tast.Closure (Tast.Flanfn n, { Tast.e = Tast.Addr (Tast.Plocal s); _ }) ->
        (match Hashtbl.find_opt makes s with
         | Some fields -> captured fields (Hashtbl.mem lits n) n
         | None -> ())
      | Tast.Handled (hs, _) ->
        List.iter
          (fun (h : Tast.hframe) ->
             match h.Tast.henv with
             | Some { Tast.e = Tast.Addr (Tast.Plocal s); _ } ->
               (match Hashtbl.find_opt makes s with
                | Some fields -> captured fields false h.Tast.hfn
                | None -> ())
             | _ -> ())
          hs
      | _ -> ()
    in
    (* Twice over the body: an environment struct is bound around the form
       that names it, and a slot's sources are complete before a use of it
       elsewhere in a loop is asked about. *)
    for _ = 1 to 2 do
      List.iter (Tast.walk go) fn.Tast.body;
      List.iter (Tast.walk go) fn.Tast.fdefers
    done;
    if is_fn fn.Tast.ret then
      match List.rev fn.Tast.body with x :: _ -> escape (roots x) | [] -> ()
  in
  while !changed do
    changed := false;
    List.iter pass fns
  done;
  if Hashtbl.length lits = 0 then fns
  else begin
    let rec rw (e : Tast.expr) : Tast.expr =
      let r = rw and rs = List.map rw in
      let e' =
        match e.Tast.e with
        | Tast.Int _ | Tast.Float _ | Tast.Bool _ | Tast.Str _ | Tast.Unit
        | Tast.Zero _ | Tast.Uninit _ | Tast.Local _ | Tast.Global _
        | Tast.None_ | Tast.FnAddr _ | Tast.Break _ | Tast.Continue _ -> e.Tast.e
        | Tast.Fill (t, b) -> Tast.Fill (t, r b)
        | Tast.DeadBeef (t, b) -> Tast.DeadBeef (t, r b)
        | Tast.Prim (p, es) -> Tast.Prim (p, rs es)
        | Tast.Call (n, es) -> Tast.Call (n, rs es)
        | Tast.Do es -> Tast.Do (rs es)
        | Tast.Make (n, es) -> Tast.Make (n, rs es)
        | Tast.MakeCase (a, b, es) -> Tast.MakeCase (a, b, rs es)
        | Tast.Arr es -> Tast.Arr (rs es)
        | Tast.InvokeRestart (a, b, es, c, d, l) ->
          Tast.InvokeRestart (a, b, rs es, c, d, l)
        | Tast.CallPtr (c, es) -> Tast.CallPtr (r c, rs es)
        (* The rewrite itself: the store of the copies into this frame and
           the value carrying their address become the value carrying the
           copies, which the backend stores into a collector allocation. *)
        | Tast.Let
            ([ (_, make) ], [ { Tast.e = Tast.Closure ((Tast.Flanfn n as fr), _); _ } ])
          when Hashtbl.mem lits n ->
          Tast.Closure (fr, r make)
        | Tast.Let (bs, body) ->
          Tast.Let (List.map (fun (s, v) -> (s, r v)) bs, rs body)
        | Tast.If (a, b, c) -> Tast.If (r a, r b, r c)
        | Tast.While (c, body, latch) -> Tast.While (r c, rs body, rs latch)
        | Tast.Return v -> Tast.Return (Option.map r v)
        | Tast.Set (p, v) -> Tast.Set (rp p, r v)
        | Tast.Addr p -> Tast.Addr (rp p)
        | Tast.Field (t, i) -> Tast.Field (r t, i)
        | Tast.Deref t -> Tast.Deref (r t)
        | Tast.CaseField (t, c, i) -> Tast.CaseField (r t, c, i)
        | Tast.Some_ t -> Tast.Some_ (r t)
        | Tast.UnwrapSome t -> Tast.UnwrapSome (r t)
        | Tast.Signal (k, i, t) -> Tast.Signal (k, i, r t)
        | Tast.Closure (f, t) -> Tast.Closure (f, r t)
        | Tast.Thicken (n, t) -> Tast.Thicken (n, r t)
        | Tast.Match (sc, arms) ->
          Tast.Match
            (r sc,
             List.map (fun (a : Tast.arm) -> { a with Tast.abody = rs a.Tast.abody }) arms)
        | Tast.Handled (hs, body) ->
          Tast.Handled
            (List.map
               (fun (h : Tast.hframe) -> { h with Tast.henv = Option.map r h.Tast.henv })
               hs,
             rs body)
        | Tast.RestartCase (cs, body) ->
          Tast.RestartCase
            (List.map (fun (c : Tast.rclause) -> { c with Tast.rbody = rs c.Tast.rbody }) cs,
             r body)
        | Tast.WithAlloc (a, body) -> Tast.WithAlloc (r a, rs body)
      in
      if e' == e.Tast.e then e else { e with Tast.e = e' }
    and rp (p : Tast.place) : Tast.place =
      match p with
      | Tast.Plocal _ | Tast.Pglobal _ -> p
      | Tast.Pfield (t, i) -> Tast.Pfield (rw t, i)
      | Tast.Pderef t -> Tast.Pderef (rw t)
      | Tast.Pindex (t, idx) -> Tast.Pindex (rw t, List.map rw idx)
    in
    List.map
      (fun (f : Tast.fn) ->
         { f with Tast.body = List.map rw f.Tast.body;
                  fdefers = List.map rw f.Tast.fdefers })
      fns
  end


let dev_program (p : Tast.program) =
  { p with Tast.fns = place ~dev:true p.Tast.fns }