Add some interface for wasm to deal with the breaking DB format changes

src/distances/distance_matrix.rs

@@ -312,6 +312,11 @@ impl<'a> SparseDistanceMatrix<'a> {
     pub fn dists_mut(&mut self) -> &mut DistVec {
         &mut self.distances
     }
+
+    /// Reference to the completed sparse distance entries.
+    pub fn dists_as_ref(&self) -> &DistVec {
+        &self.distances
+    }
 }
 
 impl<'a> Distances<'a> for SparseDistanceMatrix<'a> {

src/distances/jaccard.rs

@@ -10,11 +10,23 @@ pub fn jaccard_index(
     c1: Option<f64>,
     c2: Option<f64>,
     completeness_cutoff: f64,
+) -> f64 {
+    jaccard_index_generic::<BBITS>(sketch1, sketch2, sketchsize64, c1, c2, completeness_cutoff)
+}
+
+/// Returns the Jaccard index using the selected sketch bin width.
+pub fn jaccard_index_generic<const BITS: u64>(
+    sketch1: &[u64],
+    sketch2: &[u64],
+    sketchsize64: u64,
+    c1: Option<f64>,
+    c2: Option<f64>,
+    completeness_cutoff: f64,
 ) -> f64 {
     let unionsize = (u64::BITS as u64 * sketchsize64) as f64;
     let samebits: u32 = sketch1
-        .chunks_exact(BBITS as usize)
-        .zip(sketch2.chunks_exact(BBITS as usize))
+        .chunks_exact(BITS as usize)
+        .zip(sketch2.chunks_exact(BITS as usize))
         .map(|(chunk1, chunk2)| {
             let mut bits: u64 = !0;
             chunk1.iter().zip(chunk2.iter()).for_each(|(&s1, &s2)| {
@@ -24,7 +36,7 @@ pub fn jaccard_index(
         })
         .sum();
     let maxnbits = sketchsize64 as u32 * u64::BITS;
-    let expected_samebits = maxnbits >> BBITS;
+    let expected_samebits = maxnbits >> BITS;
 
     log::trace!("samebits:{samebits} expected_samebits:{expected_samebits} maxnbits:{maxnbits}");
     let diff = samebits.saturating_sub(expected_samebits);
@@ -65,6 +77,25 @@ pub fn core_acc_dist(
     query_sketch_idx: usize,
     completeness_vec: Option<&Vec<f64>>,
     completeness_cutoff: f64,
+) -> (f32, f32) {
+    core_acc_dist_generic::<BBITS>(
+        ref_sketches,
+        query_sketches,
+        ref_sketch_idx,
+        query_sketch_idx,
+        completeness_vec,
+        completeness_cutoff,
+    )
+}
+
+/// Calculates core/accessory distances using the selected sketch bin width.
+pub fn core_acc_dist_generic<const BITS: u64>(
+    ref_sketches: &MultiSketch,
+    query_sketches: &MultiSketch,
+    ref_sketch_idx: usize,
+    query_sketch_idx: usize,
+    completeness_vec: Option<&Vec<f64>>,
+    completeness_cutoff: f64,
 ) -> (f32, f32) {
     if ref_sketches.kmer_lengths().len() < 2 {
         panic!("Need at least two k-mer lengths to calculate core/accessory distances");
@@ -76,7 +107,7 @@ pub fn core_acc_dist(
     for (k_idx, k) in ref_sketches.kmer_lengths().iter().enumerate() {
         let c1 = completeness_vec.map(|cv| cv[ref_sketch_idx]);
         let c2 = completeness_vec.map(|cv| cv[query_sketch_idx]);
-        let y = jaccard_index(
+        let y = jaccard_index_generic::<BITS>(
             ref_sketches.get_sketch_slice(ref_sketch_idx, k_idx),
             query_sketches.get_sketch_slice(query_sketch_idx, k_idx),
             ref_sketches.sketchsize64,

src/distances/mod.rs

@@ -10,6 +10,7 @@ use crate::cli::RetainUnmatched;
 use crate::get_progress_bar;
 use crate::inverted::Inverted;
 use crate::sketch::multisketch::MultiSketch;
+use crate::sketch::BBITS;
 
 pub mod distance_matrix;
 use self::distance_matrix::*;
@@ -138,6 +139,31 @@ pub fn self_dists_knn<'a>(
     quiet: bool,
     completeness_vec: Option<&Vec<f64>>,
     completeness_cutoff: f64,
+) -> SparseDistanceMatrix<'a> {
+    self_dists_knn_generic::<BBITS>(
+        sketches,
+        n,
+        knn,
+        dist_type,
+        quiet,
+        completeness_vec,
+        completeness_cutoff,
+    )
+}
+
+/// Self kNN distances using a fixed sketch bin width.
+///
+/// Mandrake calls this with [`crate::sketch::CURRENT_BBITS`] for current
+/// databases; legacy 14-bit databases are intentionally outside the browser
+/// contract.
+pub fn self_dists_knn_generic<'a, const BITS: u64>(
+    sketches: &'a MultiSketch,
+    n: usize,
+    knn: usize,
+    dist_type: DistType,
+    quiet: bool,
+    completeness_vec: Option<&Vec<f64>>,
+    completeness_cutoff: f64,
 ) -> SparseDistanceMatrix<'a> {
     let mut sp_distances = SparseDistanceMatrix::new(sketches, knn, dist_type);
     let k_vals = sp_distances.k_vals();
@@ -162,7 +188,7 @@ pub fn self_dists_knn<'a>(
                         // This uses Option::map to safely access the completeness value for each sample.
                         let c1 = completeness_vec.map(|cv| cv[i]);
                         let c2 = completeness_vec.map(|cv| cv[j]);
-                        let dist = jaccard_index(
+                        let dist = jaccard_index_generic::<BITS>(
                             i_sketch,
                             sketches.get_sketch_slice(j, k_idx),
                             sketches.sketchsize64,
@@ -203,7 +229,7 @@ pub fn self_dists_knn<'a>(
                         if i == j {
                             continue;
                         }
-                        let dists = core_acc_dist(
+                        let dists = core_acc_dist_generic::<BITS>(
                             sketches,
                             sketches,
                             i,

src/sketch/mod.rs

@@ -30,8 +30,13 @@ pub mod sketch_datafile;
 #[cfg(not(target_family = "wasm"))]
 use self::sketch_datafile::SketchArrayWriter;
 
-/// Bin bits (lowest of 64-bits to keep)
+/// Bin bits (lowest of 64-bits to keep) used by the legacy sketch API.
+///
+/// Current databases use [`CURRENT_BBITS`]; the generic distance functions
+/// take the bin width explicitly so both formats can be read by the library.
 pub const BBITS: u64 = 14;
+/// Bin width used by current (version 0.4+) sketch databases.
+pub const CURRENT_BBITS: u64 = 16;
 /// Total width of all bins (used as sign % sign_mod)
 pub const SIGN_MOD: u64 = (1 << 61) - 1;
 

src/sketch/multisketch.rs

@@ -4,7 +4,7 @@
 use core::panic;
 use std::fmt;
 use std::fs::File;
-use std::io::{BufReader, BufWriter};
+use std::io::{BufReader, BufWriter, Cursor};
 use std::mem;
 
 use hashbrown::{HashMap, HashSet};
@@ -15,6 +15,7 @@
 use crate::sketch::sketch_datafile::SketchArrayReader;
 use crate::sketch::sketch_datafile::SketchArrayWriter;
 use crate::sketch::Sketch;
+use crate::sketch::CURRENT_BBITS;
 
 use super::sketch_datafile::append_batch;
 
@@ -102,6 +103,40 @@
         Ok(skm_obj)
     }
 
+    /// Loads metadata from an already materialised `.skm` file.
+    pub fn load_metadata_bytes(bytes: &[u8]) -> Result<Self, Error> {
+        let skm_file = BufReader::new(Cursor::new(bytes));
+        let decompress_reader = snap::read::FrameDecoder::new(skm_file);
+        let mut skm_obj: Self = ciborium::de::from_reader(decompress_reader)?;
+        if skm_obj.sketchsize64 == 0 {
+            skm_obj.sketchsize64 = skm_obj.sketch_size;
+            skm_obj.sketch_size *= 64;
+        }
+        Ok(skm_obj)
+    }
+
+    /// Loads a paired `.skm`/`.skd` database from bytes.
+    pub fn load_bytes(metadata: &[u8], data: &[u8]) -> Result<Self, Error> {
+        let mut sketches = Self::load_metadata_bytes(metadata)?;
+        sketches.read_sketch_data_bytes(data)?;
+        Ok(sketches)
+    }
+
+    /// Returns whether this database predates the 16-bit current format.
+    pub fn is_legacy_format(&self) -> bool {
+        let mut version = self.sketch_version.split('.');
+        let major = version.next().and_then(|part| part.parse::<u64>().ok());
+        let minor = version.next().and_then(|part| part.parse::<u64>().ok());
+        let current_version =
+            matches!((major, minor), (Some(major), Some(minor)) if major > 0 || minor >= 4);
+        let current_stride = self
+            .sketchsize64
+            .checked_mul(CURRENT_BBITS)
+            .map(|stride| stride as usize)
+            .is_some_and(|stride| self.kmer_stride == stride);
+        !(current_version && current_stride)
+    }
+
     /// Number of samples loaded from the .skm/.skd
     pub fn number_samples_loaded(&self) -> usize {
         match &self.block_reindex {
@@ -183,6 +218,30 @@
             sketch_reader.read_all_from_skd(self.sample_stride * self.sketch_metadata.len());
     }
 
+    /// Reads all sketch bins from an already materialised `.skd` file.
+    pub fn read_sketch_data_bytes(&mut self, bytes: &[u8]) -> Result<(), Error> {
+        if !bytes.len().is_multiple_of(std::mem::size_of::<u64>()) {
+            bail!("sketch data length is not a multiple of eight bytes");
+        }
+        let expected = self
+            .sample_stride
+            .checked_mul(self.sketch_metadata.len())
+            .ok_or_else(|| anyhow::anyhow!("sketch data size overflows memory limits"))?;
+        let values = bytes
+            .chunks_exact(8)
+            .map(|chunk| u64::from_le_bytes(chunk.try_into().expect("exact chunk size")))
+            .collect::<Vec<_>>();
+        if values.len() != expected {
+            bail!(
+                "sketch data contains {} bins but metadata requires {}",
+                values.len(),
+                expected
+            );
+        }
+        self.sketch_bins = values;
+        Ok(())
+    }
+
     /// Read a subset of the bins from an .skd file, in a memory efficient manner
     pub fn read_sketch_data_block(&mut self, file_prefix: &str, names: &[String]) {
         // Find the given names in the sketch metadata