flan/lib/closures.ml

243 lines
11 KiB
OCaml

(* 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 }