Fixes #31 - initialization failure on Firefox (#35)

justinwilaby · May 13, 2020 · 57c5ea1 · 57c5ea1
1 parent 99cd444
commit 57c5ea1
Show file tree

Hide file tree

Showing 9 changed files with 62 additions and 26 deletions.
diff --git a/lib/SpellCheckerBase.js b/lib/SpellCheckerBase.js
@@ -106,6 +106,10 @@ class SpellcheckerBase {
      */
     writeToBuffer(chunk, memory) {
         if (!this.writeBuffer || this.writeBuffer.buffer !== memory.buffer || this.writeBuffer.byteLength < chunk.byteLength) {
+            if (memory.buffer.byteLength < chunk.byteLength) {
+                const delta = Math.ceil((chunk.byteLength - memory.buffer.byteLength) / 65536); // 64 * 1024;
+                memory.grow(delta);
+            }
             this.writeBuffer = new Uint8Array(memory.buffer, 0, chunk.byteLength);
         }
         this.writeBuffer.set(chunk, 0);

diff --git a/lib/browser/SpellcheckerWasm.js b/lib/browser/SpellcheckerWasm.js
@@ -56,8 +56,17 @@ class SpellcheckerWasm extends SpellCheckerBase_1.SpellcheckerBase {
                     write_to_dictionary(0, 1, false);
                     return;
                 }
-                this.writeToBuffer(readResult.value, memory);
-                write_to_dictionary(0, readResult.value.length, isBigram);
+                let p = 0;
+                // Enforce max chunk size to prevent overwriting memory
+                // allocated for the SymSpell struct in wasm.
+                // https://github.com/justinwilaby/spellchecker-wasm/issues/31
+                const chunkSize = 1024 * 1024;
+                while (p < readResult.value.length) {
+                    const slice = readResult.value.slice(p, p + chunkSize);
+                    p += chunkSize;
+                    this.writeToBuffer(slice, memory);
+                    write_to_dictionary(0, slice.byteLength, isBigram);
+                }
             }
         };
         await readStreamIntoDictionary(dictionaryFetchResponse.body.getReader(), false);

