Add copy annot for captured values to closure

lib/ast.ml

@@ -43,13 +43,14 @@ and decl = {
 }
 
 and decl_attr = Dmut | Dmove | Dnorm | Dset
+and func_attr = Fa_single of ident | Fa_param of ident * ident list
 
 and func = {
   name : ident;
   params : decl list;
   return_annot : type_spec option;
   body : block;
-  attr : ident option;
+  attr : func_attr list;
 }
 
 and argument = { apass : decl_attr; aloc : loc; aexpr : expr }
@@ -62,7 +63,7 @@ and expr =
   | Unop of loc * unop * expr
   | If of loc * expr * expr * expr option
   | Let_e of loc * decl * passed_expr * expr
-  | Lambda of loc * decl list * block
+  | Lambda of loc * decl list * func_attr list * block
   | App of loc * expr * argument list
   | Record of loc * (string * expr) list
   | Tuple of loc * expr list

lib/cleaned_types.ml

@@ -18,7 +18,14 @@ and fun_kind = Simple | Closure of closed list
 and param = { pt : typ; pmut : bool; pmoved : bool }
 and field = { ftyp : typ; mut : bool }
 and ctor = { cname : string; ctyp : typ option; index : int }
-and closed = { clname : string; clmut : bool; cltyp : typ; clparam : bool }
+
+and closed = {
+  clname : string;
+  clmut : bool;
+  cltyp : typ;
+  clparam : bool;
+  clcopy : bool;
+}
 
 let is_type_polymorphic typ =
   let rec inner acc = function

lib/cleaned_types.mli

@@ -18,7 +18,14 @@ and fun_kind = Simple | Closure of closed list
 and param = { pt : typ; pmut : bool; pmoved : bool }
 and field = { ftyp : typ; mut : bool }
 and ctor = { cname : string; ctyp : typ option; index : int }
-and closed = { clname : string; clmut : bool; cltyp : typ; clparam : bool }
+
+and closed = {
+  clname : string;
+  clmut : bool;
+  cltyp : typ;
+  clparam : bool;
+  clcopy : bool;
+}
 
 val is_type_polymorphic : typ -> bool
 val string_of_type : typ -> string

lib/monomorph_tree.ml

@@ -985,7 +985,13 @@ and cln_kind p = function
                 extract_callname modded_name p.vars (Mvar (cl.clname, Vnorm))
               else modded_name
             in
-            { clname; cltyp = typ; clmut = cl.clmut; clparam = cl.clparam })
+            {
+              clname;
+              cltyp = typ;
+              clmut = cl.clmut;
+              clparam = cl.clparam;
+              clcopy = cl.clcopy;
+            })
           vals
       in
       Closure vals

lib/parser.mly

@@ -307,11 +307,13 @@ param:
   | Ampersand { Dmut } | Exclamation { Dmove }
 
 %inline sexp_fun:
-  | Defn; name = ident; attr = option(attr); option(String_lit); params = maybe_bracks(list(param)); body = list(stmt)
+  | Defn; name = ident; attr = list(attr); option(String_lit);
+      params = maybe_bracks(list(param)); body = list(stmt)
     { ($loc, { name; params; return_annot = None; body; attr }) }
 
 %inline attr:
-  | kw = Keyword { $loc, kw }
+  | kw = Keyword { Fa_single ($loc, kw) }
+  | kw = Keyword; lst = nonempty_list(ident) { Fa_param (($loc(kw), kw), lst) }
 
 %inline sexp_rec:
   | Rec; fst = parens(sexp_fun); tl = nonempty_list(parens(sexp_fun)) { Rec ($loc, fst :: tl) }
@@ -400,8 +402,8 @@ sexp_cond:
   | else_ = option(parens(cond_else)) { [$loc, Lit($loc, Unit), else_] }
 
 %inline sexp_lambda:
