experiment: attach optional migration expression to actor (class) as actor [exp]?

doc/md/examples/grammar.txt

@@ -14,7 +14,7 @@
 
 <obj_sort> ::= 
     'object'
-    'persistent'? 'actor'
+    'persistent'? 'actor' ('[' <exp> ']')?
     'module'
 
 <query> ::= 

rts/motoko-rts/src/persistence.rs

@@ -206,6 +206,22 @@ pub unsafe fn register_stable_type<M: Memory>(
     (*metadata).stable_type.assign(mem, &new_type);
 }
 
+/// Register the stable actor type on canister initialization and upgrade.
+/// The type is stored in the persistent metadata memory for later retrieval on canister upgrades.
+/// The `new_type` value points to a blob encoding the new stable actor type.
+#[ic_mem_fn]
+pub unsafe fn assign_stable_type<M: Memory>(
+    mem: &mut M,
+    new_candid_data: Value,
+    new_type_offsets: Value,
+) {
+    assert_eq!(new_candid_data.tag(), TAG_BLOB_B);
+    assert_eq!(new_type_offsets.tag(), TAG_BLOB_B);
+    let new_type = TypeDescriptor::new(new_candid_data, new_type_offsets);
+    let metadata = PersistentMetadata::get();
+    (*metadata).stable_type.assign(mem, &new_type);
+}
+
 pub(crate) unsafe fn stable_type_descriptor() -> &'static mut TypeDescriptor {
     let metadata = PersistentMetadata::get();
     &mut (*metadata).stable_type

src/codegen/compile_classical.ml

@@ -11616,6 +11616,15 @@ and compile_prim_invocation (env : E.t) ae p es at =
     SR.Vanilla,
     StableMem.get_mem_size env ^^ BigNum.from_word64 env
 
+  | OtherPrim "rts_in_install", [] -> (* classical specific *)
+    assert (not !Flags.enhanced_orthogonal_persistence);
+    SR.Vanilla,
+    StableMem.stable64_size env ^^
+    G.i (Test (Wasm_exts.Values.I64 I64Op.Eqz)) ^^
+    G.if1 I32Type
+      (Bool.lit true)
+      (Bool.lit false)
+
   (* Regions *)
 
   | OtherPrim "regionNew", [] ->

src/codegen/compile_enhanced.ml

@@ -1131,6 +1131,7 @@ module RTS = struct
     E.add_func_import env "rts" "allocation_barrier" [I64Type] [I64Type];
     E.add_func_import env "rts" "running_gc" [] [I32Type];
     E.add_func_import env "rts" "register_stable_type" [I64Type; I64Type] [];
+    E.add_func_import env "rts" "assign_stable_type" [I64Type; I64Type] [];
     E.add_func_import env "rts" "load_stable_actor" [] [I64Type];
     E.add_func_import env "rts" "save_stable_actor" [I64Type] [];
     E.add_func_import env "rts" "free_stable_actor" [] [];
@@ -8694,6 +8695,10 @@ module EnhancedOrthogonalPersistence = struct
     create_type_descriptor env actor_type ^^
     E.call_import env "rts" "register_stable_type"
 
+  let assign_stable_type env actor_type =
+    create_type_descriptor env actor_type ^^
+    E.call_import env "rts" "assign_stable_type"
+
   let load_old_field env field get_old_actor =
     if field.Type.typ = Type.(Opt Any) then
       (* A stable variable may have been promoted to type `Any`: Therefore, drop its former content. *)
@@ -8733,6 +8738,7 @@ module EnhancedOrthogonalPersistence = struct
     free_stable_actor env
 
   let save env actor_type =
+    assign_stable_type env actor_type ^^
     IC.get_actor_to_persist env ^^
     save_stable_actor env ^^
     NewStableMemory.backup env ^^
@@ -8751,6 +8757,7 @@ module EnhancedOrthogonalPersistence = struct
 
   let initialize env actor_type =
     register_stable_type env actor_type
+
 end (* EnhancedOrthogonalPersistence *)
 
 (* As fallback when doing persistent memory layout changes. *)
@@ -9981,12 +9988,12 @@ module IncrementalGraphStabilization = struct
   let partial_destabilization_on_upgrade env actor_type =
     (* TODO: Verify that the post_upgrade hook cannot be directly called by the IC *)
     (* Garbage collection is disabled in `start_graph_destabilization` until destabilization has completed. *)
