Backend: Compile captured actor class parameters statically (#2022)

This fixes #1847.

This went wrong before:

  * All top-level functions must be static (i.e not represented by a dynamic
    closure)
  * This means that everything they may close over, even values computed in the
    top-level scope (effectively by the init function) must be in static memory,
    at a well-known location. For the top-level `decs`, this is done by the
    `AllocHow` analysis.
  * Top-level parameters were added later. They can _also_ be captured, but were
    previously only allocated as locals. This was #1847.
  * This did not manifest when these parameters (patterns) were “ɩ-expanded”,
    e.g. if we replace the parameter `p` with `p'` and `let p = p'` in the top-level
    scope. This was effectively Claudio’s work-around, that we remove here.

The fix is to choose the location of the variable of a parameter based on
whether it is captured anywhere. This is in a way a simplistic version of the
`AllocHow` analysis.

As a side-effect, this simplifies the code (no need to pass around `setters`,
we can use `Var.set_val` these days.

Author: nomeata

src/codegen/compile.ml

@@ -1233,6 +1233,16 @@ module MutBox = struct
   (* Mutable heap objects *)
 
   let field = Tagged.header_size
+
+  let alloc env =
+    Tagged.obj env Tagged.MutBox [ compile_unboxed_zero ]
+
+  let static env =
+    let tag = bytes_of_int32 (Tagged.int_of_tag Tagged.MutBox) in
+    let zero = bytes_of_int32 0l in
+    let ptr = E.add_mutable_static_bytes env (tag ^ zero) in
+    E.add_static_root env ptr;
+    ptr
 end
 
 
@@ -4966,14 +4976,21 @@ module VarEnv = struct
       (add_local_local env ae name i, i)
 
   (* Adds the names to the environment and returns a list of setters *)
-  let rec add_argument_locals env (ae : t) = function
-    | [] -> (ae, [])
+  let rec add_arguments env (ae : t) as_local = function
+    | [] -> ae
     | (name :: names) ->
-      let i = E.add_anon_local env I32Type in
-      E.add_local_name env i name;
-      let ae' = { ae with vars = NameEnv.add name (Local i) ae.vars } in
-      let (ae_final, setters) = add_argument_locals env ae' names
-      in (ae_final, G.i (LocalSet (nr i)) :: setters)
+      if as_local name then
+        let i = E.add_anon_local env I32Type in
+        E.add_local_name env i name;
+        let ae' = { ae with vars = NameEnv.add name (Local i) ae.vars } in
+        add_arguments env ae' as_local names
+      else (* needs to go to static memory *)
+        let ptr = MutBox.static env in
+        let ae' = add_local_heap_static ae name ptr in
+        add_arguments env ae' as_local names
+
+  let add_argument_locals env (ae : t) =
+    add_arguments env ae (fun _ -> true)
 
   let add_label (ae : t) name (d : G.depth) =
       { ae with labels = NameEnv.add name d ae.labels }
@@ -5149,9 +5166,9 @@ module FuncDec = struct
       (* Deserialize argument and add params to the environment *)
       let arg_names = List.map (fun a -> a.it) args in
       let arg_tys = List.map (fun a -> a.note) args in
-      let (ae1, setters) = VarEnv.add_argument_locals env ae0 arg_names in
+      let ae1 = VarEnv.add_argument_locals env ae0 arg_names in
       Serialization.deserialize env arg_tys ^^
-      G.concat (List.rev setters) ^^
+      G.concat_map (Var.set_val env ae1) (List.rev arg_names) ^^
       mk_body env ae1 ^^
       message_cleanup env sort
     ))
@@ -5640,15 +5657,10 @@ module AllocHow = struct
       (ae1, G.nop)
     | StoreHeap ->
       let (ae1, i) = VarEnv.add_local_with_offset env ae name 1l in
-      let alloc_code =
-        Tagged.obj env Tagged.MutBox [ compile_unboxed_zero ] ^^
-        G.i (LocalSet (nr i)) in
+      let alloc_code = MutBox.alloc env ^^ G.i (LocalSet (nr i)) in
       (ae1, alloc_code)
     | StoreStatic ->
