dfir_lang/graph/

flat_to_partitioned.rs

//! Subgraph partioning algorithm

use std::collections::BTreeSet;

use itertools::Itertools;
use slotmap::{SecondaryMap, SparseSecondaryMap};

use super::meta_graph::DfirGraph;
use super::ops::DelayType;
use super::{
    Color, GraphEdgeId, GraphLoopId, GraphNode, GraphNodeId, GraphSubgraphId, HandoffKind,
};
use crate::diagnostic::{Diagnostic, Level};
use crate::graph::graph_algorithms::{SubgraphMerge, validate_topo_sort};

/// Find edge barriers: edges whose destination operator declares an input delay type.
///
/// Returns:
/// - Tick/TickLazy/Loop/LoopLazy edges keyed by edge ID (for topo-sort exclusion and handoff marking).
/// - All barrier (src, dst) node pairs (for the enemies set).
fn find_edge_barriers(
    partitioned_graph: &DfirGraph,
) -> (
    SecondaryMap<GraphEdgeId, DelayType>,
    Vec<(GraphNodeId, GraphNodeId)>,
) {
    let mut tick_edges = SecondaryMap::new();
    let mut barrier_pairs = Vec::new();

    for (edge_id, (src, dst)) in partitioned_graph.edges() {
        let Some(op_inst) = partitioned_graph.node_op_inst(dst) else {
            continue;
        };
        let (_src_port, dst_port) = partitioned_graph.edge_ports(edge_id);
        let Some(delay_type) = (op_inst.op_constraints.input_delaytype_fn)(dst_port) else {
            continue;
        };

        barrier_pairs.push((src, dst));
        tick_edges.insert(edge_id, delay_type);
    }

    (tick_edges, barrier_pairs)
}

/// Find handoff reference access group ordering constraints: for the same handoff target, operators in
/// lower access groups must run before operators in higher access groups.
fn find_access_group_ordering(partitioned_graph: &DfirGraph) -> Vec<(GraphNodeId, GraphNodeId)> {
    let mut pairs = Vec::new();
    let refs_by_target = partitioned_graph.node_handoff_reference_groups();
    for (_handoff, groups) in refs_by_target {
        for (group_a, group_b) in groups.values().tuple_windows() {
            for &(node_a, _, _) in group_a {
                for &(node_b, _, _) in group_b {
                    // TODO(mingwei): handle with diagnostics.
                    assert_ne!(
                        node_a, node_b,
                        "encounted conflicted or cyclical handoff references\n{:?}\n{:?}",
                        group_a, group_b,
                    );
                    pairs.push((node_a, node_b));
                }
            }
        }
    }
    pairs
}

