experiment: lazy migration for stable memory for regions (#4171)

* lazy migration; still need to reqwrite stable mem ops

* dynamically choose stablemem operations, based on current version - uses ifs consuming multiple values that trip up wasmopt (sigh)

* format region.rs

* update bench numbers

* update heap-32 test output

* workaround wasmopt bug

* remove dead code

* update test output

* refactor (use constants for versions)

* comment

* refactor stablememory api and prims to simplify; delete most of region0.rs

* Update test/run-drun/regions-pay-as-you-go.mo

* Update test/run-drun/stable-mem-pay-as-you-go.mo

* comment compile.ml/rename StableMemory->StableMemoryInterface

* remove dangerous stableMemoryRegion prim (that breaks aliasing)

* adjust test for removal of Regions.region0

* nuke region0.rs; change type of use_stable_regions arg to u32 (bool, really); change assert in get_region0 to debug assert for perf

rts/motoko-rts/src/gc/copying.rs

@@ -36,7 +36,7 @@ unsafe fn copying_gc<M: Memory>(mem: &mut M) {
         |hp| linear_memory::set_hp_unskewed(hp),
         ic::get_static_roots(),
         crate::continuation_table::continuation_table_loc(),
-        crate::region0::region0_get_ptr_loc(),
+        crate::region::region0_get_ptr_loc(),
         // note_live_size
         |live_size| ic::MAX_LIVE = ::core::cmp::max(ic::MAX_LIVE, live_size),
         // note_reclaimed

rts/motoko-rts/src/gc/generational.rs

@@ -48,7 +48,7 @@ unsafe fn generational_gc<M: Memory>(mem: &mut M) {
     let old_limits = get_limits();
     let roots = Roots {
         static_roots: ic::get_static_roots(),
-        region0_ptr_loc: crate::region0::region0_get_ptr_loc(),
+        region0_ptr_loc: crate::region::region0_get_ptr_loc(),
         continuation_table_ptr_loc: crate::continuation_table::continuation_table_loc(),
     };
     let heap = Heap {

rts/motoko-rts/src/gc/incremental/roots.rs

@@ -15,7 +15,7 @@ pub unsafe fn root_set() -> Roots {
     Roots {
         static_roots: ic::get_static_roots(),
         continuation_table_location: crate::continuation_table::continuation_table_loc(),
-        region0_ptr_location: crate::region0::region0_get_ptr_loc(),
+        region0_ptr_location: crate::region::region0_get_ptr_loc(),
     }
 }
 

rts/motoko-rts/src/gc/mark_compact.rs

@@ -51,7 +51,7 @@ unsafe fn compacting_gc<M: Memory>(mem: &mut M) {
         |hp| linear_memory::set_hp_unskewed(hp),
         ic::get_static_roots(),
         crate::continuation_table::continuation_table_loc(),
-        crate::region0::region0_get_ptr_loc(),
+        crate::region::region0_get_ptr_loc(),
         // note_live_size
         |live_size| ic::MAX_LIVE = ::core::cmp::max(ic::MAX_LIVE, live_size),
         // note_reclaimed

rts/motoko-rts/src/lib.rs

@@ -35,8 +35,6 @@ pub mod principal_id;
 pub mod region;
 //#[cfg(feature = "ic")]
 mod stable_mem;
-//#[cfg(feature = "ic")]
-pub mod region0;
 mod static_checks;
 pub mod stream;
 pub mod text;

rts/motoko-rts/src/region.rs

@@ -4,6 +4,12 @@ use crate::rts_trap_with;
 use crate::trap_with_prefix;
 use crate::types::{size_of, Blob, Bytes, Region, Value, TAG_REGION};
 
+// Versions
+// Should agree with constants StableMem.version_no_stable_memory etc. in compile.ml
+const VERSION_NO_STABLE_MEMORY: u32 = 0; // never manifest in serialized form
+const VERSION_SOME_STABLE_MEMORY: u32 = 1;
+const VERSION_REGIONS: u32 = 2;
+
 use motoko_rts_macros::ic_mem_fn;
 
 unsafe fn region_trap_with(msg: &str) -> ! {
@@ -418,8 +424,35 @@ unsafe fn region_reserve_id_span<M: Memory>(
     }
 }
 
+#[ic_mem_fn]
+pub unsafe fn region0_get<M: Memory>(_mem: &mut M) -> Value {
+    debug_assert_ne!(REGION_0, NO_REGION);
+    REGION_0
+}
+
+// Expose Region0 object to GC algorithms as root
+#[allow(dead_code)]
+#[cfg(feature = "ic")]
+pub(crate) unsafe fn region0_get_ptr_loc() -> *mut Value {
+    &mut REGION_0
+}
+
 #[ic_mem_fn]
 pub unsafe fn region_new<M: Memory>(mem: &mut M) -> Value {
+    match crate::stable_mem::get_version() {
+        VERSION_NO_STABLE_MEMORY => {
+            assert_eq!(crate::stable_mem::size(), 0);
+            region_migration_from_no_stable_memory(mem);
+        }
+        VERSION_SOME_STABLE_MEMORY => {
+            region_migration_from_some_stable_memory(mem);
+        }
+        VERSION_REGIONS => {}
+        _ => {
+            assert!(false);
+        }
+    };
+
     let next_id = meta_data::total_allocated_regions::get() as u16;
 
     if next_id == meta_data::max::REGIONS {
@@ -483,11 +516,12 @@ pub unsafe fn region_recover<M: Memory>(mem: &mut M, rid: &RegionId) -> Value {
     r_ptr
 }
 
-pub(crate) unsafe fn region_migration_from_v0_into_v2<M: Memory>(mem: &mut M) {
-    use crate::stable_mem::{grow, size, write};
+pub(crate) unsafe fn region_migration_from_no_stable_memory<M: Memory>(mem: &mut M) {
+    use crate::stable_mem::{get_version, grow, size, write};
     use meta_data::size::{PAGES_IN_BLOCK, PAGE_IN_BYTES};
 
-    assert!(size() == 0);
+    assert!(get_version() == VERSION_NO_STABLE_MEMORY);
+    assert_eq!(size(), 0);
 
     // pages required for meta_data (9/ 960KiB), much less than PAGES_IN_BLOCK (128/ 8MB) for a full block
     let meta_data_pages =
@@ -514,6 +548,8 @@ pub(crate) unsafe fn region_migration_from_v0_into_v2<M: Memory>(mem: &mut M) {
     // Write magic header
     write_magic();
 
+    crate::stable_mem::set_version(VERSION_REGIONS);
+
     // Region 0 -- classic API for stable memory, as a dedicated region.
     REGION_0 = region_new(mem);
 
@@ -527,7 +563,7 @@ pub(crate) unsafe fn region_migration_from_v0_into_v2<M: Memory>(mem: &mut M) {
 // region manager migration/initialization, with pre-existing stable data.
 // Case: Version 1 into version 2.
 //
-pub(crate) unsafe fn region_migration_from_v1_into_v2<M: Memory>(mem: &mut M) {
+pub(crate) unsafe fn region_migration_from_some_stable_memory<M: Memory>(mem: &mut M) {
     // Grow region0 to nearest block boundary, and add a block to fit a region manager meta data.
     //
     // Existing stable data becomes region 0, with first block relocated for region manager meta data.
@@ -611,6 +647,9 @@ pub(crate) unsafe fn region_migration_from_v1_into_v2<M: Memory>(mem: &mut M) {
             Some((RegionId(0), (i + 1) as u16)),
         );
     }
+
+    crate::stable_mem::set_version(VERSION_REGIONS);
+
     /* "Recover" the region data into a heap object. */
     REGION_0 = region_recover(mem, &RegionId(0));
 
@@ -623,7 +662,7 @@ pub(crate) unsafe fn region_migration_from_v1_into_v2<M: Memory>(mem: &mut M) {
 // region manager migration/initialization, with pre-existing stable data.
 // Case: Version 2 into version 2 ("Trivial migration" case).
 //
-pub(crate) unsafe fn region_migration_from_v2plus_into_v2<M: Memory>(mem: &mut M) {
+pub(crate) unsafe fn region_migration_from_regions_plus<M: Memory>(mem: &mut M) {
     use crate::stable_mem::{read, read_u32, size};
 
     // Check if the magic in the memory corresponds to this object.
@@ -648,18 +687,23 @@ pub(crate) unsafe fn region_migration_from_v2plus_into_v2<M: Memory>(mem: &mut M
 //
 // region manager migration/initialization, with pre-existing stable data.
 //
-// Valid cases:
-//  - from version 1 into version 2.
-//  - from version 2 into version 2.
-//
 #[ic_mem_fn]
-pub(crate) unsafe fn region_init<M: Memory>(mem: &mut M, from_version: i32) {
+pub(crate) unsafe fn region_init<M: Memory>(mem: &mut M, use_stable_regions: u32) {
     debug_assert!(meta_data::offset::FREE < meta_data::offset::BLOCK_ZERO);
-    match from_version {
-        0 => region_migration_from_v0_into_v2(mem),
-        1 => region_migration_from_v1_into_v2(mem),
-        2 => region_migration_from_v2plus_into_v2(mem),
-        _ => unreachable!(),
+    match crate::stable_mem::get_version() {
+        VERSION_NO_STABLE_MEMORY => {
+            assert!(crate::stable_mem::size() == 0);
+            if use_stable_regions != 0 {
+                region_migration_from_no_stable_memory(mem);
+            };
+        }
+        VERSION_SOME_STABLE_MEMORY => {
+            assert!(crate::stable_mem::size() > 0);
+            if use_stable_regions != 0 {
+                region_migration_from_some_stable_memory(mem);
+            };
+        }
+        _ => region_migration_from_regions_plus(mem), //check format & recover region0
     }
 }
 

rts/motoko-rts/src/stable_mem.rs

@@ -3,10 +3,20 @@ extern "C" {
     pub fn ic0_stable64_write(offset: u64, src: u64, size: u64);
     pub fn ic0_stable64_read(dst: u64, offset: u64, size: u64);
     // (virtual) stable_mem operations implemented by moc
+    pub fn moc_stable_mem_get_version() -> u32;
+    pub fn moc_stable_mem_set_version(version: u32);
     pub fn moc_stable_mem_size() -> u64;
     pub fn moc_stable_mem_grow(additional_pages: u64) -> u64;
 }
 
+pub fn get_version() -> u32 {
+    unsafe { moc_stable_mem_get_version() }
+}
+
+pub fn set_version(version: u32) {
+    unsafe { moc_stable_mem_set_version(version) }
+}
+
 pub fn size() -> u64 {
     // SAFETY: This is safe because of the ic0 api guarantees.
     unsafe { moc_stable_mem_size() }