bloq_circuit/
pauli_map.rs

1use std::ops::BitXor;
2
3use bloq_utils::Pauli;
4use glam::IVec2;
5use rustc_hash::FxHashMap;
6use smallvec::SmallVec;
7
8use crate::{CircuitError, CoordinateOverflowError, checked_translate_coordinate};
9
10const PAULI_MAP_INLINE_CAP: usize = 4;
11type PauliMapEntries = SmallVec<[(IVec2, Pauli); PAULI_MAP_INLINE_CAP]>;
12
13/// A sparse, phaseless map from qubit coordinate to non-identity [`Pauli`].
14/// Compose maps with `^`; exact product phases are discarded.
15#[derive(Debug, PartialEq, Eq, Default, serde::Serialize)]
16pub struct PauliMap {
17    entries: PauliMapEntries,
18}
19
20// Manual impl: `find_index`'s binary search silently misbehaves when the
21// entries are not strictly sorted by `(x, y)`, and every constructor upholds
22// that plus the no-identity rule — so decoded (untrusted) entries must be
23// checked here, not trusted from the wire. The wire struct mirrors the
24// derived `Serialize` shape exactly.
25impl<'de> serde::Deserialize<'de> for PauliMap {
26    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
27        #[derive(serde::Deserialize)]
28        #[serde(rename = "PauliMap")]
29        struct Wire {
30            entries: PauliMapEntries,
31        }
32
33        let Wire { entries } = Wire::deserialize(deserializer)?;
34        if !entries
35            .windows(2)
36            .all(|pair| (pair[0].0.x, pair[0].0.y) < (pair[1].0.x, pair[1].0.y))
37        {
38            return Err(serde::de::Error::custom(
39                "PauliMap entries must be strictly sorted by (x, y)",
40            ));
41        }
42        if entries.iter().any(|&(_, pauli)| pauli == Pauli::I) {
43            return Err(serde::de::Error::custom(
44                "PauliMap entries must not contain identity Paulis",
45            ));
46        }
47        Ok(Self { entries })
48    }
49}
50
51// Manual impl: `SmallVec`'s derived clone goes through a per-element iterator;
52// `from_slice` copies the `Copy` entries in one shot. Cloning is hot when
53// template chunks are translated to node coordinates.
54impl Clone for PauliMap {
55    fn clone(&self) -> Self {
56        Self {
57            entries: PauliMapEntries::from_slice(&self.entries),
58        }
59    }
60}
61
62// Manual impl: hashing is hot in flow matching, where maps key open-flow
63// tables. Each entry is folded into a single `u64` write: x in the high half,
64// y in the low half, and the 2-bit Pauli xor-mixed into y's top bits (real
65// coordinates are far too small to reach bit 30, and even an aliased entry
66// only weakens the hash, never its correctness).
67impl std::hash::Hash for PauliMap {
68    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
69        self.hash_translated(IVec2::ZERO, state);
70    }
71}
72
73impl PauliMap {
74    /// Hash the map as if translated by `offset`, without materializing the
75    /// translated copy. Equal absolute maps hash equally regardless of how
76    /// they are split into (entries, offset).
77    pub(crate) fn hash_translated<H: std::hash::Hasher>(&self, offset: IVec2, state: &mut H) {
78        state.write_usize(self.entries.len());
79        for &(coord, pauli) in &self.entries {
80            let coord = coord + offset;
81            let packed = ((coord.x as u32 as u64) << 32) | (coord.y as u32 as u64);
82            state.write_u64(packed ^ ((pauli as u64) << 30));
83        }
84    }
85}
86
87impl FromIterator<(IVec2, Pauli)> for PauliMap {
88    fn from_iter<T: IntoIterator<Item = (IVec2, Pauli)>>(iter: T) -> Self {
89        let mut map = Self::empty();
90        for (coord, pauli) in iter {
91            map.insert(coord, pauli);
92        }
93        map
94    }
95}
96
97impl std::fmt::Display for PauliMap {
98    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
99        if self.is_empty() {
100            return write!(f, "I");
101        }
102        let mut entries = self.iter().collect::<Vec<_>>();
103        entries.sort_unstable_by_key(|&(coord, _)| (coord.y, coord.x));
104        for (index, (coord, pauli)) in entries.into_iter().enumerate() {
105            if index > 0 {
106                write!(f, "*")?;
107            }
108            write!(f, "{pauli}({},{})", coord.x, coord.y)?;
109        }
110        Ok(())
111    }
112}
113
114impl PauliMap {
115    /// Last entry strictly before `coord` in the map's (x, y) order.
116    pub(crate) fn entry_before(&self, coord: IVec2) -> Option<(IVec2, Pauli)> {
117        let end = self
118            .entries
119            .partition_point(|&(entry, _)| (entry.x, entry.y) < (coord.x, coord.y));
120        end.checked_sub(1).map(|index| self.entries[index])
121    }
122
123    /// Exact power of `i` contributed by a product; the maps themselves remain phaseless.
124    pub fn product_phase(&self, other: &Self) -> u8 {
125        self.iter().fold(0, |phase, (coord, &left)| {
126            (phase
127                + match (left, other.get(coord).copied().unwrap_or(Pauli::I)) {
128                    (Pauli::X, Pauli::Y) | (Pauli::Y, Pauli::Z) | (Pauli::Z, Pauli::X) => 1,
129                    (Pauli::Y, Pauli::X) | (Pauli::Z, Pauli::Y) | (Pauli::X, Pauli::Z) => 3,
130                    _ => 0,
131                })
132                % 4
133        })
134    }
135
136    /// Returns a copy with every coordinate shifted by `offset`.
137    ///
138    /// # Panics
139    ///
140    /// Panics if any translated coordinate exceeds the signed 32-bit lattice.
141    /// Use [`Self::try_translated`] for untrusted offsets.
142    #[must_use]
143    pub fn translated(&self, offset: IVec2) -> Self {
144        self.try_translated(offset)
145            .expect("translated PauliMap coordinates must fit in the i32 lattice")
146    }
147
148    /// Returns a copy with every coordinate shifted by `offset`.
149    ///
150    /// # Errors
151    ///
152    /// Returns [`CoordinateOverflowError`] if any translated coordinate exceeds
153    /// the signed 32-bit lattice.
154    pub fn try_translated(&self, offset: IVec2) -> Result<Self, CoordinateOverflowError> {
155        // Translation by a constant offset preserves the coordinate sort
156        // order, so the entries can be copied and shifted in place.
157        let mut entries = PauliMapEntries::from_slice(&self.entries);
158        for (coord, _) in &mut entries {
159            *coord = checked_translate_coordinate(*coord, offset)?;
160        }
161        Ok(Self { entries })
162    }
163
164    /// Build a map from entries with pairwise-distinct coordinates and
165    /// non-identity Paulis.
166    ///
167    /// Skips the per-entry binary search and shifting that `FromIterator`
168    /// performs, sorting the collected entries once instead. The caller must
169    /// guarantee the entries need no deduplication or identity filtering.
170    ///
171    /// # Panics
172    ///
173    /// Panics if coordinates repeat or an entry is identity.
174    #[must_use]
175    pub fn from_unique_entries(entries: impl IntoIterator<Item = (IVec2, Pauli)>) -> Self {
176        let mut entries: PauliMapEntries = entries.into_iter().collect();
177        entries.sort_unstable_by_key(|&(coord, _)| (coord.x, coord.y));
178        assert!(
179            entries
180                .windows(2)
181                .all(|pair| (pair[0].0.x, pair[0].0.y) < (pair[1].0.x, pair[1].0.y)),
182            "entries must have pairwise-distinct coordinates"
183        );
184        assert!(
185            entries.iter().all(|&(_, pauli)| pauli != Pauli::I),
186            "entries must not contain identity Paulis"
187        );
188        Self { entries }
189    }
190
191    /// Creates an empty Pauli map.
192    #[must_use]
193    pub fn empty() -> Self {
194        Self {
195            entries: SmallVec::new(),
196        }
197    }
198
199    /// Returns the number of non-identity entries.
200    #[must_use]
201    pub fn len(&self) -> usize {
202        self.entries.len()
203    }
204
205    /// Returns whether the map contains no non-identity entries.
206    #[must_use]
207    pub fn is_empty(&self) -> bool {
208        self.entries.is_empty()
209    }
210
211    /// Iterates over the `(coordinate, pauli)` entries in sorted `(x, y)` order.
212    pub fn iter(&self) -> PauliMapIter<'_> {
213        PauliMapIter {
214            inner: self.entries.iter(),
215        }
216    }
217
218    /// The product's representative coordinate: its minimum data qubit by
219    /// `(x, y)`, used as the [`crate::Op::MPP`] measurement-record key. Entries
220    /// are kept sorted, so this is the first one. Recovering it from the product
221    /// alone lets the construct/clone/translate/merge stages agree without a
222    /// shared registry. `None` for an empty product.
223    #[must_use]
224    pub fn representative_coord(&self) -> Option<IVec2> {
225        self.iter().next().map(|(coord, _)| *coord)
226    }
227
228    /// The Pauli at `coord`, or `None` if the coordinate carries identity.
229    pub fn get(&self, coord: &IVec2) -> Option<&Pauli> {
230        self.find_index(*coord)
231            .ok()
232            .map(|index| &self.entries[index].1)
233    }
234
235    /// Sets `coord` to `pauli`, returning the Pauli previously there.
236    ///
237    /// Inserting [`Pauli::I`] removes the entry, so the no-identity invariant
238    /// holds.
239    pub fn insert(&mut self, coord: IVec2, pauli: Pauli) -> Option<Pauli> {
240        match self.find_index(coord) {
241            Ok(index) => {
242                if pauli == Pauli::I {
243                    Some(self.entries.remove(index).1)
244                } else {
245                    Some(std::mem::replace(&mut self.entries[index].1, pauli))
246                }
247            }
248            Err(index) => {
249                if pauli != Pauli::I {
250                    self.entries.insert(index, (coord, pauli));
251                }
252                None
253            }
254        }
255    }
256
257    /// Renders the map as a dense Pauli string of length `num_qubits`, placing
258    /// each entry at its `layout` index and filling gaps with `_`.
259    ///
260    /// # Errors
261    ///
262    /// Returns [`CircuitError::QubitNotFoundInLayout`] if an entry's coordinate
263    /// is absent from `layout`, or [`CircuitError::QubitIndexOutOfRange`] if its
264    /// index is at least `num_qubits`.
265    pub fn to_pauli_string(
266        &self,
267        layout: &FxHashMap<IVec2, u32>,
268        num_qubits: usize,
269    ) -> Result<String, CircuitError> {
270        let mut pauli_string = "_".repeat(num_qubits);
271        for (coord, pauli) in self {
272            let index = *layout
273                .get(coord)
274                .ok_or(CircuitError::QubitNotFoundInLayout(*coord))?
275                as usize;
276            if index >= num_qubits {
277                return Err(CircuitError::QubitIndexOutOfRange { index, num_qubits });
278            }
279            pauli_string.replace_range(
280                index..index + 1,
281                match pauli {
282                    Pauli::I => "I",
283                    Pauli::X => "X",
284                    Pauli::Y => "Y",
285                    Pauli::Z => "Z",
286                },
287            );
288        }
289        Ok(pauli_string)
290    }
291
292    fn find_index(&self, coord: IVec2) -> Result<usize, usize> {
293        self.entries
294            .binary_search_by_key(&(coord.x, coord.y), |(entry, _)| (entry.x, entry.y))
295    }
296}
297
298impl BitXor for &PauliMap {
299    type Output = PauliMap;
300
301    // Both maps are sorted, so composition is one linear merge.
302    fn bitxor(self, rhs: Self) -> Self::Output {
303        let (a, b) = (&self.entries, &rhs.entries);
304        let mut entries = PauliMapEntries::with_capacity(a.len() + b.len());
305        let (mut i, mut j) = (0, 0);
306        while i < a.len() && j < b.len() {
307            let ((a_coord, a_pauli), (b_coord, b_pauli)) = (a[i], b[j]);
308            match (a_coord.x, a_coord.y).cmp(&(b_coord.x, b_coord.y)) {
309                std::cmp::Ordering::Less => {
310                    entries.push((a_coord, a_pauli));
311                    i += 1;
312                }
313                std::cmp::Ordering::Greater => {
314                    entries.push((b_coord, b_pauli));
315                    j += 1;
316                }
317                std::cmp::Ordering::Equal => {
318                    let pauli = a_pauli ^ b_pauli;
319                    if pauli != Pauli::I {
320                        entries.push((a_coord, pauli));
321                    }
322                    i += 1;
323                    j += 1;
324                }
325            }
326        }
327        entries.extend_from_slice(&a[i..]);
328        entries.extend_from_slice(&b[j..]);
329        PauliMap { entries }
330    }
331}
332
333/// Borrowing iterator over a [`PauliMap`]'s sorted `(coordinate, pauli)`
334/// entries.
335#[derive(Debug)]
336pub struct PauliMapIter<'a> {
337    inner: std::slice::Iter<'a, (IVec2, Pauli)>,
338}
339
340impl<'a> Iterator for PauliMapIter<'a> {
341    type Item = (&'a IVec2, &'a Pauli);
342
343    fn next(&mut self) -> Option<Self::Item> {
344        self.inner.next().map(|(coord, pauli)| (coord, pauli))
345    }
346
347    fn size_hint(&self) -> (usize, Option<usize>) {
348        self.inner.size_hint()
349    }
350}
351
352impl ExactSizeIterator for PauliMapIter<'_> {}
353
354impl std::iter::FusedIterator for PauliMapIter<'_> {}
355
356impl DoubleEndedIterator for PauliMapIter<'_> {
357    fn next_back(&mut self) -> Option<Self::Item> {
358        self.inner.next_back().map(|(coord, pauli)| (coord, pauli))
359    }
360}
361
362impl IntoIterator for PauliMap {
363    type Item = (IVec2, Pauli);
364    type IntoIter = smallvec::IntoIter<[(IVec2, Pauli); PAULI_MAP_INLINE_CAP]>;
365
366    fn into_iter(self) -> Self::IntoIter {
367        self.entries.into_iter()
368    }
369}
370
371impl<'a> IntoIterator for &'a PauliMap {
372    type Item = (&'a IVec2, &'a Pauli);
373    type IntoIter = PauliMapIter<'a>;
374
375    fn into_iter(self) -> Self::IntoIter {
376        self.iter()
377    }
378}
379
380#[cfg(test)]
381mod tests {
382    use super::*;
383
384    fn map(entries: impl IntoIterator<Item = (IVec2, Pauli)>) -> PauliMap {
385        entries.into_iter().collect()
386    }
387
388    /// Decoded entries are untrusted: unsorted or identity entries would
389    /// silently break `find_index`'s binary search, so they must be rejected.
390    #[test]
391    fn deserialize_rejects_unsorted_and_identity_entries() {
392        #[derive(serde::Serialize)]
393        #[serde(rename = "PauliMap")]
394        struct RawMap {
395            entries: Vec<(IVec2, Pauli)>,
396        }
397        let encode = |entries: Vec<(IVec2, Pauli)>| {
398            postcard::to_allocvec(&RawMap { entries }).expect("test wire encodes into a Vec")
399        };
400
401        let unsorted = encode(vec![
402            (IVec2::new(1, 0), Pauli::X),
403            (IVec2::new(0, 0), Pauli::Z),
404        ]);
405        postcard::from_bytes::<PauliMap>(&unsorted).unwrap_err();
406
407        let identity = encode(vec![(IVec2::new(0, 0), Pauli::I)]);
408        postcard::from_bytes::<PauliMap>(&identity).unwrap_err();
409
410        let valid = map([(IVec2::new(0, 0), Pauli::X), (IVec2::new(1, 0), Pauli::Z)]);
411        let bytes = postcard::to_allocvec(&valid).expect("map encodes into a Vec");
412        let restored: PauliMap = postcard::from_bytes(&bytes).expect("sorted entries decode");
413        assert_eq!(restored, valid);
414    }
415
416    #[test]
417    fn unique_entries_preserve_map_invariants() {
418        for entries in [
419            [(IVec2::ZERO, Pauli::X), (IVec2::ZERO, Pauli::Z)],
420            [(IVec2::ZERO, Pauli::X), (IVec2::X, Pauli::I)],
421        ] {
422            std::panic::catch_unwind(|| PauliMap::from_unique_entries(entries))
423                .expect_err("duplicate coordinates and identity entries must panic");
424        }
425        let map = PauliMap::from_unique_entries([(IVec2::X, Pauli::Z), (IVec2::ZERO, Pauli::X)]);
426        let mut entries: crate::PauliMapIter<'_> = map.iter();
427        assert_eq!(entries.len(), 2);
428        assert_eq!(entries.next(), Some((&IVec2::ZERO, &Pauli::X)));
429        assert_eq!(entries.len(), 1);
430        assert_eq!(entries.next_back(), Some((&IVec2::X, &Pauli::Z)));
431        assert_eq!(entries.len(), 0);
432        assert_eq!(entries.next(), None);
433    }
434
435    #[test]
436    fn xor_composes_paulis() {
437        let x = map([(IVec2::ZERO, Pauli::X)]);
438        let z = map([(IVec2::ZERO, Pauli::Z)]);
439
440        assert_eq!(&x ^ &z, map([(IVec2::ZERO, Pauli::Y)]));
441        assert!((&x ^ &x).is_empty());
442    }
443
444    #[test]
445    fn pauli_string_error_reports_empty_layout_size() {
446        let coord = IVec2::ZERO;
447        let map = map([(coord, Pauli::X)]);
448        let layout: FxHashMap<_, _> = [(coord, 0)].into_iter().collect();
449
450        assert_eq!(
451            map.to_pauli_string(&layout, 0),
452            Err(CircuitError::QubitIndexOutOfRange {
453                index: 0,
454                num_qubits: 0,
455            })
456        );
457    }
458
459    #[test]
460    fn pauli_string_uses_layout_indices_and_identity_gaps() {
461        let map = map([
462            (IVec2::ZERO, Pauli::Z),
463            (IVec2::X, Pauli::X),
464            (IVec2::Y, Pauli::Y),
465        ]);
466        let layout = FxHashMap::from_iter([(IVec2::ZERO, 3), (IVec2::X, 0), (IVec2::Y, 2)]);
467        assert_eq!(map.to_pauli_string(&layout, 5), Ok("X_YZ_".into()));
468        assert_eq!(
469            map.to_pauli_string(&FxHashMap::default(), 5),
470            Err(CircuitError::QubitNotFoundInLayout(IVec2::ZERO))
471        );
472        assert_eq!(
473            PauliMap::empty().to_pauli_string(&layout, 0),
474            Ok(String::new())
475        );
476    }
477
478    #[test]
479    fn checked_translation_rejects_coordinate_overflow() {
480        let map = map([(IVec2::new(i32::MAX, 0), Pauli::X)]);
481
482        assert_eq!(
483            map.try_translated(IVec2::X),
484            Err(CoordinateOverflowError {
485                coordinate: IVec2::new(i32::MAX, 0),
486                offset: IVec2::X,
487            })
488        );
489    }
490
491    #[test]
492    fn pauli_map_deduplicates_and_sorts_entries() {
493        let map = vec![
494            (IVec2::new(2, 0), Pauli::Z),
495            (IVec2::new(0, 1), Pauli::X),
496            (IVec2::new(2, 0), Pauli::I),
497            (IVec2::new(1, 0), Pauli::Y),
498        ]
499        .into_iter()
500        .collect::<PauliMap>();
501
502        assert_eq!(
503            map.into_iter().collect::<Vec<_>>(),
504            vec![(IVec2::new(0, 1), Pauli::X), (IVec2::new(1, 0), Pauli::Y)]
505        );
506    }
507}
bloq_circuit/pauli_map.rs

bloq_circuit/
pauli_map.rs