fn find_subgraph_unionfind(
    partitioned_graph: &DfirGraph,
    tick_edges: &SecondaryMap<GraphEdgeId, DelayType>,
    edge_barrier_pairs: &[(GraphNodeId, GraphNodeId)],
    access_group_pairs: &[(GraphNodeId, GraphNodeId)],
) -> Result<(SubgraphMerge<GraphNodeId>, BTreeSet<GraphEdgeId>), Diagnostic> {
    // Modality (color) of nodes, push or pull.
    // TODO(mingwei)? This does NOT consider `DelayType` barriers (which generally imply `Pull`),
    // which makes it inconsistant with the final output in `as_code()`. But this doesn't create
    // any bugs since we exclude `DelayType` edges from joining subgraphs anyway.
    let mut node_color = partitioned_graph
        .node_ids()
        .filter_map(|node_id| {
            let op_color = partitioned_graph.node_color(node_id)?;
            Some((node_id, op_color))
        })
        .collect::<SparseSecondaryMap<_, _>>();

    // Pre-compute all predecessor edges for the topological sort.
    let mut all_preds: SecondaryMap<GraphNodeId, Vec<GraphNodeId>> = SecondaryMap::new();

    // Pipe predecessors (excluding tick edges which are cross-tick).
    for (edge_id, (src, dst)) in partitioned_graph.edges() {
        if !tick_edges.contains_key(edge_id) {
            all_preds.entry(dst).unwrap().or_default().push(src);
        }
    }

    // Handoff references: producer must run before consumer.
    for node_id in partitioned_graph.node_ids() {
        for handoff_ref in partitioned_graph.node_handoff_references(node_id).iter() {
            if let Some(src) = handoff_ref.node_id {
                all_preds.entry(node_id).unwrap().or_default().push(src);
                // Extra ordering: if the ref target is a handoff, its pipe consumers
                // depend on the borrower (borrower runs before consumer).
                if let GraphNode::Handoff { .. } = partitioned_graph.node(src) {
                    for (_edge, consumer) in partitioned_graph.node_successors(src) {
                        all_preds
                            .entry(consumer)
                            .unwrap()
                            .or_default()
                            .push(node_id);
                    }
                }
            }
        }
    }

    // Access group ordering.
    for &(src, dst) in access_group_pairs {
        all_preds.entry(dst).unwrap().or_default().push(src);
    }

    // Loop-ingress ordering constraints.
    // Fix https://github.com/hydro-project/hydro/issues/3048
    //
    // `make_loops_contiguous` (run later) gathers *all* of a loop's subgraphs to the
    // flat-order position of the loop's *first* subgraph. That is only sound if every
    // subgraph feeding into the loop from *outside* is already ordered before the loop's
    // earliest subgraph. Otherwise an ingress "receiver" inside the loop can be hoisted
    // ahead of its external "sender", violating topological order (which then produces a
    // handoff-buffer "use before declaration" compile error downstream).
    //
    // To guarantee this, we require: for each forward (non-tick) edge `src -> dst` where
    // `dst` lies inside a loop that does not contain `src`, `src` must precede *every*
    // node directly inside such loop.
    {
        for (edge_id, (src_id, dst_id)) in partitioned_graph.edges() {
            if tick_edges.contains_key(edge_id) {
                continue;
            }
            let src_loop = partitioned_graph.node_loop(src_id);
            let dst_loop = partitioned_graph.node_loop(dst_id);
            if let Some(dst_loop_some) = dst_loop
                && src_loop == partitioned_graph.loop_parent(dst_loop_some)
            {
                for &inside_dst_id in partitioned_graph.loop_nodes(dst_loop_some) {
                    all_preds
                        .entry(inside_dst_id)
                        .unwrap()
                        .or_default()
                        .push(src_id);
                }
            }
        }
    }

    // Build enemies: all node pairs that must not be in the same subgraph.
    let enemies = edge_barrier_pairs
        .iter()
        .copied()
        .chain(access_group_pairs.iter().copied())
        .chain(partitioned_graph.node_ids().flat_map(|dst| {
            partitioned_graph
                .node_handoff_references(dst)
                .iter()
                .filter_map(|r| r.node_id)
                .map(move |src| (src, dst))
        }));

    let mut subgraph_unionfind = SubgraphMerge::<GraphNodeId>::new(
        partitioned_graph.node_ids(),
        |node_id| all_preds.get(node_id).into_iter().flatten().copied(),
        enemies,
    )
    .map_err(|cycle| {
        let span = cycle
            .first()
            .map(|&node_id| partitioned_graph.node(node_id).span())
            .unwrap_or_else(proc_macro2::Span::call_site);
        let node_cycle = cycle
            .iter()
            .map(|&node_id| partitioned_graph.node(node_id).to_pretty_string())
            .collect::<Vec<_>>();
        Diagnostic::spanned(
            span,
            Level::Error,
            format!(
                "Cyclical dataflow within a tick is not supported. Use `defer_tick()` or `defer_tick_lazy()` to break the cycle across ticks. \
                Cycle: {:?}",
                node_cycle,
            ),
        )
    })?;

    // Will contain all edges which need handoffs added. Starts out with all edges and
    // we remove from this set as we combine nodes into subgraphs.
    let mut handoff_edges: BTreeSet<GraphEdgeId> = partitioned_graph.edge_ids().collect();
    // Would sort edges here for priority (for now, no sort/priority).

    // Each edge gets looked at in order. However we may not know if a linear
    // chain of operators is PUSH vs PULL until we look at the ends. A fancier
    // algorithm would know to handle linear chains from the outside inward.
    // But instead we just run through the edges in a loop until no more
    // progress is made. Could have some sort of O(N^2) pathological worst
    // case.
    let mut progress = true;
    while progress {
        progress = false;
        // TODO(mingwei): Could this iterate `handoff_edges` instead? (Modulo ownership). Then no case (1) below.
        for (edge_id, (src, dst)) in partitioned_graph.edges().collect::<Vec<_>>() {
            // Ignore existing handoffs, remove `edge_id` from `handoff_edges`.
            if matches!(partitioned_graph.node(src), GraphNode::Handoff { .. })
                || matches!(partitioned_graph.node(dst), GraphNode::Handoff { .. })
            {
                handoff_edges.remove(&edge_id);
                continue;
            }

            // Ignore (1) already added edges as well as (2) new self-cycles. (Unless reference edge).
            if subgraph_unionfind.same_set(src, dst) {
                // Note that the _edge_ `edge_id` might not be in the subgraph even when both `src` and `dst` are. This prevents case 2.
                // Handoffs will be inserted later for this self-loop.
                continue;
            }

            // Do not connect across loop contexts.
            if partitioned_graph.node_loop(src) != partitioned_graph.node_loop(dst) {
                continue;
            }

            if can_connect_colorize(&mut node_color, src, dst) {
                // At this point we have selected this edge and its src & dst to be
                // within a single subgraph.
                let ok = subgraph_unionfind.try_merge(src, dst);
                if ok {
                    assert!(handoff_edges.remove(&edge_id));
                    progress = true;
                }
            }
        }
    }

    Ok((subgraph_unionfind, handoff_edges))
}

