performance improvements

Author: holmrenser

.github/workflows/release.yml

@@ -63,12 +63,23 @@ jobs:
       - uses: PyO3/maturin-action@v1
         with:
           command: build
+          # Nightly toolchain: the published wheels enable `nj/simd`, which uses
+          # the unstable `portable_simd` feature for the explicit core::simd
+          # distance kernel (~5-7x faster distance computation than scalar).
+          # End users install the prebuilt wheel and never need nightly themselves.
+          # Pinned for reproducible releases — keep in sync with the `simd-nightly`
+          # job in test.yml; bump both together. The sdist (built on stable) keeps
+          # the autovectorized scalar kernel, so source installs need no nightly.
+          rust-toolchain: nightly-2026-06-02
+          # `--features` is listed in full (not just the simd delta) so the build is
+          # correct whether maturin merges with or replaces `[tool.maturin] features`.
           # The crate is built with pyo3's `abi3-py310` feature, so a single
           # stable-ABI wheel per platform covers CPython 3.10+ — no per-minor
           # interpreter list needed.
           args: >
             --release
             --out dist
+            --features pyo3/extension-module,nj/parallel,nj/simd
           target: ${{ matrix.target }}
           manylinux: ${{ matrix.manylinux }}
           working-directory: python

.github/workflows/test.yml

@@ -32,6 +32,25 @@ jobs:
       - run: make test
         working-directory: nj
 
+  simd-nightly:
+    name: SIMD kernel (nightly)
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      # The published PyPI wheels enable `nj/simd` (the explicit core::simd
+      # `portable_simd` kernel). This job guards that path against bit-rot and
+      # confirms it stays bit-identical to the scalar kernel (same unit/proptest/
+      # golden assertions). Pinned to match the wheel build in release.yml — bump
+      # both together.
+      - uses: dtolnay/rust-toolchain@master
+        with:
+          toolchain: nightly-2026-06-02
+      - uses: Swatinem/rust-cache@v2
+        with:
+          key: nightly-simd
+      - run: cargo test --features simd
+        working-directory: nj
+
   wasm:
     name: WASM tests
     runs-on: ubuntu-latest

README.md

@@ -20,6 +20,10 @@ cargo install nj --features cli
 
 nj sequences.fasta
 nj --substitution-model kimura2-p --n-bootstrap-samples 100 sequences.fasta > tree.nwk
+
+# Bootstrap and distance computation run in parallel by default (the `cli`
+# feature enables threading). Cap the worker count with -t/--num-threads:
+nj -t 4 --n-bootstrap-samples 1000 sequences.fasta > tree.nwk
 ```
 
 A progress bar is shown on stderr when bootstrapping.

nj/Cargo.toml

@@ -58,8 +58,15 @@ optional = true
 # default build only builds the library
 default = []
 
-# feature to build CLI
-cli = ["clap", "indicatif", "serde_json"]
+# feature to build CLI. Implies `parallel` so the shipped binary
+# (`cargo install nj --features cli`) is multi-threaded by default; the bare
+# library stays lightweight and Rayon-free unless `parallel` is requested.
+cli = ["clap", "indicatif", "serde_json", "parallel"]
 
 # feature to enable parallel bootstrap and distance matrix computation via Rayon
 parallel = ["rayon"]
+
+# feature to enable the explicit core::simd distance kernel. Requires a nightly
+# toolchain (uses the unstable `portable_simd` feature). Without it, the kernel
+# uses the branchless scalar path, which the compiler still autovectorizes.
+simd = []

nj/benches/nj_bench.rs

@@ -73,6 +73,17 @@ fn bench_distance_matrix(c: &mut Criterion) {
     group.finish();
 }
 
+/// Wide-alignment regime (large L, modest n) where the O(n²·L) distance kernel
+/// dominates — the case explicit/auto SIMD on the kernel targets.
+fn bench_distance_matrix_wide(c: &mut Criterion) {
+    let mut group = c.benchmark_group("distance_matrix_wide");
+    let msa = synthetic_dna(100, 5000);
+    group.bench_with_input(BenchmarkId::from_parameter(100), &msa, |b, msa| {
+        b.iter(|| distance_matrix(black_box(dist_config(msa.clone()))).unwrap());
+    });
+    group.finish();
+}
+
 fn bench_nj(c: &mut Criterion) {
     let mut group = c.benchmark_group("nj");
     for &n in &[25usize, 100, 250] {
@@ -95,5 +106,11 @@ fn bench_bootstrap(c: &mut Criterion) {
     group.finish();
 }
 
-criterion_group!(benches, bench_distance_matrix, bench_nj, bench_bootstrap);
+criterion_group!(
+    benches,
+    bench_distance_matrix,
+    bench_distance_matrix_wide,
+    bench_nj,
+    bench_bootstrap
+);
 criterion_main!(benches);

nj/src/alphabet.rs

@@ -18,8 +18,15 @@ use serde::{Deserialize, Serialize};
 /// [`crate::msa::MSA`] can be parameterised over any alphabet without
 /// duplicating logic.
 pub trait AlphabetEncoding {
+    /// Encoded symbol type. Implementors must be `#[repr(u8)]` (size 1, align 1)
+    /// so that a `&[Self::Symbol]` can be reinterpreted as `&[u8]` by
+    /// [`as_bytes`](AlphabetEncoding::as_bytes) for the byte-level distance kernel.
     type Symbol: Copy;
 
