stabilizer_qec 1.2.0

//! A detector error model as Stim writes one: errors, detector and observable declarations,
//! `shift_detectors`, and `repeat` blocks, kept as a program rather than unrolled.
//!
//! WHY THIS EXISTS
//! ---------------
//! A long memory's model repeats itself round after round. Stim folds the repetition into a
//! `repeat` block, so a d = 11 memory of 10,000 rounds is a few hundred lines; unrolled it is
//! millions of faults. This holds the folded form, reads and prints it, and unrolls it (into
//! `Dem`, the flat model the decoders take) only when a decoder needs it.

use std::fmt::Write as _;

use crate::circuit::{fmt_args, split_instruction, strip_comment};
use crate::dem::{xor_prob, Dem, Mechanism, Piece};

/// An observable's target: its index with the top bit set, so that within a sorted list of
/// targets the detectors come first, as in Stim.
pub const OBS: u64 = 1 << 63;

/// The largest detector index a flat model may name.
pub const MAX_DETECTOR: u64 = (1 << 24) - 1;

/// The most errors and declarations a model may unroll into.
pub const MAX_UNROLLED_LINES: u64 = 1 << 24;

#[derive(Clone, Debug, PartialEq)]
pub enum DemInstr {
    /// `error[tag](p) D0 D1 ^ D2 L0`: each piece's targets sorted (detectors relative to the
    /// current shift), pieces in order.
    Error { p: f64, pieces: Vec<Vec<u64>>, tag: String },
    /// `detector[tag](coords) D3`: detectors, relative.
    Detector { coords: Vec<f64>, targets: Vec<u64>, tag: String },
    /// `logical_observable[tag] L1`.
    Observable { targets: Vec<u32>, tag: String },
    /// `shift_detectors[tag](coords) n`.
    Shift { coords: Vec<f64>, by: u64, tag: String },
    /// `repeat[tag] count { body }`.
    Repeat { count: u64, body: Vec<DemInstr>, tag: String },
}

/// A model as a program.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct DemProgram {
    pub instrs: Vec<DemInstr>,
}

/// What a program holds, counted through its loops without unrolling them.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Stats {
    /// One more than the largest detector index named, saturating.
    pub num_detectors: u64,
    /// One more than the largest observable index named.
    pub num_observables: usize,
    /// Error instructions, counted through loops, saturating.
    pub num_errors: u64,
}

/// A body's counts: its largest detector index plus one relative to its start (0 for none),
/// how far it shifts, its errors, and its observables.
#[derive(Clone, Copy, Default)]
struct BlockStats {
    end: u64,
    shift: u64,
    errors: u64,
    observables: usize,
}

fn block_stats(instrs: &[DemInstr]) -> BlockStats {
    let mut s = BlockStats::default();
    for ins in instrs {
        match ins {
            DemInstr::Error { pieces, .. } => {
                s.errors = s.errors.saturating_add(1);
                for t in pieces.iter().flatten() {
                    if t & OBS == 0 {
                        s.end = s.end.max(s.shift.saturating_add(t + 1));
                    } else {
                        s.observables = s.observables.max((t & !OBS) as usize + 1);
                    }
                }
            }
            DemInstr::Detector { targets, .. } => {
                for t in targets {
                    s.end = s.end.max(s.shift.saturating_add(t + 1));
                }
            }
            DemInstr::Observable { targets, .. } => {
                for &t in targets {
                    s.observables = s.observables.max(t as usize + 1);
                }
            }
            DemInstr::Shift { by, .. } => s.shift = s.shift.saturating_add(*by),
            DemInstr::Repeat { count, body, .. } => {
                if *count == 0 {
                    continue;
                }
                let b = block_stats(body);
                if b.end > 0 {
                    let last = s.shift.saturating_add(b.shift.saturating_mul(count - 1));
                    s.end = s.end.max(last.saturating_add(b.end));
                }
                s.shift = s.shift.saturating_add(b.shift.saturating_mul(*count));
                s.errors = s.errors.saturating_add(b.errors.saturating_mul(*count));
                s.observables = s.observables.max(b.observables);
            }
        }
    }
    s
}

fn tagged(tag: &str) -> String {
    if tag.is_empty() {
        String::new()
    } else {
        format!("[{tag}]")
    }
}

fn with_args(name: &str, tag: &str, args: &[f64]) -> String {
    let mut s = format!("{name}{}", tagged(tag));
    if !args.is_empty() {
        let _ = write!(s, "({})", fmt_args(args));
    }
    s
}

