Polkavm heap (#2)

Use PolkaVM heap

crates/llvm-context/src/eravm/context/address_space.rs

@@ -5,9 +5,10 @@
 ///
 /// The address space aliases.
 ///
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+#[derive(Default, Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub enum AddressSpace {
     /// The stack memory.
+    #[default]
     Stack,
     /// The heap memory.
     Heap,

crates/llvm-context/src/eravm/context/function/runtime/entry.rs

@@ -30,7 +30,7 @@ impl Entry {
 
     /// Initializes the global variables.
     /// The pointers are not initialized, because it's not possible to create a null pointer.
-    pub fn initialize_globals<D>(context: &mut Context<D>)
+    pub fn initialize_globals<D>(context: &mut Context<D>) -> anyhow::Result<()>
     where
         D: Dependency + Clone,
     {
@@ -42,13 +42,18 @@ impl Entry {
             calldata_type.get_undef(),
         );
 
-        let heap_memory_type = context.array_type(context.byte_type(), 1024 * 1024);
         context.set_global(
             crate::eravm::GLOBAL_HEAP_MEMORY_POINTER,
-            heap_memory_type,
+            context.byte_type().ptr_type(AddressSpace::Generic.into()),
             AddressSpace::Stack,
-            heap_memory_type.get_undef(),
+            context.integer_type(32).get_undef(),
         );
+        context.build_store(
+            context
+                .get_global(crate::eravm::GLOBAL_HEAP_MEMORY_POINTER)?
+                .into(),
+            context.build_sbrk(context.integer_const(32, 0))?,
+        )?;
 
         context.set_global(
             crate::eravm::GLOBAL_CALLDATA_SIZE,
@@ -80,6 +85,8 @@ impl Entry {
             AddressSpace::Stack,
             extra_abi_data_type.const_zero(),
         );
+
+        Ok(())
     }
 
     /// Load the calldata via seal `input` and initialize the calldata end
@@ -256,7 +263,7 @@ where
         context.set_current_function(Runtime::FUNCTION_ENTRY)?;
         context.set_basic_block(context.current_function().borrow().entry_block());
 
-        Self::initialize_globals(context);
+        Self::initialize_globals(context)?;
         Self::load_calldata(context)?;
         Self::leave_entry(context)?;
 

crates/llvm-context/src/eravm/context/mod.rs

@@ -662,26 +662,22 @@ where
     ) -> anyhow::Result<inkwell::values::BasicValueEnum<'ctx>> {
         match pointer.address_space {
             AddressSpace::Heap => {
-                let heap_pointer = self
-                    .get_global(crate::eravm::GLOBAL_HEAP_MEMORY_POINTER)?
-                    .value
-                    .as_pointer_value();
-
-                // TODO: Ensure safe casts somehow
-                let offset = self.builder().build_ptr_to_int(
-                    pointer.value,
-                    self.integer_type(32),
-                    "offset_ptrtoint",
+                let heap_pointer = self.build_heap_gep(
+                    self.builder().build_ptr_to_int(
+                        pointer.value,
+                        self.integer_type(32),
+                        "offset_ptrtoint",
+                    )?,
+                    pointer
+                        .r#type
+                        .size_of()
+                        .expect("should be IntValue")
+                        .const_truncate(self.integer_type(32)),
                 )?;
-                let pointer_value = self.build_gep(
-                    Pointer::new(self.byte_type(), AddressSpace::Stack, heap_pointer),
-                    &[offset],
-                    self.byte_type(),
-                    "heap_offset_via_gep",
-                );
+
                 let value = self
                     .builder()
-                    .build_load(pointer.r#type, pointer_value.value, name)?;
+                    .build_load(pointer.r#type, heap_pointer.value, name)?;
                 self.basic_block()
                     .get_last_instruction()
                     .expect("Always exists")
@@ -758,17 +754,12 @@ where
                     .builder()
                     .build_load(pointer.r#type, pointer.value, name)?;
 
-                let alignment = if AddressSpace::Stack == pointer.address_space {
-                    era_compiler_common::BYTE_LENGTH_FIELD
-                } else {
-                    era_compiler_common::BYTE_LENGTH_BYTE
-                };
-
                 self.basic_block()
                     .get_last_instruction()
                     .expect("Always exists")
-                    .set_alignment(alignment as u32)
+                    .set_alignment(era_compiler_common::BYTE_LENGTH_FIELD as u32)
                     .expect("Alignment is valid");
+
                 Ok(value)
             }
         }
@@ -785,32 +776,23 @@ where
     {
         match pointer.address_space {
             AddressSpace::Heap => {
-                let heap_pointer = self
-                    .get_global(crate::eravm::GLOBAL_HEAP_MEMORY_POINTER)
-                    .unwrap()
-                    .value
-                    .as_pointer_value();
-
-                // TODO: Ensure safe casts somehow
-                let offset = self.builder().build_ptr_to_int(
-                    pointer.value,
-                    self.integer_type(32),
-                    "offset_ptrtoint",
+                let heap_pointer = self.build_heap_gep(
+                    self.builder().build_ptr_to_int(
+                        pointer.value,
+                        self.integer_type(32),
+                        "offset_ptrtoint",
+                    )?,
+                    value
+                        .as_basic_value_enum()
+                        .get_type()
+                        .size_of()
+                        .expect("should be IntValue")
+                        .const_truncate(self.integer_type(32)),
                 )?;
-                let pointer_value = self.build_gep(
-                    Pointer::new(self.byte_type(), AddressSpace::Stack, heap_pointer),
-                    &[offset],
-                    self.byte_type(),
-                    "heap_offset_via_gep",
-                );
-
                 let value = self.build_byte_swap(value.as_basic_value_enum());
 
-                let instruction = self
-                    .builder
-                    .build_store(pointer_value.value, value)
-                    .unwrap();
-                instruction
+                self.builder
+                    .build_store(heap_pointer.value, value)?
                     .set_alignment(era_compiler_common::BYTE_LENGTH_BYTE as u32)
                     .expect("Alignment is valid");
             }
@@ -875,17 +857,6 @@ where
         };
 
         Ok(())
-        // let instruction = self.builder.build_store(pointer.value, value).unwrap();
-
-        // let alignment = if AddressSpace::Stack == pointer.address_space {
-        //     era_compiler_common::BYTE_LENGTH_FIELD
-        // } else {
-        //     era_compiler_common::BYTE_LENGTH_BYTE
-        // };
-
-        // instruction
-        //     .set_alignment(alignment as u32)
-        //     .expect("Alignment is valid");
     }
 
     /// Swap the endianness of an intvalue
@@ -1212,17 +1183,9 @@ where
         //    zkevm_opcode_defs::RetForwardPageType::UseHeap
         //};
 
-        let heap_pointer = self
-            .get_global(crate::eravm::GLOBAL_HEAP_MEMORY_POINTER)?
-            .value
-            .as_pointer_value();
         let offset_truncated = self.safe_truncate_int_to_i32(offset)?;
-        let offset_into_heap = self.build_gep(
-            Pointer::new(self.byte_type(), AddressSpace::Stack, heap_pointer),
-            &[offset_truncated],
-            self.byte_type(),
-            "heap_offset_via_gep",
-        );
+        let length_truncated = self.safe_truncate_int_to_i32(length)?;
+        let offset_into_heap = self.build_heap_gep(offset_truncated, length_truncated)?;
 
         let length_pointer = self.safe_truncate_int_to_i32(length)?;
         let offset_pointer = self.builder().build_ptr_to_int(
@@ -1279,6 +1242,99 @@ where
         Ok(truncated)
     }
 
+    /// Build a call to PolkaVM `sbrk` for extending the heap by `size`.
+    pub fn build_sbrk(
+        &self,
+        size: inkwell::values::IntValue<'ctx>,
+    ) -> anyhow::Result<inkwell::values::PointerValue<'ctx>> {
+        Ok(self
+            .builder()
+            .build_call(
+                self.module().get_function("__sbrk").expect("is declared"),
+                &[size.into()],
+                "call_sbrk",
+            )?
+            .try_as_basic_value()
+            .left()
+            .expect("sbrk returns a pointer")
+            .into_pointer_value())
+    }
+
+    /// Call PolkaVM `sbrk` for extending the heap by `size`,
+    /// trapping the contract if the call failed.
+    pub fn build_heap_alloc(
+        &self,
+        size: inkwell::values::IntValue<'ctx>,
+    ) -> anyhow::Result<inkwell::values::PointerValue<'ctx>> {
+        let end_of_memory = self.build_sbrk(size)?;
+        let return_is_nil = self.builder().build_int_compare(
+            inkwell::IntPredicate::EQ,
+            end_of_memory,
+            self.byte_type().ptr_type(Default::default()).const_null(),
+            "compare_end_of_memory_nil",
+        )?;
+
+        let continue_block = self.append_basic_block("sbrk_not_nil");
+        let trap_block = self.append_basic_block("sbrk_nil");
+        self.build_conditional_branch(return_is_nil, trap_block, continue_block)?;
+
+        self.set_basic_block(trap_block);
+        self.build_call(self.intrinsics().trap, &[], "invalid_trap");
+        self.build_unreachable();
+
+        self.set_basic_block(continue_block);
+
+        Ok(end_of_memory)
+    }
+
+    /// Returns a pointer to `offset` into the heap, allocating `length`
+    /// bytes more memory if `offset + length` would be out of bounds.
+    ///
+    /// # Panics
+    /// Assumes `offset` and `length` to be an i32 value.
+    pub fn build_heap_gep(
+        &self,
+        offset: inkwell::values::IntValue<'ctx>,
+        length: inkwell::values::IntValue<'ctx>,
+    ) -> anyhow::Result<Pointer<'ctx>> {
+        assert_eq!(offset.get_type(), self.integer_type(32));
+        assert_eq!(length.get_type(), self.integer_type(32));
+
+        let heap_start = self
+            .get_global(crate::eravm::GLOBAL_HEAP_MEMORY_POINTER)?
+            .value
+            .as_pointer_value();
+        let heap_end = self.build_sbrk(self.integer_const(32, 0))?;
+        let value_end = self.build_gep(
+            Pointer::new(self.byte_type(), AddressSpace::Stack, heap_start),
+            &[self.builder().build_int_nuw_add(offset, length, "end")?],
+            self.byte_type(),
+            "heap_end_gep",
+        );
+        let is_out_of_bounds = self.builder().build_int_compare(
+            inkwell::IntPredicate::UGT,
+            value_end.value,
+            heap_end,
+            "is_value_overflowing_heap",
+        )?;
+
+        let out_of_bounds_block = self.append_basic_block("heap_offset_out_of_bounds");
+        let heap_offset_block = self.append_basic_block("build_heap_pointer");
+        self.build_conditional_branch(is_out_of_bounds, out_of_bounds_block, heap_offset_block)?;
+
+        self.set_basic_block(out_of_bounds_block);
+        self.build_heap_alloc(length)?;
+        self.build_unconditional_branch(heap_offset_block);
+
+        self.set_basic_block(heap_offset_block);
+        Ok(self.build_gep(
+            Pointer::new(self.byte_type(), AddressSpace::Stack, heap_start),
+            &[offset],
+            self.byte_type(),
+            "heap_offset_via_gep",
+        ))
+    }
+
     ///
     /// Writes the ABI pointer to the global variable.
     ///

crates/llvm-context/src/eravm/evm/calldata.rs

@@ -59,16 +59,9 @@ pub fn copy<'ctx, D>(
 where
     D: Dependency + Clone,
 {
-    let heap_pointer = context
-        .get_global(crate::eravm::GLOBAL_HEAP_MEMORY_POINTER)?
-        .value
-        .as_pointer_value();
-    let destination = context.build_gep(
-        Pointer::new(context.byte_type(), AddressSpace::Stack, heap_pointer),
-        &[context.safe_truncate_int_to_i32(destination_offset)?],
-        context.byte_type(),
-        "heap_pointer_with_offset",
-    );
+    let offset = context.safe_truncate_int_to_i32(destination_offset)?;
+    let size = context.safe_truncate_int_to_i32(size)?;
+    let destination = context.build_heap_gep(offset, size)?;
 
     let calldata_pointer = context
         .get_global(crate::eravm::GLOBAL_CALLDATA_POINTER)?

crates/llvm-context/src/eravm/evm/crypto.rs

@@ -17,23 +17,14 @@ where
     D: Dependency + Clone,
 {
     let offset_casted = context.safe_truncate_int_to_i32(offset)?;
-    let heap_pointer = context.get_global(crate::eravm::GLOBAL_HEAP_MEMORY_POINTER)?;
-    let input_pointer = unsafe {
-        context.builder().build_gep(
-            context.byte_type(),
-            heap_pointer.value.as_pointer_value(),
-            &[offset_casted],
-            "heap_offset_via_gep",
-        )
-    }?;
+    let length_casted = context.safe_truncate_int_to_i32(length)?;
+    let input_pointer = context.build_heap_gep(offset_casted, length_casted)?;
     let input_pointer_casted = context.builder().build_ptr_to_int(
-        input_pointer,
+        input_pointer.value,
         context.integer_type(32),
         "input_pointer_casted",
     )?;
 
-    let length_casted = context.safe_truncate_int_to_i32(length)?;
-
     let output_pointer = context.build_alloca(context.field_type(), "output_pointer");
     let output_pointer_casted = context.builder().build_ptr_to_int(
         output_pointer.value,