Use custom allocator for the program (#801)

To allow making use of extended heap sizes.

programs/mango-v4/Cargo.toml

@@ -14,11 +14,12 @@ no-entrypoint = []
 no-idl = []
 no-log-ix-name = []
 cpi = ["no-entrypoint"]
-default = []
+default = ["custom-heap"]
 test-bpf = ["client"]
 client = ["solana-sdk", "no-entrypoint"]
 # Enables GPL-licensed parts of the code. See LICENSE file.
 enable-gpl = ["openbook-v2/enable-gpl"]
+custom-heap = []
 
 [dependencies]
 # todo: when to fix, when to use caret? need a regular chore to bump dependencies

programs/mango-v4/src/allocator.rs

@@ -0,0 +1,77 @@
+#![allow(dead_code)]
+
+use std::alloc::{GlobalAlloc, Layout};
+
+#[cfg(not(feature = "no-entrypoint"))]
+#[global_allocator]
+pub static ALLOCATOR: BumpAllocator = BumpAllocator {};
+
+pub fn heap_used() -> usize {
+    #[cfg(not(feature = "no-entrypoint"))]
+    return ALLOCATOR.used();
+
+    #[cfg(feature = "no-entrypoint")]
+    return 0;
+}
+
+/// Custom bump allocator for on-chain operations
+///
+/// The default allocator is also a bump one, but grows from a fixed
+/// HEAP_START + 32kb downwards and has no way of making use of extra
+/// heap space requested for the transaction.
+///
+/// This implementation starts at HEAP_START and grows upward, producing
+/// a segfault once out of available heap memory.
+pub struct BumpAllocator {}
+
+unsafe impl GlobalAlloc for BumpAllocator {
+    #[inline]
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let heap_start = solana_program::entrypoint::HEAP_START_ADDRESS as usize;
+        let pos_ptr = heap_start as *mut usize;
+
+        let mut pos = *pos_ptr;
+        if pos == 0 {
+            // First time, override the current position to be just past the location
+            // where the current heap position is stored.
+            pos = heap_start + 8;
+        }
+
+        // The result address needs to be aligned to layout.align(),
+        // which is guaranteed to be a power of two.
+        // Find the first address >=pos that has the required alignment.
+        // Wrapping ops are used for performance.
+        let mask = layout.align().wrapping_sub(1);
+        let begin = pos.wrapping_add(mask) & (!mask);
+
+        // Update allocator state
+        let end = begin.checked_add(layout.size()).unwrap();
+        *pos_ptr = end;
+
+        // Write a byte to trigger heap overflow errors early
+        let end_ptr = end as *mut u8;
+        *end_ptr = 0;
+
+        begin as *mut u8
+    }
+    #[inline]
+    unsafe fn dealloc(&self, _: *mut u8, _: Layout) {
+        // I'm a bump allocator, I don't free
+    }
+}
+
+impl BumpAllocator {
+    #[inline]
+    pub fn used(&self) -> usize {
+        let heap_start = solana_program::entrypoint::HEAP_START_ADDRESS as usize;
+        unsafe {
+            let pos_ptr = heap_start as *mut usize;
+
+            let pos = *pos_ptr;
+            if pos == 0 {
+                return 0;
+            }
+            return pos - heap_start;
+        }
+    }
+}

programs/mango-v4/src/lib.rs

@@ -15,6 +15,7 @@ use accounts_ix::*;
 pub mod accounts_ix;
 pub mod accounts_zerocopy;
 pub mod address_lookup_table_program;
+mod allocator;
 pub mod error;
 pub mod events;
 pub mod health;