-    GraphCopyStabilization.start_graph_destabilization env actor_type ^^
+    GraphCopyStabilization.start_graph_destabilization env actor_type.Ir.pre ^^
     get_destabilized_actor env ^^
     compile_test I64Op.Eqz ^^
     E.if0
       begin
-        destabilization_increment env actor_type ^^
+        destabilization_increment env actor_type.Ir.pre ^^
         get_destabilized_actor env ^^
         (E.if0
           G.nop
@@ -10018,22 +10025,22 @@ module IncrementalGraphStabilization = struct
       })
     | _ -> ()
     end
-  
+
   let load env =
     get_destabilized_actor env ^^
     compile_test I64Op.Eqz ^^
     E.then_trap_with env "Destabilization is not yet completed: Call __motoko_destabilize_after_upgrade" ^^
     get_destabilized_actor env
     (* Upgrade costs are already record in RTS for graph-copy-based (de-)stabilization. *)
 
-  let define_methods env actor_type =
+  let define_methods env (actor_type : Ir.stable_actor_typ) =
     define_async_stabilization_reply_callback env;
     define_async_stabilization_reject_callback env;
     export_async_stabilization_method env;
-    export_stabilize_before_upgrade_method env actor_type;
+    export_stabilize_before_upgrade_method env actor_type.Ir.post;
     define_async_destabilization_reply_callback env;
     define_async_destabilization_reject_callback env;
-    export_async_destabilization_method env actor_type;
+    export_async_destabilization_method env actor_type.Ir.pre;
     export_destabilize_after_upgrade_method env;
 
 end (* IncrementalGraphStabilization *)
@@ -11678,6 +11685,15 @@ and compile_prim_invocation (env : E.t) ae p es at =
     SR.Vanilla,
     StableMem.get_mem_size env ^^ BigNum.from_word64 env
 
+  | OtherPrim "rts_in_install", [] -> (* EOP specific *)
+    assert (!Flags.enhanced_orthogonal_persistence);
+    SR.Vanilla,
+    EnhancedOrthogonalPersistence.load_stable_actor env ^^
+    compile_test I64Op.Eqz ^^
+    E.if1 I64Type
+      (Bool.lit true)
+      (Bool.lit false)
+
   (* Regions *)
 
   | OtherPrim "regionNew", [] ->
@@ -12252,9 +12268,11 @@ and compile_prim_invocation (env : E.t) ae p es at =
 
   | ICStableRead ty, [] ->
     SR.Vanilla,
+    (*    IC.compile_static_print env ("ICStableRead" ^ Type.string_of_typ ty) ^^ *)
     Persistence.load env ty
   | ICStableWrite ty, [] ->
     SR.unit,
+    (*    IC.compile_static_print env ("ICStableWrite" ^ Type.string_of_typ ty) ^^ *)
     Persistence.save env ty
 
   (* Cycles *)

src/docs/extract.ml

