Refactored flexible memory implementation (#378)

* Refactor dymnaic memory

* Refactor memory flags

* Refactor frames flags

* Small optimization

* Use less instructions for _CHECK_READ_FRAMES on aarch64

* Small optimization

* Small optimization

* Use alloc instead of vec

* Add cast_ptr_to_slice and cast_ptr_to_slice_mut methods

* Fix clippy

* Remove memory size limit

* Let AsmCoreMachine be responsible for memory allocation and release

Author: mohanson

Makefile

@@ -28,7 +28,8 @@ clippy_rule = -D warnings \
 	-A clippy::inconsistent_digit_grouping \
 	-A clippy::large_digit_groups \
 	-A clippy::suspicious_operation_groupings \
-	-A clippy::unnecessary_cast
+	-A clippy::unnecessary_cast \
+	-A clippy::mut_from_ref
 clippy:
 	cargo clippy --all --features=asm -- $(clippy_rule)
 	cd definitions && cargo clippy --all -- $(clippy_rule)

definitions/src/asm.rs

@@ -1,8 +1,8 @@
 use crate::{
-    instructions::Instruction, MEMORY_FRAMES, MEMORY_FRAMESIZE, MEMORY_FRAME_SHIFTS,
-    RISCV_GENERAL_REGISTER_NUMBER, RISCV_MAX_MEMORY, RISCV_PAGES, RISCV_PAGESIZE,
+    instructions::Instruction, DEFAULT_MEMORY_SIZE, MEMORY_FRAMESIZE, MEMORY_FRAME_SHIFTS,
+    RISCV_GENERAL_REGISTER_NUMBER, RISCV_PAGESIZE,
 };
-use std::alloc::{alloc, Layout};
+use std::alloc::{alloc, alloc_zeroed, dealloc, Layout};
 
 // The number of trace items to keep
 pub const TRACE_SIZE: usize = 8192;
@@ -101,15 +101,25 @@ pub struct AsmCoreMachine {
     pub last_read_frame: u64,
     pub last_write_page: u64,
 
-    pub flags: [u8; RISCV_PAGES],
-    pub frames: [u8; MEMORY_FRAMES],
+    pub memory_ptr: u64,
+    pub flags_ptr: u64,
+    pub frames_ptr: u64,
+}
 
-    pub memory: [u8; RISCV_MAX_MEMORY],
+impl Drop for AsmCoreMachine {
+    fn drop(&mut self) {
+        let memory_layout = Layout::array::<u8>(self.memory_size as usize).unwrap();
+        unsafe { dealloc(self.memory_ptr as *mut u8, memory_layout) };
+        let flags_layout = Layout::array::<u8>(self.flags_size as usize).unwrap();
+        unsafe { dealloc(self.flags_ptr as *mut u8, flags_layout) };
+        let frames_layout = Layout::array::<u8>(self.frames_size as usize).unwrap();
+        unsafe { dealloc(self.frames_ptr as *mut u8, frames_layout) };
+    }
 }
 
 impl AsmCoreMachine {
     pub fn new(isa: u8, version: u32, max_cycles: u64) -> Box<AsmCoreMachine> {
-        Self::new_with_memory(isa, version, max_cycles, RISCV_MAX_MEMORY)
+        Self::new_with_memory(isa, version, max_cycles, DEFAULT_MEMORY_SIZE)
     }
 
     pub fn new_with_memory(
@@ -119,15 +129,10 @@ impl AsmCoreMachine {
         memory_size: usize,
     ) -> Box<AsmCoreMachine> {
         assert_ne!(memory_size, 0);
-        assert!(memory_size <= RISCV_MAX_MEMORY);
         assert_eq!(memory_size % RISCV_PAGESIZE, 0);
         assert_eq!(memory_size % (1 << MEMORY_FRAME_SHIFTS), 0);
-
         let mut machine = unsafe {
-            let machine_size =
-                std::mem::size_of::<AsmCoreMachine>() - RISCV_MAX_MEMORY + memory_size;
-
-            let layout = Layout::array::<u8>(machine_size).unwrap();
+            let layout = Layout::new::<AsmCoreMachine>();
             let raw_allocation = alloc(layout) as *mut AsmCoreMachine;
             Box::from_raw(raw_allocation)
         };
@@ -146,8 +151,6 @@ impl AsmCoreMachine {
         machine.reset_signal = 0;
         machine.version = version;
         machine.isa = isa;
-        machine.flags = [0; RISCV_PAGES];
-        machine.frames = [0; MEMORY_FRAMES];
 
         machine.memory_size = memory_size as u64;
         machine.frames_size = (memory_size / MEMORY_FRAMESIZE) as u64;
@@ -156,10 +159,35 @@ impl AsmCoreMachine {
         machine.last_read_frame = u64::max_value();
         machine.last_write_page = u64::max_value();
 
+        let memory_layout = Layout::array::<u8>(machine.memory_size as usize).unwrap();
+        machine.memory_ptr = unsafe { alloc(memory_layout) } as u64;
+        let flags_layout = Layout::array::<u8>(machine.flags_size as usize).unwrap();
+        machine.flags_ptr = unsafe { alloc_zeroed(flags_layout) } as u64;
+        let frames_layout = Layout::array::<u8>(machine.frames_size as usize).unwrap();
+        machine.frames_ptr = unsafe { alloc_zeroed(frames_layout) } as u64;
+
         machine
     }
 
     pub fn set_max_cycles(&mut self, cycles: u64) {
         self.max_cycles = cycles;
     }
 }
+
+impl AsmCoreMachine {
+    pub fn cast_ptr_to_slice(&self, ptr: u64, offset: usize, size: usize) -> &[u8] {
+        unsafe {
+            let ptr = ptr as *mut u8;
+            let ptr = ptr.add(offset);
+            std::slice::from_raw_parts(ptr, size)
+        }
+    }
+
+    pub fn cast_ptr_to_slice_mut(&self, ptr: u64, offset: usize, size: usize) -> &mut [u8] {
+        unsafe {
+            let ptr = ptr as *mut u8;
+            let ptr = ptr.add(offset);
+            std::slice::from_raw_parts_mut(ptr, size)
+        }
+    }
+}

definitions/src/generate_asm_constants.rs

@@ -8,8 +8,8 @@ use ckb_vm_definitions::{
     instructions::{instruction_opcode_name, MAXIMUM_OPCODE, MINIMAL_OPCODE},
     memory::{FLAG_DIRTY, FLAG_EXECUTABLE, FLAG_FREEZED, FLAG_WRITABLE, FLAG_WXORX_BIT},
     registers::{RA, SP},
-    MEMORY_FRAMES, MEMORY_FRAMESIZE, MEMORY_FRAME_PAGE_SHIFTS, MEMORY_FRAME_SHIFTS,
-    RISCV_MAX_MEMORY, RISCV_PAGES, RISCV_PAGESIZE, RISCV_PAGE_SHIFTS,
+    MEMORY_FRAMESIZE, MEMORY_FRAME_PAGE_SHIFTS, MEMORY_FRAME_SHIFTS, RISCV_PAGESIZE,
+    RISCV_PAGE_SHIFTS,
 };
 use std::mem::{size_of, zeroed};
 
@@ -30,17 +30,14 @@ macro_rules! print_inst_label {
 // of this as a workaround to the problem that build.rs cannot depend on any
 // of its crate contents.
 fn main() {
-    println!("#define CKB_VM_ASM_RISCV_MAX_MEMORY {}", RISCV_MAX_MEMORY);
     println!("#define CKB_VM_ASM_RISCV_PAGE_SHIFTS {}", RISCV_PAGE_SHIFTS);
     println!("#define CKB_VM_ASM_RISCV_PAGE_SIZE {}", RISCV_PAGESIZE);
     println!("#define CKB_VM_ASM_RISCV_PAGE_MASK {}", RISCV_PAGESIZE - 1);
-    println!("#define CKB_VM_ASM_RISCV_PAGES {}", RISCV_PAGES);
     println!(
         "#define CKB_VM_ASM_MEMORY_FRAME_SHIFTS {}",
         MEMORY_FRAME_SHIFTS
     );
     println!("#define CKB_VM_ASM_MEMORY_FRAMESIZE {}", MEMORY_FRAMESIZE);
-    println!("#define CKB_VM_ASM_MEMORY_FRAMES {}", MEMORY_FRAMES);
     println!(
         "#define CKB_VM_ASM_MEMORY_FRAME_PAGE_SHIFTS {}",
         MEMORY_FRAME_PAGE_SHIFTS
@@ -183,31 +180,18 @@ fn main() {
         "#define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_LAST_WRITE_PAGE {}",
         (&m.last_write_page as *const u64 as usize) - m_address
     );
-
-    println!(
-        "#define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_FLAGS {}",
-        (&m.flags as *const u8 as usize) - m_address
-    );
-    let memory_offset_address = (&m.memory as *const u8 as usize) - m_address;
-    println!(
-        "#define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_MEMORY {}",
-        memory_offset_address
-    );
     println!(
-        "#define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_FRAMES {}",
-        (&m.frames as *const u8 as usize) - m_address
+        "#define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_MEMORY_PTR {}",
+        (&m.memory_ptr as *const u64 as usize) - m_address
     );
-    println!();
-
     println!(
-        "#define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_MEMORY_H {}",
-        memory_offset_address.wrapping_shr(12).wrapping_shl(12)
+        "#define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_FLAGS_PTR {}",
+        (&m.flags_ptr as *const u64 as usize) - m_address
     );
     println!(
-        "#define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_MEMORY_L {}",
-        memory_offset_address & 0xFFF
+        "#define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_FRAMES_PTR {}",
+        (&m.frames_ptr as *const u64 as usize) - m_address
     );
-
     println!();
 
     for op in MINIMAL_OPCODE..MAXIMUM_OPCODE {

definitions/src/lib.rs

@@ -6,17 +6,12 @@ pub mod registers;
 pub const RISCV_PAGE_SHIFTS: usize = 12;
 pub const RISCV_PAGESIZE: usize = 1 << RISCV_PAGE_SHIFTS;
 pub const RISCV_GENERAL_REGISTER_NUMBER: usize = 32;
-// 4 MB
-pub const RISCV_MAX_MEMORY: usize = 4 << 20;
-// 1 MB
-pub const DEFAULT_STACK_SIZE: usize = 1 << 20;
-pub const RISCV_PAGES: usize = RISCV_MAX_MEMORY / RISCV_PAGESIZE;
-// 256 KB
 pub const MEMORY_FRAME_SHIFTS: usize = 18;
-pub const MEMORY_FRAMESIZE: usize = 1 << MEMORY_FRAME_SHIFTS;
-pub const MEMORY_FRAMES: usize = RISCV_MAX_MEMORY / MEMORY_FRAMESIZE;
+pub const MEMORY_FRAMESIZE: usize = 1 << MEMORY_FRAME_SHIFTS; // 256 KB
 pub const MEMORY_FRAME_PAGE_SHIFTS: usize = MEMORY_FRAME_SHIFTS - RISCV_PAGE_SHIFTS;
 
+pub const DEFAULT_MEMORY_SIZE: usize = 4 << 20; // 4 MB
+
 pub const ISA_IMC: u8 = 0b0000_0000;
 pub const ISA_B: u8 = 0b0000_0001;
 pub const ISA_MOP: u8 = 0b0000_0010;

fuzz/fuzz_targets/isa_a.rs

@@ -43,7 +43,6 @@ fuzz_target!(|data: [u8; 512]| {
         ckb_vm::machine::asm::AsmCoreMachine::new(ckb_vm_isa, ckb_vm_version, u64::MAX),
     )
     .build();
-
     let insts: [u32; 18] = [
         0b00001_00_00000_00000_010_00000_0101111, // AMOSWAP.W
         0b00000_00_00000_00000_010_00000_0101111, // AMOADD.W

src/lib.rs

@@ -30,9 +30,8 @@ pub use crate::{
 pub use bytes::Bytes;
 
 pub use ckb_vm_definitions::{
-    registers, DEFAULT_STACK_SIZE, ISA_A, ISA_B, ISA_IMC, ISA_MOP, MEMORY_FRAMES, MEMORY_FRAMESIZE,
-    MEMORY_FRAME_SHIFTS, RISCV_GENERAL_REGISTER_NUMBER, RISCV_MAX_MEMORY, RISCV_PAGES,
-    RISCV_PAGESIZE, RISCV_PAGE_SHIFTS,
+    registers, DEFAULT_MEMORY_SIZE, ISA_A, ISA_B, ISA_IMC, ISA_MOP, MEMORY_FRAMESIZE,
+    MEMORY_FRAME_SHIFTS, RISCV_GENERAL_REGISTER_NUMBER, RISCV_PAGESIZE, RISCV_PAGE_SHIFTS,
 };
 
 pub use error::Error;
@@ -74,7 +73,7 @@ mod tests {
 
     #[test]
     fn test_max_memory_must_be_multiple_of_pages() {
-        assert_eq!(RISCV_MAX_MEMORY % RISCV_PAGESIZE, 0);
+        assert_eq!(DEFAULT_MEMORY_SIZE % RISCV_PAGESIZE, 0);
     }
 
     #[test]

src/machine/asm/cdefinitions_generated.h

@@ -1,11 +1,8 @@
-#define CKB_VM_ASM_RISCV_MAX_MEMORY 4194304
 #define CKB_VM_ASM_RISCV_PAGE_SHIFTS 12
 #define CKB_VM_ASM_RISCV_PAGE_SIZE 4096
 #define CKB_VM_ASM_RISCV_PAGE_MASK 4095
-#define CKB_VM_ASM_RISCV_PAGES 1024
 #define CKB_VM_ASM_MEMORY_FRAME_SHIFTS 18
 #define CKB_VM_ASM_MEMORY_FRAMESIZE 262144
-#define CKB_VM_ASM_MEMORY_FRAMES 16
 #define CKB_VM_ASM_MEMORY_FRAME_PAGE_SHIFTS 6
 
 #define CKB_VM_ASM_MAXIMUM_TRACE_ADDRESS_LENGTH 64
@@ -52,12 +49,9 @@
 #define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_FLAGS_SIZE 336
 #define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_LAST_READ_FRAME 344
 #define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_LAST_WRITE_PAGE 352
-#define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_FLAGS 360
-#define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_MEMORY 1400
-#define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_FRAMES 1384
-
-#define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_MEMORY_H 0
-#define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_MEMORY_L 1400
+#define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_MEMORY_PTR 360
+#define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_FLAGS_PTR 368
+#define CKB_VM_ASM_ASM_CORE_MACHINE_OFFSET_FRAMES_PTR 376
 
 #define CKB_VM_ASM_OP_UNLOADED 16
 #define CKB_VM_ASM_OP_ADD 17