-      let tag = bytes_of_int32 (Tagged.int_of_tag Tagged.MutBox) in
-      let zero = bytes_of_int32 0l in
-      let ptr = E.add_mutable_static_bytes env (tag ^ zero) in
-      E.add_static_root env ptr;
+      let ptr = MutBox.static env in
       let ae1 = VarEnv.add_local_heap_static ae name ptr in
       (ae1, G.nop)
 
@@ -7066,8 +7078,7 @@ and compile_exp (env : E.t) ae exp =
     let (ae1, i) = VarEnv.add_local_with_offset env ae name 1l in
     let sr, code = compile_exp env ae1 e in
     sr,
-    Tagged.obj env Tagged.MutBox [ compile_unboxed_zero ] ^^
-    G.i (LocalSet (nr i)) ^^
+    MutBox.alloc env ^^ G.i (LocalSet (nr i)) ^^
     code
   | DefineE (name, _, e) ->
     SR.unit,
@@ -7598,11 +7609,15 @@ and main_actor as_opt mod_env ds fs up =
   Func.define_built_in mod_env "init" [] [] (fun env ->
     let ae0 = VarEnv.empty_ae in
 
+    let captured = Freevars.captured_vars (Freevars.actor ds fs up) in
+
     (* Add any params to the environment *)
+    (* Captured ones need to go into static memory, the rest into locals *)
     let args = match as_opt with None -> [] | Some as_ -> as_ in
     let arg_names = List.map (fun a -> a.it) args in
     let arg_tys = List.map (fun a -> a.note) args in
-    let (ae1, setters) = VarEnv.add_argument_locals env ae0 arg_names in
+    let as_local n = not (Freevars.S.mem n captured) in
+    let ae1 = VarEnv.add_arguments env ae0 as_local arg_names in
 
     (* Reverse the fs, to a map from variable to exported name *)
     let v2en = E.NameEnv.from_list (List.map (fun f -> (f.it.var, f.it.name)) fs) in
@@ -7623,7 +7638,7 @@ and main_actor as_opt mod_env ds fs up =
     Dfinity.export_upgrade_methods env;
 
     (* Deserialize any arguments *)
-    (match as_opt with
+    begin match as_opt with
      | None
      | Some [] ->
        (* Liberally accept empty as well as unit argument *)
@@ -7632,7 +7647,8 @@ and main_actor as_opt mod_env ds fs up =
        G.if_ [] (Serialization.deserialize env arg_tys) G.nop
      | Some (_ :: _) ->
        Serialization.deserialize env arg_tys ^^
-       G.concat (List.rev setters)) ^^
+       G.concat_map (Var.set_val env ae1) (List.rev arg_names)
+    end ^^
     (* Continue with decls *)
     decls_codeW G.nop
   )

src/ir_def/check_ir.ml

@@ -840,7 +840,7 @@ and check_args env args =
       check env a.at (not (T.Env.mem a.it ve))
         "duplicate binding for %s in argument list" a.it;
       check_typ env a.note;
-      let val_info = {typ = a.note; const = false; loc_known = false} in
+      let val_info = {typ = a.note; const = false; loc_known = env.lvl = TopLvl } in
       let env' = T.Env.add a.it val_info ve in
       go env' as_
   in go T.Env.empty args

src/ir_def/freevars.ml