+    /// Byte value of the gap symbol (`Self::Symbol::Gap as u8`). Used by the
+    /// branchless / SIMD distance kernel for gap detection without enum matching.
+    const GAP_BYTE: u8;
+
     /// Encodes a byte into the corresponding symbol in the alphabet.
     fn encode(symbol: u8) -> Self::Symbol;
 
@@ -28,6 +35,18 @@ pub trait AlphabetEncoding {
 
     /// Returns `true` if `symbol` represents a gap (`-`).
     fn is_gap(symbol: Self::Symbol) -> bool;
+
+    /// Reinterprets an encoded sequence as a raw byte slice.
+    ///
+    /// The default implementation relies on the trait invariant that
+    /// `Self::Symbol` is `#[repr(u8)]`, so a slice of symbols has the exact same
+    /// layout as a `&[u8]` of equal length. All built-in alphabets uphold this.
+    fn as_bytes(seq: &[Self::Symbol]) -> &[u8] {
+        debug_assert_eq!(core::mem::size_of::<Self::Symbol>(), 1);
+        // SAFETY: `Self::Symbol` is `#[repr(u8)]` (1 byte, align 1), so the slice
+        // has identical layout to `&[u8]` of the same length.
+        unsafe { core::slice::from_raw_parts(seq.as_ptr() as *const u8, seq.len()) }
+    }
 }
 
 /// A single nucleotide in the DNA alphabet.
