From 8e2bd8a4a9ae54fc707943ced7ba625fc3295a11 Mon Sep 17 00:00:00 2001
From: Roman Kashitsyn <roman@dfinity.org>
Date: Thu, 10 Mar 2022 09:54:17 +0100
Subject: [PATCH] feat: add an impl of Memory that uses ic0 API

---
 src/ic0_api.rs | 35 +++++++++++++++++++++++++++++++----
 src/log.rs     | 12 +++++-------
 2 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/src/ic0_api.rs b/src/ic0_api.rs
index 679c6990..0ef149dc 100644
--- a/src/ic0_api.rs
+++ b/src/ic0_api.rs
@@ -1,7 +1,34 @@
+use crate::Memory;
+
 #[link(wasm_import_module = "ic0")]
 extern "C" {
-    pub fn stable_size() -> u32;
-    pub fn stable_grow(additional_pages: u32) -> i32;
-    pub fn stable_read(dst: u32, offset: u32, size: u32);
-    pub fn stable_write(offset: u32, src: u32, size: u32);
+    pub fn stable64_size() -> u64;
+    pub fn stable64_grow(additional_pages: u64) -> i64;
+    pub fn stable64_read(dst: u64, offset: u64, size: u64);
+    pub fn stable64_write(offset: u64, src: u64, size: u64);
+}
+
+#[derive(Clone, Copy)]
+pub struct Ic0StableMemory;
+
+impl Memory for Ic0StableMemory {
+    fn size(&self) -> u64 {
+        // SAFETY: This is safe because of the ic0 api guarantees.
+        unsafe { stable64_size() }
+    }
+
+    fn grow(&self, pages: u64) -> i64 {
+        // SAFETY: This is safe because of the ic0 api guarantees.
+        unsafe { stable64_grow(pages) }
+    }
+
+    fn read(&self, offset: u64, dst: &mut [u8]) {
+        // SAFETY: This is safe because of the ic0 api guarantees.
+        unsafe { stable64_read(dst.as_ptr() as u64, offset, dst.len() as u64) }
+    }
+
+    fn write(&self, offset: u64, src: &[u8]) {
+        // SAFETY: This is safe because of the ic0 api guarantees.
+        unsafe { stable64_write(offset, src.as_ptr() as u64, src.len() as u64) }
+    }
 }
diff --git a/src/log.rs b/src/log.rs
index 6ce90eec..ff64f2e2 100644
--- a/src/log.rs
+++ b/src/log.rs
@@ -64,13 +64,11 @@ impl<M: Memory> StableLog<M> {
             )
         };
 
-        if memory.size() < 1 {
-            if memory.grow(1) == -1 {
-                return Err(AllocError::GrowFailed {
-                    current: 0,
-                    delta: 1,
-                });
-            }
+        if memory.size() < 1 && memory.grow(1) == -1 {
+            return Err(AllocError::GrowFailed {
+                current: 0,
+                delta: 1,
+            });
         }
         memory.write(0, header_slice);