/// Find subgraphs and insert handoffs.
fn make_subgraphs(
    partitioned_graph: &mut DfirGraph,
    tick_edges: &mut SecondaryMap<GraphEdgeId, DelayType>,
    edge_barrier_pairs: &[(GraphNodeId, GraphNodeId)],
    access_group_pairs: &[(GraphNodeId, GraphNodeId)],
) -> Result<(), Diagnostic> {
    let (subgraph_merge, handoff_edges) = find_subgraph_unionfind(
        partitioned_graph,
        tick_edges,
        edge_barrier_pairs,
        access_group_pairs,
    )?;

    // Insert handoffs between subgraphs (or on subgraph self-loop edges)
    for edge_id in handoff_edges {
        let (src_id, dst_id) = partitioned_graph.edge(edge_id);

        // Already has a handoff, no need to insert one.
        let src_node = partitioned_graph.node(src_id);
        let dst_node = partitioned_graph.node(dst_id);
        if matches!(src_node, GraphNode::Handoff { .. })
            || matches!(dst_node, GraphNode::Handoff { .. })
        {
            continue;
        }

        let hoff = GraphNode::Handoff {
            kind: HandoffKind::Vec,
            src_span: src_node.span(),
            dst_span: dst_node.span(),
        };
        let (_node_id, out_edge_id) = partitioned_graph.insert_intermediate_node(edge_id, hoff);

        // Update tick_edges for inserted node.
        if let Some(delay_type) = tick_edges.remove(edge_id) {
            tick_edges.insert(out_edge_id, delay_type);
        }
    }

    // Register subgraphs. `SubgraphMerge` maintains operators in topo-sorted order per subgraph.
    // Filter out handoff nodes — they are not part of any subgraph.
    // Collect subgraph IDs in flat topological order.
    let mut flat_toposort = Vec::new();
    for nodes in subgraph_merge.subgraphs() {
        if nodes.is_empty() {
            continue;
        }
        // Skip single-node "subgraphs" that are handoff nodes.
        if nodes
            .iter()
            .any(|&n| matches!(partitioned_graph.node(n), GraphNode::Handoff { .. }))
        {
            continue;
        }
        let sg_id = partitioned_graph.insert_subgraph(nodes.to_vec()).unwrap();
        flat_toposort.push(sg_id);
    }

    // Rearrange the flat toposort to make loop subgraphs contiguous.
    let subgraph_toposort = make_loops_contiguous(partitioned_graph, &flat_toposort);

    // Defensive check: the toposort must still be valid after making loops contiguous.
    // A violation here indicates a bug in the ordering/contiguity logic.
    assert!(
        validate_topo_sort(
            subgraph_toposort
                .iter()
                .flat_map(|&sg_id| partitioned_graph.subgraph(sg_id))
                .copied(),
            |succ_id| {
                partitioned_graph
                    .node_predecessors(succ_id)
                    .filter_map(|(edge_id, pred_id)| {
                        (!tick_edges.contains_key(edge_id)).then_some(pred_id)
                    })
                    .map(|pred_id| {
                        // Jump over handoff nodes.
                        if matches!(partitioned_graph.node(pred_id), GraphNode::Handoff { .. }) {
                            partitioned_graph
                                .node_predecessor_nodes(pred_id)
                                .next()
                                .unwrap()
                        } else {
                            pred_id
                        }
                    })
            }
        )
        .is_ok(),
        "bug: toposort is invalid after `make_loops_contiguous`.",
    );

    partitioned_graph.set_subgraph_toposort(subgraph_toposort);

    Ok(())
}

