Record Upgrade Instructions for Classical Persistence (#4676)

Also support `Prim.rts_upgrade_instructions()` on classical persistence, to inspect the instructions consumed by the last upgrade (including stabilization plus destabilization) for performance measurements. This functionality is already supported with enhanced orthogonal persistence.

## Mechanism
1. The pre-upgrade records its consumed instructions, in particular for stabilization, in the stable memory metadata. The location depends on the stable memory version.
2. The post-upgrade loads this information, if present, and adds its consumed instructions, in particular for destabilization.
3. `Prim.rts_upgrade_instructions()` returns the sum of those costs.

## Backwards Compatibility
The record of upgrade instructions is optional in the stable memory. This is to ensure backwards compatibility to older Motoko programs that do not record this information. When upgrading from such an older Motoko program,  `Prim.rts_upgrade_instructions()` returns `Nat64.maximumValue`.

## Notes
* `Prim.rts_upgrade_instructions()` currently returns >0 after installation because some upgrade check needs to be performed there too.
* This function is to be exposed to the Motoko base library together with the other diagnostic runtime information, since `Prim` is only intended for internal use.

design/OldStableMemory.md

@@ -191,13 +191,15 @@ NOTE: A program with no stable variables still writes an empty record value `v =
 (case !size == 0) // hence N = 0
   [0..3] StableVariable data len
   [4..4+len-1] StableVariable data
-  [4+len-1,..M-1] 0...0 // zero padding
+  [len+4..len+12-1] instruction costs of stabilization (optional)
+  [len+12,..M-1] 0...0 // zero padding
 (case !size > 0)
 [0..3]  0...0
 [4..N-1]  StableMemory bytes
 [N..N+3]  StableVariable data len
 [N+4..(N+4)+len-1] StableVariable data
-[(N+4)+len..M-13] 0...0 // zero padding
+[(N+4)+len..M-20] 0...0 // zero padding
+[M-20..M-13] instruction costs of stabilization (optional)
 [M-12..M-9] value N/64Ki = !size
 [M-8..M-5] saved StableMemory bytes
 [M-4..M-1]  version word

src/codegen/compile_classical.ml

@@ -5853,31 +5853,40 @@ module StableMem = struct
     read env false "word32" I32Type 4l load_unskewed_ptr
   let write_word32 env =
     write env false "word32" I32Type 4l store_unskewed_ptr
+  let write_word64 env =
+    write env false "word64" I64Type 8l (G.i (Store {ty = I64Type; align = 2; offset = 0L; sz = None}))
+
+  let read_and_clear env name typ bytes zero load store =
+    Func.share_code1 Func.Always env (Printf.sprintf "__stablemem_read_and_clear_%s" name)
+      ("offset", I64Type) [typ]
+      (fun env get_offset ->
+        let words = Int32.div (Int32.add bytes 3l) 4l in
+        Stack.with_words env "temp_ptr" words (fun get_temp_ptr ->
+          let (set_word, get_word, _) = new_local_ env typ "word" in
+          (* read *)
+          get_temp_ptr ^^ G.i (Convert (Wasm.Values.I64 I64Op.ExtendUI32)) ^^
+          get_offset ^^
+          compile_const_64 (Int64.of_int32 bytes) ^^
+          stable64_read env ^^
+          get_temp_ptr ^^ load ^^
+          set_word ^^
+          (* write 0 *)
+          get_temp_ptr ^^ zero ^^ store ^^
+          get_offset ^^
+          get_temp_ptr ^^ G.i (Convert (Wasm.Values.I64 I64Op.ExtendUI32)) ^^
+          compile_const_64 (Int64.of_int32 bytes) ^^
+          stable64_write env ^^
+          (* return word *)
+          get_word
+      ))
 
-
-  (* read and clear word32 from stable mem offset on stack *)
   let read_and_clear_word32 env =
-      Func.share_code1 Func.Always env "__stablemem_read_and_clear_word32"
-        ("offset", I64Type) [I32Type]
-        (fun env get_offset ->
-          Stack.with_words env "temp_ptr" 1l (fun get_temp_ptr ->
-            let (set_word, get_word) = new_local env "word" in
-            (* read word *)
-            get_temp_ptr ^^ G.i (Convert (Wasm.Values.I64 I64Op.ExtendUI32)) ^^
-            get_offset ^^
-            compile_const_64 4L ^^
-            stable64_read env ^^
-            get_temp_ptr ^^ load_unskewed_ptr ^^
-            set_word ^^
-            (* write 0 *)
-            get_temp_ptr ^^ compile_unboxed_const 0l ^^ store_unskewed_ptr ^^
-            get_offset ^^
-            get_temp_ptr ^^ G.i (Convert (Wasm.Values.I64 I64Op.ExtendUI32)) ^^
-            compile_const_64 4L ^^
-            stable64_write env ^^
-            (* return word *)
-            get_word
-        ))
+      read_and_clear env "word32" I32Type 4l (compile_unboxed_const 0l) 
+      load_unskewed_ptr store_unskewed_ptr
+  let read_and_clear_word64 env =
+    read_and_clear env "word64" I64Type 8l (compile_const_64 0L)
+      (G.i (Load {ty = I64Type; align = 2; offset = 0L; sz = None}))
+      (G.i (Store {ty = I64Type; align = 2; offset = 0L; sz = None}))
 
   (* ensure_pages : ensure at least num pages allocated,
      growing (real) stable memory if needed *)
@@ -6229,6 +6238,16 @@ module StableMemoryInterface = struct
 
 end
 
+module UpgradeStatistics = struct
+  let register_globals env =
+    E.add_global64 env "__upgrade_instructions" Mutable 0L
+
+  let get_upgrade_instructions env =
+    G.i (GlobalGet (nr (E.get_global env "__upgrade_instructions")))
+  let set_upgrade_instructions env =
+    G.i (GlobalSet (nr (E.get_global env "__upgrade_instructions")))
+end
+
 module RTS_Exports = struct
   (* Must be called late, after main codegen, to ensure correct generation of
      of functioning or unused-but-trapping stable memory exports (as required)
@@ -8574,10 +8593,10 @@ module Stabilization = struct
 
         (* Case-true: Stable variables only --
            no use of either regions or experimental API. *)
-        (* ensure [0,..,3,...len+4) *)
+        (* ensure [0,..,3,...len+4, .. len+4+8 ) *)
         compile_const_64 0L ^^
         extend64 get_len ^^
-        compile_add64_const 4L ^^  (* reserve one word for size *)
+        compile_add64_const 12L ^^  (* reserve 4 bytes for size, and 8 bytes for upgrade instructions *)
         StableMem.ensure env ^^
 
         (* write len to initial word of stable memory*)
@@ -8592,7 +8611,12 @@ module Stabilization = struct
             extend64 get_dst ^^
             extend64 get_len ^^
             StableMem.stable64_write env
-          end
+          end ^^
+
+          (* store stabilization instructions at len + 4 *)
+          extend64 get_len ^^ compile_add64_const 4L ^^
+          GC.instruction_counter env ^^
+          StableMem.write_word64 env
       end
       begin
         (* Case-false: Either regions or experimental API. *)
@@ -8604,11 +8628,11 @@ module Stabilization = struct
         set_N ^^
 
         (* grow mem to page including address
-           N + 4 + len + 4 + 4 + 4 = N + len + 16
+           N + 4 + len + 4 + 4 + 4 + 8 = N + len + 24
         *)
         get_N ^^
         extend64 get_len ^^
-        compile_add64_const 16L ^^
+        compile_add64_const 24L ^^
         StableMem.ensure env ^^
 
         get_N ^^
@@ -8633,6 +8657,12 @@ module Stabilization = struct
         compile_shl64_const (Int64.of_int page_size_bits) ^^
         set_M ^^
 
+        (* store stabilization instructions at M + (pagesize - 20) *)
+        get_M ^^
+        compile_add64_const (Int64.sub page_size64 20L) ^^
+        GC.instruction_counter env ^^
+        StableMem.write_word64 env ^^
+
         (* store mem_size at M + (pagesize - 12) *)
         get_M ^^
         compile_add64_const (Int64.sub page_size64 12L) ^^
@@ -8672,6 +8702,19 @@ module Stabilization = struct
   let destabilize env ty save_version =
     match E.mode env with
     | Flags.ICMode | Flags.RefMode ->
+      let (set_instructions, get_instructions) = new_local64 env "instructions" in
+      let handle_missing_instructions = 
+        get_instructions ^^
+        compile_eq64_const 0L ^^
+        (G.if0
+        begin
+          (* Default to -1 if no upgrade instructions were recorded, i.e. 
+             because the record space was lacking or was zero padding. *)
+          compile_const_64 (-1L) ^^
+          set_instructions
+        end
+        G.nop) in
+      compile_const_64 0L ^^ set_instructions ^^
       let (set_pages, get_pages) = new_local64 env "pages" in
       StableMem.stable64_size env ^^
       set_pages ^^
@@ -8759,16 +8802,37 @@ module Stabilization = struct
               (* set offset *)
               get_N ^^
               compile_add64_const 4L ^^
-              set_offset
+              set_offset ^^
+
+              (* Backwards compatibility: Check if upgrade instructions have space in the last page. *)
+              get_offset ^^ extend64 get_len ^^ G.i (Binary (Wasm.Values.I64 I64Op.Add)) ^^
+              get_M ^^ compile_add64_const (Int64.sub page_size64 20L) ^^
+              G.i (Compare (Wasm.Values.I64 I64Op.LeU)) ^^
+              (G.if0
+              begin
+                  (* Load stabilization instructions if defined, otherwise zero padding. *)
+                  get_M ^^
+                  compile_add64_const (Int64.sub page_size64 20L) ^^
+                  StableMem.read_and_clear_word64 env ^^
+                  set_instructions
+              end
+              G.nop) ^^
+              handle_missing_instructions
             end
             begin
               (* Sub-Case: Version 0.
                  Stable vars with NO Regions/Experimental API. *)
               (* assert mem_size == 0 *)
+              let (set_M, get_M) = new_local64 env "M" in
+
               StableMem.get_mem_size env ^^
               G.i (Test (Wasm.Values.I64 I64Op.Eqz)) ^^
               E.else_trap_with env "unexpected, non-zero stable memory size" ^^
 
+              StableMem.stable64_size env ^^
+              compile_shl64_const (Int64.of_int page_size_bits) ^^
+              set_M ^^
+
               (* set len *)
               get_marker ^^
               set_len ^^
@@ -8778,7 +8842,21 @@ module Stabilization = struct
               set_offset ^^
 
               compile_unboxed_const (Int32.of_int 0) ^^
-              save_version
+              save_version ^^
+
+              (* Backwards compatibility: Check if stabilization instructions have space in the last page. *)
+              get_offset ^^ extend64 get_len ^^ G.i (Binary (Wasm.Values.I64 I64Op.Add)) ^^
+              get_M ^^ compile_sub64_const 8L ^^
+              G.i (Compare (Wasm.Values.I64 I64Op.LeU)) ^^
+              (G.if0
+              begin
+                  (* Load stabilization instructions if defined, otherwise zero padding *)
+                  get_offset ^^ extend64 get_len ^^ G.i (Binary (Wasm.Values.I64 I64Op.Add)) ^^
+                  StableMem.read_and_clear_word64 env ^^
+                  set_instructions
+              end
+              G.nop) ^^
+              handle_missing_instructions
             end ^^ (* if_ *)
 
           let (set_blob, get_blob) = new_local env "blob" in
@@ -8808,7 +8886,20 @@ module Stabilization = struct
 
           (* return val *)
           get_val
-        end
+        end ^^
+        (* Record the total upgrade instructions if defined. 
+           If stabilization costs were missing due to upgrades from old Motoko programs,
+           the costs are defaulted to 0xFFFF_FFFF_FFFF_FFFF. *)
+        get_instructions ^^
+        compile_eq64_const (-1L) ^^
+        (G.if1 I64Type
+          get_instructions
+          begin
+            get_instructions ^^
+            GC.instruction_counter env ^^
+            G.i (Binary (Wasm.Values.I64 I64Op.Add))
+          end) ^^
+          UpgradeStatistics.set_upgrade_instructions env
     | _ -> assert false
 end
 
@@ -11505,6 +11596,10 @@ and compile_prim_invocation (env : E.t) ae p es at =
     SR.Vanilla,
     GC.get_collector_instructions env ^^ BigNum.from_word64 env
 
+  | OtherPrim "rts_upgrade_instructions", [] ->
+    SR.Vanilla,
+    UpgradeStatistics.get_upgrade_instructions env ^^ BigNum.from_word64 env
+
   | OtherPrim "rts_stable_memory_size", [] ->
     SR.Vanilla,
     StableMem.stable64_size env ^^ BigNum.from_word64 env
@@ -13061,6 +13156,7 @@ let compile mode rts (prog : Ir.prog) : Wasm_exts.CustomModule.extended_module =
   StableMem.register_globals env;
   Serialization.Registers.register_globals env;
   Serialization.Registers.define_idl_limit_check env;
+  UpgradeStatistics.register_globals env;
 
   (* See Note [Candid subtype checks] *)
   let set_serialization_globals = Serialization.register_delayed_globals env in

test/run-drun/migrate-upgrade-instr-stable-mem.drun

@@ -0,0 +1,5 @@
+# CLASSICAL-PERSISTENCE-ONLY
+# SKIP ic-ref-run
+install $ID migrate-upgrade-instr-stable-mem/old.wasm ""
+upgrade $ID migrate-upgrade-instr-stable-mem/test.mo ""
+ingress $ID test "DIDL\x00\x00"

test/run-drun/migrate-upgrade-instr-stable-mem/note.txt

@@ -0,0 +1 @@
+`old.wasm` has been generated by `moc -o old.wasm test.mo` using classical persistence without `rts_upgrade_instructions()` support (e.g. commit 247aa05).

test/run-drun/migrate-upgrade-instr-stable-mem/test.mo

@@ -0,0 +1,13 @@
+import Prim "mo:prim";
+
+import StableMemory "../stable-mem/StableMemory";
+
+
+actor {
+    ignore StableMemory.grow(10);
+    Prim.debugPrint("Stable memory size: " # debug_show(StableMemory.size()));
+
+    public func test() : async () {
+        Prim.debugPrint("Upgrade instructions: " # debug_show (Prim.rts_upgrade_instructions()));
+    };
+};

test/run-drun/migrate-upgrade-instr-stable-region.drun

@@ -0,0 +1,5 @@
+# CLASSICAL-PERSISTENCE-ONLY
+# SKIP ic-ref-run
+install $ID migrate-upgrade-instr-stable-region/old.wasm ""
+upgrade $ID migrate-upgrade-instr-stable-region/test.mo ""
+ingress $ID test "DIDL\x00\x00"

test/run-drun/migrate-upgrade-instr-stable-region/note.txt

@@ -0,0 +1 @@
+`old.wasm` has been generated by `moc -o old.wasm test.mo` using classical persistence without `rts_upgrade_instructions()` support (e.g. commit 247aa05).

test/run-drun/migrate-upgrade-instr-stable-region/test.mo

@@ -0,0 +1,12 @@
+import Prim "mo:prim";
+import Region "../stable-region/Region";
+
+actor {
+    stable var region = Region.new();
+    ignore Region.grow(region, 1);
+    Prim.debugPrint("Region size: " # debug_show(Region.size(region)));
+
+    public func test() : async () {
+        Prim.debugPrint("Upgrade instructions: " # debug_show (Prim.rts_upgrade_instructions()));
+    };
+};

test/run-drun/migrate-upgrade-instr-stable-var.drun

@@ -0,0 +1,5 @@
+# CLASSICAL-PERSISTENCE-ONLY
+# SKIP ic-ref-run
+install $ID migrate-upgrade-instr-stable-var/old.wasm ""
+upgrade $ID migrate-upgrade-instr-stable-var/test.mo ""
+ingress $ID test "DIDL\x00\x00"

test/run-drun/migrate-upgrade-instr-stable-var/note.txt

@@ -0,0 +1 @@
+`old.wasm` has been generated by `moc -o old.wasm test.mo` using classical persistence without `rts_upgrade_instructions()` support (e.g. commit 247aa05).

test/run-drun/migrate-upgrade-instr-stable-var/test.mo

@@ -0,0 +1,15 @@
+import Prim "mo:prim";
+
+actor {
+    stable let state = {
+        var number = 0;
+        var text = "Test";
+    };
+    state.number += 1;
+    state.text #= "Test";
+    Prim.debugPrint(debug_show(state));
+
+    public func test() : async () {
+        Prim.debugPrint("Upgrade instructions: " # debug_show (Prim.rts_upgrade_instructions()));
+    };
+};

test/run-drun/ok/migrate-upgrade-instr-stable-mem.drun.ok

@@ -0,0 +1,7 @@
+ingress Completed: Reply: 0x4449444c016c01b3c4b1f204680100010a00000000000000000101
+debug.print: Stable memory size: 10
+ingress Completed: Reply: 0x4449444c0000
+debug.print: Stable memory size: 20
+ingress Completed: Reply: 0x4449444c0000
+debug.print: Upgrade instructions: 18_446_744_073_709_551_615
+ingress Completed: Reply: 0x4449444c0000

test/run-drun/ok/migrate-upgrade-instr-stable-region.drun.ok

@@ -0,0 +1,7 @@
+ingress Completed: Reply: 0x4449444c016c01b3c4b1f204680100010a00000000000000000101
+debug.print: Region size: 1
+ingress Completed: Reply: 0x4449444c0000
+debug.print: Region size: 2
+ingress Completed: Reply: 0x4449444c0000
+debug.print: Upgrade instructions: 18_446_744_073_709_551_615
+ingress Completed: Reply: 0x4449444c0000

test/run-drun/ok/migrate-upgrade-instr-stable-var.drun.ok

@@ -0,0 +1,7 @@
+ingress Completed: Reply: 0x4449444c016c01b3c4b1f204680100010a00000000000000000101
+debug.print: {number = 1; text = "TestTest"}
+ingress Completed: Reply: 0x4449444c0000
+debug.print: {number = 2; text = "TestTestTest"}
+ingress Completed: Reply: 0x4449444c0000
+debug.print: Upgrade instructions: 18_446_744_073_709_551_615
+ingress Completed: Reply: 0x4449444c0000

test/run-drun/upgrade-instructions.mo

@@ -1,4 +1,3 @@
-//ENHANCED-ORTHOGONAL-PERSISTENCE-ONLY
 import Prim "mo:prim";
 
 actor {
@@ -11,6 +10,8 @@ actor {
     lastInstructions := Prim.rts_upgrade_instructions();
     assert(lastInstructions > 0);
   };
+  // Assert that the upgrade instructions are defined, i.e. no backwards compatibility to old Motoko programs.
+  assert(Prim.rts_upgrade_instructions() != 18_446_744_073_709_551_615);
   Prim.debugPrint("Ignore Diff: Upgrade instructions: " # debug_show (Prim.rts_upgrade_instructions()));
 
   public func increase() : async () {