perf(runtime): faster map_insert + specialized diagonal 2-qubit rotations

crates/ppvm-runtime/Cargo.toml

@@ -10,6 +10,7 @@ approx = { version = "0.5.1", optional = true }
 bincode = { version = "2.0.1", optional = true }
 bitvec = "1.0.1"
 bon = "3.7.2"
+bytemuck = { version = "1", features = ["min_const_generics"] }
 dashmap = { version = "6.1.0", features = ["rayon"], optional = true }
 fxhash = "0.2.1"
 gxhash = { version = "3.5.0", features = ["deterministic"], optional = true }

crates/ppvm-runtime/src/map/dashmap.rs

@@ -168,6 +168,17 @@ where
         })
     }
 
+    fn map_insert_vec<F>(&mut self, dest: &mut Vec<(W, C)>, f: F)
+    where
+        F: Fn(&W, &mut C) -> Option<(W, C)> + Sync + Send,
+    {
+        use rayon::iter::ParallelExtend;
+        dest.par_extend(self.par_iter_mut().filter_map(|mut entry| {
+            let (k, v) = entry.pair_mut();
+            f(k, v)
+        }));
+    }
+
     fn map_insert_multiple<F>(&mut self, dest: &mut Self, f: F)
     where
         F: Fn(&W, &mut C) -> Option<Vec<(W, C)>> + Sync + Send,

crates/ppvm-runtime/src/map/hashmap.rs

@@ -167,6 +167,17 @@ macro_rules! impl_acmap_insert {
                 }
             }
 
+            fn map_insert_vec<F>(&mut self, dest: &mut Vec<(W, C)>, f: F)
+            where
+                F: Fn(&W, &mut C) -> Option<(W, C)> + Sync + Send,
+            {
+                for (k, v) in self.iter_mut() {
+                    if let Some(entry) = f(k, v) {
+                        dest.push(entry);
+                    }
+                }
+            }
+
             fn map_insert_multiple<F>(&mut self, dest: &mut Self, f: F)
             where
                 F: Fn(&W, &mut C) -> Option<Vec<(W, C)>> + Sync + Send,

crates/ppvm-runtime/src/sum/data.rs

@@ -47,6 +47,10 @@ pub struct PauliSum<T: Config> {
     /// Keep-set: strings [`PauliSum::truncate`] must always re-insert
     /// after the strategy runs. Empty by default.
     preserve_strings: HashSet<T::PauliWordType>,
+    /// Reusable scratch buffer for [`map_insert`](Self::map_insert): produced
+    /// entries are pushed here (no hashing) then merged into the primary map,
+    /// avoiding a second hashmap probe per entry.
+    scratch: Vec<(T::PauliWordType, T::Coeff)>,
 }
 
 #[bon::bon]
@@ -82,6 +86,7 @@ impl<T: Config> PauliSum<T> {
             capacity,
             strategy,
             preserve_strings,
+            scratch: Vec::new(),
         }
     }
 }
@@ -205,10 +210,19 @@ impl<T: Config> PauliSum<T> {
             + Sync
             + Send,
     {
-        let (data, aux) = self.data_aux_mut();
-        aux.clear();
-        data.map_insert(aux, f);
-        self.consume();
+        // Collect produced terms into a reusable scratch Vec (cheap pushes,
+        // no hashing), then merge them into the primary map in one pass. This
+        // avoids the extra hashmap probe + insert the old aux-map buffer paid
+        // for every produced term, and makes the per-call clear O(1).
+        let mut scratch = std::mem::take(&mut self.scratch);
+        scratch.clear();
+        self.data_mut().map_insert_vec(&mut scratch, f);
+
+        let data = self.data_mut();
+        for (k, v) in scratch.drain(..) {
+            data.add_assign(k, v);
+        }
+        self.scratch = scratch;
     }
 
     /// modify in place existing entries and insert some new entries

crates/ppvm-runtime/src/sum/rot2.rs

@@ -56,6 +56,130 @@ where
             }
         });
     }