/// Rearranges a flat topological order of subgraphs so that all subgraphs within a loop are contiguous, while
/// preserving the relative topological order.
///
/// PRECONDITION: The first member of a loop MUST COME AFTER all outer subgraphs which send into the loop. Subgraphs
/// from different loop contexts that appear interleaved in the flat order have no dependency between them **EXCEPT**
/// ingress (from outer to inner loop) or egress (from inner to outer). This method hoists all members of a loop to the
/// **first** subgraph of the loop in the flat order, hence the correctness requirement only on the ingress side.
///
/// https://github.com/hydro-project/hydro/issues/3048
fn make_loops_contiguous(
    graph: &DfirGraph,
    flat_order: &[GraphSubgraphId],
) -> Vec<GraphSubgraphId> {
    use std::collections::HashMap;

    // Pre-compute: for each loop, collect *all* descendant subgraphs in flat-order (not just direct).
    // Each sg_id is inserted into every ancestor loop.
    let mut loop_descendants = HashMap::<GraphLoopId, Vec<GraphSubgraphId>>::new();
    for &sg_id in flat_order {
        let mut current = graph.subgraph_loop(sg_id);
        while let Some(loop_id) = current {
            loop_descendants.entry(loop_id).or_default().push(sg_id);
            current = graph.loop_parent(loop_id);
        }
    }

    fn helper(
        graph: &DfirGraph,
        flat_order: &[GraphSubgraphId],
        current_loop: Option<GraphLoopId>,
        loop_descendants: &mut HashMap<GraphLoopId, Vec<GraphSubgraphId>>,
        output: &mut Vec<GraphSubgraphId>,
    ) {
        for &sg_id in flat_order {
            let sg_loop = graph.subgraph_loop(sg_id);
            if current_loop == sg_loop {
                // Directly at this level — emit in place.
                output.push(sg_id);
                continue;
            }
            let sg_loop = sg_loop.expect("root-level subgraph cannot be within a loop: `sg_loop == None` implies `current_loop == None`");
            if current_loop == graph.loop_parent(sg_loop) {
                // In a direct child loop of current_loop, recurse if we haven't yet (tracked by `loop_descendant`
                // entry removal).
                if let Some(inner_order) = loop_descendants.remove(&sg_loop) {
                    helper(graph, &inner_order, Some(sg_loop), loop_descendants, output);
                }
            }
            // else: Deeper nested — skip, will be handled by recursion into its ancestor.
        }
    }

    let mut output = Vec::with_capacity(flat_order.len());
    helper(graph, flat_order, None, &mut loop_descendants, &mut output);
    output
}

