Skip to main content

bloqade_lanes_bytecode/policy/
trace.rs

1//! `trace-policy` subcommand: per-event verbose trace output.
2
3use std::fs::File;
4use std::io::{self, BufWriter, Write};
5use std::path::Path;
6use std::sync::Arc;
7
8use bloqade_lanes_dsl_core::sandbox::SandboxConfig;
9use bloqade_lanes_search::dsl::fixture::{self, Problem};
10use bloqade_lanes_search::dsl::move_policy_dsl::{
11    GraphDelta, JsonMoveTraceObserver, MoveKernelObserver, PolicyGraphSnapshot, PolicyOptions,
12    PolicyStatus, solve_with_policy,
13};
14use bloqade_lanes_search::dsl::target_generator_dsl::{
15    CandidateSummary, JsonTargetTraceObserver, TargetContextSnapshot, TargetKernelObserver,
16    run_target_policy,
17};
18use bloqade_lanes_search::primitives::lane_index::LaneIndex;
19
20#[allow(clippy::too_many_arguments)]
21pub fn run_trace_policy(
22    policy: &Path,
23    problem: &Path,
24    params: Option<&Path>,
25    max_expansions: Option<u64>,
26    timeout_s: Option<f64>,
27    json: bool,
28    _seed: Option<u64>,
29    out: Option<&Path>,
30) -> Result<(), String> {
31    let (parsed, arch_path) = fixture::load(problem).map_err(|e| format!("error: {e}"))?;
32    let arch_json = std::fs::read_to_string(&arch_path)
33        .map_err(|e| format!("error: reading arch {}: {e}", arch_path.display()))?;
34    let arch = bloqade_lanes_bytecode_core::arch::ArchSpec::from_json(&arch_json)
35        .map_err(|e| format!("error: parsing arch {}: {e}", arch_path.display()))?;
36
37    let writer: Box<dyn Write> = match out {
38        Some(p) => Box::new(BufWriter::new(
39            File::create(p).map_err(|e| format!("error: writing {}: {e}", p.display()))?,
40        )),
41        None => Box::new(io::stdout()),
42    };
43
44    match parsed {
45        Problem::Move(mp) => trace_move(
46            policy,
47            mp,
48            arch,
49            params,
50            max_expansions,
51            timeout_s,
52            json,
53            writer,
54        ),
55        Problem::Target(tp) => trace_target(policy, tp, arch, params, json, writer),
56    }
57}
58
59#[allow(clippy::too_many_arguments)]
60fn trace_move(
61    policy: &Path,
62    mp: bloqade_lanes_search::dsl::fixture::MoveProblem,
63    arch: bloqade_lanes_bytecode_core::arch::ArchSpec,
64    params: Option<&Path>,
65    max_expansions: Option<u64>,
66    timeout_s: Option<f64>,
67    json: bool,
68    writer: Box<dyn Write>,
69) -> Result<(), String> {
70    let index = Arc::new(LaneIndex::new(arch));
71    let initial = mp.initial_locations();
72    let target = mp.target_locations();
73    let blocked = mp.blocked_locations();
74
75    let opts = PolicyOptions {
76        policy_path: policy.display().to_string(),
77        sandbox: SandboxConfig::default(),
78        policy_params: super::eval::load_params(params, &mp.policy_params)?,
79        max_expansions: max_expansions
80            .or(mp.budget.as_ref().map(|b| b.max_expansions))
81            .unwrap_or(5_000),
82        timeout_s: Some(
83            timeout_s
84                .or(mp.budget.as_ref().map(|b| b.timeout_s))
85                .unwrap_or(10.0),
86        ),
87    };
88
89    if json {
90        let mut obs = JsonMoveTraceObserver::new(writer);
91        let _ = solve_with_policy(initial, target, blocked, index, opts, &mut obs)
92            .map_err(|e| format!("error: {e}"))?;
93    } else {
94        let mut obs = HumanMoveTraceObserver::new(writer);
95        let _ = solve_with_policy(initial, target, blocked, index, opts, &mut obs)
96            .map_err(|e| format!("error: {e}"))?;
97    }
98    Ok(())
99}
100
101fn trace_target(
102    policy: &Path,
103    tp: bloqade_lanes_search::dsl::fixture::TargetProblem,
104    arch: bloqade_lanes_bytecode_core::arch::ArchSpec,
105    params: Option<&Path>,
106    json: bool,
107    writer: Box<dyn Write>,
108) -> Result<(), String> {
109    let index = Arc::new(LaneIndex::new(arch));
110    let placement = tp.current_placement_locations();
111    let cfg = SandboxConfig::default();
112    let params_value = super::eval::load_params(params, &tp.policy_params)?;
113    let controls = tp.controls.clone();
114    let targets = tp.targets.clone();
115    let lookahead = tp.lookahead_cz_layers.clone();
116    let stage_idx = tp.cz_stage_index;
117
118    if json {
119        let mut obs = JsonTargetTraceObserver::new(writer);
120        let _ = run_target_policy(
121            policy,
122            index,
123            placement,
124            controls,
125            targets,
126            lookahead,
127            stage_idx,
128            params_value,
129            &cfg,
130            &mut obs,
131        );
132    } else {
133        let mut obs = HumanTargetTraceObserver::new(writer);
134        let _ = run_target_policy(
135            policy,
136            index,
137            placement,
138            controls,
139            targets,
140            lookahead,
141            stage_idx,
142            params_value,
143            &cfg,
144            &mut obs,
145        );
146    }
147    Ok(())
148}
149
150/// Human-readable observer for Move policies. One line per event,
151/// terse fields, designed for interactive debugging.
152struct HumanMoveTraceObserver<W: Write> {
153    w: W,
154}
155
156impl<W: Write> HumanMoveTraceObserver<W> {
157    fn new(w: W) -> Self {
158        Self { w }
159    }
160}
161
162impl<W: Write> MoveKernelObserver for HumanMoveTraceObserver<W> {
163    fn on_init(&mut self, root: &PolicyGraphSnapshot) {
164        let _ = writeln!(
165            self.w,
166            "init   qubits={} target={} blocked={}",
167            root.root_qubits.len(),
168            root.target_qubits.len(),
169            root.blocked_count
170        );
171    }
172
173    fn on_step(
174        &mut self,
175        step: u64,
176        depth: u32,
177        action: &bloqade_lanes_search::dsl::move_policy_dsl::actions::MoveAction,
178        _delta: &GraphDelta,
179    ) {
180        let _ = writeln!(self.w, "step   #{step:04} depth={depth} action={action:?}");
181    }
182
183    fn on_builtin(&mut self, step: u64, name: &str, ok: bool) {
184        let _ = writeln!(self.w, "builtin #{step:04} {name} ok={ok}");
185    }
186
187    fn on_halt(&mut self, status: &PolicyStatus) {
188        let _ = writeln!(self.w, "halt   status={status:?}");
189    }
190}
191
192/// Human-readable observer for Target policies.
193struct HumanTargetTraceObserver<W: Write> {
194    w: W,
195}
196
197impl<W: Write> HumanTargetTraceObserver<W> {
198    fn new(w: W) -> Self {
199        Self { w }
200    }
201}
202
203impl<W: Write> TargetKernelObserver for HumanTargetTraceObserver<W> {
204    fn on_invoke(&mut self, stage: u64, ctx: &TargetContextSnapshot) {
205        let _ = writeln!(
206            self.w,
207            "invoke stage={stage} qubits={} controls={} targets={} lookahead={}",
208            ctx.current_qubit_count, ctx.controls_len, ctx.targets_len, ctx.lookahead_layers
209        );
210    }
211
212    fn on_result(&mut self, stage: u64, s: &CandidateSummary, ok: bool) {
213        let _ = writeln!(
214            self.w,
215            "result stage={stage} ok={ok} candidates={} first_size={}",
216            s.num_candidates, s.first_candidate_size
217        );
218    }
219}