+
+    #[inline]
+    fn rzz(&mut self, a: usize, b: usize, theta: impl Into<T::Coeff>) {
+        let (sin, cos) = theta.into().sin_cos();
+        self.map_insert(|k, v| {
+            // Loss fallbacks — identical to the generic `rotate_2` path
+            // (axis on the surviving qubit is Z for a ZZ rotation). These
+            // branches are dead-code-eliminated for non-lossy PauliWord,
+            // since `get_lbit` is a const `false`.
+            if k.get_lbit(a) {
+                return rotate_1_map_insert_closure::<T>(k, v, Pauli::Z, b, &sin, &cos);
+            }
+            if k.get_lbit(b) {
+                return rotate_1_map_insert_closure::<T>(k, v, Pauli::Z, a, &sin, &cos);
+            }
+            let xa = k.get_xbit(a);
+            let xb = k.get_xbit(b);
+            // ZZ commutes iff both qubits agree on having an X-component.
+            if xa == xb {
+                return None;
+            }
+            let za = k.get_zbit(a);
+            let zb = k.get_zbit(b);
+            // The anticommuting qubit is the one with xbit == 1.
+            // sign = +1 if it is Y (zbit set), -1 if it is X (zbit clear).
+            let z_anti = if xa { za } else { zb };
+            let eps: i8 = if z_anti { 1 } else { -1 };
+
+            let mut coeff = v.clone();
+            *v *= cos.clone();
+
+            let mut new_word = k.clone();
+            new_word.set_zbit(a, !za);
+            new_word.set_zbit(b, !zb);
+            new_word.rehash();
+
+            coeff *= sin.mul_sign(eps);
+            Some((new_word, coeff))
+        });
+    }
+
+    #[inline]
+    fn rxx(&mut self, a: usize, b: usize, theta: impl Into<T::Coeff>) {
+        let (sin, cos) = theta.into().sin_cos();
+        self.map_insert(|k, v| {
+            // Loss fallback: a lost qubit leaves a single-qubit X rotation on
+            // the surviving partner (axis on the surviving qubit is X for an
+            // XX rotation), matching the generic `rotate_2` path.
+            if k.get_lbit(a) {
+                return rotate_1_map_insert_closure::<T>(k, v, Pauli::X, b, &sin, &cos);
+            }
+            if k.get_lbit(b) {
+                return rotate_1_map_insert_closure::<T>(k, v, Pauli::X, a, &sin, &cos);
+            }
+            let za = k.get_zbit(a);
+            let zb = k.get_zbit(b);
+            // X anticommutes with a qubit's Pauli iff that Pauli carries a
+            // Z-component (Z or Y). XX commutes iff both qubits agree on that.
+            if za == zb {
+                return None;
+            }
+            let xa = k.get_xbit(a);
+            let xb = k.get_xbit(b);
+            // The anticommuting qubit is the one carrying Z (its z-bit set).
+            // sign = +1 if it is Z (x-bit clear), -1 if it is Y (x-bit set).
+            let x_anti = if za { xa } else { xb };
+            let eps: i8 = if x_anti { -1 } else { 1 };
+
+            let mut coeff = v.clone();
+            *v *= cos.clone();
+
+            let mut new_word = k.clone();
+            new_word.set_xbit(a, !xa);
+            new_word.set_xbit(b, !xb);
+            new_word.rehash();
+
+            coeff *= sin.mul_sign(eps);
+            Some((new_word, coeff))
+        });
+    }
+
+    #[inline]
+    fn ryy(&mut self, a: usize, b: usize, theta: impl Into<T::Coeff>) {
+        let (sin, cos) = theta.into().sin_cos();
+        self.map_insert(|k, v| {
+            // Loss fallback: a lost qubit leaves a single-qubit Y rotation on
+            // the surviving partner (axis on the surviving qubit is Y for a
+            // YY rotation), matching the generic `rotate_2` path.
+            if k.get_lbit(a) {
+                return rotate_1_map_insert_closure::<T>(k, v, Pauli::Y, b, &sin, &cos);
+            }
+            if k.get_lbit(b) {
+                return rotate_1_map_insert_closure::<T>(k, v, Pauli::Y, a, &sin, &cos);
+            }
+            let xa = k.get_xbit(a);
+            let za = k.get_zbit(a);
+            let xb = k.get_xbit(b);
+            let zb = k.get_zbit(b);
+            // Y anticommutes with a qubit's Pauli iff it is X or Z, i.e. iff
+            // its x-bit and z-bit differ. YY commutes iff both qubits agree.
+            let pa = xa ^ za;
+            let pb = xb ^ zb;
+            if pa == pb {
+                return None;
+            }
+            // The anticommuting qubit is X or Z (x-bit != z-bit).
+            // sign = +1 if it is X (x-bit set), -1 if it is Z (x-bit clear).
+            let x_anti = if pa { xa } else { xb };
+            let eps: i8 = if x_anti { 1 } else { -1 };
+
+            let mut coeff = v.clone();
+            *v *= cos.clone();
+
+            let mut new_word = k.clone();
+            new_word.set_xbit(a, !xa);
+            new_word.set_zbit(a, !za);
+            new_word.set_xbit(b, !xb);
+            new_word.set_zbit(b, !zb);
+            new_word.rehash();
+
+            coeff *= sin.mul_sign(eps);
+            Some((new_word, coeff))
+        });
+    }
 }
 
 // R_{G}[P] = cos(theta) P - sin(theta) [G, P]/2i