/// Set `src` or `dst` color if `None` based on the other (if possible):
/// `None` indicates an op could be pull or push i.e. unary-in & unary-out.
/// So in that case we color `src` or `dst` based on its newfound neighbor (the other one).
///
/// Returns if `src` and `dst` can be in the same subgraph.
fn can_connect_colorize(
    node_color: &mut SparseSecondaryMap<GraphNodeId, Color>,
    src: GraphNodeId,
    dst: GraphNodeId,
) -> bool {
    // Pull -> Pull
    // Push -> Push
    // Pull -> [Computation] -> Push
    // Push -> [Handoff] -> Pull
    let can_connect = match (node_color.get(src), node_color.get(dst)) {
        // Linear chain, can't connect because it may cause future conflicts.
        // But if it doesn't in the _future_ we can connect it (once either/both ends are determined).
        (None, None) => false,

        // Infer left side.
        (None, Some(Color::Pull | Color::Comp)) => {
            node_color.insert(src, Color::Pull);
            true
        }
        (None, Some(Color::Push | Color::Hoff)) => {
            node_color.insert(src, Color::Push);
            true
        }

        // Infer right side.
        (Some(Color::Pull | Color::Hoff), None) => {
            node_color.insert(dst, Color::Pull);
            true
        }
        (Some(Color::Comp | Color::Push), None) => {
            node_color.insert(dst, Color::Push);
            true
        }

        // Both sides already specified.
        (Some(Color::Pull), Some(Color::Pull)) => true,
        (Some(Color::Pull), Some(Color::Comp)) => true,
        (Some(Color::Pull), Some(Color::Push)) => true,

        (Some(Color::Comp), Some(Color::Pull)) => false,
        (Some(Color::Comp), Some(Color::Comp)) => false,
        (Some(Color::Comp), Some(Color::Push)) => true,

        (Some(Color::Push), Some(Color::Pull)) => false,
        (Some(Color::Push), Some(Color::Comp)) => false,
        (Some(Color::Push), Some(Color::Push)) => true,

        // Handoffs are not part of subgraphs.
        (Some(Color::Hoff), Some(_)) => false,
        (Some(_), Some(Color::Hoff)) => false,
    };
    can_connect
}

/// Marks tick-boundary (`defer_tick` / `defer_tick_lazy`) handoffs with their delay type
/// for double-buffered codegen in `as_code`.
///
/// Also remaps `DelayType::Tick` → `DelayType::Loop` (and `TickLazy` → `LoopLazy`) for
/// handoffs whose consumer is inside a nested loop. This allows a single `defer_tick()`
/// operator to work in both tick-level and loop-iteration contexts.
fn mark_tick_boundary_handoffs(
    partitioned_graph: &mut DfirGraph,
    tick_edges: &SecondaryMap<GraphEdgeId, DelayType>,
) {
    let tick_handoffs: Vec<_> = partitioned_graph
        .nodes()
        .filter_map(|(hoff_id, hoff)| {
            if !matches!(hoff, GraphNode::Handoff { .. }) {
                return None;
            }
            if partitioned_graph.node_degree_out(hoff_id) == 0 {
                return None;
            }
            let (succ_edge, _) = partitioned_graph.node_successors(hoff_id).next().unwrap();
            let &delay_type = tick_edges.get(succ_edge)?;

            // Remap Tick/TickLazy → Loop/LoopLazy if the consumer is in a nested loop.
            let consumer_loop = partitioned_graph
                .node_successors(hoff_id)
                .next()
                .and_then(|(_, succ)| partitioned_graph.node_loop(succ));
            let effective_delay_type = if let Some(loop_id) = consumer_loop {
                if partitioned_graph.loop_parent(loop_id).is_some() {
                    // Nested loop: remap to loop-level delay.
                    match delay_type {
                        DelayType::Tick => DelayType::Loop,
                        DelayType::TickLazy => DelayType::LoopLazy,
                        other => other,
                    }
                } else {
                    delay_type
                }
            } else {
                delay_type
            };

            Some((hoff_id, effective_delay_type))
        })
        .collect();

    for (hoff_id, delay_type) in tick_handoffs {
        partitioned_graph.set_handoff_delay_type(hoff_id, delay_type);
    }
}