-  | Fn; params = maybe_bracks(list(param)); body = list(stmt)
-    { Lambda ($loc, params, body) }
+  | Fn; attr = list(attr); params = maybe_bracks(list(param)); body = list(stmt)
+    { Lambda ($loc, params, attr, body) }
 
 %inline sexp_field_set:
   | Set; Ampersand; var = sexp_expr; Exclamation; value = sexp_expr
@@ -466,7 +468,7 @@ let with_loc(x) :=
     { Por ($loc, (head :: tail)) }
 
 %inline record_item_pattern:
-  | attr = attr; p = option(sexp_pattern) { attr, p }
+  | attr = Keyword; p = option(sexp_pattern) { ($loc(attr), attr), p }
 
 %inline ctor_pattern_item:
   | sexp_ctor; sexp_pattern { Pctor ($1, Some $2) }

lib/syntax_errors.messages

@@ -574,10 +574,14 @@ prog: Lpar Defn Kebab_id Wildcard
 
 <YOUR SYNTAX ERROR MESSAGE HERE>
 
-prog: Lpar Defn Kebab_id Keyword Wildcard
+prog: Lpar Fn Keyword Wildcard
 
 <YOUR SYNTAX ERROR MESSAGE HERE>
 
+prog: Lpar Fn Keyword Kebab_id Wildcard
+
+Expecting parameter list
+
 prog: Lpar Defn Kebab_id String_lit Wildcard
 
 <YOUR SYNTAX ERROR MESSAGE HERE>
@@ -758,6 +762,10 @@ prog: Lpar Fn Wildcard
 
 <YOUR SYNTAX ERROR MESSAGE HERE>
 
+prog: Lpar Fn Keyword String_lit
+
+Expecting parameter list
+
 prog: Lpar Fn Lbrack Rbrack Wildcard
 
 <YOUR SYNTAX ERROR MESSAGE HERE>

lib/typing/types.ml

@@ -41,6 +41,7 @@ and closed = {
   cltyp : typ;
   clparam : bool;
   clmname : Path.t option;
+  clcopy : bool; (* otherwise move *)
 }
 
 and dattr = Ast.decl_attr = Dmut | Dmove | Dnorm | Dset
@@ -223,3 +224,13 @@ let rec contains_allocation = function
       true
 
 let mut_of_pattr = function Dmut | Dset -> true | Dnorm | Dmove -> false
+
+let add_closure_copy clsd id =
+  let changed, clsd =
+    List.fold_left_map
+      (fun changed c ->
+        if String.equal c.clname id then (true, { c with clcopy = true })
+        else (changed, c))
+      false clsd
+  in
+  if changed then Some clsd else None

lib/typing/types.mli

@@ -33,6 +33,7 @@ and closed = {
   cltyp : typ;
   clparam : bool;
   clmname : Path.t option;
+  clcopy : bool; (* otherwise move *)
 }
 
 val clean : typ -> typ
@@ -52,3 +53,4 @@ val is_weak : sub:Sset.t -> typ -> bool
 val extract_name_path : typ -> Path.t option
 val contains_allocation : typ -> bool
 val mut_of_pattr : Ast.decl_attr -> bool
+val add_closure_copy : closed list -> string -> closed list option

lib/typing/typing.ml

@@ -619,7 +619,7 @@ end = struct
     | Lit (loc, Unit) ->
         let attr = { no_attr with const = true } in
         { typ = Tunit; expr = Const Unit; attr; loc }
-    | Lambda (loc, id, e) -> convert_lambda env loc id e
+    | Lambda (loc, id, attr, e) -> convert_lambda env loc id attr e
     | Let_e (loc, decl, expr, cont) -> convert_let_e env loc decl expr cont
     | App (loc, e1, e2) -> convert_app ~switch_uni:false env loc e1 e2
     | Bop (loc, bop, es) -> convert_bop env loc bop es
