From 74812a6a0177db5bbb83cff61161b1282e0d2851 Mon Sep 17 00:00:00 2001
From: Kelvin Ly <kelvin.ly1618@gmail.com>
Date: Fri, 12 May 2023 12:20:11 -0400
Subject: [PATCH] Looks like we've got meaningful communication between Rust
 and JavaScript, time for some refactoring

---
 .gitignore            |   1 +
 Makefile              |   4 +-
 wordle.htm            |   3 +
 wordle_opt/src/lib.rs | 125 +++++++++++++++++++++++++++++++++---------
 wordle_shim.js        |  22 ++++++++
 5 files changed, 126 insertions(+), 29 deletions(-)
 create mode 100644 .gitignore
 create mode 100644 wordle_shim.js

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..31114b3
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+wordle_opt.wasm
diff --git a/Makefile b/Makefile
index 3ffccf2..b56f13a 100644
--- a/Makefile
+++ b/Makefile
@@ -1,8 +1,8 @@
 all: wordle_opt.wasm
 
 wordle_opt.wasm: wordle_opt/target/wasm32-unknown-unknown/release/wordle_opt.wasm
-	#wasm-snip $< -o $@ 
-	cp $< $@
+	wasm-snip $< -o $@  --snip-rust-fmt-code --snip-rust-panicking-code
+	#cp $< $@
 
 wordle_opt/target/wasm32-unknown-unknown/release/wordle_opt.wasm: wordle_opt/src/*.rs
 	cd wordle_opt; cargo build --target wasm32-unknown-unknown --release
diff --git a/wordle.htm b/wordle.htm
index 322c3da..d443cfe 100644
--- a/wordle.htm
+++ b/wordle.htm
@@ -27,6 +27,7 @@
 </style>
 
 <script type="text/javascript" src="wordle.js"></script>
+<script type="text/javascript" src="wordle_shim.js"></script>
 <script type="text/javascript">
 
 var words = []
@@ -52,6 +53,8 @@ var valid_words = null;
     console.log("found " + words.length+ " words")
     status("found " + words.length + " words")
     valid_words = gen_all_words_list()
+
+    init_wasm_solver(text)
   })
   req.send()
 
diff --git a/wordle_opt/src/lib.rs b/wordle_opt/src/lib.rs
index e333eb4..660298e 100644
--- a/wordle_opt/src/lib.rs
+++ b/wordle_opt/src/lib.rs
@@ -7,11 +7,20 @@ static ALLOC: wee_alloc::WeeAlloc = wee_alloc::WeeAlloc::INIT;
 
 // webassembly version of all the wordle logic
 // because the javascript version is slow as hell on firefox
-use std::os::raw::c_char;
 
-static mut LOOKUP: Option<&'static mut [u8]> = None;
-static mut STRINGS: Option<&'static mut [u8]> = None;
-static mut STR_IDXS: Option<&'static mut [*const c_char]> = None;
+// TODO wrap all this in a struct to reduce the insane amount of
+// unsafe in the code
+pub static mut LOOKUP: &'static mut [u8] = &mut [];
+pub static mut STRINGS: &'static mut [u8] = &mut [];
+pub static mut STR_IDXS: &'static mut [*const u8] = &mut [];
+pub static mut VALID_STRS: &'static mut [usize] = &mut [];
+
+#[no_mangle]
+pub fn idxs_offset() -> *const *const u8 {
+    unsafe {
+        STR_IDXS.as_ptr()
+    }
+}
 
 fn unwrap_or_abort<T>(v: Option<T>) -> T {
     match v {
@@ -37,42 +46,104 @@ unsafe fn dealloc_ary<T>(v: &'static mut [T]) {
     dealloc(v.as_ptr() as *mut u8, layout);
 }
 
+fn realloc_ary<T>(v: &'static mut [T], sz: usize) -> &'static mut [T] {
+    unsafe {
+        if v.len() > 0 {
+            dealloc_ary(v);
+        }
+        alloc_ary(sz)
+    }
+}
+
+
+extern {
+    fn fill_string(p: *mut u8);
+    fn log_num_idxs(i: usize);
+}
 
 #[no_mangle]
-pub extern fn init_strings(sz: usize) -> *mut c_char {
+pub extern fn init(str_sz: usize) {
     unsafe {
-        if let Some(v) = &mut STRINGS {
-            dealloc_ary(v);
-            STRINGS = None;
+        STRINGS = realloc_ary(STRINGS, str_sz);
+        fill_string(STRINGS.as_mut_ptr());
+    }
+
+    let num_strs = init_idx();
+    unsafe {
+        log_num_idxs(num_strs);
+
+        let num_strings = STR_IDXS.len();
+        LOOKUP = realloc_ary(LOOKUP, num_strings*num_strings);
+    }
+}
+
+fn init_idx() -> usize {
+    let strings = unsafe { &*STRINGS };
+
+    let mut last_alpha = false;
+    let mut num_words = 0;
+    for v in strings.iter() {
+        let cur_alpha = *v >= ('a' as u8) && *v <= ('z' as u8);
+        if cur_alpha && !last_alpha {
+            num_words += 1;
         }
-        STRINGS = Some(alloc_ary(sz));
-        unwrap_or_abort(STRINGS.as_deref_mut().map(|v| v.as_mut_ptr() as *mut c_char))
+        last_alpha = cur_alpha;
+    }
+
+    unsafe {
+        if STR_IDXS.len() > 0 {
+            dealloc_ary(STR_IDXS);
+        }
+        STR_IDXS = alloc_ary(num_words);
+    }
+
+    let idxs = unsafe { &mut *STR_IDXS };
+    let mut idx = 0;
+    for (i, v) in strings.iter().enumerate() {
+        let cur_alpha = *v >= ('a' as u8) && *v <= ('z' as u8);
+        if cur_alpha && !last_alpha {
+            idxs[idx] = &strings[i] as *const u8;
+            idx += 1;
+        }
+        last_alpha = cur_alpha;
+    }
+
+    num_words
+}
+
+static mut IDX: usize = 0;
+#[no_mangle]
+fn precalc(num_words: usize) {
+    let mut i = unsafe { IDX };
+    let idxs = unsafe { &*STR_IDXS };
+    let lookup = unsafe { &mut *LOOKUP };
+    let l = idxs.len();
+
+    for _ in 0..num_words {
+        if i >= idxs.len() {
+            break;
+        }
+        let guess = idxs[i];
+        for (j, reference) in idxs.iter().enumerate() {
+            lookup[i*l + j] = calc_match(guess, *reference);
+        }
+        i += 1;
+    }
+
+    unsafe {
+        IDX = i;
     }
 }
 
 #[no_mangle]
-pub extern fn init_idx() -> usize {
-    // TODO
-    0
-}
-
-#[no_mangle]
-pub extern fn init_lookup() -> *mut c_char {
+fn precalc_done() -> bool {
     unsafe {
-        let num_strings = STR_IDXS.as_ref().map(|idxs| idxs.len()).unwrap_or(0);
-        if let Some(v) = &mut LOOKUP {
-            dealloc_ary(v);
-            STRINGS = None;
-        }
-
-        LOOKUP = Some(alloc_ary(num_strings*num_strings));
-        unwrap_or_abort(
-            LOOKUP.as_deref_mut().map(|v| v.as_mut_ptr() as *mut c_char))
+        IDX >= STR_IDXS.len()
     }
 }
 
 #[no_mangle]
-pub extern fn calc_match(guess: *const c_char, reference: *const c_char) -> u8 {
+fn calc_match(guess: *const u8, reference: *const u8) -> u8 {
     let mut unmatched = vec![0u8; 26];
     let mut ret = 0;
     let mut matched = 0;
diff --git a/wordle_shim.js b/wordle_shim.js
new file mode 100644
index 0000000..d21bbcf
--- /dev/null
+++ b/wordle_shim.js
@@ -0,0 +1,22 @@
+var wasm_solver = null
+
+function init_wasm_solver(words_str) {
+  const fill_string = (offset) => {
+    console.log("fill strings called " + offset)
+    const str_buf = new Uint8Array(wasm_solver.exports.memory.buffer, offset)
+    const buf = new TextEncoder().encode(words_str, "utf8")
+    str_buf.set(buf, 0)
+  }
+  const log_num_idxs = i => {
+    console.log("wasm: found " + i + " words")
+  }
+  WebAssembly.instantiateStreaming(fetch("wordle_opt.wasm"), {
+    env: {
+      fill_string: fill_string,
+      log_num_idxs: log_num_idxs,
+    }
+  }).then(wm => {
+    wasm_solver = wm.instance
+    wasm_solver.exports.init(words_str.length)
+  })
+}