bloqade_lanes_bytecode/policy/
trace.rs1use 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
150struct 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
192struct 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}