Make the compilation truly single pass, increase guard page granularity to 64k

crates/polkavm-assembler/src/amd64.rs

@@ -2157,7 +2157,7 @@ mod tests {
             }
 
             self.disassembly_1.pop();
-            disassemble_into(code, &mut self.disassembly_2);
+            disassemble_into(&code, &mut self.disassembly_2);
             assert_eq!(self.disassembly_1, self.disassembly_2, "broken encoding for: {inst:?}");
         }
     }
@@ -2244,7 +2244,7 @@ mod tests {
         let mut asm = crate::Assembler::new();
         let label = asm.forward_declare_label();
         asm.push_with_label(label, jmp_label8(label));
-        let disassembly = disassemble(asm.finalize());
+        let disassembly = disassemble(&asm.finalize());
         assert_eq!(disassembly, "00000000 ebfe jmp short 0x0");
     }
 
@@ -2254,7 +2254,7 @@ mod tests {
         let mut asm = crate::Assembler::new();
         let label = asm.forward_declare_label();
         asm.push_with_label(label, jmp_label32(label));
-        let disassembly = disassemble(asm.finalize());
+        let disassembly = disassemble(&asm.finalize());
         assert_eq!(disassembly, "00000000 e9fbffffff jmp 0x0");
     }
 
@@ -2264,7 +2264,7 @@ mod tests {
         let mut asm = crate::Assembler::new();
         let label = asm.forward_declare_label();
         asm.push_with_label(label, call_label32(label));
-        let disassembly = disassemble(asm.finalize());
+        let disassembly = disassemble(&asm.finalize());
         assert_eq!(disassembly, "00000000 e8fbffffff call 0x0");
     }
 
