From b59716600ab78fe53d71d3954cd98d446ea22a73 Mon Sep 17 00:00:00 2001
From: Max Brunsfeld <maxbrunsfeld@gmail.com>
Date: Fri, 24 Nov 2023 20:03:54 -0800
Subject: [PATCH] Add a test that loads multiple languages in a wasm store

---
 cli/src/tests/helpers/dirs.rs       |  1 +
 cli/src/tests/mod.rs                |  3 ++
 cli/src/tests/wasm_language_test.rs | 46 +++++++++++++++++++++++++++++
 lib/src/wasm.c                      |  6 ++++
 4 files changed, 56 insertions(+)
 create mode 100644 cli/src/tests/wasm_language_test.rs

diff --git a/cli/src/tests/helpers/dirs.rs b/cli/src/tests/helpers/dirs.rs
index 7ff48ab7..2eac45d5 100644
--- a/cli/src/tests/helpers/dirs.rs
+++ b/cli/src/tests/helpers/dirs.rs
@@ -8,6 +8,7 @@ lazy_static! {
         fs::create_dir_all(&result).unwrap();
         result
     };
+    pub static ref WASM_DIR: PathBuf = ROOT_DIR.join("target").join("release");
     pub static ref SCRATCH_DIR: PathBuf = {
         // https://doc.rust-lang.org/reference/conditional-compilation.html
         let vendor = if cfg!(target_vendor = "apple") {
diff --git a/cli/src/tests/mod.rs b/cli/src/tests/mod.rs
index 308fc2c5..e09dc838 100644
--- a/cli/src/tests/mod.rs
+++ b/cli/src/tests/mod.rs
@@ -14,3 +14,6 @@ mod test_highlight_test;
 mod test_tags_test;
 mod text_provider_test;
 mod tree_test;
+
+#[cfg(feature = "wasm")]
+mod wasm_language_test;
diff --git a/cli/src/tests/wasm_language_test.rs b/cli/src/tests/wasm_language_test.rs
new file mode 100644
index 00000000..ec64066f
--- /dev/null
+++ b/cli/src/tests/wasm_language_test.rs
@@ -0,0 +1,46 @@
+use crate::tests::helpers::fixtures::WASM_DIR;
+use lazy_static::lazy_static;
+use std::fs;
+use tree_sitter::{wasmtime::Engine, Parser, WasmStore};
+
+lazy_static! {
+    static ref ENGINE: Engine = Engine::default();
+}
+
+#[test]
+fn test_wasm_store() {
+    let mut store = WasmStore::new(ENGINE.clone());
+    let mut parser = Parser::new();
+
+    let wasm_cpp = fs::read(&WASM_DIR.join(format!("tree-sitter-cpp.wasm"))).unwrap();
+    let wasm_rs = fs::read(&WASM_DIR.join(format!("tree-sitter-rust.wasm"))).unwrap();
+    let wasm_rb = fs::read(&WASM_DIR.join(format!("tree-sitter-ruby.wasm"))).unwrap();
+
+    let language_rust = store.load_language("rust", &wasm_rs);
+    let language_cpp = store.load_language("cpp", &wasm_cpp);
+    let language_ruby = store.load_language("ruby", &wasm_rb);
+    parser.set_wasm_store(store).unwrap();
+
+    for _ in 0..2 {
+        parser.set_language(language_cpp).unwrap();
+        let tree = parser.parse("A<B> c = d();", None).unwrap();
+        assert_eq!(
+            tree.root_node().to_sexp(),
+            "(translation_unit (declaration type: (template_type name: (type_identifier) arguments: (template_argument_list (type_descriptor type: (type_identifier)))) declarator: (init_declarator declarator: (identifier) value: (call_expression function: (identifier) arguments: (argument_list)))))"
+        );
+
+        parser.set_language(language_rust).unwrap();
+        let tree = parser.parse("const A: B = c();", None).unwrap();
+        assert_eq!(
+            tree.root_node().to_sexp(),
+            "(source_file (const_item name: (identifier) type: (type_identifier) value: (call_expression function: (identifier) arguments: (arguments))))"
+        );
+
+        parser.set_language(language_ruby).unwrap();
+        let tree = parser.parse("class A; end", None).unwrap();
+        assert_eq!(
+            tree.root_node().to_sexp(),
+            "(program (class name: (constant)))"
+        );
+    }
+}
diff --git a/lib/src/wasm.c b/lib/src/wasm.c
index ccfdc9c4..a5538da7 100644
--- a/lib/src/wasm.c
+++ b/lib/src/wasm.c
@@ -704,6 +704,12 @@ static bool ts_wasm_store__instantiate(
   wasmtime_error_t *error = NULL;
   wasm_trap_t *trap = NULL;
 
+  // Grow the function table to make room for the new functions.
+  wasmtime_val_t initializer = {.kind = WASMTIME_FUNCREF};
+  uint32_t prev_size;
+  error = wasmtime_table_grow(context, &self->function_table, dylink_info->table_size, &initializer, &prev_size);
+  assert(!error);
+
   // Construct the language function name as string.
   unsigned prefix_len = strlen("tree_sitter_");
   size_t name_len = strlen(language_name);