From 781c5279b5b907ee7483be7cbf6b5a75b16100f5 Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Thu, 1 Oct 2026 09:36:40 +0000
Subject: [PATCH] [qec] Share Stim observable assembly in the emitter module

Both Stim builders call stim_emit::emit_observable_include for OBSERVABLE_INCLUDE record offsets. Headers stay in the builders. Checked-in Stim gold stays byte-identical.

Issue 473 stays open.

Co-authored-by: Arnab Ghosh <ag2583@cornell.edu>
---
 quon_qec/src/experiment.rs           | 28 ++++--------
 quon_qec/src/experiment/stim_emit.rs | 66 +++++++++++++++++++++++++---
 2 files changed, 68 insertions(+), 26 deletions(-)

diff --git a/quon_qec/src/experiment.rs b/quon_qec/src/experiment.rs
index 4447190..3fd29b6 100644
--- a/quon_qec/src/experiment.rs
+++ b/quon_qec/src/experiment.rs
@@ -594,16 +594,12 @@ fn emit_stim_single_block_memory(expanded: &ExpandedWorkload) -> Result<String,
         }
     }
 
-    out.push_str("OBSERVABLE_INCLUDE(0)");
     let obs_atoms = logical_observable_atoms(block, measure_basis);
-    for atom in &obs_atoms {
-        let pos = data_atoms
-            .iter()
-            .position(|a| a.0 == *atom)
-            .ok_or(ExperimentError::MissingDataMeasurement { atom: *atom })?;
-        let rec = -(d - pos as i32);
-        out.push_str(&format!(" rec[{rec}]"));
+    let mut data_ids = Vec::with_capacity(data_atoms.len());
+    for atom in &data_atoms {
+        data_ids.push(atom.0);
     }
+    stim_emit::emit_observable_include(&mut out, 0, &data_ids, &obs_atoms)?;
     out.push('\n');
 
     Ok(out)
@@ -676,8 +672,9 @@ fn emit_stim_lattice_surgery_cx(expanded: &ExpandedWorkload) -> Result<String, E
     stim_emit::emit_reset_tick(&mut out, &all_atoms);
 
     // Round kinds append through [`stim_emit::StimRoundEmitter`]. Byproduct
-    // handles and deferred measure-logical rounds come back out for observable
-    // assembly below.
+    // handles and deferred measure-logical rounds come back out so the
+    // builder can write measure lines and frame records around
+    // [`stim_emit::emit_observable_include`].
     let mut ctx = stim_emit::LatticeSurgeryCtx::new(expanded, control, target, ancilla, out);
     for round in &expanded.rounds {
         stim_emit::lattice_round_emitter(round.kind).emit(&mut ctx, round)?;
@@ -718,16 +715,7 @@ fn emit_stim_lattice_surgery_cx(expanded: &ExpandedWorkload) -> Result<String, E
                 block_count: expanded.blocks.len(),
             })?;
         let obs_atoms = logical_observable_atoms(block, LogicalBasis::Z);
-        let d = data_atoms.len() as i32;
-        out.push_str(&format!("OBSERVABLE_INCLUDE({obs_id})"));
-        for atom in &obs_atoms {
-            let pos = data_atoms
-                .iter()
-                .position(|a| a == atom)
-                .ok_or(ExperimentError::MissingDataMeasurement { atom: *atom })?;
-            let rec = -(d - pos as i32);
-            out.push_str(&format!(" rec[{rec}]"));
-        }
+        stim_emit::emit_observable_include(&mut out, obs_id, &data_atoms, &obs_atoms)?;
         // Horsman on Z readout: X-frame byproducts flip Z (include m_zz, m_a on
         // target). Z-frame byproducts (smooth/ancilla → Z on control) commute
         // with Z and are listed below as OBSERVABLE_INCLUDE frame witnesses so
diff --git a/quon_qec/src/experiment/stim_emit.rs b/quon_qec/src/experiment/stim_emit.rs
index 3968259..81bf7cf 100644
--- a/quon_qec/src/experiment/stim_emit.rs
+++ b/quon_qec/src/experiment/stim_emit.rs
@@ -9,8 +9,10 @@
 //! [`StimRoundEmitter::emit_single_block`] on the same registry. A new round
 //! kind is a new impl plus one arm in [`lattice_round_emitter`].
 //!