diff --git a/lib/browser/index.js b/lib/browser/index.js
diff --git a/lib/spellchecker-wasm.js b/lib/spellchecker-wasm.js
@@ -1,2 +1 @@
-import * as wasm from './spellchecker-wasm.wasm_bg.wasm';
-
+export * from "./spellchecker-wasm.wasm_bg.js";
diff --git a/lib/spellchecker-wasm.wasm b/lib/spellchecker-wasm.wasm
diff --git a/src/js/SpellCheckerBase.ts b/src/js/SpellCheckerBase.ts
@@ -184,6 +184,10 @@ export abstract class SpellcheckerBase {
      */
     protected writeToBuffer(chunk: Uint8Array, memory: WebAssembly.Memory): void {
         if (!this.writeBuffer || this.writeBuffer.buffer !== memory.buffer || this.writeBuffer.byteLength < chunk.byteLength) {
+            if (memory.buffer.byteLength < chunk.byteLength) {
+                const delta = Math.ceil((chunk.byteLength - memory.buffer.byteLength) / 65536) // 64 * 1024;
+                memory.grow(delta);
+            }
             this.writeBuffer = new Uint8Array(memory.buffer, 0, chunk.byteLength);
         }
         this.writeBuffer.set(chunk, 0);

diff --git a/src/js/__tests__/index.html b/src/js/__tests__/index.html
@@ -22,7 +22,8 @@
     const spellchecker = new SpellcheckerWasm();
     const {notEqual} = window.chai.assert;
 
-    describe('The SpellcheckerWasm', () => {
+    describe('The SpellcheckerWasm', function() {
+        this.timeout(30000);
         it('should read from the supplied dictionary and perform lookups', async () => {
             const wasm = await fetch('../../../lib/spellchecker-wasm.wasm');
             const dictionary = await fetch('../../../lib/frequency_dictionary_en_82_765.txt');

diff --git a/src/js/browser/SpellcheckerWasm.ts b/src/js/browser/SpellcheckerWasm.ts
@@ -68,8 +68,17 @@ export class SpellcheckerWasm extends SpellcheckerBase {
                     write_to_dictionary(0, 1, false);
                     return
                 }
-                this.writeToBuffer(readResult.value, memory);
-                write_to_dictionary(0, readResult.value.length, isBigram);
+                let p = 0;
+                // Enforce max chunk size to prevent overwriting memory
+                // allocated for the SymSpell struct in wasm.
+                // https://github.com/justinwilaby/spellchecker-wasm/issues/31
+                const chunkSize = 1024 * 1024;
+                while(p < readResult.value.length) {
+                    const slice = readResult.value.slice(p, p + chunkSize);
+                    p += chunkSize;
+                    this.writeToBuffer(slice, memory);
+                    write_to_dictionary(0, slice.byteLength, isBigram);
+                }
             }
         };
 

diff --git a/src/spellchecker_wasm.rs b/src/spellchecker_wasm.rs
@@ -1,59 +1,69 @@
+use std::cell::RefCell;
 use std::mem::transmute;
 use std::slice;
 use std::str;
+use std::mem;
 
-use crate::sym_spell::sym_spell::SymSpell;
-use crate::sym_spell::verbosity::Verbosity;
 use crate::sym_spell::Encode;
 use crate::sym_spell::suggested_item::SuggestItem;
+use crate::sym_spell::sym_spell::SymSpell;
+use crate::sym_spell::verbosity::Verbosity;
 
-static mut SYM: *mut SymSpell = 0 as *mut SymSpell;
-static mut BUFFER: *mut Vec<u8> = 0 as *mut Vec<u8>;
+static mut BUFFER: Option<RefCell<Vec<u8>>> = None;
+static mut SYM: Option<RefCell<SymSpell>> = None;
 
 #[no_mangle]
 pub unsafe extern fn symspell(max_dictionary_edit_distance: usize, count_threshold: usize) {
-    if SYM == 0 as *mut SymSpell {
-        let sym_spell: SymSpell = SymSpell::new(Some(max_dictionary_edit_distance), Some(7), Some(count_threshold));
-        SYM = transmute(Box::new(sym_spell));
-        BUFFER = transmute(Box::new(vec![] as Vec<u8>));
-    }
+    let sym = SymSpell::new(Some(max_dictionary_edit_distance), Some(7), Some(count_threshold));
+
+    SYM = Some(RefCell::new(sym));
+    BUFFER = Some(RefCell::new(Vec::new()));
 }
 
 #[no_mangle]
 pub unsafe extern fn write_to_dictionary(ptr: *const u8, length: usize, is_bigram: bool) {
-    (*BUFFER).extend_from_slice(slice::from_raw_parts(ptr, length));
-    let len = (*BUFFER).len();
+    let buffer_cell = BUFFER.as_ref().unwrap();
+    let sym_cell = SYM.as_ref().unwrap();
+    let mut sym = sym_cell.borrow_mut();
+    let mut buffer = buffer_cell.borrow_mut();
+    buffer.extend_from_slice(slice::from_raw_parts(ptr, length));
+
+    let len = buffer.len();
     let mut cursor: usize = 0;
     for i in 0..len {
-        let ch = (*BUFFER)[i];
+        let ch = buffer[i];
         if ch == b'\n' {
             if i > 1 {
-                let chunk = str::from_utf8_unchecked(&(*BUFFER)[cursor..i - 1]);  // do not write the '\n' char
+                let chunk = str::from_utf8_unchecked(&buffer[cursor..i - 1]);  // do not write the '\n' char
                 if is_bigram {
-                    (*SYM).write_line_to_bigram_dictionary(chunk, " ");
+                    sym.write_line_to_bigram_dictionary(chunk, " ");
                 } else {
-                    (*SYM).write_line_to_dictionary(chunk, " ");
+                    sym.write_line_to_dictionary(chunk, " ");
                 }
             }
             cursor = i + 1; // skip the '\n' char for the next iteration
         }
     }
 
-    (*BUFFER).drain(0..cursor);
+    buffer.drain(0..cursor);
 }
 
 #[no_mangle]
 pub unsafe extern fn lookup(ptr: *mut u8, length: usize, verbosity: Verbosity, max_edit_distance: usize, include_unknown: bool, include_self: bool) {
+    let sym_cell = SYM.as_ref().unwrap();
+    let sym = sym_cell.borrow();
     let bytes = slice::from_raw_parts(ptr, length);
-    let results = (*SYM).lookup(str::from_utf8_unchecked(bytes), verbosity, max_edit_distance, include_unknown, include_self);
+    let results = sym.lookup(str::from_utf8_unchecked(bytes), verbosity, max_edit_distance, include_unknown, include_self);
 
     emit_results(results)
 }
 
 #[no_mangle]
 pub unsafe extern fn lookup_compound(ptr: *mut u8, length: usize, max_edit_distance: usize) {
+    let sym_cell = SYM.as_ref().unwrap();
+    let sym = sym_cell.borrow();
     let bytes = slice::from_raw_parts(ptr, length);
-    let results = (*SYM).lookup_compound(str::from_utf8_unchecked(bytes), max_edit_distance);
+    let results = sym.lookup_compound(str::from_utf8_unchecked(bytes), max_edit_distance);
 
     emit_results(results);
 }