@@ -118,3 +242,108 @@ pub fn comm_2(q0: [u8; 2], q1: [u8; 2], p0: [u8; 2], p1: [u8; 2]) -> (i8, u8, u8
 
     (coeff, x_out0, z_out0, x_out1, z_out1)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::config::fxhash::ByteF64;
+
+    type C = ByteF64<2>;
+    const PAULI_CHARS: [char; 4] = ['I', 'X', 'Y', 'Z'];
+
+    fn sum_with(word: &str) -> PauliSum<C> {
+        let mut s: PauliSum<C> = PauliSum::builder().n_qubits(2).build();
+        s += (word, 1.0);
+        s
+    }
+
+    /// Every diagonal two-qubit rotation has a hand-written fast path in this
+    /// file that bypasses `comm_2`. Assert it produces exactly the same
+    /// `PauliSum` as the generic `rotate_2` (the pre-existing `comm_2`-based
+    /// path) for every two-qubit input Pauli, at a representative angle.
+    fn assert_matches_generic(
+        axis: [u8; 2],
+        special: impl Fn(&mut PauliSum<C>, usize, usize, f64),
+    ) {
+        let theta = 0.7_f64;
+        for &p in &PAULI_CHARS {
+            for &q in &PAULI_CHARS {
+                let word = format!("{p}{q}");
+
+                let mut got = sum_with(&word);
+                special(&mut got, 0, 1, theta);
+
+                let mut want = sum_with(&word);
+                want.rotate_2(axis, axis, 0, 1, theta);
+
+                assert_eq!(got, want, "fast path disagrees with rotate_2 for {word}");
+            }
+        }
+    }
+
+    #[test]
+    fn rxx_matches_generic() {
+        assert_matches_generic([1, 0], |s, a, b, t| s.rxx(a, b, t));
+    }
+
+    #[test]
+    fn ryy_matches_generic() {
+        assert_matches_generic([1, 1], |s, a, b, t| s.ryy(a, b, t));
+    }
+
+    #[test]
+    fn rzz_matches_generic() {
+        assert_matches_generic([0, 1], |s, a, b, t| s.rzz(a, b, t));
+    }
+
+    /// Explicit hand-computed values, independent of `rotate_2`/`comm_2`, so a
+    /// bug shared by both the fast path and the generic path can't hide.
+    #[test]
+    fn rzz_explicit_values() {
+        let t = 0.7_f64;
+
+        // X_0 I_1 --rzz--> cos·XI − sin·YZ  (anticommutes; X-carrier → −1)
+        let mut s = sum_with("XI");
+        s.rzz(0, 1, t);
+        let mut want: PauliSum<C> = PauliSum::builder().n_qubits(2).build();
+        want += ("XI", t.cos());
+        want += ("YZ", -t.sin());
+        assert_eq!(s, want);
+
+        // Y_0 I_1 --rzz--> cos·YI + sin·XZ  (anticommutes; Y-carrier → +1)
+        let mut s = sum_with("YI");
+        s.rzz(0, 1, t);
+        let mut want: PauliSum<C> = PauliSum::builder().n_qubits(2).build();
+        want += ("YI", t.cos());
+        want += ("XZ", t.sin());
+        assert_eq!(s, want);
+
+        // ZZ commutes with the ZZ generator → unchanged.
+        let mut s = sum_with("ZZ");
+        s.rzz(0, 1, t);
+        assert_eq!(s, sum_with("ZZ"));
+    }
+
+    /// A diagonal rotation on non-adjacent qubits must address the right
+    /// slots (exercises `a`/`b` other than `0,1`).
+    #[test]
+    fn rzz_non_adjacent_addressing() {
+        let theta = 0.7_f64;
+        for &p in &PAULI_CHARS {
+            for &q in &PAULI_CHARS {
+                // 3-qubit word with the acted-on Paulis on qubits 0 and 2.
+                let word = format!("{p}I{q}");
+
+                let mut got: PauliSum<C> = PauliSum::builder().n_qubits(3).build();
+                got += (word.as_str(), 1.0);
+                got.rzz(0, 2, theta);
+
+                let mut want: PauliSum<C> = PauliSum::builder().n_qubits(3).build();
+                want += (word.as_str(), 1.0);
+                want.rotate_2([0, 1], [0, 1], 0, 2, theta);
+
+                assert_eq!(got, want, "rzz(0,2) disagrees with rotate_2 for {word}");
+            }
+        }
+    }
+}