-//! Both builders share [`emit_reset_tick`] for the prepare-`R` / `TICK` line.
-//! Headers and observable assembly stay in the builders.
+//! Both builders share [`emit_reset_tick`] for the prepare-`R` / `TICK` line
+//! and [`emit_observable_include`] for `OBSERVABLE_INCLUDE` record offsets.
+//! Headers stay in the builders. Lattice surgery still appends frame-byproduct
+//! records after the shared observable line.
 
 use std::collections::HashMap;
 
@@ -34,10 +36,41 @@ pub(crate) fn emit_reset_tick(out: &mut String, atom_ids: &[u32]) {
     out.push_str("\nTICK\n");
 }
 
+/// Write `OBSERVABLE_INCLUDE({obs_id})` and one `rec[-(d - pos)]` per
+/// observable atom.
+///
+/// `data_atoms` is the final measurement window in Stim order. Each
+/// `obs_atoms` id must appear in that window; the first match wins. The
+/// trailing newline is left to the caller so lattice surgery can append
+/// frame-byproduct `rec` targets on the same line.
+pub(crate) fn emit_observable_include(
+    out: &mut String,
+    obs_id: u32,
+    data_atoms: &[u32],
+    obs_atoms: &[u32],
+) -> Result<(), ExperimentError> {
+    let d = data_atoms.len() as i32;
+    out.push_str(&format!("OBSERVABLE_INCLUDE({obs_id})"));
+    for atom in obs_atoms {
+        let mut found = None;
+        for (pos, id) in data_atoms.iter().enumerate() {
+            if id == atom {
+                found = Some(pos);
+                break;
+            }
+        }
+        let pos = found.ok_or(ExperimentError::MissingDataMeasurement { atom: *atom })?;
+        let rec = -(d - pos as i32);
+        out.push_str(&format!(" rec[{rec}]"));
+    }
+    Ok(())
+}
+
 /// Shared Stim state for one lattice-surgery circuit.
 ///
-/// Round impls append instructions and record byproduct handles. Observable
-/// assembly after the schedule still lives in the circuit builder.
+/// Round impls append instructions and record byproduct handles. Measure
+/// lines, frame-byproduct records, and frame comments stay in the builder.
+/// [`emit_observable_include`] writes the shared observable prefix.
 pub(crate) struct LatticeSurgeryCtx<'a> {
     pub(crate) expanded: &'a ExpandedWorkload,
     pub(crate) control: &'a ExpandedBlock,
@@ -76,8 +109,9 @@ impl<'a> LatticeSurgeryCtx<'a> {
 
 /// Stim state for one single-block memory circuit.
 ///
-/// The builder still writes the header, reset, and final observable. Round
-/// impls append construct locals, memory rounds, and the measure-logical record.
+/// The builder still writes the header, the final measure line, and closing
+/// detectors. Round impls append construct locals, memory rounds, and the
+/// measure-logical record. [`emit_observable_include`] writes the observable.
 pub(crate) struct SingleBlockCtx<'a> {
     pub(crate) out: String,
     pub(crate) n_checks: usize,
@@ -696,4 +730,24 @@ mod tests {
         emit_reset_tick(&mut empty, &[]);
         assert_eq!(empty, "R\nTICK\n");
     }
+
+    #[test]
+    fn observable_include_writes_record_offsets() {
+        let mut out = String::new();
+        emit_observable_include(&mut out, 0, &[10, 20, 30], &[30, 10]).expect("present");
+        assert_eq!(out, "OBSERVABLE_INCLUDE(0) rec[-1] rec[-3]");
+        let mut numbered = String::new();
+        emit_observable_include(&mut numbered, 1, &[7], &[7]).expect("present");
+        assert_eq!(numbered, "OBSERVABLE_INCLUDE(1) rec[-1]");
+    }
+
+    #[test]
+    fn observable_include_missing_atom_errors() {
+        let mut out = String::new();
+        let err = emit_observable_include(&mut out, 1, &[10], &[99]).expect_err("missing");
+        assert!(matches!(
+            err,
+            ExperimentError::MissingDataMeasurement { atom: 99 }
+        ));
+    }
 }