@@ -2275,7 +2275,7 @@ mod tests {
             let mut asm = crate::Assembler::new();
             let label = asm.forward_declare_label();
             asm.push_with_label(label, jcc_label8(cond, label));
-            let disassembly = disassemble(asm.finalize());
+            let disassembly = disassemble(&asm.finalize());
             assert_eq!(
                 disassembly,
                 format!("00000000 {:02x}fe j{} short 0x0", 0x70 + cond as u8, cond.suffix())
@@ -2291,7 +2291,7 @@ mod tests {
             let label = asm.forward_declare_label();
             asm.push(jcc_label8(cond, label));
             asm.push_with_label(label, nop());
-            let disassembly = disassemble(asm.finalize());
+            let disassembly = disassemble(&asm.finalize());
             assert_eq!(
                 disassembly,
                 format!(
@@ -2311,7 +2311,7 @@ mod tests {
             let label = asm.forward_declare_label();
             asm.push_with_label(label, nop());
             asm.push(jcc_label8(cond, label));
-            let disassembly = disassemble(asm.finalize());
+            let disassembly = disassemble(&asm.finalize());
             assert_eq!(
                 disassembly,
                 format!(
@@ -2331,7 +2331,7 @@ mod tests {
             let label = asm.forward_declare_label();
             asm.push(jcc_label32(cond, label));
             asm.push_with_label(label, nop());
-            let disassembly = disassemble(asm.finalize());
+            let disassembly = disassemble(&asm.finalize());
             assert_eq!(
                 disassembly,
                 format!(
@@ -2349,7 +2349,7 @@ mod tests {
         let mut asm = crate::Assembler::new();
         let label = asm.forward_declare_label();
         asm.push_with_label(label, lea_rip_label(super::Reg::rax, label));
-        let disassembly = disassemble(asm.finalize());
+        let disassembly = disassemble(&asm.finalize());
         assert_eq!(disassembly, "00000000 488d05f9ffffff lea rax, [rip-0x7]");
     }
 
@@ -2360,7 +2360,7 @@ mod tests {
         let label = asm.forward_declare_label();
         asm.push(lea_rip_label(super::Reg::rax, label));
         asm.push_with_label(label, nop());
-        let disassembly = disassemble(asm.finalize());
+        let disassembly = disassemble(&asm.finalize());
         assert_eq!(disassembly, "00000000 488d0500000000 lea rax, [rip]\n00000007 90 nop");
     }
 }

crates/polkavm-assembler/src/assembler.rs

@@ -23,6 +23,30 @@ impl Default for Assembler {
     }
 }
 
+#[repr(transparent)]
+pub struct AssembledCode<'a>(&'a mut Assembler);
+
+impl<'a> core::ops::Deref for AssembledCode<'a> {
+    type Target = [u8];
+
+    #[inline]
+    fn deref(&self) -> &Self::Target {
+        &self.0.code
+    }
+}
+
+impl<'a> From<AssembledCode<'a>> for Vec<u8> {
+    fn from(code: AssembledCode<'a>) -> Vec<u8> {
+        core::mem::take(&mut code.0.code)
+    }
+}
+
+impl<'a> Drop for AssembledCode<'a> {
+    fn drop(&mut self) {
+        self.0.clear();
+    }
+}
+
 impl Assembler {
     pub const fn new() -> Self {
         Assembler {
@@ -136,7 +160,7 @@ impl Assembler {
         self
     }
 
-    pub fn finalize(&mut self) -> &[u8] {
+    pub fn finalize(&mut self) -> AssembledCode {
         for fixup in self.fixups.drain(..) {
             let origin = fixup.instruction_offset + fixup.instruction_length as usize;
             let target_absolute = self.labels[fixup.target_label.0 as usize];
@@ -157,7 +181,8 @@ impl Assembler {
                 unreachable!()
             }
         }
-        &self.code
+
+        AssembledCode(self)
     }
 
     pub fn is_empty(&self) -> bool {
@@ -168,6 +193,10 @@ impl Assembler {
         self.code.len()
     }
 
+    pub fn code_mut(&mut self) -> &mut [u8] {
+        &mut self.code
+    }
+
     pub fn spare_capacity(&self) -> usize {
         self.code.capacity() - self.code.len()
     }

crates/polkavm-common/src/abi.rs

@@ -1,40 +1,46 @@
 //! Everything in this module affects the ABI of the guest programs, either by affecting
 //! their observable behavior (no matter how obscure), or changing which programs are accepted by the VM.
 
-use crate::utils::align_to_next_page_u64;
+use crate::utils::{align_to_next_page_u32, align_to_next_page_u64};
 use core::ops::Range;
 
+const ADDRESS_SPACE_SIZE: u64 = 0x100000000_u64;
+
 /// The page size of the VM.
 ///
 /// This is the minimum granularity with which the VM can allocate memory.
 pub const VM_PAGE_SIZE: u32 = 0x4000;
 
+/// The maximum page size of the VM.
+pub const VM_MAX_PAGE_SIZE: u32 = 0x10000;
+
+static_assert!(VM_PAGE_SIZE <= VM_MAX_PAGE_SIZE);
+static_assert!(VM_MAX_PAGE_SIZE % VM_PAGE_SIZE == 0);
+
 /// The address at which the program's memory starts inside of the VM.
 ///
 /// This is directly accessible by the program running inside of the VM.
-pub const VM_ADDR_USER_MEMORY: u32 = 0x00010000;
+pub const VM_ADDR_USER_MEMORY: u32 = VM_MAX_PAGE_SIZE;
 
 /// The address at which the program's stack starts inside of the VM.
 ///
 /// This is directly accessible by the program running inside of the VM.
-pub const VM_ADDR_USER_STACK_HIGH: u32 = 0xffffc000;
-static_assert!(0xffffffff - VM_PAGE_SIZE as u64 + 1 == VM_ADDR_USER_STACK_HIGH as u64);
+pub const VM_ADDR_USER_STACK_HIGH: u32 = (ADDRESS_SPACE_SIZE - VM_MAX_PAGE_SIZE as u64) as u32;
 
 /// The address which, when jumped to, will return to the host.
 ///
 /// There isn't actually anything there; it's just a virtual address.
-pub const VM_ADDR_RETURN_TO_HOST: u32 = 0xffffc000;
+pub const VM_ADDR_RETURN_TO_HOST: u32 = 0xffff0000;
 static_assert!(VM_ADDR_RETURN_TO_HOST & 0b11 == 0);
 
 /// The total maximum amount of memory a program can use.
 ///
 /// This is the whole 32-bit address space, except:
-///   * the guard pages at the start,
+///   * the guard page at the start,
 ///   * the guard page between read-only data and read-write data
 ///   * the guard page between the heap and the stack,
 ///   * and the guard page at the end.
-pub const VM_MAXIMUM_MEMORY_SIZE: u32 = 0xfffe4000;
-static_assert!(VM_MAXIMUM_MEMORY_SIZE as u64 == (1_u64 << 32) - VM_ADDR_USER_MEMORY as u64 - VM_PAGE_SIZE as u64 * 3);
+pub const VM_MAXIMUM_MEMORY_SIZE: u32 = (ADDRESS_SPACE_SIZE - VM_MAX_PAGE_SIZE as u64 * 4) as u32;
 
 /// The maximum number of VM instructions a program can be composed of.
 pub const VM_MAXIMUM_INSTRUCTION_COUNT: u32 = 2 * 1024 * 1024;
@@ -94,7 +100,7 @@ impl GuestMemoryConfig {
         // We already checked that these are less than the maximum memory size, so these cannot fail
         // because the maximum memory size is going to be vastly smaller than what an u64 can hold.
         const _: () = {
-            assert!(VM_MAXIMUM_MEMORY_SIZE as u64 + VM_PAGE_SIZE as u64 <= u32::MAX as u64);
+            assert!(VM_MAXIMUM_MEMORY_SIZE as u64 + VM_MAX_PAGE_SIZE as u64 <= u32::MAX as u64);
         };
 
         let ro_data_size = match align_to_next_page_u64(VM_PAGE_SIZE as u64, ro_data_size) {
@@ -204,7 +210,10 @@ impl GuestMemoryConfig {
         if self.ro_data_size == 0 {
             self.user_memory_region_address()
         } else {
-            self.ro_data_address() + self.ro_data_size + VM_PAGE_SIZE
+            match align_to_next_page_u32(VM_MAX_PAGE_SIZE, self.ro_data_address() + self.ro_data_size) {
+                Some(offset) => offset + VM_MAX_PAGE_SIZE,
+                None => unreachable!(),
+            }
         }
     }