crates/ppvm-runtime/src/sum/snapshots/ppvm_runtime__sum__data__tests__pauli_sum_creation-5.snap

@@ -2,4 +2,4 @@
 source: crates/ppvm-runtime/src/sum/data.rs
 expression: sum.to_string()
 ---
-3.000 * IIII + 3.000 * XIII + 1.000 * IYII
+3.000 * IIII + 1.000 * IYII + 3.000 * XIII

crates/ppvm-runtime/src/traits/map.rs

@@ -74,6 +74,14 @@ pub trait ACMapInsert<
     where
         F: Fn(&W, &mut V) -> Option<(W, V)> + Sync + Send;
 
+    /// Like [`map_insert`](Self::map_insert) but appends produced entries
+    /// into a plain `Vec` (a push per entry, no hashing) instead of a map.
+    /// The caller merges the buffer into the destination map afterwards,
+    /// avoiding a second hashmap probe per produced entry.
+    fn map_insert_vec<F>(&mut self, dest: &mut Vec<(W, V)>, f: F)
+    where
+        F: Fn(&W, &mut V) -> Option<(W, V)> + Sync + Send;
+
     /// Same as [`map_insert`](Self::map_insert) but `f` may return a
     /// `Vec` of new entries to be inserted into `dest`.
     fn map_insert_multiple<F>(&mut self, dest: &mut Self, f: F)

crates/ppvm-runtime/src/traits/storage.rs

@@ -7,6 +7,12 @@ use std::hash::Hash;
 /// Backing storage for a [`PauliWord`](crate::word::PauliWord) — a
 /// fixed-size, `Copy`-able block of bits (typically `[u8; N]` or
 /// `[u64; N]`).
+///
+/// The [`bytemuck::Pod`] bound guarantees the storage is a plain-old-data
+/// type with no padding and all bit patterns valid, which lets
+/// [`PauliWord::rehash`](crate::word::PauliWord) view it as a `&[u8]` for
+/// fast byte-slice hashing *without* `unsafe`. Every concrete storage used
+/// here (`[u8; N]`, `[u64; N]`) is already `Pod`.
 pub trait PauliStorage:
     BitViewSized
     + Clone
@@ -19,6 +25,7 @@ pub trait PauliStorage:
     + std::fmt::Debug
     + Ord
     + PartialOrd
+    + bytemuck::Pod
 {
 }
 impl<
@@ -32,7 +39,8 @@ impl<
         + Sync
         + std::fmt::Debug
         + Ord
-        + PartialOrd,
+        + PartialOrd
+        + bytemuck::Pod,
 > PauliStorage for A
 {
 }