piecrust: distinguish 32bit and 64bit memories

Two different trees are created - one for 32-bit and another for 64-bit
- and ensure they're properly distinguished.

piecrust/src/instance.rs

@@ -13,7 +13,7 @@ use piecrust_uplink::{ContractId, Event, ARGBUF_LEN};
 use crate::contract::WrappedContract;
 use crate::imports::Imports;
 use crate::session::Session;
-use crate::store::{Memory, MAX_MEM_SIZE};
+use crate::store::Memory;
 use crate::Error;
 
 pub struct WrappedInstance {
@@ -169,7 +169,7 @@ impl WrappedInstance {
             _ => return Err(Error::InvalidArgumentBuffer),
         };
 
-        if arg_buf_ofs + ARGBUF_LEN >= MAX_MEM_SIZE {
+        if arg_buf_ofs + ARGBUF_LEN >= memory.len() {
             return Err(Error::InvalidArgumentBuffer);
         }
 

piecrust/src/store.rs

@@ -22,7 +22,7 @@ use std::{fs, io, thread};
 
 pub use bytecode::Bytecode;
 use dusk_wasmtime::Engine;
-pub use memory::{Memory, MAX_MEM_SIZE, PAGE_SIZE};
+pub use memory::{Memory, PAGE_SIZE};
 pub use metadata::Metadata;
 pub use module::Module;
 use piecrust_uplink::ContractId;

piecrust/src/store/memory.rs

@@ -16,18 +16,15 @@ use crumbles::{LocateFile, Mmap};
 use dusk_wasmtime::LinearMemory;
 
 pub const PAGE_SIZE: usize = 0x10000;
-const WASM_MAX_PAGES: u32 = 0x4000000;
 
-const MIN_PAGES: usize = 4;
-const MIN_MEM_SIZE: usize = MIN_PAGES * PAGE_SIZE;
-const MAX_PAGES: usize = WASM_MAX_PAGES as usize;
-
-pub const MAX_MEM_SIZE: usize = MAX_PAGES * PAGE_SIZE;
+const WASM32_MAX_PAGES: usize = 0x10000;
+const WASM64_MAX_PAGES: usize = 0x4000000;
 
 pub struct MemoryInner {
     pub mmap: Mmap,
     pub current_len: usize,
     pub is_new: bool,
+    is_64: bool,
     ref_count: AtomicUsize,
 }
 
@@ -63,32 +60,54 @@ pub struct Memory {
 }
 
 impl Memory {
-    pub fn new() -> io::Result<Self> {
+    pub fn new(is_64: bool) -> io::Result<Self> {
+        let max_pages = if is_64 {
+            WASM64_MAX_PAGES
+        } else {
+            WASM32_MAX_PAGES
+        };
+
         Ok(Self {
             inner: Box::leak(Box::new(MemoryInner {
-                mmap: Mmap::new(MAX_PAGES, PAGE_SIZE)?,
-                current_len: MIN_MEM_SIZE,
+                mmap: Mmap::new(max_pages, PAGE_SIZE)?,
+                current_len: 0,
                 is_new: true,
+                is_64,
                 ref_count: AtomicUsize::new(1),
             })),
         })
     }
 
-    pub fn from_files<FL>(file_locator: FL, len: usize) -> io::Result<Self>
+    pub fn from_files<FL>(
+        is_64: bool,
+        file_locator: FL,
+        len: usize,
+    ) -> io::Result<Self>
     where
         FL: 'static + LocateFile,
     {
+        let max_pages = if is_64 {
+            WASM64_MAX_PAGES
+        } else {
+            WASM32_MAX_PAGES
+        };
+
         Ok(Self {
             inner: Box::leak(Box::new(MemoryInner {
                 mmap: unsafe {
-                    Mmap::with_files(MAX_PAGES, PAGE_SIZE, file_locator)?
+                    Mmap::with_files(max_pages, PAGE_SIZE, file_locator)?
                 },
                 current_len: len,
                 is_new: false,
+                is_64,
                 ref_count: AtomicUsize::new(1),
             })),
         })
     }
+
+    pub fn is_64(&self) -> bool {
+        self.inner.is_64
+    }
 }
 
 /// This implementation of clone is dangerous, and must be accompanied by the
@@ -143,7 +162,7 @@ unsafe impl LinearMemory for Memory {
     }
 
     fn maximum_byte_size(&self) -> Option<usize> {
-        Some(MAX_MEM_SIZE)
+        Some(self.inner.len())
     }
 
     fn grow_to(&mut self, new_size: usize) -> Result<(), dusk_wasmtime::Error> {

piecrust/src/store/session.rs

@@ -182,6 +182,7 @@ impl ContractSession {
                                     let memory_path = memory_path.clone();
 
                                     Memory::from_files(
+                                        module.is_64(),
                                         move |page_index: usize| {
                                             match page_indices
                                                 .contains(&page_index)
@@ -199,7 +200,7 @@ impl ContractSession {
                                         elem.len,
                                     )?
                                 }
-                                None => Memory::new()?,
+                                None => Memory::new(module.is_64())?,
                             };
 
                             let contract = entry
@@ -250,10 +251,10 @@ impl ContractSession {
         metadata: ContractMetadata,
         metadata_bytes: B,
     ) -> io::Result<()> {
-        let memory = Memory::new()?;
         let bytecode = Bytecode::new(bytecode)?;
         let module = Module::new(&self.engine, objectcode)?;
         let metadata = Metadata::new(metadata_bytes, metadata)?;
+        let memory = Memory::new(module.is_64())?;
 
         // If the position is already filled in the tree, the contract cannot be
         // inserted.

piecrust/src/store/tree.rs

@@ -15,16 +15,53 @@ use rkyv::{Archive, Deserialize, Serialize};
 
 use crate::store::memory::Memory;
 
-// There are max `2^26` pages in a memory
-const P_HEIGHT: usize = 13;
-const P_ARITY: usize = 4;
+// There are max `2^16` pages in a 32-bit memory
+const P32_HEIGHT: usize = 8;
+const P32_ARITY: usize = 4;
 
-pub type PageTree = dusk_merkle::Tree<Hash, P_HEIGHT, P_ARITY>;
+type PageTree32 = dusk_merkle::Tree<Hash, P32_HEIGHT, P32_ARITY>;
+
+// There are max `2^26` pages in a 64-bit memory
+const P64_HEIGHT: usize = 13;
+const P64_ARITY: usize = 4;
+
+type PageTree64 = dusk_merkle::Tree<Hash, P64_HEIGHT, P64_ARITY>;
 
 // This means we have max `2^32` contracts
 const C_HEIGHT: usize = 32;
 const C_ARITY: usize = 2;
 
+#[derive(Debug, Clone, Archive, Deserialize, Serialize)]
+#[archive_attr(derive(CheckBytes))]
+pub enum PageTree {
+    Wasm32(PageTree32),
+    Wasm64(PageTree64),
+}
+
+impl PageTree {
+    pub fn new(is_64: bool) -> Self {
+        if is_64 {
+            Self::Wasm64(PageTree64::new())
+        } else {
+            Self::Wasm32(PageTree32::new())
+        }
+    }
+
+    pub fn insert(&mut self, position: u64, item: impl Into<Hash>) {
+        match self {
+            Self::Wasm32(tree) => tree.insert(position, item),
+            Self::Wasm64(tree) => tree.insert(position, item),
+        }
+    }
+
+    pub fn root(&self) -> Ref<Hash> {
+        match self {
+            Self::Wasm32(tree) => tree.root(),
+            Self::Wasm64(tree) => tree.root(),
+        }
+    }
+}
+
 pub type Tree = dusk_merkle::Tree<Hash, C_HEIGHT, C_ARITY>;
 
 #[derive(Debug, Clone, Archive, Deserialize, Serialize)]
@@ -57,7 +94,7 @@ impl ContractIndex {
             self.contracts.insert(
                 contract,
                 ContractIndexElement {
-                    tree: PageTree::new(),
+                    tree: PageTree::new(memory.is_64()),
                     len: 0,
                     page_indices: BTreeSet::new(),
                 },
@@ -155,10 +192,10 @@ impl dusk_merkle::Aggregate<C_ARITY> for Hash {
     }
 }
 
-impl dusk_merkle::Aggregate<P_ARITY> for Hash {
+impl dusk_merkle::Aggregate<P32_ARITY> for Hash {
     const EMPTY_SUBTREE: Self = Hash([0; blake3::OUT_LEN]);
 
-    fn aggregate(items: [&Self; P_ARITY]) -> Self {
+    fn aggregate(items: [&Self; P32_ARITY]) -> Self {
         let mut hasher = Hasher::new();
         for item in items {
             hasher.update(item.as_bytes());