fix(wasm): when reallocating the last allocated region, properly grow

the region in place.

Previous changes to `malloc` caused `realloc` to sometimes pull regions
off of the free list during this optimization. Because no `memcpy` is
performed, this resulted in corrupted data returning to the caller.

Co-authored-by: trim21 <i@trim21.me>

crates/cli/src/tests/wasm_language_test.rs

@@ -131,6 +131,19 @@ fn test_wasm_realloc_smaller_size() {
     });
 }
 
+#[test]
+fn test_wasm_realloc_clobber_region() {
+    allocations::record(|| {
+        let store = WasmStore::new(&ENGINE).unwrap();
+        let mut parser = Parser::new();
+        let language = get_test_fixture_language_wasm("wasm_realloc_clobber_region");
+        parser.set_wasm_store(store).unwrap();
+        parser.set_language(&language).unwrap();
+        let tree = parser.parse("hello", None).unwrap();
+        assert_eq!(tree.root_node().to_sexp(), "(document (zero_width))");
+    });
+}
+
 #[test]
 fn test_load_multiple_wasm_languages() {
     allocations::record(|| {

crates/language/wasm/src/stdlib.c

@@ -48,6 +48,19 @@ static int grow_heap(size_t size) {
   return __builtin_wasm_memory_grow(0, new_page_count) != SIZE_MAX;
 }
 
+// Grows the heap if necessary to fit a region at the _end_ of the heap
+// ending at `region_end` by `size` bytes.
+//
+// Returns 0 if the heap could not be grown, 1 otherwise.
+static inline int grow_heap_for_region(Region *region_end, size_t size) {
+  if (region_end > heap_end) {
+    if ((char *)region_end - (char *)heap_start > MAX_HEAP_SIZE) return 0;
+    if (!grow_heap(size)) return 0;
+    heap_end = get_heap_end();
+  }
+  return 1;
+}
+
 // Clear out the heap, and move it to the given address.
 void reset_heap(void *new_heap_start) {
   heap_start = new_heap_start;
@@ -76,13 +89,7 @@ void *malloc(size_t size) {
 
   Region *region_end = region_after(next, size);
 
-  if (region_end > heap_end) {
-    if ((char *)region_end - (char *)heap_start > MAX_HEAP_SIZE) {
-      return NULL;
-    }
-    if (!grow_heap(size)) return NULL;
-    heap_end = get_heap_end();
-  }
+  if (!grow_heap_for_region(region_end, size)) return NULL;
 
   void *result = &next->data;
   next->size = size;
@@ -127,11 +134,18 @@ void *realloc(void *ptr, size_t new_size) {
   Region *region = region_for_ptr(ptr);
   Region *region_end = region_after(region, region->size);
 
-  // When reallocating the last allocated region, return
-  // the same pointer, and skip copying the data.
+  // When reallocating the last allocated region, resize
+  // in place if possible, return the same pointer, and
+  // skip copying the data.
   if (region_end == next) {
-    next = region;
-    return malloc(new_size);
+    Region *new_region_end = region_after(region, new_size);
+
+    size_t additional_size = (char *)new_region_end - (char *)heap_end;
+    if (!grow_heap_for_region(new_region_end, additional_size)) return NULL;
+
+    region->size = new_size;
+    next = new_region_end;
+    return &region->data;
   }
 
   void *result = malloc(new_size);

test/fixtures/test_grammars/wasm_realloc_clobber_region/corpus.txt

@@ -0,0 +1,9 @@
+==========================
+A leading zero-width token
+==========================
+
+hello
+
+---
+
+(document (zero_width))

test/fixtures/test_grammars/wasm_realloc_clobber_region/grammar.js

@@ -0,0 +1,10 @@
+export default grammar({
+    name: 'wasm_realloc_clobber_region',
+
+    extras: _ => [/\s/],
+    externals: $ => [$.zero_width],
+
+    rules: {
+        document: $ => seq($.zero_width, 'hello'),
+    }
+});

test/fixtures/test_grammars/wasm_realloc_clobber_region/scanner.c

@@ -0,0 +1,68 @@
+#include "tree_sitter/parser.h"
+#include <string.h>
+
+#ifdef __wasm__
+  #define panic(x) __builtin_trap()
+#else
+  #include <assert.h>
+  #define panic(x) assert(0)
+#endif
+
+enum TokenType {
+  ZERO_WIDTH_TOKEN
+};
+
+void *tree_sitter_wasm_realloc_clobber_region_external_scanner_create(void) {
+  size_t a_count = 32;
+  void *a = malloc(a_count);
+  memset(a, 0x41, a_count); // fill `a` with 'A's
+  for (int i = 0; i < 32; i++) {
+    // This condition is unreachable, but it prevents wasi-sdk clang from optimizing
+    // out the allocation to `a`.
+    if (((char *)a)[i] != 0x41) panic();
+  }
+
+  size_t b_count = 16;
+  void *b = malloc(b_count);
+  memset(b, 0x42, b_count); // fill `b` with 'B's
+
+  // `a` is now first on the free list
+  free(a);
+  // We're re-allocating the last region, so `realloc`'s optimization to re-use the region
+  // by altering its metadata's size (thus avoiding a memcpy) should be used.
+  void *c = realloc(b, 24); // realloc `b` should still contain 'B's
+
+  // ensure the contents of `b` were not clobbered. A previous bug caused `a`s region
+  // to be used for `c`.
+  char *p = (char *)c;
+  for (int i = 0; i < (int)b_count; i++) {
+    if (p[i] != 0x42) panic();
+  }
+
+  free(c);
+  return NULL;
+}
+
+bool tree_sitter_wasm_realloc_clobber_region_external_scanner_scan(
+  void *payload,
+  TSLexer *lexer,
+  const bool *valid_symbols
+) {
+  lexer->result_symbol = ZERO_WIDTH_TOKEN;
+  return true;
+}
+
+unsigned tree_sitter_wasm_realloc_clobber_region_external_scanner_serialize(
+  void *payload,
+  char *buffer
+) {
+  return 0;
+}
+
+void tree_sitter_wasm_realloc_clobber_region_external_scanner_deserialize(
+  void *payload,
+  const char *buffer,
+  unsigned length
+) {}
+
+void tree_sitter_wasm_realloc_clobber_region_external_scanner_destroy(void *payload) {}