eop changes (WIP)

rts/motoko-rts/src/persistence.rs

@@ -206,6 +206,23 @@ 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_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 ^^
@@ -9981,12 +9987,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 +10024,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 *)

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

@@ -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/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/lowering/desugar.ml

@@ -550,18 +550,16 @@ and build_actor at ts exp_opt self_id es obj_typ =
   let state = fresh_var "state" (T.Mut (T.Opt ty)) in
   let get_state = fresh_var "getState" (T.Func(T.Local, T.Returns, [], [], [ty])) in
   let ds = List.map (fun mk_d -> mk_d get_state) mk_ds in
-  let migration = match exp_opt with
-    | None -> primE (I.ICStableRead ty) [] (* as before *)
+  let stable_type, migration = match exp_opt with
+    | None ->
+      I.{pre = ty; post = ty},
+      primE (I.ICStableRead ty) [] (* as before *)
     | Some exp0 ->
       let e = exp exp0 in
       let [@warning "-8"] (_s,_c, [], [dom], [rng]) = T.as_func (exp0.note.S.note_typ) in
       let [@warning "-8"] (T.Object, dom_fields) = T.as_obj dom in
       let [@warning "-8"] (T.Object, rng_fields) = T.as_obj rng in
-      ifE (primE (Ir.RelPrim (T.nat, Operator.EqOp)) [
-               primE (I.OtherPrim "rts_stable_memory_size") [];
-               natE Numerics.Nat.zero])
-        (primE (I.ICStableRead ty) [])
-        (let fields' =
+      let fields' =
            List.map
              (fun (i,t) ->
                 T.{lab = i; typ = T.Opt (T.as_immut t); src = T.empty_src})
@@ -572,11 +570,16 @@ and build_actor at ts exp_opt self_id es obj_typ =
                      | Some t -> None (* ignore overriden *)
                      | None -> Some (i, t) (* retain others *))
                    ids)) in
-         let ty' = T.Obj (T.Memory, List.sort T.compare_field fields') in
-         let v = fresh_var "v" ty' in
-         let v_dom = fresh_var "v_dom" dom in
-         let v_rng = fresh_var "v_rng" rng in
-         letE v (primE (I.ICStableRead ty') [])
+      let ty' = T.Obj (T.Memory, List.sort T.compare_field fields') in
+      let v = fresh_var "v" ty' in
+      let v_dom = fresh_var "v_dom" dom in
+      let v_rng = fresh_var "v_rng" rng in
+      I.{pre = ty'; post = ty},
+      ifE (primE (Ir.RelPrim (T.nat, Operator.EqOp)) [
+               primE (I.OtherPrim "rts_stable_memory_size") [];
+               natE Numerics.Nat.zero])
+        (primE (I.ICStableRead ty) [])
+        (letE v (primE (I.ICStableRead ty') [])
            (letE v_dom
               (objectE T.Object
                  (List.map (fun T.{lab=i;typ=t;_} ->
@@ -665,7 +668,7 @@ and build_actor at ts exp_opt self_id es obj_typ =
           | Some call -> call
           | None -> tupE []);
        stable_record = with_stable_vars (fun e -> e);
-       stable_type = ty;
+       stable_type = stable_type
      },
      obj_typ))