/// Main method for this module. Partitions a flat [`DfirGraph`] into one with subgraphs.
///
/// Returns an error if an intra-tick cycle exists in the graph.
pub fn partition_graph(flat_graph: DfirGraph) -> Result<DfirGraph, Diagnostic> {
    let (mut tick_edges, edge_barrier_pairs) = find_edge_barriers(&flat_graph);
    let access_group_pairs = find_access_group_ordering(&flat_graph);
    let mut partitioned_graph = flat_graph;

    // Partition into subgraphs and insert handoffs.
    make_subgraphs(
        &mut partitioned_graph,
        &mut tick_edges,
        &edge_barrier_pairs,
        &access_group_pairs,
    )?;

    // Mark tick-boundary handoffs for double-buffering.
    mark_tick_boundary_handoffs(&mut partitioned_graph, &tick_edges);

    Ok(partitioned_graph)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::graph::GraphNode;

    /// Regression test for a loop-contiguity toposort bug.
    /// https://github.com/hydro-project/hydro/issues/3048
    ///
    /// [`make_loops_contiguous`] gathers every subgraph of a loop at the position of the
    /// loop's *first* subgraph in the flat order. If an external subgraph that *feeds into*
    /// the loop (loop ingress) happens to sit — in the flat topological order — after the
    /// loop's first subgraph but before a later loop subgraph, then gathering the loop
    /// pulls that later (ingress-receiving) subgraph *ahead* of its producer.
    ///
    /// Concretely, the flat order `[S1, loopA, S2, loopB]` (valid, since `S2 -> loopB`)
    /// would become `[S1, loopA, loopB, S2]`, so the loop-ingress *receiver* `loopB` would
    /// be emitted before its *sender* `S2`. Downstream this produces a handoff-buffer "use
    /// before declaration" compile error (and, logically, a drained-but-never-filled
    /// buffer).
    ///
    /// The fix adds loop-ingress ordering constraints so every ingress sender precedes the
    /// entire loop block, keeping gather-at-first sound.
    #[test]
    fn test_make_loops_contiguous_ingress_toposort() {
        use crate::graph::{FlatGraphBuilder, FlatGraphBuilderOutput};

        let mut builder = FlatGraphBuilder::new();
        // External source #1 (feeds the loop's first ingress).
        builder.add_dfir(
            syn::parse_quote! {
                inp1 = source_iter([1, 2, 3]);
            },
            None,
            None,
        );
        let loop_id = builder.insert_loop(None);
        // Loop ingress A — independent of `inp2`, so the flat toposort can place it before
        // `inp2`.
        builder.add_dfir(
            syn::parse_quote! {
                inp1 -> batch() -> for_each(std::mem::drop);
            },
            Some(loop_id),
            None,
        );
        // External source #2 — the loop-ingress "sender", defined after loop ingress A so
        // the flat toposort interleaves it: `[inp1, loopA, inp2, loopB]`.
        builder.add_dfir(
            syn::parse_quote! {
                inp2 = source_iter([4, 5, 6]);
            },
            None,
            None,
        );
        // Loop ingress B — the "receiver" of `inp2`.
        builder.add_dfir(
            syn::parse_quote! {
                inp2 -> batch() -> for_each(std::mem::drop);
            },
            Some(loop_id),
            None,
        );

        let FlatGraphBuilderOutput { flat_graph, .. } =
            builder.build().expect("should build without errors");
        let partitioned = partition_graph(flat_graph).expect("should partition without errors");

        // Verify: every forward (non-delay) handoff has its producer subgraph *before* its
        // consumer subgraph in the final `subgraph_toposort`.
        let order = partitioned.subgraph_toposort();
        let pos: std::collections::HashMap<_, _> =
            order.iter().enumerate().map(|(i, &sg)| (sg, i)).collect();

        for (hoff_id, node) in partitioned.nodes() {
            if !matches!(node, GraphNode::Handoff { .. }) {
                continue;
            }
            // Skip delay (back-edge) handoffs, which are allowed to point "backwards".
            if partitioned.handoff_delay_type(hoff_id).is_some() {
                continue;
            }
            let Some((_, pred)) = partitioned.node_predecessors(hoff_id).next() else {
                continue;
            };
            let Some((_, succ)) = partitioned.node_successors(hoff_id).next() else {
                continue;
            };
            let (Some(pred_sg), Some(succ_sg)) = (
                partitioned.node_subgraph(pred),
                partitioned.node_subgraph(succ),
            ) else {
                continue;
            };
            let (Some(&pp), Some(&sp)) = (pos.get(&pred_sg), pos.get(&succ_sg)) else {
                continue;
            };
            assert!(
                pp < sp,
                "toposort violated: producer subgraph {pred_sg:?} (pos {pp}) emitted after \
                 consumer subgraph {succ_sg:?} (pos {sp}) via handoff {hoff_id:?}; order = {order:?}",
            );
        }
    }

    fn make_op_node(graph: &mut DfirGraph, loop_ctx: Option<GraphLoopId>) -> GraphNodeId {
        let operator: crate::parse::Operator = syn::parse_quote! { identity() };
        graph.insert_node(GraphNode::Operator(operator), None, loop_ctx)
    }

    /// Test that subgraphs within a loop are contiguous after rearrangement,
    /// even when the flat input order interleaves them with unrelated subgraphs.
    #[test]
    fn test_make_loops_contiguous_basic() {
        let mut graph = DfirGraph::default();

        // Create a loop.
        let loop_id = graph.insert_loop(None);

        // Create operator nodes: two outside the loop, two inside.
        let op_a = make_op_node(&mut graph, None);
        let op_b = make_op_node(&mut graph, None);
        let op_c = make_op_node(&mut graph, Some(loop_id));
        let op_d = make_op_node(&mut graph, Some(loop_id));

        // Create subgraphs (one per operator node).
        let sg_a = graph.insert_subgraph(vec![op_a]).unwrap();
        let sg_b = graph.insert_subgraph(vec![op_b]).unwrap();
        let sg_c = graph.insert_subgraph(vec![op_c]).unwrap();
        let sg_d = graph.insert_subgraph(vec![op_d]).unwrap();

        // Flat order: A, C, B, D — the loop subgraphs (C, D) are not contiguous.
        let flat_order = vec![sg_a, sg_c, sg_b, sg_d];
        let result = make_loops_contiguous(&graph, &flat_order);

        // After rearrangement: A, C, D, B — loop subgraphs are now contiguous.
        assert_eq!(result, vec![sg_a, sg_c, sg_d, sg_b]);
    }

    /// Test that two independent sibling loops each have their subgraphs contiguous.
    #[test]
    fn test_make_loops_contiguous_independent_loops() {
        let mut graph = DfirGraph::default();

        let loop1 = graph.insert_loop(None);
        let loop2 = graph.insert_loop(None);

        // Loop1: C, D. Loop2: E, F.
        let op_c = make_op_node(&mut graph, Some(loop1));
        let op_d = make_op_node(&mut graph, Some(loop1));
        let op_e = make_op_node(&mut graph, Some(loop2));
        let op_f = make_op_node(&mut graph, Some(loop2));

        let sg_c = graph.insert_subgraph(vec![op_c]).unwrap();
        let sg_d = graph.insert_subgraph(vec![op_d]).unwrap();
        let sg_e = graph.insert_subgraph(vec![op_e]).unwrap();
        let sg_f = graph.insert_subgraph(vec![op_f]).unwrap();

        // Flat order interleaves the two loops: C, E, D, F.
        let flat_order = vec![sg_c, sg_e, sg_d, sg_f];
        let result = make_loops_contiguous(&graph, &flat_order);

        // Both loops should be contiguous. Loop1 appears first (C seen first).
        let pos_c = result.iter().position(|&s| s == sg_c).unwrap();
        let pos_d = result.iter().position(|&s| s == sg_d).unwrap();
        let pos_e = result.iter().position(|&s| s == sg_e).unwrap();
        let pos_f = result.iter().position(|&s| s == sg_f).unwrap();

        // C before D, E before F (relative order preserved).
        assert!(pos_c < pos_d);
        assert!(pos_e < pos_f);

        // Contiguity: C and D are adjacent, E and F are adjacent.
        assert_eq!(pos_d - pos_c, 1, "Loop1 subgraphs must be contiguous");
        assert_eq!(pos_f - pos_e, 1, "Loop2 subgraphs must be contiguous");
    }

    /// Test nested loops: inner loop subgraphs are contiguous within the outer loop block.
    #[test]
    fn test_make_loops_contiguous_nested() {
        let mut graph = DfirGraph::default();

        let outer = graph.insert_loop(None);
        let inner = graph.insert_loop(Some(outer));

        // Outer loop has: op_a, then inner loop (op_b, op_c), then op_d.
        let op_a = make_op_node(&mut graph, Some(outer));
        let op_b = make_op_node(&mut graph, Some(inner));
        let op_c = make_op_node(&mut graph, Some(inner));
        let op_d = make_op_node(&mut graph, Some(outer));

        let sg_a = graph.insert_subgraph(vec![op_a]).unwrap();
        let sg_b = graph.insert_subgraph(vec![op_b]).unwrap();
        let sg_c = graph.insert_subgraph(vec![op_c]).unwrap();
        let sg_d = graph.insert_subgraph(vec![op_d]).unwrap();

        // Flat order: A, B, C, D (already valid).
        let flat_order = vec![sg_a, sg_b, sg_c, sg_d];
        let result = make_loops_contiguous(&graph, &flat_order);

        // Expected: A, B, C, D (all contiguous within outer, B/C contiguous within inner).
        assert_eq!(result, vec![sg_a, sg_b, sg_c, sg_d]);
    }

    /// Test nested loops with interleaving: outer-level subgraph between inner loop items.
    #[test]
    fn test_make_loops_contiguous_nested_interleaved() {
        let mut graph = DfirGraph::default();

        let outer = graph.insert_loop(None);
        let inner = graph.insert_loop(Some(outer));

        let op_a = make_op_node(&mut graph, Some(outer));
        let op_b = make_op_node(&mut graph, Some(inner));
        let op_c = make_op_node(&mut graph, Some(inner));
        let op_d = make_op_node(&mut graph, Some(outer));

        let sg_a = graph.insert_subgraph(vec![op_a]).unwrap();
        let sg_b = graph.insert_subgraph(vec![op_b]).unwrap();
        let sg_c = graph.insert_subgraph(vec![op_c]).unwrap();
        let sg_d = graph.insert_subgraph(vec![op_d]).unwrap();

        // Flat order: A, B, D, C — D (outer) is between B and C (inner).
        let flat_order = vec![sg_a, sg_b, sg_d, sg_c];
        let result = make_loops_contiguous(&graph, &flat_order);

        // After rearrangement within outer: A, B, C, D — inner loop (B, C) contiguous.
        assert_eq!(result, vec![sg_a, sg_b, sg_c, sg_d]);
    }
}