1use std::collections::{HashMap, HashSet};
4
5use ascent::Lattice;
6use ascent::lattice::BoundedLattice;
7use itertools::Itertools;
8
9use hugr_core::extension::prelude::{MakeTuple, UnpackTuple};
10use hugr_core::ops::{DataflowOpTrait, OpTag, OpTrait, OpType, TailLoop};
11use hugr_core::{HugrView, IncomingPort, OutgoingPort, PortIndex as _, Wire};
12
13use super::value_row::ValueRow;
14use super::{
15 AbstractValue, AnalysisResults, DFContext, LoadedFunction, PartialValue, partial_from_const,
16 row_contains_bottom,
17};
18
19type PV<V, N> = PartialValue<V, N>;
20
21type NodeInputs<V, N> = Vec<(IncomingPort, PV<V, N>)>;
22type NodeOutputs<V, N> = Vec<(OutgoingPort, PV<V, N>)>;
23
24#[deprecated(
33 note = "`hugr-passes` is deprecated. Use tket::passes instead",
34 since = "0.26.2"
35)]
36pub struct Machine<H: HugrView, V: AbstractValue> {
37 hugr: H,
38 in_wire_proto: HashMap<H::Node, NodeInputs<V, H::Node>>,
39 out_wire_proto: HashMap<H::Node, NodeOutputs<V, H::Node>>,
40}
41
42impl<H: HugrView, V: AbstractValue> Machine<H, V> {
43 pub fn new(hugr: H) -> Self {
45 Self {
46 hugr,
47 in_wire_proto: Default::default(),
48 out_wire_proto: Default::default(),
49 }
50 }
51}
52
53impl<H: HugrView, V: AbstractValue> Machine<H, V> {
54 pub fn prepopulate_wire(&mut self, w: Wire<H::Node>, v: PartialValue<V, H::Node>) {
57 self.out_wire_proto
58 .entry(w.node())
59 .or_default()
60 .push((w.source(), v));
61 }
62
63 #[expect(
69 clippy::result_large_err,
70 reason = "Not called recursively and not a performance bottleneck"
71 )]
72 #[inline]
73 pub fn prepopulate_inputs(
74 &mut self,
75 parent: H::Node,
76 in_values: impl IntoIterator<Item = (IncomingPort, PartialValue<V, H::Node>)>,
77 ) -> Result<(), OpType> {
78 if !self.hugr.contains_node(parent) {
79 return Ok(());
80 }
81 match self.hugr.get_optype(parent) {
82 OpType::DataflowBlock(_) | OpType::Case(_) | OpType::FuncDefn(_) => {
83 let [inp, _] = self.hugr.get_io(parent).unwrap();
85 let mut vals =
86 vec![PartialValue::Top; self.hugr.signature(inp).unwrap().output_types().len()];
87 for (ip, v) in in_values {
88 vals[ip.index()] = v;
89 }
90 for (i, v) in vals.into_iter().enumerate() {
91 self.prepopulate_wire(Wire::new(inp, i), v);
92 }
93 }
94 OpType::DFG(_) | OpType::TailLoop(_) | OpType::CFG(_) | OpType::Conditional(_) => {
95 let mut vals = vec![
97 PartialValue::Top;
98 self.hugr.signature(parent).unwrap().input_types().len()
99 ];
100 for (ip, v) in in_values {
101 vals[ip.index()] = v;
102 }
103 self.in_wire_proto
104 .entry(parent)
105 .or_default()
106 .extend(vals.into_iter().enumerate().map(|(i, v)| (i.into(), v)));
107 }
108 op => return Err(op.clone()),
109 }
110 Ok(())
111 }
112
113 pub fn run(
128 mut self,
129 context: impl DFContext<V, Node = H::Node>,
130 in_values: impl IntoIterator<Item = (IncomingPort, PartialValue<V, H::Node>)>,
131 ) -> AnalysisResults<V, H> {
132 if self.hugr.entrypoint_optype().is_module() {
133 assert!(
134 in_values.into_iter().next().is_none(),
135 "No inputs possible for Module"
136 );
137 } else {
138 let ep = self.hugr.entrypoint();
139 let have_value_for_entry = self.in_wire_proto.contains_key(&ep)
140 || (self.hugr.entrypoint_optype().tag() <= OpTag::DataflowParent
141 && self.out_wire_proto.contains_key(&ep));
142 let mut p = in_values.into_iter().peekable();
143 if p.peek().is_some() || !have_value_for_entry {
145 self.prepopulate_inputs(ep, p).unwrap();
146 }
147 }
148 run_datalog(
149 context,
150 self.hugr,
151 self.in_wire_proto
152 .into_iter()
153 .flat_map(|(n, vals)| vals.into_iter().map(move |(ip, v)| (n, ip, v)))
154 .collect(),
155 self.out_wire_proto
156 .into_iter()
157 .flat_map(|(n, vals)| vals.into_iter().map(move |(op, v)| (n, op, v)))
158 .collect(),
159 )
160 }
161}
162
163pub(super) type InWire<V, N> = (N, IncomingPort, PartialValue<V, N>);
164type OutWire<V, N> = (N, OutgoingPort, PartialValue<V, N>);
165
166fn run_datalog<V: AbstractValue, H: HugrView>(
167 mut ctx: impl DFContext<V, Node = H::Node>,
168 hugr: H,
169 in_wire_value_proto: Vec<InWire<V, H::Node>>,
170 out_wire_value_proto: Vec<OutWire<V, H::Node>>,
171) -> AnalysisResults<V, H> {
172 #![allow(
175 clippy::clone_on_copy,
176 clippy::unused_enumerate_index,
177 clippy::collapsible_if
178 )]
179 let all_results = ascent::ascent_run! {
180 pub(super) struct AscentProgram<V: AbstractValue, H: HugrView>;
181 relation node(H::Node); relation in_wire(H::Node, IncomingPort); relation out_wire(H::Node, OutgoingPort); relation parent_of_node(H::Node, H::Node); relation input_child(H::Node, H::Node); relation output_child(H::Node, H::Node); lattice out_wire_value(H::Node, OutgoingPort, PV<V, H::Node>); lattice in_wire_value(H::Node, IncomingPort, PV<V, H::Node>); lattice node_in_value_row(H::Node, ValueRow<V, H::Node>); node(n) <-- for n in hugr.nodes();
198
199 in_wire(n, p) <-- node(n), for (p,_) in hugr.in_value_types(*n); out_wire(n, p) <-- node(n), for (p,_) in hugr.out_value_types(*n); parent_of_node(parent, child) <--
203 node(child), if let Some(parent) = hugr.get_parent(*child);
204
205 input_child(parent, input) <-- node(parent), if let Some([input, _output]) = hugr.get_io(*parent);
206 output_child(parent, output) <-- node(parent), if let Some([_input, output]) = hugr.get_io(*parent);
207
208 out_wire_value(n, p, PV::bottom()) <-- out_wire(n, p);
210 in_wire_value(n, p, PV::bottom()) <-- in_wire(n, p);
211
212 in_wire_value(n, ip, v) <-- in_wire(n, ip),
214 if let Some((m, op)) = hugr.single_linked_output(*n, *ip),
215 out_wire_value(m, op, v);
216
217 in_wire_value(n, p, v) <-- for (n, p, v) in &in_wire_value_proto,
220 node(n),
221 if let Some(sig) = hugr.signature(*n),
222 if sig.input_ports().contains(p);
223
224 out_wire_value(n, p, v) <-- for (n, p, v) in &out_wire_value_proto,
226 node(n),
227 if let Some(sig) = hugr.signature(*n),
228 if sig.output_ports().contains(p);
229
230 node_in_value_row(n, ValueRow::new(sig.input_count())) <-- node(n), if let Some(sig) = hugr.signature(*n);
232 node_in_value_row(n, ValueRow::new(hugr.signature(*n).unwrap().input_count()).set(p.index(), v.clone())) <-- in_wire_value(n, p, v);
233
234 out_wire_value(n, p, v) <--
236 node(n),
237 let op_t = hugr.get_optype(*n),
238 if !op_t.is_container(),
239 if let Some(sig) = op_t.dataflow_signature(),
240 node_in_value_row(n, vs),
241 if let Some(outs) = propagate_leaf_op(&mut ctx, &hugr, *n, &vs[..], sig.output_count()),
242 for (p, v) in (0..).map(OutgoingPort::from).zip(outs);
243
244 relation dfg_node(H::Node); dfg_node(n) <-- node(n), if hugr.get_optype(*n).is_dfg();
247
248 out_wire_value(i, OutgoingPort::from(p.index()), v) <-- dfg_node(dfg),
249 input_child(dfg, i), in_wire_value(dfg, p, v);
250
251 out_wire_value(dfg, OutgoingPort::from(p.index()), v) <-- dfg_node(dfg),
252 output_child(dfg, o), in_wire_value(o, p, v);
253
254 out_wire_value(i, OutgoingPort::from(p.index()), v) <-- node(tl),
257 if hugr.get_optype(*tl).is_tail_loop(),
258 input_child(tl, i),
259 in_wire_value(tl, p, v);
260
261 out_wire_value(in_n, OutgoingPort::from(out_p), v) <-- node(tl),
263 if let Some(tailloop) = hugr.get_optype(*tl).as_tail_loop(),
264 input_child(tl, in_n),
265 output_child(tl, out_n),
266 node_in_value_row(out_n, out_in_row), if let Some(fields) = out_in_row.unpack_first(TailLoop::CONTINUE_TAG, tailloop.just_inputs.len()),
269 for (out_p, v) in fields.enumerate();
270
271 out_wire_value(tl, OutgoingPort::from(out_p), v) <-- node(tl),
273 if let Some(tailloop) = hugr.get_optype(*tl).as_tail_loop(),
274 output_child(tl, out_n),
275 node_in_value_row(out_n, out_in_row), if let Some(fields) = out_in_row.unpack_first(TailLoop::BREAK_TAG, tailloop.just_outputs.len()),
278 for (out_p, v) in fields.enumerate();
279
280 relation case_node(H::Node, usize, H::Node);
283 case_node(cond, i, case) <-- node(cond),
284 if hugr.get_optype(*cond).is_conditional(),
285 for (i, case) in hugr.children(*cond).enumerate(),
286 if hugr.get_optype(case).is_case();
287
288 out_wire_value(i_node, OutgoingPort::from(out_p), v) <--
290 case_node(cond, case_index, case),
291 input_child(case, i_node),
292 node_in_value_row(cond, in_row),
293 let conditional = hugr.get_optype(*cond).as_conditional().unwrap(),
294 if let Some(fields) = in_row.unpack_first(*case_index, conditional.sum_rows[*case_index].len()),
295 for (out_p, v) in fields.enumerate();
296
297 out_wire_value(cond, OutgoingPort::from(o_p.index()), v) <--
299 case_node(cond, _i, case),
300 case_reachable(cond, case),
301 output_child(case, o),
302 in_wire_value(o, o_p, v);
303
304 relation case_reachable(H::Node, H::Node);
306 case_reachable(cond, case) <-- case_node(cond, i, case),
307 in_wire_value(cond, IncomingPort::from(0), v),
308 if v.supports_tag(*i);
309
310 relation cfg_node(H::Node); cfg_node(n) <-- node(n), if hugr.get_optype(*n).is_cfg();
313
314 relation bb_reachable(H::Node, H::Node);
316 bb_reachable(cfg, entry) <-- cfg_node(cfg), if let Some(entry) = hugr.children(*cfg).next();
317 bb_reachable(cfg, bb) <-- cfg_node(cfg),
318 bb_reachable(cfg, pred),
319 output_child(pred, pred_out),
320 in_wire_value(pred_out, IncomingPort::from(0), predicate),
321 for (tag, bb) in hugr.output_neighbours(*pred).enumerate(),
322 if predicate.supports_tag(tag);
323
324 out_wire_value(i_node, OutgoingPort::from(p.index()), v) <--
326 cfg_node(cfg),
327 if let Some(entry) = hugr.children(*cfg).next(),
328 input_child(entry, i_node),
329 in_wire_value(cfg, p, v);
330
331 relation _cfg_succ_dest(H::Node, H::Node, H::Node);
334 _cfg_succ_dest(cfg, exit, cfg) <-- cfg_node(cfg), if let Some(exit) = hugr.children(*cfg).nth(1);
335 _cfg_succ_dest(cfg, blk, inp) <-- cfg_node(cfg),
336 for blk in hugr.children(*cfg),
337 if hugr.get_optype(blk).is_dataflow_block(),
338 input_child(blk, inp);
339
340 out_wire_value(dest, OutgoingPort::from(out_p), v) <--
342 bb_reachable(cfg, pred),
343 if let Some(df_block) = hugr.get_optype(*pred).as_dataflow_block(),
344 for (succ_n, succ) in hugr.output_neighbours(*pred).enumerate(),
345 output_child(pred, out_n),
346 _cfg_succ_dest(cfg, succ, dest),
347 node_in_value_row(out_n, out_in_row),
348 if let Some(fields) = out_in_row.unpack_first(succ_n, df_block.sum_rows.get(succ_n).unwrap().len()),
349 for (out_p, v) in fields.enumerate();
350
351 relation func_call(H::Node, H::Node); func_call(call, func_defn) <--
354 node(call),
355 if hugr.get_optype(*call).is_call(),
356 if let Some(func_defn) = hugr.static_source(*call);
357
358 out_wire_value(inp, OutgoingPort::from(p.index()), v) <--
359 func_call(call, func),
360 input_child(func, inp),
361 in_wire_value(call, p, v);
362
363 out_wire_value(call, OutgoingPort::from(p.index()), v) <--
364 func_call(call, func),
365 output_child(func, outp),
366 in_wire_value(outp, p, v);
367
368 lattice indirect_call(H::Node, LatticeWrapper<H::Node>); indirect_call(call, tgt) <--
371 node(call),
372 if let OpType::CallIndirect(_) = hugr.get_optype(*call),
373 in_wire_value(call, IncomingPort::from(0), v),
374 let tgt = load_func(v);
375
376 out_wire_value(inp, OutgoingPort::from(p.index()-1), v) <--
377 indirect_call(call, lv),
378 if let LatticeWrapper::Value(func) = lv,
379 input_child(func, inp),
380 in_wire_value(call, p, v)
381 if p.index() > 0;
382
383 out_wire_value(call, OutgoingPort::from(p.index()), v) <--
384 indirect_call(call, lv),
385 if let LatticeWrapper::Value(func) = lv,
386 output_child(func, outp),
387 in_wire_value(outp, p, v);
388
389 out_wire_value(call, p, PV::Top) <--
392 node(call),
393 if let OpType::CallIndirect(ci) = hugr.get_optype(*call),
394 in_wire_value(call, IncomingPort::from(0), v),
395 if matches!(v, PartialValue::Top | PartialValue::Value(_)),
397 for p in ci.signature().output_ports();
398 };
399 let entry_descs = hugr.entry_descendants().collect::<HashSet<_>>();
400 let out_wire_values = all_results
401 .out_wire_value
402 .iter()
403 .filter(|(n, _, _)| entry_descs.contains(n))
404 .map(|(n, p, v)| (Wire::new(*n, *p), v.clone()))
405 .collect();
406 AnalysisResults {
407 hugr,
408 out_wire_values,
409 in_wire_value: all_results
410 .in_wire_value
411 .into_iter()
412 .filter(|(n, _, _)| entry_descs.contains(n))
413 .collect(),
414 case_reachable: all_results
415 .case_reachable
416 .into_iter()
417 .filter(|(_, n)| entry_descs.contains(n))
418 .collect(),
419 bb_reachable: all_results
420 .bb_reachable
421 .into_iter()
422 .filter(|(_, n)| entry_descs.contains(n))
423 .collect(),
424 }
425}
426
427#[derive(Debug, PartialEq, Eq, Hash, Clone, PartialOrd)]
428enum LatticeWrapper<T> {
429 Bottom,
430 Value(T),
431 Top,
432}
433
434impl<N: PartialEq + PartialOrd> Lattice for LatticeWrapper<N> {
435 fn meet_mut(&mut self, other: Self) -> bool {
436 if *self == other || *self == LatticeWrapper::Bottom || other == LatticeWrapper::Top {
437 return false;
438 }
439 if *self == LatticeWrapper::Top || other == LatticeWrapper::Bottom {
440 *self = other;
441 return true;
442 }
443 *self = LatticeWrapper::Bottom;
445 true
446 }
447
448 fn join_mut(&mut self, other: Self) -> bool {
449 if *self == other || *self == LatticeWrapper::Top || other == LatticeWrapper::Bottom {
450 return false;
451 }
452 if *self == LatticeWrapper::Bottom || other == LatticeWrapper::Top {
453 *self = other;
454 return true;
455 }
456 *self = LatticeWrapper::Top;
458 true
459 }
460}
461
462fn load_func<V, N: Copy>(v: &PV<V, N>) -> LatticeWrapper<N> {
463 match v {
464 PartialValue::Bottom | PartialValue::PartialSum(_) => LatticeWrapper::Bottom,
465 PartialValue::LoadedFunction(LoadedFunction { func_node, .. }) => {
466 LatticeWrapper::Value(*func_node)
467 }
468 PartialValue::Value(_) | PartialValue::Top => LatticeWrapper::Top,
469 }
470}
471
472fn propagate_leaf_op<V: AbstractValue, H: HugrView>(
473 ctx: &mut impl DFContext<V, Node = H::Node>,
474 hugr: &H,
475 n: H::Node,
476 ins: &[PV<V, H::Node>],
477 num_outs: usize,
478) -> Option<ValueRow<V, H::Node>> {
479 match hugr.get_optype(n) {
480 op if op.cast::<MakeTuple>().is_some() => Some(ValueRow::from_iter([PV::new_variant(
483 0,
484 ins.iter().cloned(),
485 )])),
486 op if op.cast::<UnpackTuple>().is_some() => {
487 let elem_tys = op.cast::<UnpackTuple>().unwrap().0;
488 let tup = ins.iter().exactly_one().unwrap();
489 tup.variant_values(0, elem_tys.len())
490 .map(ValueRow::from_iter)
491 }
492 OpType::Tag(t) => Some(ValueRow::from_iter([PV::new_variant(
493 t.tag,
494 ins.iter().cloned(),
495 )])),
496 OpType::Input(_) | OpType::Output(_) | OpType::ExitBlock(_) => None, OpType::Call(_) | OpType::CallIndirect(_) => None, OpType::LoadConstant(load_op) => {
499 assert!(ins.is_empty()); let const_node = hugr
501 .single_linked_output(n, load_op.constant_port())
502 .unwrap()
503 .0;
504 let const_val = hugr.get_optype(const_node).as_const().unwrap().value();
505 Some(ValueRow::singleton(partial_from_const(ctx, n, const_val)))
506 }
507 OpType::LoadFunction(load_op) => {
508 assert!(ins.is_empty()); let func_node = hugr
510 .single_linked_output(n, load_op.function_port())
511 .unwrap()
512 .0;
513 Some(ValueRow::singleton(PartialValue::new_load(
515 func_node,
516 load_op.type_args.clone(),
517 )))
518 }
519 OpType::ExtensionOp(e) => {
520 Some(ValueRow::from_iter(if row_contains_bottom(ins) {
521 vec![PartialValue::Bottom; num_outs]
524 } else {
525 let mut outs = vec![PartialValue::Top; num_outs];
528 ctx.interpret_leaf_op(n, e, ins, &mut outs[..]);
532 outs
533 }))
534 }
535 o => todo!("Unhandled: {:?}", o), }
538}
539
540#[cfg(test)]
541mod test {
542 use ascent::Lattice;
543
544 use super::LatticeWrapper;
545
546 #[test]
547 fn latwrap_join() {
548 for lv in [
549 LatticeWrapper::Value(3),
550 LatticeWrapper::Value(5),
551 LatticeWrapper::Top,
552 ] {
553 let mut subject = LatticeWrapper::Bottom;
554 assert!(subject.join_mut(lv.clone()));
555 assert_eq!(subject, lv);
556 assert!(!subject.join_mut(lv.clone()));
557 assert_eq!(subject, lv);
558 assert_eq!(
559 subject.join_mut(LatticeWrapper::Value(11)),
560 lv != LatticeWrapper::Top
561 );
562 assert_eq!(subject, LatticeWrapper::Top);
563 }
564 }
565
566 #[test]
567 fn latwrap_meet() {
568 for lv in [
569 LatticeWrapper::Bottom,
570 LatticeWrapper::Value(3),
571 LatticeWrapper::Value(5),
572 ] {
573 let mut subject = LatticeWrapper::Top;
574 assert!(subject.meet_mut(lv.clone()));
575 assert_eq!(subject, lv);
576 assert!(!subject.meet_mut(lv.clone()));
577 assert_eq!(subject, lv);
578 assert_eq!(
579 subject.meet_mut(LatticeWrapper::Value(11)),
580 lv != LatticeWrapper::Bottom
581 );
582 assert_eq!(subject, LatticeWrapper::Bottom);
583 }
584 }
585}