cleanup(kernel): only try to re-use free blocks before memory.grow (#765)

Updates the kernel to only scan for free blocks before a `memory.grow`
operation, this should improve performance of `alloc` by not iterating
over every block for each allocation.

kernel/src/lib.rs

@@ -255,15 +255,6 @@ impl MemoryRoot {
     pub unsafe fn alloc(&mut self, length: u64) -> Option<&'static mut MemoryBlock> {
         let self_position = self.position.load(Ordering::Acquire);
         let self_length = self.length.load(Ordering::Acquire);
-        let b = self.find_free_block(length, self_position);
-
-        // If there's a free block then re-use it
-        if let Some(b) = b {
-            b.used = length as usize;
-            b.status
-                .store(MemoryStatus::Active as u8, Ordering::Release);
-            return Some(b);
-        }
 
         // Get the current index for a new block
         let curr = self.blocks.as_ptr() as u64 + self_position;
@@ -275,6 +266,21 @@ impl MemoryRoot {
         // When the allocation is larger than the number of bytes available
         // we will need to try to grow the memory
         if length_with_block >= mem_left {
+            // If the current position is large enough to hold the length of the block being
+            // allocated then check for existing free blocks that can be re-used before
+            // growing memory
+            if length_with_block <= self_position {
+                let b = self.find_free_block(length, self_position);
+
+                // If there's a free block then re-use it
+                if let Some(b) = b {
+                    b.used = length as usize;
+                    b.status
+                        .store(MemoryStatus::Active as u8, Ordering::Release);
+                    return Some(b);
+                }
+            }
+
             // Calculate the number of pages needed to cover the remaining bytes
             let npages = num_pages(length_with_block - mem_left);
             let x = core::arch::wasm32::memory_grow(0, npages);
@@ -582,45 +588,6 @@ mod test {
     use crate::*;
     use wasm_bindgen_test::*;
 
-    // See https://github.com/extism/extism/pull/659
-    #[wasm_bindgen_test]
-    fn test_659() {
-        unsafe {
-            // Warning: These offsets will need to change if we adjust the kernel memory layout at all
-            reset();
-            assert_eq!(alloc(1065), 77);
-            assert_eq!(alloc(288), 1154);
-            assert_eq!(alloc(128), 1454);
-            assert_eq!(length(1154), 288);
-            assert_eq!(length(1454), 128);
-            free(1454);
-            assert_eq!(alloc(213), 1594);
-            length_unsafe(1594);
-            assert_eq!(alloc(511), 1819);
-            assert_eq!(alloc(4), 1454);
-            assert_eq!(length(1454), 4);
-            assert_eq!(length(1819), 511);
-            assert_eq!(alloc(13), 2342);
-            assert_eq!(length(2342), 13);
-            assert_eq!(alloc(336), 2367);
-            assert_eq!(alloc(1077), 2715);
-            assert_eq!(length(2367), 336);
-            assert_eq!(length(2715), 1077);
-            free(2715);
-            assert_eq!(alloc(1094), 3804);
-            length_unsafe(3804);
-
-            // Allocate 4 bytes, expect to receive address 3788
-            assert_eq!(alloc(4), 3788);
-
-            assert_eq!(alloc(4), 3772);
-            assert_eq!(length(3772), 4);
-
-            // Address 3788 has not been freed yet, so expect it to have 4 bytes allocated
-            assert_eq!(length(3788), 4);
-        }
-    }
-
     #[wasm_bindgen_test]
     fn test_oom() {
         let size = 1024 * 1024 * 5;