fix(wasm): correct several bugs in realloc

- free memory if 0 size is passed in
- Don't `memcpy` contents if new pointer is `NULL`
- Copy old region's contents only up to size of new region
- free old region

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

crates/cli/src/tests/wasm_language_test.rs

@@ -118,6 +118,19 @@ fn test_load_fixture_language_wasm() {
     });
 }
 
+#[test]
+fn test_wasm_realloc_smaller_size() {
+    allocations::record(|| {
+        let store = WasmStore::new(&ENGINE).unwrap();
+        let mut parser = Parser::new();
+        let language = get_test_fixture_language_wasm("wasm_realloc");
+        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

@@ -117,6 +117,11 @@ void *realloc(void *ptr, size_t new_size) {
   if (ptr == NULL) {
     return malloc(new_size);
   }
+  if (new_size == 0) {
+    free(ptr);
+    return NULL;
+  }
+
 
   Region *region = region_for_ptr(ptr);
   Region *region_end = region_after(region, region->size);
@@ -129,7 +134,12 @@ void *realloc(void *ptr, size_t new_size) {
   }
 
   void *result = malloc(new_size);
-  memcpy(result, &region->data, region->size);
+  if (!result) return NULL;
+
+  size_t copy_size = region->size < new_size ? region->size : new_size;
+  memcpy(result, &region->data, copy_size);
+
+  free(ptr);
   return result;
 }
 

lib/src/wasm_store.c

@@ -255,7 +255,7 @@ static bool wasm_dylink_info__parse(
  * Native callbacks exposed to Wasm modules
  *******************************************/
 
- static wasm_trap_t *callback__abort(
+static wasm_trap_t *callback__abort(
   void *env,
   wasmtime_caller_t* caller,
   wasmtime_val_raw_t *args_and_results,

test/fixtures/test_grammars/wasm_realloc/corpus.txt

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

test/fixtures/test_grammars/wasm_realloc/grammar.js

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

test/fixtures/test_grammars/wasm_realloc/scanner.c

@@ -0,0 +1,43 @@
+#include "tree_sitter/parser.h"
+
+// Constant copied from `crates/language/wasm/src/stdlib.c`,
+// Must be kept in sync for a reliable repro. 
+#define MAX_HEAP_SIZE (4 * 1024 * 1024)
+
+enum TokenType {
+  ZERO_WIDTH_TOKEN
+};
+
+void *tree_sitter_wasm_realloc_external_scanner_create(void) {
+  size_t large_size = MAX_HEAP_SIZE - 64;
+  void *p = malloc(large_size);
+  void *q = malloc(4);
+  p = realloc(p, 4);
+  free(p);
+  free(q);
+  return NULL;
+}
+
+bool tree_sitter_wasm_realloc_external_scanner_scan(
+  void *payload,
+  TSLexer *lexer,
+  const bool *valid_symbols
+) {
+  lexer->result_symbol = ZERO_WIDTH_TOKEN;
+  return true;
+}
+
+unsigned tree_sitter_wasm_realloc_external_scanner_serialize(
+  void *payload,
+  char *buffer
+) {
+  return 0;
+}
+
+void tree_sitter_wasm_realloc_external_scanner_deserialize(
+  void *payload,
+  const char *buffer,
+  unsigned length
+) {}
+
+void tree_sitter_wasm_realloc_external_scanner_destroy(void *payload) {}