fn push_target(s: &mut String, t: u64) {
    if t & OBS == 0 {
        let _ = write!(s, " D{t}");
    } else {
        let _ = write!(s, " L{}", t & !OBS);
    }
}

/// XOR `t` into a sorted target list.
fn toggle(v: &mut Vec<u64>, t: u64) {
    match v.binary_search(&t) {
        Ok(i) => {
            v.remove(i);
        }
        Err(i) => v.insert(i, t),
    }
}

impl DemProgram {
    pub fn stats(&self) -> Stats {
        let b = block_stats(&self.instrs);
        Stats { num_detectors: b.end, num_observables: b.observables, num_errors: b.errors }
    }

    /// Stim's text.
    pub fn to_stim(&self) -> String {
        let mut s = String::new();
        emit(&self.instrs, "", &mut s);
        s
    }

    pub fn parse(text: &str) -> Result<DemProgram, String> {
        let lines: Vec<&str> = text.lines().collect();
        let mut pos = 0usize;
        let instrs = parse_block(&lines, &mut pos, false)?;
        Ok(DemProgram { instrs })
    }

    /// The program unrolled: no `repeat` and no `shift_detectors`, detectors absolute and
    /// coordinates shifted, in the program's order. Refuses past `MAX_UNROLLED_LINES` errors and
    /// declarations, or past detector `MAX_DETECTOR`.
    pub fn flattened(&self) -> Result<DemProgram, String> {
        let mut out = Vec::new();
        let mut st = Unroll { offset: 0, coords: Vec::new(), budget: MAX_UNROLLED_LINES };
        unroll(&self.instrs, &mut st, &mut |ins| out.push(ins))?;
        Ok(DemProgram { instrs: out })
    }

    /// The flat model the decoders take, in the program's order.
    pub fn to_dem(&self) -> Result<Dem, String> {
        let stats = self.stats();
        if stats.num_detectors > MAX_DETECTOR + 1 {
            return Err(format!("the model names detector {}, past the largest index, {MAX_DETECTOR}", stats.num_detectors - 1));
        }
        let mut dem = Dem { num_detectors: stats.num_detectors as usize, num_observables: stats.num_observables, ..Dem::default() };
        dem.detector_coords = vec![Vec::new(); dem.num_detectors];
        let mut st = Unroll { offset: 0, coords: Vec::new(), budget: MAX_UNROLLED_LINES };
        unroll(&self.instrs, &mut st, &mut |ins| match ins {
            DemInstr::Error { p, pieces, tag } => dem.mechanisms.push(mechanism(p, &pieces, tag)),
            DemInstr::Detector { coords, targets, tag } => {
                for d in targets {
                    let d = d as usize;
                    dem.detector_coords[d] = coords.clone();
                    if !tag.is_empty() {
                        if dem.detector_tags.len() <= d {
                            dem.detector_tags.resize(d + 1, String::new());
                        }
                        dem.detector_tags[d] = tag.clone();
                    }
                }
            }
            DemInstr::Observable { targets, tag } => {
                for l in targets {
                    let l = l as usize;
                    if !tag.is_empty() {
                        if dem.observable_tags.len() <= l {
                            dem.observable_tags.resize(l + 1, String::new());
                        }
                        dem.observable_tags[l] = tag.clone();
                    }
                }
            }
            DemInstr::Shift { .. } | DemInstr::Repeat { .. } => {}
        })?;
        Ok(dem)
    }

    /// The flat model of a model this engine built: `to_dem`, then faults that are the same in
    /// every way merged and all sorted, as the decoders have always been given it.
    pub fn to_dem_merged(&self) -> Result<Dem, String> {
        let mut dem = self.to_dem()?;
        let key = |m: &Mechanism| (m.detectors.clone(), m.observables, m.tag.clone(), m.pieces.clone());
        dem.mechanisms.sort_by(|a, b| {
            a.detectors
                .cmp(&b.detectors)
                .then(a.observables.cmp(&b.observables))
                .then(a.tag.cmp(&b.tag))
                .then(a.pieces.cmp(&b.pieces))
        });
        let mut merged: Vec<Mechanism> = Vec::with_capacity(dem.mechanisms.len());
        for m in dem.mechanisms {
            match merged.last_mut() {
                Some(last) if key(last) == key(&m) => last.p = xor_prob(last.p, m.p),
                _ => merged.push(m),
            }
        }
        dem.mechanisms = merged;
        Ok(dem)
    }
}