@@ -140,7 +140,7 @@ struct
           _;
         } -> (
         match rhs with
-        | Source.{ it = Syntax.ObjBlockE (sort, _, fields); _ } ->
+        | Source.{ it = Syntax.ObjBlockE (sort, _, _, fields); _ } ->
             let mk_field_xref xref = mk_xref (Xref.XClass (name, xref)) in
             Some
               ( mk_xref (Xref.XType name),
@@ -155,7 +155,7 @@ struct
         )
     | Source.{ it = Syntax.VarD ({ it = name; _ }, rhs); _ } -> (
         match rhs with
-        | Source.{ it = Syntax.ObjBlockE (sort, _, fields); _ } ->
+        | Source.{ it = Syntax.ObjBlockE (sort, _, _, fields); _ } ->
             let mk_field_xref xref = mk_xref (Xref.XClass (name, xref)) in
             Some
               ( mk_xref (Xref.XType name),
@@ -184,7 +184,15 @@ struct
         {
           it =
             Syntax.ClassD
-              (shared_pat, name, type_args, ctor, _, obj_sort, _, fields);
+              ( shared_pat,
+                exp_opt,
+                name,
+                type_args,
+                ctor,
+                _,
+                obj_sort,
+                _,
+                fields );
           _;
         } ->
         let mk_field_xref xref = mk_xref (Xref.XClass (name.it, xref)) in

src/docs/namespace.ml

@@ -32,7 +32,7 @@ let from_module =
         | Syntax.ExpD _ -> acc
         | Syntax.LetD
             ( { it = Syntax.VarP id; _ },
-              { it = Syntax.ObjBlockE (_, _, decs); _ },
+              { it = Syntax.ObjBlockE (_, _, _, decs); _ },
               _ ) ->
             let mk_nested x = mk_xref (Xref.XNested (id.it, x)) in
             {
@@ -69,7 +69,7 @@ let from_module =
                   (mk_xref (Xref.XValue id.it), None)
                   acc.values;
             }
-        | Syntax.ClassD (_, id, _, _, _, _, _, _) ->
+        | Syntax.ClassD (_, _, id, _, _, _, _, _, _) ->
             {
               acc with
               types = StringMap.add id.it (mk_xref (Xref.XType id.it)) acc.types;

src/gen-grammar/grammar.sed

@@ -12,7 +12,7 @@ s/<id>/ID/g
 /^<parse_module_header> ::=/,+2d
 /^<stab_field> ::=/,+2d
 /^<typ_dec> ::=/,+2d
-/^<parse_stab_sig> ::=/,+2d
+/^<parse_stab_sig> ::=/,+5d
 /.*PRIM.*/d
 /^<bl> ::=/,+2d
 /^<ob> ::=/,+2d

src/ir_def/arrange_ir.ml

@@ -43,7 +43,7 @@ and system { meta; preupgrade; postupgrade; heartbeat; timer; inspect; stable_re
       "Timer" $$ [exp timer];
       "Inspect" $$ [exp inspect];
       "StableRecord" $$ [exp stable_record];
-      "StableType" $$ [typ stable_type]
+      "StableType" $$ [typ stable_type.pre; typ stable_type.post]
     ]
 
 and lexp le = match le.it with

src/ir_def/check_ir.ml

@@ -241,7 +241,7 @@ let rec check_typ env typ : unit =
     if not (Lib.List.is_strictly_ordered T.compare_field fields) then
       error env no_region "variant type's fields are not distinct and sorted %s" (T.string_of_typ typ)
   | T.Mut typ ->
-    error env no_region "unexpected T.Mut"
+    error env no_region "unexpected T.Mut %s" (T.string_of_typ typ)
   | T.Typ c ->
     error env no_region "unexpected T.Typ"
 
@@ -834,7 +834,7 @@ let rec check_exp env (exp:Ir.exp) : unit =
     typ heartbeat <: T.unit;
     typ timer <: T.unit;
     typ inspect <: T.unit;
-    typ stable_record <: stable_type;
+    typ stable_record <: stable_type.post;
     check (T.is_obj t0) "bad annotation (object type expected)";
     let (s0, tfs0) = T.as_obj t0 in
     let val_tfs0 = List.filter (fun tf -> not (T.is_typ tf.T.typ)) tfs0 in
@@ -1184,7 +1184,7 @@ let check_comp_unit env = function
     typ heartbeat <: T.unit;
     typ timer <: T.unit;
     typ inspect <: T.unit;
-    typ stable_record <: stable_type;
+    typ stable_record <: stable_type.post;
     check (T.is_obj t0) "bad annotation (object type expected)";
     let (s0, tfs0) = T.as_obj t0 in
     let val_tfs0 = List.filter (fun tf -> not (T.is_typ tf.T.typ)) tfs0 in

src/ir_def/construct.ml

@@ -129,6 +129,7 @@ let primE prim es =
     | OtherPrim "rts_max_stack_size" -> T.nat
     | OtherPrim "rts_callback_table_count" -> T.nat
     | OtherPrim "rts_callback_table_size" -> T.nat
+    | OtherPrim "rts_in_install" -> T.bool
     | _ -> assert false (* implement more as needed *)
   in
   let eff = map_max_effs eff es in
@@ -790,8 +791,42 @@ let objE sort typ_flds flds =
   in
   go [] [] [] flds
 
+
 let recordE flds = objE T.Object [] flds
 
+let objectE sort flds (tfs : T.field list) =
+  let rec go ds fields = function
+    | [] ->
+      blockE
+        (List.rev ds)
+        (newObjE sort fields
+          (T.Obj (sort, List.sort T.compare_field tfs)))
+    | (lab, exp)::flds ->
+       let v, typ, ds =
+         match T.lookup_val_field_opt lab tfs with
+         | None -> assert false
+         | Some typ ->
+           if T.is_mut typ
+           then
+             let v = fresh_var lab typ in
+             v, typ, varD v exp :: ds
+           else
+             match exp.it with
+             | VarE (Const, v) ->
+                var v typ, typ, ds
+             | _ ->
+                let v = fresh_var lab typ in
+                v, typ, letD v exp :: ds
+       in
+       let field = {
+         it = {name = lab; var = id_of_var v};
+         at = no_region;
+         note = typ
+       } in
+      go ds (field::fields) flds
+  in
+  go [] [] flds
+
 let check_call_perform_status success mk_failure =
   ifE
     (callE

src/ir_def/construct.mli

@@ -147,6 +147,8 @@ val (-*-) : exp -> exp -> exp (* application *)
 
 val objE : obj_sort -> (lab * con) list -> (lab * exp) list -> exp
 
+val objectE : obj_sort -> (lab * exp) list -> field list -> exp
+
 (* Records *)
 
 val recordE : (lab * exp) list -> exp

src/ir_def/ir.ml

@@ -77,6 +77,8 @@ and exp' =
   | NewObjE of Type.obj_sort * field list * Type.typ     (* make an object *)
   | TryE of exp * case list * (id * Type.typ) option (* try/catch/cleanup *)
 
+and stable_actor_typ = { pre: Type.typ; post: Type.typ }
+
 and system = {
   meta : meta;
   (* TODO: use option expressions for (some or all of) these *)
@@ -86,7 +88,7 @@ and system = {
   timer : exp; (* TODO: use an option type: (Default of exp | UserDefined of exp) option *)
   inspect : exp;
   stable_record: exp;
-  stable_type: Type.typ;
+  stable_type: stable_actor_typ;
 }
 
 and candid = {
@@ -242,7 +244,7 @@ type actor_type = {
   transient_actor_type: Type.typ;
   (* record of stable actor fields used for persistence,
      the fields are without mutability distinctions *)
-  stable_actor_type: Type.typ
+  stable_actor_type: stable_actor_typ
 }
 
 (* Program *)

src/ir_passes/async.ml

@@ -451,7 +451,7 @@ let transform prog =
          timer = t_exp timer;
          inspect = t_exp inspect;
          stable_record = t_exp stable_record;
-         stable_type = t_typ stable_type;
+         stable_type = {pre = t_typ stable_type.pre; post = t_typ stable_type.post};
         },
         t_typ typ)
     | NewObjE (sort, ids, t) ->
@@ -532,7 +532,10 @@ let transform prog =
           timer = t_exp timer;
           inspect = t_exp inspect;
           stable_record = t_exp stable_record;
-          stable_type = t_typ stable_type;
+          stable_type = {
+            pre = t_typ stable_type.pre;
+            post = t_typ stable_type.post
+          }
         },
         t_typ t)
 

src/ir_passes/erase_typ_field.ml

@@ -135,7 +135,10 @@ let transform prog =
         timer = t_exp timer;
         inspect = t_exp inspect;
         stable_record = t_exp stable_record;
-        stable_type = t_typ stable_type;
+        stable_type = {
+            pre = t_typ stable_type.pre;
+            post = t_typ stable_type.post
+        }
        },
        t_typ typ)
 
@@ -220,7 +223,10 @@ let transform prog =
           timer = t_exp timer;
           inspect = t_exp inspect;
           stable_record = t_exp stable_record;
-          stable_type = t_typ stable_type;
+          stable_type = {
+            pre = t_typ stable_type.pre;
+            post = t_typ stable_type.post
+          }
         },
         t_typ t)
   and t_prog (cu, flavor) = (t_comp_unit cu, { flavor with has_typ_field = false } )

src/languageServer/declaration_index.ml

@@ -251,7 +251,7 @@ let populate_definitions (project_root : string) (libs : Syntax.lib list)
   let is_type_def dec_field =
     match dec_field.it.Syntax.dec.it with
     | Syntax.TypD (typ_id, _, _) -> Some typ_id
-    | Syntax.ClassD (_, typ_id, _, _, _, _, _, _) -> Some typ_id
+    | Syntax.ClassD (_, _, typ_id, _, _, _, _, _, _) -> Some typ_id
     | _ -> None
   in
   let extract_binders env (pat : Syntax.pat) = gather_pat env pat in