use wasm_bindgen::prelude::*;
use crate::config::{Config, Metric};
use crate::hnsw::{Hnsw, Neighbor};
use crate::Error;
#[derive(serde::Serialize)]
struct Hit {
id: String,
distance: f32,
}
impl From<Neighbor> for Hit {
fn from(n: Neighbor) -> Self {
Hit {
id: n.id,
distance: n.distance,
}
}
}
fn metric_from_str(s: &str) -> Result<Metric, JsValue> {
match s.to_ascii_lowercase().as_str() {
"cosine" => Ok(Metric::Cosine),
"dot" => Ok(Metric::Dot),
"l2" | "euclidean" => Ok(Metric::L2),
other => Err(JsValue::from_str(&format!(
"unknown metric {other:?} (expected \"cosine\", \"dot\", or \"l2\")"
))),
}
}
fn err_to_js(e: Error) -> JsValue {
js_sys::Error::new(&e.to_string()).into()
}
#[wasm_bindgen]
pub struct FerrovecCore {
inner: Hnsw,
}
#[wasm_bindgen]
impl FerrovecCore {
#[wasm_bindgen(constructor)]
pub fn new(dims: usize) -> FerrovecCore {
console_error_panic_hook::set_once();
FerrovecCore {
inner: Hnsw::new(dims),
}
}
#[wasm_bindgen(js_name = withConfig)]
pub fn with_config(
dims: usize,
max_connections: usize,
ef_construction: usize,
ef_search: usize,
metric: &str,
seed: f64,
) -> Result<FerrovecCore, JsValue> {
console_error_panic_hook::set_once();
let metric = metric_from_str(metric)?;
let config = Config {
max_connections,
ef_construction,
ef_search,
metric,
seed: seed.to_bits(),
};
Ok(FerrovecCore {
inner: Hnsw::with_config(dims, config),
})
}
pub fn insert(&mut self, id: String, vector: &[f32]) -> Result<(), JsValue> {
self.inner.insert(id, vector).map_err(err_to_js)
}
pub fn search(&self, query: &[f32], k: usize) -> Result<JsValue, JsValue> {
let hits: Vec<Hit> = self
.inner
.search(query, k)
.map_err(err_to_js)?
.into_iter()
.map(Hit::from)
.collect();
serde_wasm_bindgen::to_value(&hits).map_err(|e| JsValue::from_str(&e.to_string()))
}
pub fn remove(&mut self, id: &str) -> bool {
self.inner.remove(id)
}
pub fn contains(&self, id: &str) -> bool {
self.inner.contains(id)
}
pub fn compact(&mut self) {
self.inner.compact();
}
pub fn clear(&mut self) {
self.inner.clear();
}
#[wasm_bindgen(js_name = len)]
pub fn len(&self) -> usize {
self.inner.len()
}
#[wasm_bindgen(js_name = isEmpty)]
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
pub fn dims(&self) -> usize {
self.inner.dims()
}
#[wasm_bindgen(js_name = toBytes)]
pub fn to_bytes(&self) -> Result<Vec<u8>, JsValue> {
self.inner.to_bytes().map_err(err_to_js)
}
#[wasm_bindgen(js_name = fromBytes)]
pub fn from_bytes(data: &[u8]) -> Result<FerrovecCore, JsValue> {
Hnsw::from_bytes(data)
.map(|inner| FerrovecCore { inner })
.map_err(err_to_js)
}
}