/// A flat model's mechanism from an error's absolute pieces, as the `.dem` reader has always
/// made one: the pieces' XOR, and the pieces kept when there are several or the one is
/// graph-like.
fn mechanism(p: f64, pieces: &[Vec<u64>], tag: String) -> Mechanism {
    let mut all = Vec::new();
    let mut kept = Vec::new();
    for piece in pieces {
        let (dets, obs) = split(piece);
        kept.push(Piece { detectors: dets, observables: obs });
        for &t in piece {
            toggle(&mut all, t);
        }
    }
    let (detectors, observables) = split(&all);
    let pieces = if kept.len() > 1 {
        kept.into_iter().filter(|p| !p.detectors.is_empty() || p.observables != 0).collect()
    } else if (1..=2).contains(&detectors.len()) {
        vec![Piece { detectors: detectors.clone(), observables }]
    } else {
        Vec::new()
    };
    Mechanism { p, detectors, observables, pieces, tag }
}

/// A sorted target list as detectors (`u32`) and an observable mask.
pub fn split(targets: &[u64]) -> (Vec<u32>, u64) {
    let mut dets = Vec::new();
    let mut obs = 0u64;
    for &t in targets {
        if t & OBS == 0 {
            dets.push(t as u32);
        } else {
            obs ^= 1u64 << (t & !OBS);
        }
    }
    (dets, obs)
}

struct Unroll {
    offset: u64,
    coords: Vec<f64>,
    budget: u64,
}

fn unroll(instrs: &[DemInstr], st: &mut Unroll, out: &mut dyn FnMut(DemInstr)) -> Result<(), String> {
    let absolute = |t: u64, offset: u64| -> Result<u64, String> {
        if t & OBS != 0 {
            return Ok(t);
        }
        t.checked_add(offset)
            .filter(|&d| d <= MAX_DETECTOR)
            .ok_or_else(|| format!("the model names a detector past the largest index, {MAX_DETECTOR}"))
    };
    let spend = |st: &mut Unroll| -> Result<(), String> {
        if st.budget == 0 {
            return Err(format!("the model unrolls past {MAX_UNROLLED_LINES} errors and declarations"));
        }
        st.budget -= 1;
        Ok(())
    };
    for ins in instrs {
        match ins {
            DemInstr::Error { p, pieces, tag } => {
                spend(st)?;
                let mut abs = Vec::with_capacity(pieces.len());
                for piece in pieces {
                    let mut v = piece.iter().map(|&t| absolute(t, st.offset)).collect::<Result<Vec<_>, _>>()?;
                    v.sort_unstable();
                    abs.push(v);
                }
                out(DemInstr::Error { p: *p, pieces: abs, tag: tag.clone() });
            }
            DemInstr::Detector { coords, targets, tag } => {
                spend(st)?;
                let coords = coords.iter().enumerate().map(|(i, v)| v + st.coords.get(i).copied().unwrap_or(0.0)).collect();
                let targets = targets.iter().map(|&t| absolute(t, st.offset)).collect::<Result<Vec<_>, _>>()?;
                out(DemInstr::Detector { coords, targets, tag: tag.clone() });
            }
            DemInstr::Observable { .. } => {
                spend(st)?;
                out(ins.clone());
            }
            DemInstr::Shift { coords, by, .. } => {
                if st.coords.len() < coords.len() {
                    st.coords.resize(coords.len(), 0.0);
                }
                for (a, b) in st.coords.iter_mut().zip(coords) {
                    *a += b;
                }
                st.offset = st.offset.checked_add(*by).filter(|&o| o <= MAX_DETECTOR + 1).ok_or_else(|| {
                    format!("shift_detectors moves past detector {MAX_DETECTOR}, the largest index")
                })?;
            }
            DemInstr::Repeat { count, body, .. } => {
                for _ in 0..*count {
                    // Every pass costs at least one unit, so an empty body cannot spin forever.
                    spend(st)?;
                    unroll(body, st, out)?;
                }
            }
        }
    }
    Ok(())
}