@@ -725,7 +725,7 @@ end = struct
     let expr = Let { id; uniq; rmut; pass = expr.pattr; rhs; cont } in
     { typ = cont.typ; expr; attr = cont.attr; loc }
 
-  and convert_lambda env loc params body =
+  and convert_lambda env loc params attr body =
     let env = Env.open_function env in
     enter_level ();
     let env, params_t, qparams, ret_annot =
@@ -753,6 +753,28 @@ end = struct
         touched body
     in
 
+    (* Copied from function below *)
+    let closed_vars =
+      List.fold_left
+        (fun clsd -> function
+          | Ast.Fa_single (loc, attr) ->
+              raise (Error (loc, "Unknown attribute: " ^ attr))
+          | Fa_param ((_, "copy"), lst) ->
+              List.fold_left
+                (fun clsd (loc, id) ->
+                  match add_closure_copy clsd id with
+                  | Some c -> c
+                  | None ->
+                      let msg =
+                        "Value " ^ id ^ " is not captured, cannot copy"
+                      in
+                      raise (Error (loc, msg)))
+                clsd lst
+          | Fa_param ((loc, attr), _) ->
+              raise (Error (loc, "Unknown attribute: " ^ attr)))
+        closed_vars attr
+    in
+
     let kind = match closed_vars with [] -> Simple | lst -> Closure lst in
     check_unused env unused;
 
@@ -782,13 +804,6 @@ end = struct
        and use it in the function body *)
     let unique = uniq_name name in
 
-    let inline =
-      match attr with
-      | Some (_, "inline") -> true
-      | Some (loc, attr) -> raise (Error (loc, "Unknown attribute: " ^ attr))
-      | None -> false
-    in
-
     enter_level ();
     let env =
       if inrec then
@@ -834,6 +849,29 @@ end = struct
         touched body
     in
 
+    let inline, closed_vars =
+      List.fold_left
+        (fun (inl, clsd) -> function
+          | Ast.Fa_single (_, "inline") -> (true, clsd)
+          | Fa_single (loc, attr) ->
+              raise (Error (loc, "Unknown attribute: " ^ attr))
+          | Fa_param ((_, "copy"), lst) ->
+              ( inl,
+                List.fold_left
+                  (fun clsd (loc, id) ->
+                    match add_closure_copy clsd id with
+                    | Some c -> c
+                    | None ->
+                        let msg =
+                          "Value " ^ id ^ " is not captured, cannot copy"
+                        in
+                        raise (Error (loc, msg)))
+                  clsd lst )
+          | Fa_param ((loc, attr), _) ->
+              raise (Error (loc, "Unknown attribute: " ^ attr)))
+        (false, closed_vars) attr
+    in
+
     let kind = match closed_vars with [] -> Simple | lst -> Closure lst in
     check_unused env unused;
 

test/typing.ml

@@ -78,6 +78,14 @@ let test_func_recursive_if () =
 let test_func_generic_return () =
   test "int" "(defn apply [f x] (f x)) (defn add1 [x] (+ x 1)) (apply add1 1)"
 
+let test_func_capture_annot () =
+  test "unit"
+    "(external somefn (fun unit int)) (defn wrapper () (def a (somefn)) (defn \
+     captured :copy a () (+ 1 a)) ()) ()"
+
+let test_func_capture_annot_wrong () =
+  test_exn "Value a is not captured, cannot copy" "(defn somefn :copy a () ())"
+
 let test_record_clear () = test "t" "(type t { :x int :y int }) { :x 2 :y 2 }"
 
 let test_record_false () =
@@ -1031,6 +1039,8 @@ let () =
           case "1st_stay_gen" test_func_1st_stay_general;
           case "recursive_if" test_func_recursive_if;
           case "generic_return" test_func_generic_return;
+          case "capture annot" test_func_capture_annot;
+          case "capture annot wrong" test_func_capture_annot_wrong;
         ] );
       ( "records",
         [