@@ -111,11 +111,13 @@ let rec exp e : f = match e.it with
   | DeclareE (i, t, e)  -> exp e  // i
   | DefineE (i, m, e)   -> id i ++ exp e
   | FuncE (x, s, c, tp, as_, t, e) -> under_lambda (exp e /// args as_)
-  | ActorE (ds, fs, u, _)  -> close (decs ds +++ fields fs +++ upgrade u)
+  | ActorE (ds, fs, u, _)  -> actor ds fs u
   | NewObjE (_, fs, _)  -> fields fs
   | TryE (e, cs)        -> exp e ++ cases cs
   | SelfCallE (_, e1, e2, e3) -> under_lambda (exp e1) ++ exp e2 ++ exp e3
 
+and actor ds fs u = close (decs ds +++ fields fs +++ upgrade u)
+
 and upgrade {pre; post} = under_lambda (exp pre) ++ under_lambda (exp post)
 
 and exps es : f = unions exp es

src/ir_passes/const.ml

@@ -120,8 +120,8 @@ type env = info M.t
 
 
 let no_info = { loc_known = false; const = surely_false }
-let arg env a = M.add a.it no_info env
-let args env as_ = List.fold_left arg env as_
+let arg lvl env a = M.add a.it { no_info with loc_known = lvl = TopLvl } env
+let args lvl env as_ = List.fold_left (arg lvl) env as_
 
 let rec pat env p = match p.it with
   | WildP
@@ -153,7 +153,7 @@ let rec exp lvl (env : env) e : lazy_bool =
     match e.it with
     | VarE v -> (find v env).const
     | FuncE (x, s, c, tp, as_ , ts, body) ->
-      exp_ NotTopLvl (args env as_) body;
+      exp_ NotTopLvl (args NotTopLvl env as_) body;
       begin match s, lvl with
       (* shared functions are not const for now *)
       | Type.Shared _, _ -> surely_false
@@ -271,7 +271,7 @@ and comp_unit = function
   | ActorU (as_opt, ds, fs, {pre; post}, typ) ->
     let env = match as_opt with
       | None -> M.empty
-      | Some as_ -> args M.empty as_
+      | Some as_ -> args TopLvl M.empty as_
     in
     let (env', _) = decs TopLvl env ds in
     exp_ TopLvl env' pre;

src/lowering/desugar.ml

@@ -452,11 +452,7 @@ and dec' at n d = match d with
     let id' = {id with note = ()} in
     let sort, _, _, _, _ = Type.as_func n.S.note_typ in
     let op = match sp.it with
-      | T.Local ->
-        if s.it = T.Actor then (* HACK: work around for issue #1847 (also below) *)
-          Some { it = S.WildP; at = no_region; note = T.ctxt }
-        else
-          None
+      | T.Local -> None
       | T.Shared (_, p) -> Some p in
     let inst = List.map
                  (fun tb ->
@@ -764,9 +760,7 @@ let transform_unit_body (u : S.comp_unit_body) : Ir.comp_unit =
   | S.ActorClassU (sp, typ_id, p, _, self_id, fields) ->
     let fun_typ = u.note.S.note_typ in
     let op = match sp.it with
-      | T.Local ->
-        (* HACK: work around for issue #1847 (also above) *)
-        Some { it = S.WildP; at = no_region; note = T.ctxt }
+      | T.Local -> None
       | T.Shared (_, p) -> Some p in
     let args, wrap, control, _n_res = to_args fun_typ op p in
     let obj_typ =

test/run-drun/issue-1847.mo

@@ -0,0 +1,5 @@
+actor class (t : Text) {
+
+  public func f () : async Text { t };
+
+};

test/run-drun/ok/issue-1847.drun-run.ok

@@ -0,0 +1,2 @@
+ingress Completed: Reply: 0x4449444c016c01b3c4b1f2046801000112010000000000000000000000000000000001
+ingress Err: IC0503: Canister 75hes-oqbaa-aaaaa-aaaaa-aaaaa-aaaaa-aaaaa-q trapped explicitly: IDL error: empty input

test/run-drun/ok/issue-1847.ic-ref-run.ok

@@ -0,0 +1,4 @@
+→ update create_canister()
+← completed: (record {1313628723 = service "cvccv-qqaaq-aaaaa-aaaaa-c"})
+→ update install_code(record {4849238 = blob ""; 265441191 = blob "\00asm\01\00\0…
+← rejected (RC_CANISTER_ERROR): Initialization trapped: EvalTrapError :0.1 "canister trapped explicitly: IDL error: empty input"