fn emit(instrs: &[DemInstr], indent: &str, s: &mut String) {
    for ins in instrs {
        match ins {
            DemInstr::Error { p, pieces, tag } => {
                let _ = write!(s, "{indent}{}", with_args("error", tag, &[*p]));
                for (k, piece) in pieces.iter().enumerate() {
                    if k > 0 {
                        s.push_str(" ^");
                    }
                    for &t in piece {
                        push_target(s, t);
                    }
                }
            }
            DemInstr::Detector { coords, targets, tag } => {
                let _ = write!(s, "{indent}{}", with_args("detector", tag, coords));
                for &t in targets {
                    push_target(s, t);
                }
            }
            DemInstr::Observable { targets, tag } => {
                let _ = write!(s, "{indent}logical_observable{}", tagged(tag));
                for &t in targets {
                    let _ = write!(s, " L{t}");
                }
            }
            DemInstr::Shift { coords, by, tag } => {
                let _ = write!(s, "{indent}{} {by}", with_args("shift_detectors", tag, coords));
            }
            DemInstr::Repeat { count, body, tag } => {
                let _ = writeln!(s, "{indent}repeat{} {count} {{", tagged(tag));
                emit(body, &format!("{indent}    "), s);
                let _ = write!(s, "{indent}}}");
            }
        }
        s.push('\n');
    }
}

fn parse_block(lines: &[&str], pos: &mut usize, nested: bool) -> Result<Vec<DemInstr>, String> {
    let mut out = Vec::new();
    while *pos < lines.len() {
        let lineno = *pos + 1;
        let line = strip_comment(lines[*pos]).trim();
        *pos += 1;
        if line.is_empty() {
            continue;
        }
        if line == "}" {
            if nested {
                return Ok(out);
            }
            return Err(format!("line {lineno}: unmatched '}}'"));
        }
        if let Some(rest) = line.strip_suffix('{') {
            let (name, tag, args, parts) = split_instruction(rest).map_err(|e| format!("line {lineno}: {e}"))?;
            if name != "REPEAT" {
                return Err(format!("line {lineno}: only repeat opens a block"));
            }
            let count: u64 = match (args.is_empty(), parts.as_slice()) {
                (true, [n]) => n.parse().ok(),
                _ => None,
            }
            .ok_or_else(|| format!("line {lineno}: repeat needs a count"))?;
            let body = parse_block(lines, pos, true)?;
            out.push(DemInstr::Repeat { count, body, tag: tag.to_string() });
            continue;
        }
        let (name, tag, args, tokens) = split_instruction(line).map_err(|e| format!("line {lineno}: {e}"))?;
        let tag = tag.to_string();
        let bad = |t: &str| format!("line {lineno}: bad target '{t}'");
        let detector = |t: &str| -> Result<u64, String> {
            t.strip_prefix('D').and_then(|d| d.parse::<u64>().ok()).filter(|&d| d <= MAX_DETECTOR).ok_or_else(|| bad(t))
        };
        let observable = |t: &str| -> Result<u32, String> {
            let l = t.strip_prefix('L').and_then(|l| l.parse::<u32>().ok()).ok_or_else(|| bad(t))?;
            if l >= 64 {
                return Err(format!("line {lineno}: at most 64 observables are supported"));
            }
            Ok(l)
        };
        out.push(match name.as_str() {
            "ERROR" => {
                if args.len() != 1 || !(0.0..=1.0).contains(&args[0]) {
                    return Err(format!("line {lineno}: error takes one probability"));
                }
                let mut pieces = vec![Vec::new()];
                for t in &tokens {
                    if *t == "^" {
                        pieces.push(Vec::new());
                    } else if t.starts_with('D') {
                        toggle(pieces.last_mut().unwrap(), detector(t)?);
                    } else if t.starts_with('L') {
                        toggle(pieces.last_mut().unwrap(), OBS | u64::from(observable(t)?));
                    } else {
                        return Err(bad(t));
                    }
                }
                DemInstr::Error { p: args[0], pieces, tag }
            }
            "DETECTOR" => {
                let targets = tokens.iter().map(|t| detector(t)).collect::<Result<Vec<_>, _>>()?;
                DemInstr::Detector { coords: args, targets, tag }
            }
            "LOGICAL_OBSERVABLE" => {
                let targets = tokens.iter().map(|t| observable(t)).collect::<Result<Vec<_>, _>>()?;
                DemInstr::Observable { targets, tag }
            }
            "SHIFT_DETECTORS" => {
                let by = match tokens.as_slice() {
                    [] => 0,
                    [n] => n.parse::<u64>().ok().filter(|&n| n <= MAX_DETECTOR + 1).ok_or_else(|| bad(n))?,
                    _ => return Err(format!("line {lineno}: shift_detectors takes one count")),
                };
                DemInstr::Shift { coords: args, by, tag }
            }
            _ => return Err(format!("line {lineno}: unsupported instruction '{name}'")),
        });
    }
    if nested {
        return Err("unterminated repeat block".into());
    }
    Ok(out)
}