@@ -36,13 +55,14 @@ pub trait AlphabetEncoding {
 /// character in aligned sequences. Any byte not matching `A/C/G/T/N/-`
 /// (case-insensitive) is mapped to `N`.
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[repr(u8)]
 pub enum DnaSymbol {
-    A,
-    C,
-    G,
-    T,
-    N,
-    Gap,
+    A = 0,
+    C = 1,
+    G = 2,
+    T = 3,
+    N = 4,
+    Gap = 5,
 }
 
 /// Marker struct for the DNA alphabet.
@@ -54,6 +74,8 @@ pub struct DNA;
 impl AlphabetEncoding for DNA {
     type Symbol = DnaSymbol;
 
+    const GAP_BYTE: u8 = DnaSymbol::Gap as u8;
+
     fn encode(symbol: u8) -> Self::Symbol {
         match symbol {
             b'A' | b'a' => DnaSymbol::A,
@@ -88,29 +110,30 @@ impl AlphabetEncoding for DNA {
 /// (the `-` alignment character). Encoding is case-insensitive; any byte that
 /// does not match a known amino acid letter is mapped to `X`.
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
+#[repr(u8)]
 pub enum ProteinSymbol {
-    A,
-    R,
-    N,
-    D,
-    C,
-    Q,
-    E,
-    G,
-    H,
-    I,
-    L,
-    K,
-    M,
-    F,
-    P,
-    S,
-    T,
-    W,
-    Y,
-    V,
-    X,
-    Gap,
+    A = 0,
+    R = 1,
+    N = 2,
+    D = 3,
+    C = 4,
+    Q = 5,
+    E = 6,
+    G = 7,
+    H = 8,
+    I = 9,
+    L = 10,
+    K = 11,
+    M = 12,
+    F = 13,
+    P = 14,
+    S = 15,
+    T = 16,
+    W = 17,
+    Y = 18,
+    V = 19,
+    X = 20,
+    Gap = 21,
 }
 
 /// Marker struct for the protein alphabet.
@@ -122,6 +145,8 @@ pub struct Protein;
 impl AlphabetEncoding for Protein {
     type Symbol = ProteinSymbol;
 
+    const GAP_BYTE: u8 = ProteinSymbol::Gap as u8;
+
     fn encode(symbol: u8) -> Self::Symbol {
         // Match on the uppercased byte so lowercase residues encode identically
         // (consistent with `DNA::encode` and `detect_alphabet`, which both

nj/src/config.rs

@@ -32,6 +32,11 @@ impl SequenceObject {
     pub fn len(&self) -> usize {
         self.sequence.len()
     }
+
+    /// Returns `true` if the sequence is empty.
+    pub fn is_empty(&self) -> bool {
+        self.sequence.is_empty()
+    }
 }
 
 /// Configuration for distance-only computation (no NJ, no bootstrap).

nj/src/distance_matrix.rs

@@ -260,7 +260,7 @@ mod tests {
     fn test_dist_from_msa_no_overlap() {
         // no non-gap mismatches are counted, so the distance remains 0.0
         let seqs = vec![("X".into(), "A--".into()), ("Y".into(), "--A".into())];
-        let msa = MSA::<DNA>::from_iter(seqs.into_iter());
+        let msa = MSA::<DNA>::from_iter(seqs);
         let mat = msa.into_dist::<PDiff>();
         assert_eq!(mat.names, vec!["X", "Y"]);
         assert!((mat.get(0, 1) - 0.0).abs() < 1e-12);
@@ -329,7 +329,7 @@ mod tests {
         // ensure non-negative branch lengths in the produced tree (clamping is applied in implementation)
         fn check_nonneg(node: &TreeNode) {
             match &node.children {
-                None => return,
+                None => (),
                 Some(children) => {
                     assert!(children[0].len.unwrap() >= -1e-12, "left_len negative");
                     assert!(children[1].len.unwrap() >= -1e-12, "right_len negative");

nj/src/lib.rs

@@ -1,3 +1,7 @@
+// The explicit `core::simd` distance kernel needs the nightly `portable_simd`
+// feature. Gated on the `simd` Cargo feature so default (stable) builds are
+// unaffected — the kernel falls back to the autovectorizable scalar path.
+#![cfg_attr(feature = "simd", feature(portable_simd))]
 //! Neighbor-Joining phylogenetic tree inference library.
 //!
 //! # Data flow
@@ -313,11 +317,12 @@ fn add_bootstrap_to_tree(
         bitset_of(node, idx, &mut bv)?;
 
         let n = bv.count_ones();
-        if n > 1 && n < n_taxa {
-            if let Some(c) = counts.get(&bv.as_raw_slice().to_vec()) {
-                let pct = c * 100 / n_bootstrap_samples;
-                node.label = Some(NameOrSupport::Support(pct));
-            }
+        if n > 1
+            && n < n_taxa
+            && let Some(c) = counts.get(&bv.as_raw_slice().to_vec())
+        {
+            let pct = c * 100 / n_bootstrap_samples;
+            node.label = Some(NameOrSupport::Support(pct));
         }
 
         if let Some([l, r]) = &mut node.children {

nj/src/main.rs

@@ -62,8 +62,9 @@ mod cli {
         #[arg(long, default_value_t = false)]
         pub average_distance: bool,
 
-        /// Number of threads to use for parallel computation (default: all available).
-        /// Only effective when built with the `parallel` feature.
+        /// Number of threads for parallel bootstrap and distance computation
+        /// (default: all available cores). The CLI is always built with
+        /// threading enabled, so this flag is always effective.
         #[arg(short = 't', long, value_name = "N")]
         pub num_threads: Option<usize>,