1#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Hash)]
28pub enum JitMode {
29 Off,
31 #[default]
34 Auto,
35 Force,
38}
39
40#[cfg(not(feature = "jit"))]
41pub(crate) fn extract_jit_cones(_dag: &mut super::assembly::ResolvedDag, _mode: JitMode) {}
42
43#[cfg(feature = "jit")]
44pub(crate) use jit_impl::extract_jit_cones;
45
46#[cfg(feature = "jit")]
47mod jit_impl {
48 use super::JitMode;
49 use crate::ast::{NodeMeta, PolydatNode, Port, PortType, Purity, Slot, SlotShape, Value};
50 use crate::compile::assembly::{PolydatAssembler, ResolvedDag};
51 use crate::compile::jit::{JitOp, classify_node_typed};
52 use crate::kernel::{InputDef, InputKind, WireSource};
53 use std::collections::HashMap;
54
55 pub(crate) struct JitConeNode {
62 meta: NodeMeta,
63 code_fn: crate::compile::jit::NativeFn,
64 total_slots: usize,
65 scratch: crate::compile::jit::ScratchPlan,
68 attribution: std::sync::Arc<crate::compile::Attribution>,
72 in_slots: Vec<usize>,
74 out_slots: Vec<usize>,
76 in_types: Vec<PortType>,
77 out_types: Vec<PortType>,
78 members: Vec<Box<dyn PolydatNode>>,
82 sub_wiring: Vec<Vec<WireSource>>,
86 out_ports: Vec<(usize, usize)>,
88 _module: crate::compile::jit::JitCode,
91 fallible: bool,
94 }
95
96 impl PolydatNode for JitConeNode {
97 fn meta(&self) -> &NodeMeta {
98 &self.meta
99 }
100
101 fn fusion_subgraph(&self) -> Option<crate::ast::FusionSubgraph<'_>> {
102 Some(crate::ast::FusionSubgraph {
103 members: &self.members,
104 wiring: &self.sub_wiring,
105 out_ports: &self.out_ports,
106 })
107 }
108
109 fn scratch_layout(&self) -> Vec<crate::ast::ScratchElem> {
114 let mut layout = vec![crate::ast::ScratchElem::Slots];
115 layout.extend(self.scratch.elems.iter().copied());
116 layout
117 }
118
119 fn eval_in(
120 &self,
121 scratch: &mut [crate::ast::ScratchBuf],
122 inputs: &[Value],
123 outputs: &mut [Value],
124 ) {
125 let (slots, members) = scratch.split_at_mut(1);
126 let crate::ast::ScratchBuf::Slots(buf) = &mut slots[0] else {
127 unreachable!("a cone's scratch is its slot buffer");
128 };
129 self.eval_with(buf, members, inputs, outputs)
130 }
131
132 fn eval(&self, inputs: &[Value], outputs: &mut [Value]) {
135 let mut buf = Vec::new();
136 let mut members: Vec<crate::ast::ScratchBuf> = self
137 .scratch
138 .elems
139 .iter()
140 .map(|e| crate::ast::ScratchBuf::new(*e))
141 .collect();
142 self.eval_with(&mut buf, &mut members, inputs, outputs)
143 }
144 }
145
146 impl JitConeNode {
147 fn eval_with(
153 &self,
154 buf: &mut Vec<u64>,
155 members: &mut [crate::ast::ScratchBuf],
156 inputs: &[Value],
157 outputs: &mut [Value],
158 ) {
159 buf.clear();
160 buf.resize(self.total_slots + 1, 0);
161 for (i, v) in inputs.iter().enumerate() {
162 let start = self.in_slots[i];
163 if crate::compile::marshal::encode_slots(v, &mut buf[start..]).is_none() {
164 panic!(
165 "cone `{}` boundary input [{i}] expected {:?}, got {:?}",
166 self.meta.name,
167 self.in_types[i],
168 v.port_type()
169 );
170 }
171 }
172 let code_fn = self.code_fn;
177 let cp = buf.as_ptr();
178 let mp = buf.as_mut_ptr();
179 let sc = members.as_mut_ptr();
180 if !self.fallible {
181 unsafe { (code_fn)(cp, mp, sc) };
183 } else {
184 buf[self.total_slots] = u64::MAX;
185 let capture = crate::kernel::engines::EvalPanicCaptureGuard::arm();
186 let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
187 crate::compile::jit::invoke_with_catch(move || unsafe {
188 (code_fn)(cp, mp, sc);
189 })
190 }));
191 drop(capture);
192 if let Err(payload) = outcome {
193 let step = buf[self.total_slots] as usize;
194 self.attribution.reraise(payload, step, buf, None);
195 }
196 }
197 #[cfg(debug_assertions)]
198 for &(slot, idx) in &self.scratch.refs {
199 let (p, l) = members[idx].ptr_len();
200 assert!(
201 buf[slot] == p && buf[slot + 1] == l,
202 "S9 ref-validator: cone `{}` slot pair ({slot}, {}) does not name \
203 scratch[{idx}]",
204 self.meta.name,
205 slot + 1
206 );
207 }
208 for (k, slot) in self.out_slots.iter().enumerate() {
209 outputs[k] = crate::compile::marshal::decode_output(buf, *slot, self.out_types[k]);
210 }
211 }
212 }
213
214 fn audit_skip(member_count: usize, reason: &str) {
218 crate::library::support::audit::debug(&format!(
219 "jit cone: leaving a {member_count}-member component on the interpreter: {reason}"
220 ));
221 }
222
223 fn scalar_ok(ty: PortType) -> bool {
231 use crate::ast::SlotColor;
232 match ty.slot_color() {
233 SlotColor::Imm1 | SlotColor::Ref2 => true,
234 SlotColor::Imm2 => false,
235 }
236 }
237
238 fn node_eligible(node: &dyn PolydatNode, wire_types: &[PortType]) -> bool {
244 matches!(node.purity(), Purity::Pure)
245 && !matches!(classify_node_typed(node, wire_types), JitOp::Fallback)
246 && node.meta().outs.iter().all(|p| scalar_ok(p.typ))
247 && wire_types.iter().all(|t| scalar_ok(*t))
248 && node.meta().wire_inputs().iter().all(|p| scalar_ok(p.typ))
249 }
250
251 #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
259 enum Lc {
260 CompileConst,
261 ScopeInit,
262 Dynamic,
263 }
264
265 fn classify_lifecycles(dag: &ResolvedDag, nodes: &[Box<dyn PolydatNode>]) -> Vec<Lc> {
266 let n = dag.wiring.len();
267 let mut lc = vec![Lc::CompileConst; n];
268 for i in 0..n {
269 for src in &dag.wiring[i] {
270 if let WireSource::Input(idx) = src {
271 let kind = dag
272 .input_defs
273 .get(*idx)
274 .map(|d| d.kind)
275 .unwrap_or(InputKind::Coordinate);
276 let seed = match kind {
277 InputKind::IterationExtern => Lc::ScopeInit,
278 InputKind::Coordinate | InputKind::ExternalWrite => Lc::Dynamic,
279 };
280 lc[i] = lc[i].max(seed);
281 }
282 }
283 if matches!(nodes[i].purity(), Purity::Nondeterministic { .. }) {
284 lc[i] = Lc::Dynamic;
285 }
286 }
287 loop {
288 let mut changed = false;
289 for i in 0..n {
290 for src in &dag.wiring[i] {
291 if let WireSource::NodeOutput(j, _) = src
292 && lc[*j] > lc[i]
293 {
294 lc[i] = lc[*j];
295 changed = true;
296 }
297 }
298 }
299 if !changed {
300 break;
301 }
302 }
303 lc
304 }
305
306 fn src_key(src: &WireSource) -> (u8, usize, usize) {
308 match src {
309 WireSource::Input(i) => (0, *i, 0),
310 WireSource::NodeOutput(j, p) => (1, *j, *p),
311 }
312 }
313
314 struct ConePlan {
315 members: Vec<usize>,
317 boundary_in: Vec<WireSource>,
319 in_types: Vec<PortType>,
320 boundary_out: Vec<(usize, usize)>,
322 out_types: Vec<PortType>,
323 }
324
325 pub(crate) fn extract_jit_cones(dag: &mut ResolvedDag, mode: JitMode) {
329 let min_members = match mode {
330 JitMode::Off => return,
331 JitMode::Auto => 2,
332 JitMode::Force => 1,
333 };
334 let n = dag.nodes.len();
335 if n == 0 {
336 return;
337 }
338
339 let lifecycles = classify_lifecycles(dag, &dag.nodes);
340 let mut eligible: Vec<bool> = vec![false; n];
349 for i in 0..n {
350 if lifecycles[i] != Lc::Dynamic {
351 continue;
352 }
353 let nd = dag.nodes[i].as_ref();
354 if !node_eligible(nd, &crate::compile::assembly::wire_types_of(dag, i)) {
355 continue;
356 }
357 if nd.accepts_none_inputs()
358 && !dag.wiring[i]
359 .iter()
360 .all(|src| matches!(src, WireSource::NodeOutput(j, _) if eligible[*j]))
361 {
362 continue;
363 }
364 eligible[i] = true;
365 }
366
367 let mut parent: Vec<usize> = (0..n).collect();
369 fn find(parent: &mut [usize], mut i: usize) -> usize {
370 while parent[i] != i {
371 parent[i] = parent[parent[i]];
372 i = parent[i];
373 }
374 i
375 }
376 for i in 0..n {
377 if !eligible[i] {
378 continue;
379 }
380 for src in &dag.wiring[i] {
381 if let WireSource::NodeOutput(j, _) = src
382 && eligible[*j]
383 {
384 let (a, b) = (find(&mut parent, i), find(&mut parent, *j));
385 parent[a] = b;
386 }
387 }
388 }
389 let mut components: HashMap<usize, Vec<usize>> = HashMap::new();
390 for (i, &is_eligible) in eligible.iter().enumerate().take(n) {
391 if is_eligible {
392 components.entry(find(&mut parent, i)).or_default().push(i);
393 }
394 }
395 let mut roots: Vec<usize> = components.keys().copied().collect();
396 roots.sort_unstable();
397
398 let mut consumers: Vec<Vec<usize>> = vec![Vec::new(); n];
401 for (i, wiring) in dag.wiring.iter().enumerate() {
402 for src in wiring {
403 if let WireSource::NodeOutput(j, _) = src {
404 consumers[*j].push(i);
405 }
406 }
407 }
408
409 let mut nodes_opt: Vec<Option<Box<dyn PolydatNode>>> = std::mem::take(&mut dag.nodes)
410 .into_iter()
411 .map(Some)
412 .collect();
413 let mut cones: Vec<(ConePlan, JitConeNode)> = Vec::new();
414
415 for root in roots {
416 let members = &components[&root];
417 if members.len() < min_members {
418 continue;
419 }
420 if !component_is_convex(members, &consumers, n) {
435 audit_skip(
436 members.len(),
437 "non-convex component (an external path re-enters the cone)",
438 );
439 continue;
440 }
441 let Some(plan) = plan_cone(dag, members, &nodes_opt) else {
442 continue;
445 };
446 match build_cone(dag, &plan, &mut nodes_opt) {
447 Ok(cone) => {
448 crate::library::support::audit::debug(&format!(
452 "jit cone: fused {} members ({} boundary in, {} out): {}",
453 plan.members.len(),
454 plan.boundary_in.len(),
455 plan.boundary_out.len(),
456 cone.meta().name,
457 ));
458 cones.push((plan, cone));
459 }
460 Err(e) => {
466 crate::library::support::audit::warn(&format!(
467 "jit cone: codegen failed for a {}-member cone — staying on the interpreter: {e}",
468 plan.members.len(),
469 ));
470 }
471 }
472 }
473
474 if cones.is_empty() {
475 dag.nodes = nodes_opt.into_iter().map(Option::unwrap).collect();
476 return;
477 }
478 rebuild(dag, nodes_opt, cones);
479 }
480
481 fn component_is_convex(members: &[usize], consumers: &[Vec<usize>], n: usize) -> bool {
487 let mut is_member = vec![false; n];
488 for &m in members {
489 is_member[m] = true;
490 }
491 let mut seen = vec![false; n];
492 let mut stack: Vec<usize> = Vec::new();
493 for &m in members {
494 for &c in &consumers[m] {
495 if !is_member[c] && !seen[c] {
496 seen[c] = true;
497 stack.push(c);
498 }
499 }
500 }
501 while let Some(x) = stack.pop() {
502 for &c in &consumers[x] {
503 if is_member[c] {
504 return false;
505 }
506 if !seen[c] {
507 seen[c] = true;
508 stack.push(c);
509 }
510 }
511 }
512 true
513 }
514
515 fn plan_cone(
518 dag: &ResolvedDag,
519 members: &[usize],
520 nodes: &[Option<Box<dyn PolydatNode>>],
521 ) -> Option<ConePlan> {
522 let is_member = |j: usize| members.binary_search(&j).is_ok();
523
524 let mut boundary_in: Vec<WireSource> = Vec::new();
525 let mut in_types: Vec<PortType> = Vec::new();
526 let mut seen_in: HashMap<(u8, usize, usize), usize> = HashMap::new();
527 for &m in members {
528 let member = nodes[m].as_ref()?;
529 let member_ports: Vec<PortType> =
530 member.meta().wire_inputs().iter().map(|p| p.typ).collect();
531 let wire_types: Vec<PortType> = dag.wiring[m]
532 .iter()
533 .map(|src| match src {
534 WireSource::Input(i) => Some(dag.input_defs[*i].port_type),
535 WireSource::NodeOutput(j, p) => Some(nodes[*j].as_ref()?.meta().outs[*p].typ),
536 })
537 .collect::<Option<_>>()?;
538 let typed_by_wires = matches!(
544 classify_node_typed(member.as_ref(), &wire_types),
545 JitOp::SlotCall { .. }
546 );
547 for (k, src) in dag.wiring[m].iter().enumerate() {
548 let ty = wire_types[k];
549 if !typed_by_wires
551 && let Some(expected) = member_ports.get(k)
552 && *expected != ty
553 {
554 audit_skip(
555 members.len(),
556 &format!(
557 "input [{k}] of `{}` is a {ty:?} wire on a {expected:?} port",
558 member.meta().name
559 ),
560 );
561 return None;
562 }
563 let intra = matches!(src, WireSource::NodeOutput(j, _) if is_member(*j));
564 if !intra && member.accepts_none_inputs() {
568 audit_skip(
569 members.len(),
570 &format!(
571 "`{}` tolerates None inputs and input [{k}] is a boundary wire",
572 member.meta().name
573 ),
574 );
575 return None;
576 }
577 if intra {
578 continue;
579 }
580 let key = src_key(src);
581 if seen_in.contains_key(&key) {
582 continue;
583 }
584 if !scalar_ok(ty) {
585 audit_skip(
586 members.len(),
587 &format!("boundary input of type {ty:?} is not marshalable"),
588 );
589 return None;
590 }
591 seen_in.insert(key, boundary_in.len());
592 boundary_in.push(src.clone());
593 in_types.push(ty);
594 }
595 }
596 if boundary_in.len() > 64 {
601 audit_skip(
602 members.len(),
603 &format!(
604 "{} boundary inputs exceeds the 64-input bound (no re-split implemented — catchup item B2)",
605 boundary_in.len()
606 ),
607 );
608 return None;
609 }
610 if boundary_in.is_empty() {
616 return None;
619 }
620
621 let mut boundary_out: Vec<(usize, usize)> = Vec::new();
622 let mut seen_out: HashMap<(usize, usize), usize> = HashMap::new();
623 let mut note_out = |j: usize, p: usize| {
624 if let std::collections::hash_map::Entry::Vacant(e) = seen_out.entry((j, p)) {
625 e.insert(boundary_out.len());
626 boundary_out.push((j, p));
627 }
628 };
629 for (i, wiring) in dag.wiring.iter().enumerate() {
630 if is_member(i) {
631 continue;
632 }
633 for src in wiring {
634 if let WireSource::NodeOutput(j, p) = src
635 && is_member(*j)
636 {
637 note_out(*j, *p);
638 }
639 }
640 }
641 for (j, p) in dag.output_map.values() {
642 if is_member(*j) {
643 note_out(*j, *p);
644 }
645 }
646 if boundary_out.is_empty() {
647 return None;
650 }
651 let out_types: Vec<PortType> = boundary_out
652 .iter()
653 .map(|(j, p)| nodes[*j].as_ref().map(|nd| nd.meta().outs[*p].typ))
654 .collect::<Option<_>>()?;
655 if out_types.iter().any(|t| !scalar_ok(*t)) {
656 audit_skip(members.len(), "a boundary output type is not marshalable");
657 return None;
658 }
659
660 Some(ConePlan {
661 members: members.to_vec(),
662 boundary_in,
663 in_types,
664 boundary_out,
665 out_types,
666 })
667 }
668
669 fn default_for(ty: PortType) -> Value {
673 match ty {
674 PortType::F64 => Value::F64(0.0),
675 PortType::Bool => Value::Bool(false),
676 PortType::Str => Value::Str("".into()),
677 PortType::Bytes => Value::Bytes(Vec::new().into()),
678 PortType::Json => Value::Json(std::sync::Arc::new(serde_json::Value::Null)),
679 PortType::U64 => Value::U64(0),
680 _ => Value::None,
681 }
682 }
683
684 fn build_cone(
688 dag: &ResolvedDag,
689 plan: &ConePlan,
690 nodes: &mut [Option<Box<dyn PolydatNode>>],
691 ) -> Result<JitConeNode, String> {
692 let local: HashMap<usize, usize> = plan
693 .members
694 .iter()
695 .enumerate()
696 .map(|(l, &g)| (g, l))
697 .collect();
698 let in_pos: HashMap<(u8, usize, usize), usize> = plan
699 .boundary_in
700 .iter()
701 .enumerate()
702 .map(|(i, s)| (src_key(s), i))
703 .collect();
704
705 let sub_wiring: Vec<Vec<WireSource>> = plan
706 .members
707 .iter()
708 .map(|&m| {
709 dag.wiring[m]
710 .iter()
711 .map(|src| match src {
712 WireSource::NodeOutput(j, p) if local.contains_key(j) => {
713 WireSource::NodeOutput(local[j], *p)
714 }
715 other => WireSource::Input(in_pos[&src_key(other)]),
716 })
717 .collect()
718 })
719 .collect();
720 let sub_input_defs: Vec<InputDef> = plan
721 .in_types
722 .iter()
723 .enumerate()
724 .map(|(i, ty)| InputDef {
725 name: format!("c{i}"),
726 default: default_for(*ty),
727 port_type: *ty,
728 kind: InputKind::Coordinate,
729 })
730 .collect();
731 let mut sub_output_map: HashMap<String, (usize, usize)> = HashMap::new();
732 let mut sub_output_order: Vec<String> = Vec::new();
733 for (k, (j, p)) in plan.boundary_out.iter().enumerate() {
734 let name = format!("o{k}");
735 sub_output_map.insert(name.clone(), (local[j], *p));
736 sub_output_order.push(name);
737 }
738
739 let taken: Vec<Box<dyn PolydatNode>> = plan
740 .members
741 .iter()
742 .map(|&m| nodes[m].take().expect("cone member present"))
743 .collect();
744 let member_label = cone_label(&taken);
745
746 let mut sub = ResolvedDag {
747 nodes: taken,
748 wiring: sub_wiring,
749 input_defs: sub_input_defs,
750 coord_count: plan.boundary_in.len(),
751 output_map: sub_output_map,
752 output_order: sub_output_order,
753 cursor_schemas: Vec::new(),
754 source: String::new(),
755 context: dag.context.clone(),
759 output_modifiers: HashMap::new(),
760 const_outputs: std::collections::HashSet::new(),
761 ledger: crate::kernel::CompileLedger::new(),
765 };
766
767 let restore = |sub_nodes: Vec<Box<dyn PolydatNode>>,
768 nodes: &mut [Option<Box<dyn PolydatNode>>]| {
769 for (&m, nd) in plan.members.iter().zip(sub_nodes) {
770 nodes[m] = Some(nd);
771 }
772 };
773
774 let layout = match PolydatAssembler::build_jit_layout(&sub) {
775 Ok(l) => l,
776 Err(e) => {
777 restore(sub.nodes, nodes);
778 return Err(e);
779 }
780 };
781 let (coord_slots, total_slots, jit_steps, jit_outputs, scratch, _volatile) = layout;
782 let mut in_slots = Vec::with_capacity(plan.in_types.len());
785 let mut next = 0usize;
786 for ty in &plan.in_types {
787 in_slots.push(next);
788 next += ty.slot_width();
789 }
790 debug_assert_eq!(coord_slots, next);
791 let compiled = crate::compile::jit::compile_jit_entry(&jit_steps, Some(total_slots));
792 let (code_fn, code) = match compiled {
793 Ok(parts) => parts,
794 Err(e) => {
795 restore(sub.nodes, nodes);
796 return Err(e);
797 }
798 };
799
800 let out_slots: Vec<usize> = (0..plan.boundary_out.len())
801 .map(|k| jit_outputs[&format!("o{k}")])
802 .collect();
803 let meta = NodeMeta {
810 name: member_label,
811 ins: plan
812 .boundary_in
813 .iter()
814 .zip(&plan.in_types)
815 .enumerate()
816 .map(|(i, (src, ty))| {
817 let mut port = match src {
821 WireSource::NodeOutput(j, p) => nodes[*j]
822 .as_ref()
823 .map(|nd| nd.meta().outs[*p].clone())
824 .unwrap_or_else(|| Port::new("", *ty)),
825 WireSource::Input(_) => Port::new("", *ty),
826 };
827 port.name = format!("c{i}");
828 port.constraint = None;
829 Slot::Wire(port)
830 })
831 .collect(),
832 outs: plan
833 .boundary_out
834 .iter()
835 .enumerate()
836 .map(|(k, (j, p))| {
837 let mut port = sub.nodes[local[j]].meta().outs[*p].clone();
838 port.name = format!("o{k}");
839 port.constraint = None;
840 port
841 })
842 .collect(),
843 };
844 let out_ports: Vec<(usize, usize)> = plan
845 .boundary_out
846 .iter()
847 .map(|(j, p)| (local[j], *p))
848 .collect();
849 let mut named = sub.output_map.clone();
853 for (k, (j, p)) in plan.boundary_out.iter().enumerate() {
854 let names: Vec<String> = dag
855 .output_map
856 .iter()
857 .filter(|(_, v)| **v == (*j, *p))
858 .map(|(n, _)| n.clone())
859 .collect();
860 if !names.is_empty()
861 && let Some(target) = named.remove(&format!("o{k}"))
862 {
863 for n in names {
864 named.insert(n, target);
865 }
866 }
867 }
868 let numbered = std::mem::replace(&mut sub.output_map, named);
869 let attribution = std::sync::Arc::new(PolydatAssembler::attribution_of(&sub));
870 sub.output_map = numbered;
871 Ok(JitConeNode {
872 attribution,
873 in_slots,
874 meta,
875 code_fn,
876 total_slots,
877 out_slots,
878 in_types: plan.in_types.clone(),
879 out_types: plan.out_types.clone(),
880 members: sub.nodes,
881 sub_wiring: sub.wiring,
882 out_ports,
883 scratch,
884 fallible: code.fallible(),
885 _module: code,
886 })
887 }
888
889 fn cone_label(members: &[Box<dyn PolydatNode>]) -> String {
892 const SHOWN: usize = 6;
893 let names: Vec<&str> = members
894 .iter()
895 .take(SHOWN)
896 .map(|n| n.meta().name.as_str())
897 .collect();
898 let suffix = if members.len() > SHOWN {
899 format!("+{} more", members.len() - SHOWN)
900 } else {
901 String::new()
902 };
903 format!("jit_cone[{}{}]", names.join("+"), suffix)
904 }
905
906 fn rebuild(
909 dag: &mut ResolvedDag,
910 nodes_opt: Vec<Option<Box<dyn PolydatNode>>>,
911 cones: Vec<(ConePlan, JitConeNode)>,
912 ) {
913 let old_n = nodes_opt.len();
914 let mut cone_port: HashMap<(usize, usize), (usize, usize)> = HashMap::new();
916 for (ci, (plan, _)) in cones.iter().enumerate() {
917 for (k, (j, p)) in plan.boundary_out.iter().enumerate() {
918 cone_port.insert((*j, *p), (ci, k));
919 }
920 }
921
922 let mut kept_map: HashMap<usize, usize> = HashMap::new();
923 let mut new_nodes: Vec<Box<dyn PolydatNode>> = Vec::new();
924 let mut new_wiring: Vec<Vec<WireSource>> = Vec::new();
925 for (old, slot) in nodes_opt.into_iter().enumerate() {
926 if let Some(node) = slot {
927 kept_map.insert(old, new_nodes.len());
928 new_nodes.push(node);
929 new_wiring.push(dag.wiring[old].clone());
930 }
931 }
932 let cone_base = new_nodes.len();
933 let mut cone_plans: Vec<ConePlan> = Vec::with_capacity(cones.len());
934 for (plan, cone) in cones {
935 new_nodes.push(Box::new(cone));
936 new_wiring.push(plan.boundary_in.clone());
937 cone_plans.push(plan);
938 }
939
940 let remap = |src: &WireSource| -> WireSource {
941 match src {
942 WireSource::Input(i) => WireSource::Input(*i),
943 WireSource::NodeOutput(j, p) => {
944 if let Some(&nj) = kept_map.get(j) {
945 WireSource::NodeOutput(nj, *p)
946 } else {
947 let (ci, k) = cone_port[&(*j, *p)];
948 WireSource::NodeOutput(cone_base + ci, k)
949 }
950 }
951 }
952 };
953 for wiring in new_wiring.iter_mut() {
954 for src in wiring.iter_mut() {
955 *src = remap(src);
956 }
957 }
958 let mut new_output_map: HashMap<String, (usize, usize)> = HashMap::new();
959 for (name, (j, p)) in dag.output_map.iter() {
960 let (nj, np) = match remap(&WireSource::NodeOutput(*j, *p)) {
961 WireSource::NodeOutput(a, b) => (a, b),
962 WireSource::Input(_) => unreachable!("outputs map to nodes"),
963 };
964 new_output_map.insert(name.clone(), (nj, np));
965 }
966
967 let m = new_nodes.len();
970 let mut indegree = vec![0usize; m];
971 let mut dependents: Vec<Vec<usize>> = vec![Vec::new(); m];
972 for (i, wiring) in new_wiring.iter().enumerate() {
973 let mut producers: Vec<usize> = wiring
974 .iter()
975 .filter_map(|s| match s {
976 WireSource::NodeOutput(j, _) => Some(*j),
977 WireSource::Input(_) => None,
978 })
979 .collect();
980 producers.sort_unstable();
981 producers.dedup();
982 indegree[i] = producers.len();
983 for j in producers {
984 dependents[j].push(i);
985 }
986 }
987 let mut order: Vec<usize> = Vec::with_capacity(m);
988 let mut ready: std::collections::BinaryHeap<std::cmp::Reverse<usize>> = (0..m)
989 .filter(|&i| indegree[i] == 0)
990 .map(std::cmp::Reverse)
991 .collect();
992 while let Some(std::cmp::Reverse(i)) = ready.pop() {
993 order.push(i);
994 for &d in &dependents[i] {
995 indegree[d] -= 1;
996 if indegree[d] == 0 {
997 ready.push(std::cmp::Reverse(d));
998 }
999 }
1000 }
1001 assert_eq!(
1002 order.len(),
1003 m,
1004 "cone splice must not introduce a cycle (old_n={old_n})"
1005 );
1006 let mut pos = vec![0usize; m];
1007 for (new_idx, &i) in order.iter().enumerate() {
1008 pos[i] = new_idx;
1009 }
1010
1011 let mut sorted_nodes: Vec<Option<Box<dyn PolydatNode>>> =
1012 new_nodes.into_iter().map(Some).collect();
1013 dag.nodes = order
1014 .iter()
1015 .map(|&i| sorted_nodes[i].take().expect("each node placed once"))
1016 .collect();
1017 dag.wiring = order
1018 .iter()
1019 .map(|&i| {
1020 new_wiring[i]
1021 .iter()
1022 .map(|s| match s {
1023 WireSource::Input(k) => WireSource::Input(*k),
1024 WireSource::NodeOutput(j, p) => WireSource::NodeOutput(pos[*j], *p),
1025 })
1026 .collect()
1027 })
1028 .collect();
1029 dag.output_map = new_output_map
1030 .into_iter()
1031 .map(|(name, (j, p))| (name, (pos[j], p)))
1032 .collect();
1033 let _ = cone_plans;
1034 }
1035}