Skip to main content

celox_slt/
scheduler.rs

1use crate::{
2    HashMap, HashSet, LogicPath, LogicPathTarget, NodeId, SLTNode, SLTNodeArena, SLTNodeFactsError,
3};
4use celox_analysis::dag_schedule::{
5    MappedGraphRows, MappedNodeRows, schedule_min_live_values_and_tokens_with_mapped_rows,
6};
7use celox_analysis::interval::{DisjointIntervalError, DisjointIntervalMap, ExactInterval};
8use celox_design::{BinaryOp, BitAccess, RuntimeErrorInfo, UnaryOp, VarAtomBase};
9use celox_sir::{
10    BlockId, ExecutionUnit, RegisterId, SIRBuilder, SIRInstruction, SIROffset, SIRTerminator,
11    SIRValue,
12};
13use std::collections::{BTreeMap, BTreeSet};
14use std::fmt::Debug;
15use std::fmt::Display;
16use std::hash::Hash;
17use thiserror::Error;
18
19/// Sparse scheduler-facing memory effects for one same-trigger FF group.
20#[derive(Debug, Clone, Default)]
21pub struct FfAccessSummary<A> {
22    pub reads: Vec<VarAtomBase<A>>,
23    pub writes: Vec<VarAtomBase<A>>,
24    pub dynamic_writes: HashSet<A>,
25}
26
27fn greedy_fas_sort(scc: &[usize], global_adj: &[Vec<usize>]) -> Vec<usize> {
28    let scc_set: HashSet<usize> = scc.iter().cloned().collect();
29    let mut local_adj: HashMap<usize, Vec<usize>> = HashMap::default();
30    let mut in_degree: HashMap<usize, usize> = HashMap::default();
31
32    for &u in scc {
33        in_degree.entry(u).or_insert(0);
34        let entries = local_adj.entry(u).or_default();
35        for &v in &global_adj[u] {
36            if scc_set.contains(&v) {
37                entries.push(v);
38                *in_degree.entry(v).or_insert(0) += 1;
39            }
40        }
41    }
42
43    let mut left = Vec::new();
44    let mut right = Vec::new();
45    let mut current_nodes: HashSet<usize> = scc.iter().cloned().collect();
46
47    while !current_nodes.is_empty() {
48        // 1. Sinks
49        while let Some(&u) = current_nodes
50            .iter()
51            .find(|&&u| local_adj.get(&u).is_none_or(|v| v.is_empty()))
52        {
53            right.push(u);
54            current_nodes.remove(&u);
55        }
56        // 2. Sources
57        while let Some(&u) = current_nodes
58            .iter()
59            .find(|&&u| in_degree.get(&u).is_none_or(|&d| d == 0))
60        {
61            left.push(u);
62            current_nodes.remove(&u);
63            if let Some(neighbors) = local_adj.remove(&u) {
64                for v in neighbors {
65                    if let Some(d) = in_degree.get_mut(&v) {
66                        *d -= 1;
67                    }
68                }
69            }
70        }
71        if current_nodes.is_empty() {
72            break;
73        }
74        // 3. Maximum Degree Difference
75        let &u = current_nodes
76            .iter()
77            .max_by_key(|&&u| {
78                let out_d = local_adj.get(&u).map_or(0, |v| v.len());
79                let in_d = in_degree.get(&u).cloned().unwrap_or(0);
80                out_d as i32 - in_d as i32
81            })
82            .unwrap();
83
84        left.push(u);
85        current_nodes.remove(&u);
86        if let Some(neighbors) = local_adj.remove(&u) {
87            for v in neighbors {
88                if let Some(d) = in_degree.get_mut(&v) {
89                    *d -= 1;
90                }
91            }
92        }
93    }
94    right.reverse();
95    left.extend(right);
96    left
97}
98fn calculate_required_iterations(adj: &[Vec<usize>], order: &[usize]) -> usize {
99    let pos: HashMap<usize, usize> = order.iter().enumerate().map(|(i, &n)| (n, i)).collect();
100    let scc_nodes: HashSet<usize> = order.iter().cloned().collect();
101
102    // Record already visited nodes to ensure a "simple path"
103    fn find_longest_backedge_path(
104        u: usize,
105        visited: &mut Vec<bool>,
106        adj: &[Vec<usize>],
107        pos: &HashMap<usize, usize>,
108        scc_nodes: &HashSet<usize>,
109    ) -> usize {
110        visited[u] = true;
111        let mut max_delay = 0;
112
113        for &v in &adj[u] {
114            if scc_nodes.contains(&v) && !visited[v] {
115                // 0 if forward direction, 1 if back-edge
116                let weight = if pos[&u] >= pos[&v] { 1 } else { 0 };
117                max_delay = max_delay
118                    .max(weight + find_longest_backedge_path(v, visited, adj, pos, scc_nodes));
119            }
120        }
121
122        visited[u] = false; // backtrack
123        max_delay
124    }
125
126    let mut overall_max_delay = 0;
127    let mut visited = vec![false; adj.len()];
128
129    // Search for the longest "waiting time (number of back-edges)" starting from each node
130    for &start_node in order {
131        overall_max_delay = overall_max_delay.max(find_longest_backedge_path(
132            start_node,
133            &mut visited,
134            adj,
135            &pos,
136            &scc_nodes,
137        ));
138    }
139
140    // Base execution (1) + number of times signals loop back (overall_max_delay)
141    overall_max_delay + 1
142}
143
144fn ranges_cover_access(ranges: &[BitAccess], access: BitAccess) -> bool {
145    let mut ranges = ranges.to_vec();
146    ranges.sort_unstable_by_key(|range| (range.lsb, range.msb));
147    let mut next = access.lsb;
148    for range in ranges {
149        if range.msb < next {
150            continue;
151        }
152        if range.lsb > next {
153            return false;
154        }
155        if range.msb >= access.msb {
156            return true;
157        }
158        let Some(after) = range.msb.checked_add(1) else {
159            return false;
160        };
161        next = after;
162    }
163    false
164}
165
166/// Conservatively check that every input of `address` reachable from `root`
167/// is supplied by the carried ranges. Dynamic inputs retain their complete
168/// declared access here; retaining an extra external dependency is safe.
169fn node_reads_only_covered_ranges<Addr: Clone + Eq + Hash>(
170    root: NodeId,
171    address: &Addr,
172    ranges: &[BitAccess],
173    arena: &SLTNodeArena<Addr>,
174) -> bool {
175    let mut visited = HashSet::default();
176    let mut work = vec![root];
177    while let Some(node) = work.pop() {
178        if !visited.insert(node) {
179            continue;
180        }
181        match arena.get(node) {
182            SLTNode::Input {
183                variable,
184                index,
185                access,
186                ..
187            } => {
188                if variable == address && !ranges_cover_access(ranges, *access) {
189                    return false;
190                }
191                work.extend(index.iter().map(|entry| entry.node));
192            }
193            SLTNode::Constant(..) => {}
194            SLTNode::Binary(lhs, _, rhs) => {
195                work.push(*lhs);
196                work.push(*rhs);
197            }
198            SLTNode::Unary(_, inner) | SLTNode::Capture { expr: inner, .. } => work.push(*inner),
199            SLTNode::Mux {
200                cond,
201                then_expr,
202                else_expr,
203            } => {
204                work.push(*cond);
205                work.push(*then_expr);
206                work.push(*else_expr);
207            }
208            SLTNode::Concat(parts) => work.extend(parts.iter().map(|(part, _)| *part)),
209            SLTNode::Slice { expr, .. } => work.push(*expr),
210            SLTNode::ForFold {
211                start,
212                end,
213                result,
214                initials,
215                updates,
216                effects,
217                continue_cond,
218                ..
219            } => {
220                if let crate::SLTLoopBound::Expr(node) = start {
221                    work.push(*node);
222                }
223                if let crate::SLTLoopBound::Expr(node) = end {
224                    work.push(*node);
225                }
226                if let crate::SLTForFoldResult::Transient { initial, update } = result {
227                    work.push(*initial);
228                    work.push(*update);
229                }
230                work.extend(initials.iter().map(|state| state.expr));
231                work.extend(updates.iter().map(|state| state.expr));
232                for effect in effects {
233                    match effect {
234                        crate::SLTForEffect::Event { guard, args, .. } => {
235                            work.extend(*guard);
236                            work.extend(args.iter().copied());
237                        }
238                        crate::SLTForEffect::Runner(runner) => work.push(*runner),
239                    }
240                }
241                work.push(*continue_cond);
242            }
243            SLTNode::ForFoldGroup {
244                entry_guard,
245                states,
246                ..
247            } => {
248                work.push(*entry_guard);
249                for state in states {
250                    work.push(state.initial);
251                    work.push(state.update);
252                }
253            }
254        }
255    }
256    true
257}
258
259fn collect_node_input_deps<Addr: Clone + Eq + Hash + Debug + Copy + Display>(
260    node: crate::NodeId,
261    arena: &SLTNodeArena<Addr>,
262    memo: &mut HashMap<crate::NodeId, HashSet<Addr>>,
263    inverse_memo: &mut HashMap<Addr, HashSet<crate::NodeId>>,
264) -> HashSet<Addr> {
265    if let Some(found) = memo.get(&node) {
266        return found.clone();
267    }
268
269    let deps = match arena.get(node) {
270        crate::SLTNode::Input {
271            variable, index, ..
272        } => {
273            let mut set = HashSet::default();
274            set.insert(*variable);
275            for idx in index {
276                set.extend(collect_node_input_deps(idx.node, arena, memo, inverse_memo));
277            }
278            set
279        }
280        crate::SLTNode::Slice { expr, .. } => {
281            collect_node_input_deps(*expr, arena, memo, inverse_memo)
282        }
283        crate::SLTNode::Concat(parts) => {
284            let mut set = HashSet::default();
285            for (part, _) in parts {
286                set.extend(collect_node_input_deps(*part, arena, memo, inverse_memo));
287            }
288            set
289        }
290        crate::SLTNode::Binary(lhs, _, rhs) => {
291            let mut set = collect_node_input_deps(*lhs, arena, memo, inverse_memo);
292            set.extend(collect_node_input_deps(*rhs, arena, memo, inverse_memo));
293            set
294        }
295        crate::SLTNode::Unary(_, inner) => {
296            collect_node_input_deps(*inner, arena, memo, inverse_memo)
297        }
298        crate::SLTNode::Capture { expr, .. } => {
299            collect_node_input_deps(*expr, arena, memo, inverse_memo)
300        }
301        crate::SLTNode::Mux {
302            cond,
303            then_expr,
304            else_expr,
305        } => {
306            let mut set = collect_node_input_deps(*cond, arena, memo, inverse_memo);
307            set.extend(collect_node_input_deps(
308                *then_expr,
309                arena,
310                memo,
311                inverse_memo,
312            ));
313            set.extend(collect_node_input_deps(
314                *else_expr,
315                arena,
316                memo,
317                inverse_memo,
318            ));
319            set
320        }
321        crate::SLTNode::ForFold {
322            loop_var,
323            start,
324            end,
325            result,
326            initials,
327            updates,
328            effects,
329            continue_cond,
330            ..
331        } => {
332            let mut set = HashSet::default();
333            match start {
334                crate::SLTLoopBound::Const(_) => {}
335                crate::SLTLoopBound::Expr(node) => {
336                    set.extend(collect_node_input_deps(*node, arena, memo, inverse_memo));
337                }
338            }
339            match end {
340                crate::SLTLoopBound::Const(_) => {}
341                crate::SLTLoopBound::Expr(node) => {
342                    set.extend(collect_node_input_deps(*node, arena, memo, inverse_memo));
343                }
344            }
345            if let crate::SLTForFoldResult::Transient { initial, update } = result {
346                set.extend(collect_node_input_deps(*initial, arena, memo, inverse_memo));
347                set.extend(collect_node_input_deps(*update, arena, memo, inverse_memo));
348            }
349            for init in initials {
350                set.extend(collect_node_input_deps(
351                    init.expr,
352                    arena,
353                    memo,
354                    inverse_memo,
355                ));
356            }
357            for update in updates {
358                set.extend(collect_node_input_deps(
359                    update.expr,
360                    arena,
361                    memo,
362                    inverse_memo,
363                ));
364            }
365            for effect in effects {
366                match effect {
367                    crate::SLTForEffect::Event { guard, args, .. } => {
368                        if let Some(guard) = guard {
369                            set.extend(collect_node_input_deps(*guard, arena, memo, inverse_memo));
370                        }
371                        for arg in args {
372                            set.extend(collect_node_input_deps(*arg, arena, memo, inverse_memo));
373                        }
374                    }
375                    crate::SLTForEffect::Runner(runner) => {
376                        set.extend(collect_node_input_deps(*runner, arena, memo, inverse_memo));
377                    }
378                }
379            }
380            set.remove(loop_var);
381            set.extend(collect_node_input_deps(
382                *continue_cond,
383                arena,
384                memo,
385                inverse_memo,
386            ));
387            set.remove(loop_var);
388            set
389        }
390        crate::SLTNode::ForFoldGroup {
391            loop_var,
392            entry_guard,
393            states,
394            ..
395        } => {
396            let mut set = collect_node_input_deps(*entry_guard, arena, memo, inverse_memo);
397            for state in states {
398                set.extend(collect_node_input_deps(
399                    state.initial,
400                    arena,
401                    memo,
402                    inverse_memo,
403                ));
404            }
405            let mut update_deps = HashSet::default();
406            for state in states {
407                update_deps.extend(collect_node_input_deps(
408                    state.update,
409                    arena,
410                    memo,
411                    inverse_memo,
412                ));
413            }
414            update_deps.remove(loop_var);
415
416            let mut state_ranges: HashMap<Addr, Vec<BitAccess>> = HashMap::default();
417            for state in states {
418                state_ranges
419                    .entry(state.target.id)
420                    .or_default()
421                    .push(state.target.access);
422            }
423            for (state_id, ranges) in state_ranges {
424                if states.iter().all(|state| {
425                    node_reads_only_covered_ranges(state.update, &state_id, &ranges, arena)
426                }) {
427                    update_deps.remove(&state_id);
428                }
429            }
430            set.extend(update_deps);
431            set
432        }
433        crate::SLTNode::Constant(_, _, _, _) => HashSet::default(),
434    };
435
436    for &addr in &deps {
437        inverse_memo.entry(addr).or_default().insert(node);
438    }
439    memo.insert(node, deps.clone());
440    deps
441}
442
443fn collect_logic_path_input_deps<Addr: Clone + Eq + Hash + Debug + Copy + Display>(
444    path: &LogicPath<Addr>,
445    arena: &SLTNodeArena<Addr>,
446    memo: &mut HashMap<NodeId, HashSet<Addr>>,
447    inverse_memo: &mut HashMap<Addr, HashSet<NodeId>>,
448) {
449    collect_node_input_deps(path.expr, arena, memo, inverse_memo);
450    for (_, node) in &path.local_inputs {
451        collect_node_input_deps(*node, arena, memo, inverse_memo);
452    }
453    for node in &path.pre_lower_nodes {
454        collect_node_input_deps(*node, arena, memo, inverse_memo);
455    }
456}
457
458struct TarjanContext {
459    index: usize,
460    stack: Vec<usize>,
461    on_stack: HashSet<usize>,
462    indices: Vec<Option<usize>>,
463    lowlink: Vec<Option<usize>>,
464    sccs: Vec<Vec<usize>>,
465}
466
467fn strong_connect(u: usize, adj: &Vec<Vec<usize>>, ctx: &mut TarjanContext) {
468    ctx.indices[u] = Some(ctx.index);
469    ctx.lowlink[u] = Some(ctx.index);
470    ctx.index += 1;
471    ctx.stack.push(u);
472    ctx.on_stack.insert(u);
473
474    for &v in &adj[u] {
475        if ctx.indices[v].is_none() {
476            strong_connect(v, adj, ctx);
477            ctx.lowlink[u] = Some(ctx.lowlink[u].unwrap().min(ctx.lowlink[v].unwrap()));
478        } else if ctx.on_stack.contains(&v) {
479            ctx.lowlink[u] = Some(ctx.lowlink[u].unwrap().min(ctx.indices[v].unwrap()));
480        }
481    }
482
483    if ctx.lowlink[u] == ctx.indices[u] {
484        let mut scc = Vec::new();
485        while let Some(w) = ctx.stack.pop() {
486            ctx.on_stack.remove(&w);
487            scc.push(w);
488            if w == u {
489                break;
490            }
491        }
492        ctx.sccs.push(scc);
493    }
494}
495
496fn component_map(adj: &Vec<Vec<usize>>) -> Vec<usize> {
497    let mut ctx = TarjanContext {
498        index: 0,
499        stack: Vec::new(),
500        on_stack: HashSet::default(),
501        indices: vec![None; adj.len()],
502        lowlink: vec![None; adj.len()],
503        sccs: Vec::new(),
504    };
505    for node in 0..adj.len() {
506        if ctx.indices[node].is_none() {
507            strong_connect(node, adj, &mut ctx);
508        }
509    }
510    let mut component_by_node = vec![usize::MAX; adj.len()];
511    for (component, nodes) in ctx.sccs.iter().enumerate() {
512        for &node in nodes {
513            component_by_node[node] = component;
514        }
515    }
516    component_by_node
517}
518
519/// Add old-state anti-dependencies which make a later FF write eligible for
520/// direct publication.  These edges are profitable ordering constraints, not
521/// RTL semantic requirements: WORKING staging is always a correct fallback.
522///
523/// Add all candidates once, identify the SCCs they create, then remove every
524/// newly-added edge internal to such an SCC.  Any remaining edge is acyclic in
525/// the resulting graph.  This keeps all old-state readers before a directly
526/// publishable writer without turning feedback pipelines into comb/FF loops.
527fn add_acyclic_ff_write_order_edges(
528    adj: &mut Vec<Vec<usize>>,
529    optional_edges: impl IntoIterator<Item = (usize, usize)>,
530) -> HashSet<(usize, usize)> {
531    let mut requested = HashSet::default();
532    let mut added = vec![Vec::<usize>::new(); adj.len()];
533    for (source, target) in optional_edges {
534        if source == target || source >= adj.len() || target >= adj.len() {
535            continue;
536        }
537        requested.insert((source, target));
538        if !adj[source].contains(&target) {
539            adj[source].push(target);
540            added[source].push(target);
541        }
542    }
543    if added.iter().all(Vec::is_empty) {
544        return requested;
545    }
546
547    let component_by_node = component_map(adj);
548    for (source, targets) in added.iter_mut().enumerate() {
549        targets.retain(|target| component_by_node[source] == component_by_node[*target]);
550        targets.sort_unstable();
551        if !targets.is_empty() {
552            adj[source].retain(|target| targets.binary_search(target).is_err());
553        }
554    }
555    requested.retain(|(source, target)| adj[*source].contains(target));
556    requested
557}
558
559/// Memory definitions and uses induced by one set of LogicPaths.
560///
561/// A variable target is one bit-range MemoryDef. A current-value source is a
562/// MemoryUse of every overlapping definition; a previous-value source is a
563/// live-on-entry use plus an anti-dependence which keeps that use before an
564/// overlapping definition. `dependencies` contains all semantic edges while
565/// `values` contains only Def-to-Use edges which can become a forwarded value
566/// during later lowering. Both relations retain their forward and reverse
567/// adjacency so regional scheduling can borrow either direction.
568struct LogicPathMemorySsa {
569    dependencies: LogicPathEdges,
570    values: LogicPathEdges,
571}
572
573struct LogicPathEdges {
574    users: Vec<Vec<usize>>,
575    predecessors: Vec<Vec<usize>>,
576}
577
578impl LogicPathEdges {
579    fn new(node_count: usize) -> Self {
580        Self {
581            users: vec![Vec::new(); node_count],
582            predecessors: vec![Vec::new(); node_count],
583        }
584    }
585
586    fn resize(&mut self, node_count: usize) {
587        self.users.resize_with(node_count, Vec::new);
588        self.predecessors.resize_with(node_count, Vec::new);
589    }
590
591    fn push(&mut self, predecessor: usize, user: usize) {
592        self.users[predecessor].push(user);
593        self.predecessors[user].push(predecessor);
594    }
595
596    fn canonicalize(&mut self) {
597        for row in self.users.iter_mut().chain(&mut self.predecessors) {
598            row.sort_unstable();
599            row.dedup();
600        }
601    }
602}
603
604#[derive(Debug, Clone, PartialEq, Eq)]
605pub(crate) struct FfCombSchedulePlan {
606    pub required_comb: Vec<bool>,
607    pub comb_value_predecessors: Vec<Vec<usize>>,
608    pub comb_before_direct_write: Vec<Vec<usize>>,
609    pub ff_before_direct_write: Vec<Vec<usize>>,
610}
611
612fn bit_interval(access: BitAccess) -> Option<(usize, usize)> {
613    Some((
614        access.lsb,
615        access.msb.checked_sub(access.lsb)?.checked_add(1)?,
616    ))
617}
618
619fn build_logic_path_memory_ssa<Addr>(
620    input: &[LogicPath<Addr>],
621) -> Result<LogicPathMemorySsa, SchedulerError<Addr>>
622where
623    Addr: Copy + Ord + Hash + Eq + Display + Debug,
624{
625    let mut definition_intervals = Vec::new();
626    for (path, logic_path) in input.iter().enumerate() {
627        let Some(target) = logic_path.target.var() else {
628            continue;
629        };
630        let Some((start, length)) = bit_interval(target.access) else {
631            return Err(SchedulerError::InvalidDependencyGraph);
632        };
633        definition_intervals.push(ExactInterval {
634            object: target.id,
635            start,
636            length,
637            value: path,
638        });
639    }
640    let definitions = match DisjointIntervalMap::try_new(definition_intervals) {
641        Ok(definitions) => definitions,
642        Err(DisjointIntervalError::Overlap { first, second }) => {
643            return Err(SchedulerError::MultipleDriver {
644                blocks: vec![input[first].clone(), input[second].clone()],
645            });
646        }
647        Err(DisjointIntervalError::Empty { .. } | DisjointIntervalError::Overflow { .. }) => {
648            return Err(SchedulerError::InvalidDependencyGraph);
649        }
650    };
651
652    let mut dependencies = LogicPathEdges::new(input.len());
653    let mut values = LogicPathEdges::new(input.len());
654    for (user, path) in input.iter().enumerate() {
655        for source in &path.sources {
656            let Some((start, length)) = bit_interval(source.access) else {
657                return Err(SchedulerError::InvalidDependencyGraph);
658            };
659            let reaching = definitions
660                .overlapping(&source.id, start, length)
661                .map_err(|_| SchedulerError::InvalidDependencyGraph)?;
662            for definition in reaching {
663                dependencies.push(definition, user);
664                values.push(definition, user);
665            }
666        }
667
668        // A previous-value source reads the live-on-entry version. Preserve
669        // that read before every overlapping definition. A self read/write is
670        // already one atomic LogicPath and needs no self anti-dependence.
671        for source in &path.previous_sources {
672            let Some((start, length)) = bit_interval(source.access) else {
673                return Err(SchedulerError::InvalidDependencyGraph);
674            };
675            let reaching = definitions
676                .overlapping(&source.id, start, length)
677                .map_err(|_| SchedulerError::InvalidDependencyGraph)?;
678            for definition in reaching.filter(|definition| *definition != user) {
679                dependencies.push(user, definition);
680            }
681        }
682        for target in &path.order_before {
683            if target.0 < input.len() && target.0 != user {
684                dependencies.push(user, target.0);
685            }
686        }
687    }
688    dependencies.canonicalize();
689    values.canonicalize();
690    Ok(LogicPathMemorySsa {
691        dependencies,
692        values,
693    })
694}
695
696pub(crate) fn plan_ff_comb_schedule<Addr>(
697    input: &[LogicPath<Addr>],
698    ff: &[FfAccessSummary<Addr>],
699) -> Result<FfCombSchedulePlan, SchedulerError<Addr>>
700where
701    Addr: Copy + Ord + Hash + Eq + Display + Debug,
702{
703    let memory = build_logic_path_memory_ssa(input)?;
704    let mut definition_intervals = Vec::new();
705    for (path, logic_path) in input.iter().enumerate() {
706        let Some(target) = logic_path.target.var() else {
707            continue;
708        };
709        let Some((start, length)) = bit_interval(target.access) else {
710            return Err(SchedulerError::InvalidDependencyGraph);
711        };
712        definition_intervals.push(ExactInterval {
713            object: target.id,
714            start,
715            length,
716            value: path,
717        });
718    }
719    let definitions = DisjointIntervalMap::try_new(definition_intervals)
720        .map_err(|_| SchedulerError::InvalidDependencyGraph)?;
721
722    let mut comb_value_predecessors = vec![Vec::new(); ff.len()];
723    let mut required_comb = vec![false; input.len()];
724    let mut work = Vec::new();
725    for (ff_index, summary) in ff.iter().enumerate() {
726        for read in &summary.reads {
727            let Some((start, length)) = bit_interval(read.access) else {
728                return Err(SchedulerError::InvalidDependencyGraph);
729            };
730            for definition in definitions
731                .overlapping(&read.id, start, length)
732                .map_err(|_| SchedulerError::InvalidDependencyGraph)?
733            {
734                comb_value_predecessors[ff_index].push(definition);
735                if !std::mem::replace(&mut required_comb[definition], true) {
736                    work.push(definition);
737                }
738            }
739        }
740    }
741    for (path, logic_path) in input.iter().enumerate() {
742        if logic_path_is_scheduling_barrier(logic_path)
743            && !std::mem::replace(&mut required_comb[path], true)
744        {
745            work.push(path);
746        }
747    }
748    for row in &mut comb_value_predecessors {
749        row.sort_unstable();
750        row.dedup();
751    }
752
753    while let Some(path) = work.pop() {
754        for &predecessor in &memory.dependencies.predecessors[path] {
755            if !std::mem::replace(&mut required_comb[predecessor], true) {
756                work.push(predecessor);
757            }
758        }
759    }
760
761    let mut comb_before_direct_write = vec![Vec::new(); ff.len()];
762    for (path_index, path) in input.iter().enumerate() {
763        if !required_comb[path_index] {
764            continue;
765        }
766        for (ff_index, summary) in ff.iter().enumerate() {
767            if path
768                .sources
769                .iter()
770                .chain(&path.previous_sources)
771                .any(|read| {
772                    summary
773                        .writes
774                        .iter()
775                        .any(|write| read.id == write.id && read.access.overlaps(&write.access))
776                })
777            {
778                comb_before_direct_write[ff_index].push(path_index);
779            }
780        }
781    }
782
783    let mut ff_before_direct_write = vec![Vec::new(); ff.len()];
784    for (writer, write_summary) in ff.iter().enumerate() {
785        for (reader, read_summary) in ff.iter().enumerate() {
786            if writer != reader
787                && read_summary.reads.iter().any(|read| {
788                    write_summary
789                        .writes
790                        .iter()
791                        .any(|write| read.id == write.id && read.access.overlaps(&write.access))
792                })
793            {
794                ff_before_direct_write[writer].push(reader);
795            }
796        }
797    }
798    for row in comb_before_direct_write
799        .iter_mut()
800        .chain(&mut ff_before_direct_write)
801    {
802        row.sort_unstable();
803        row.dedup();
804    }
805
806    Ok(FfCombSchedulePlan {
807        required_comb,
808        comb_value_predecessors,
809        comb_before_direct_write,
810        ff_before_direct_write,
811    })
812}
813
814fn lower_logic_path_expr<Addr: Clone + Eq + Ord + Hash + Debug + Copy + Display>(
815    lowerer: &crate::SLTToSIRLowerer,
816    builder: &mut SIRBuilder<Addr>,
817    path: &LogicPath<Addr>,
818    arena: &SLTNodeArena<Addr>,
819    lower_cache: &mut HashMap<NodeId, RegisterId>,
820) -> RegisterId {
821    if path.local_inputs.is_empty() {
822        return lowerer.lower(builder, path.expr, arena, lower_cache);
823    }
824
825    let mut env_inputs = HashMap::default();
826    for (addr, node) in &path.local_inputs {
827        let reg = lowerer.lower(builder, *node, arena, lower_cache);
828        let width = crate::get_width(*node, arena);
829        if width > 0 {
830            env_inputs.insert(VarAtomBase::new(*addr, 0, width - 1), reg);
831        }
832    }
833    lowerer.lower_with_inputs(builder, path.expr, arena, lower_cache, env_inputs)
834}
835
836fn lower_logic_path_node<Addr: Clone + Eq + Ord + Hash + Debug + Copy + Display>(
837    lowerer: &crate::SLTToSIRLowerer,
838    builder: &mut SIRBuilder<Addr>,
839    path: &LogicPath<Addr>,
840    node: NodeId,
841    arena: &SLTNodeArena<Addr>,
842    lower_cache: &mut HashMap<NodeId, RegisterId>,
843) -> RegisterId {
844    if path.local_inputs.is_empty() {
845        return lowerer.lower(builder, node, arena, lower_cache);
846    }
847    // Unbound captures may already have been materialized before a later
848    // write. Do not rebuild them under the observer's unrelated local inputs.
849    if matches!(arena.get(node), crate::SLTNode::Capture { .. })
850        && let Some(reg) = lower_cache.get(&node)
851    {
852        return *reg;
853    }
854    let mut env_inputs = HashMap::default();
855    for (addr, local_node) in &path.local_inputs {
856        let reg = lowerer.lower(builder, *local_node, arena, lower_cache);
857        let width = crate::get_width(*local_node, arena);
858        if width > 0 {
859            env_inputs.insert(VarAtomBase::new(*addr, 0, width - 1), reg);
860        }
861    }
862    lowerer.lower_with_inputs(builder, node, arena, lower_cache, env_inputs)
863}
864
865fn pre_lower_logic_path_node<Addr: Clone + Eq + Ord + Hash + Debug + Copy + Display>(
866    lowerer: &crate::SLTToSIRLowerer,
867    builder: &mut SIRBuilder<Addr>,
868    path: &LogicPath<Addr>,
869    node: NodeId,
870    arena: &SLTNodeArena<Addr>,
871    lower_cache: &mut HashMap<NodeId, RegisterId>,
872) {
873    if path.local_inputs.is_empty() {
874        lowerer.lower(builder, node, arena, lower_cache);
875    }
876}
877
878fn static_access_offset<Addr: Eq + Hash>(
879    addr: &Addr,
880    access: BitAccess,
881    unpacked_element_widths: &HashMap<Addr, usize>,
882) -> SIROffset {
883    let width = access.msb - access.lsb + 1;
884    match unpacked_element_widths.get(addr).copied() {
885        Some(element_width)
886            if element_width != 0
887                && width > element_width
888                && access.lsb.is_multiple_of(element_width)
889                && width.is_multiple_of(element_width) =>
890        {
891            SIROffset::PackedElements {
892                bit_offset: access.lsb,
893                element_width,
894            }
895        }
896        _ => SIROffset::Static(access.lsb),
897    }
898}
899
900fn emit_logic_path_store_with_result<Addr: Clone + Eq + Ord + Hash + Debug + Copy + Display>(
901    lowerer: &crate::SLTToSIRLowerer,
902    builder: &mut SIRBuilder<Addr>,
903    path: &LogicPath<Addr>,
904    arena: &SLTNodeArena<Addr>,
905    lower_cache: &mut HashMap<NodeId, RegisterId>,
906    unpacked_element_widths: &HashMap<Addr, usize>,
907    prepared_result: Option<RegisterId>,
908) {
909    match &path.target {
910        LogicPathTarget::Var(target) => {
911            for node in &path.pre_lower_nodes {
912                pre_lower_logic_path_node(lowerer, builder, path, *node, arena, lower_cache);
913            }
914            let result_reg = match prepared_result {
915                Some(result) => result,
916                None => lower_logic_path_expr(lowerer, builder, path, arena, lower_cache),
917            };
918            let width = 1 + target.access.msb - target.access.lsb;
919            let offset = static_access_offset(&target.id, target.access, unpacked_element_widths);
920            let old_reg =
921                if path.comb_capture_enable_sites.is_empty() || path.comb_capture_enable_always {
922                    None
923                } else {
924                    let old_reg = builder.alloc_bit(width, false);
925                    builder.emit(SIRInstruction::Load(
926                        old_reg,
927                        target.id,
928                        offset.clone(),
929                        width,
930                    ));
931                    Some(old_reg)
932                };
933            builder.emit(SIRInstruction::Store(
934                target.id,
935                offset,
936                width,
937                result_reg,
938                Vec::new(),
939                Vec::new(),
940            ));
941            if !path.comb_capture_enable_sites.is_empty() {
942                let (old, new) = if path.comb_capture_enable_always {
943                    let old = builder.alloc_bit(1, false);
944                    let new = builder.alloc_bit(1, false);
945                    builder.emit(SIRInstruction::Imm(old, SIRValue::new(0u8)));
946                    builder.emit(SIRInstruction::Imm(new, SIRValue::new(1u8)));
947                    (old, new)
948                } else {
949                    (
950                        old_reg.expect("changed capture enable loads the old value"),
951                        result_reg,
952                    )
953                };
954                builder.emit(SIRInstruction::CombCaptureEnableIfChanged {
955                    old,
956                    new,
957                    sites: path.comb_capture_enable_sites.clone(),
958                });
959            }
960        }
961        LogicPathTarget::CombCaptureEvent {
962            site_id,
963            guard,
964            emit_on_true,
965            args,
966            loop_runner,
967            fatal_error_code,
968            consume_enabled,
969        } => {
970            debug_assert!(prepared_result.is_none());
971            if let Some(loop_runner) = loop_runner {
972                lower_logic_path_node(lowerer, builder, path, *loop_runner, arena, lower_cache);
973                return;
974            }
975            let emit = |builder: &mut SIRBuilder<Addr>,
976                        lower_cache: &mut HashMap<NodeId, RegisterId>| {
977                let regs = args
978                    .iter()
979                    .map(|arg| {
980                        lower_logic_path_node(lowerer, builder, path, *arg, arena, lower_cache)
981                    })
982                    .collect();
983                builder.emit(SIRInstruction::CombCaptureEvent {
984                    site_id: *site_id,
985                    args: regs,
986                    fatal_error_code: *fatal_error_code,
987                    consume_enabled: *consume_enabled,
988                });
989            };
990            if let Some(guard) = guard {
991                let cond =
992                    lower_logic_path_node(lowerer, builder, path, *guard, arena, lower_cache);
993                let branch_cond = if *emit_on_true {
994                    cond
995                } else {
996                    let inverted = builder.alloc_bit(1, false);
997                    builder.emit(SIRInstruction::Unary(inverted, UnaryOp::LogicNot, cond));
998                    inverted
999                };
1000                let event_block = builder.new_block();
1001                let done_block = builder.new_block();
1002                builder.seal_block(SIRTerminator::Branch {
1003                    cond: branch_cond,
1004                    true_block: (event_block, vec![]),
1005                    false_block: (done_block, vec![]),
1006                });
1007                builder.switch_to_block(event_block);
1008                emit(builder, lower_cache);
1009                builder.seal_block(SIRTerminator::Jump(done_block, vec![]));
1010                builder.switch_to_block(done_block);
1011            } else {
1012                emit(builder, lower_cache);
1013            }
1014        }
1015    }
1016}
1017
1018fn emit_logic_path_store<Addr: Clone + Eq + Ord + Hash + Debug + Copy + Display>(
1019    lowerer: &crate::SLTToSIRLowerer,
1020    builder: &mut SIRBuilder<Addr>,
1021    path: &LogicPath<Addr>,
1022    arena: &SLTNodeArena<Addr>,
1023    lower_cache: &mut HashMap<NodeId, RegisterId>,
1024    unpacked_element_widths: &HashMap<Addr, usize>,
1025) {
1026    emit_logic_path_store_with_result(
1027        lowerer,
1028        builder,
1029        path,
1030        arena,
1031        lower_cache,
1032        unpacked_element_widths,
1033        None,
1034    );
1035}
1036
1037fn invalidate_logic_path_target<Addr: Clone + Eq + Ord + Hash + Debug + Copy + Display>(
1038    path: &LogicPath<Addr>,
1039    inverse_dep_memo: &HashMap<Addr, HashSet<NodeId>>,
1040    lower_cache: &mut HashMap<NodeId, RegisterId>,
1041) {
1042    let Some(target) = path.target.var() else {
1043        return;
1044    };
1045    if let Some(to_remove) = inverse_dep_memo.get(&target.id) {
1046        for node in to_remove {
1047            if !path.pre_lower_nodes.contains(node) {
1048                lower_cache.remove(node);
1049            }
1050        }
1051    }
1052}
1053
1054#[allow(clippy::too_many_arguments)]
1055fn emit_scheduled_guard_region<Addr: Clone + Eq + Ord + Hash + Debug + Copy + Display>(
1056    lowerer: &crate::SLTToSIRLowerer,
1057    builder: &mut SIRBuilder<Addr>,
1058    condition: NodeId,
1059    paths: &[usize],
1060    input: &[LogicPath<Addr>],
1061    arena: &SLTNodeArena<Addr>,
1062    lower_cache: &mut HashMap<NodeId, RegisterId>,
1063    dep_memo: &mut HashMap<NodeId, HashSet<Addr>>,
1064    inverse_dep_memo: &mut HashMap<Addr, HashSet<NodeId>>,
1065    unpacked_element_widths: &HashMap<Addr, usize>,
1066) {
1067    for &path in paths {
1068        collect_logic_path_input_deps(&input[path], arena, dep_memo, inverse_dep_memo);
1069    }
1070
1071    let condition_value = lowerer.lower(builder, condition, arena, lower_cache);
1072    let then_block = builder.new_block();
1073    let else_block = builder.new_block();
1074    let merge_block = builder.new_block();
1075    builder.seal_block(SIRTerminator::Branch {
1076        cond: condition_value,
1077        true_block: (then_block, Vec::new()),
1078        false_block: (else_block, Vec::new()),
1079    });
1080
1081    let mut affected = paths
1082        .iter()
1083        .filter_map(|path| input[*path].target.var())
1084        .filter_map(|target| inverse_dep_memo.get(&target.id))
1085        .flatten()
1086        .copied()
1087        .collect::<Vec<_>>();
1088    affected.sort_unstable();
1089    affected.dedup();
1090    let saved_cache = affected
1091        .iter()
1092        .filter_map(|node| lower_cache.get(node).copied().map(|value| (*node, value)))
1093        .collect::<Vec<_>>();
1094
1095    let emit_arm = |builder: &mut SIRBuilder<Addr>,
1096                    lower_cache: &mut HashMap<NodeId, RegisterId>,
1097                    take_true: bool| {
1098        let mut inserted = Vec::new();
1099        for &path in paths {
1100            let (_, then_expr, else_expr) = scheduled_root_mux(&input[path], arena)
1101                .expect("scheduled guard region must retain root Muxes");
1102            let root = if take_true { then_expr } else { else_expr };
1103            // Keep the existing per-LogicPath lowering model for nested
1104            // decisions. The scheduled region owns only the common outer
1105            // guard; changing fanout/cost facts for all nested Muxes at once
1106            // would silently replace the established branch decisions.
1107            let value = lowerer.lower(builder, root, arena, lower_cache);
1108            inserted.extend(lowerer.take_scheduled_region_insertions());
1109            emit_logic_path_store_with_result(
1110                lowerer,
1111                builder,
1112                &input[path],
1113                arena,
1114                lower_cache,
1115                unpacked_element_widths,
1116                Some(value),
1117            );
1118            invalidate_logic_path_target(&input[path], inverse_dep_memo, lower_cache);
1119        }
1120        builder.seal_block(SIRTerminator::Jump(merge_block, Vec::new()));
1121        for node in inserted {
1122            lower_cache.remove(&node);
1123        }
1124        for node in &affected {
1125            lower_cache.remove(node);
1126        }
1127        lower_cache.extend(saved_cache.iter().copied());
1128    };
1129
1130    builder.switch_to_block(then_block);
1131    emit_arm(builder, lower_cache, true);
1132    builder.switch_to_block(else_block);
1133    emit_arm(builder, lower_cache, false);
1134    builder.switch_to_block(merge_block);
1135
1136    // Both arms publish every path target, so cached expressions reading any
1137    // such target are invalid at the merge even when both arm caches happened
1138    // to contain the same SLT node identity.
1139    for node in affected {
1140        lower_cache.remove(&node);
1141    }
1142}
1143
1144fn projected_for_fold_group<Addr: Clone + Eq + Hash>(
1145    mut node: NodeId,
1146    arena: &SLTNodeArena<Addr>,
1147) -> Option<NodeId> {
1148    loop {
1149        match arena.get(node) {
1150            SLTNode::ForFoldGroup { .. } => return Some(node),
1151            SLTNode::Slice { expr, .. } => node = *expr,
1152            _ => return None,
1153        }
1154    }
1155}
1156
1157fn fold_group_projection_access<Addr: Clone + Eq + Hash>(
1158    node: NodeId,
1159    group: NodeId,
1160    arena: &SLTNodeArena<Addr>,
1161) -> Option<BitAccess> {
1162    if node == group {
1163        let width = crate::get_width(group, arena);
1164        return (width != 0).then(|| BitAccess::new(0, width - 1));
1165    }
1166    let SLTNode::Slice { expr, access } = arena.get(node) else {
1167        return None;
1168    };
1169    let parent = fold_group_projection_access(*expr, group, arena)?;
1170    let parent_width = parent.msb.checked_sub(parent.lsb)?.checked_add(1)?;
1171    if access.msb >= parent_width {
1172        return None;
1173    }
1174    Some(BitAccess::new(
1175        parent.lsb.checked_add(access.lsb)?,
1176        parent.lsb.checked_add(access.msb)?,
1177    ))
1178}
1179
1180fn packed_fold_group_state_accesses<Addr: Clone + Eq + Hash>(
1181    states: &[crate::SLTForFoldGroupState<Addr>],
1182) -> Option<Vec<BitAccess>> {
1183    let total_width = states.iter().try_fold(0usize, |total, state| {
1184        let width = state
1185            .target
1186            .access
1187            .msb
1188            .checked_sub(state.target.access.lsb)?
1189            .checked_add(1)?;
1190        total.checked_add(width)
1191    })?;
1192    if total_width == 0 {
1193        return None;
1194    }
1195
1196    let mut next_msb = total_width;
1197    let mut result = Vec::with_capacity(states.len());
1198    for state in states {
1199        let width = state.target.access.msb - state.target.access.lsb + 1;
1200        next_msb = next_msb.checked_sub(width)?;
1201        result.push(BitAccess::new(next_msb, next_msb + width - 1));
1202    }
1203    Some(result)
1204}
1205
1206fn push_scheduler_node_children<Addr: Clone + Eq + Hash>(
1207    node: NodeId,
1208    arena: &SLTNodeArena<Addr>,
1209    work: &mut Vec<NodeId>,
1210) {
1211    match arena.get(node) {
1212        SLTNode::Input { index, .. } => work.extend(index.iter().map(|entry| entry.node)),
1213        SLTNode::Constant(..) => {}
1214        SLTNode::Binary(lhs, _, rhs) => {
1215            work.push(*lhs);
1216            work.push(*rhs);
1217        }
1218        SLTNode::Unary(_, inner) | SLTNode::Capture { expr: inner, .. } => work.push(*inner),
1219        SLTNode::Mux {
1220            cond,
1221            then_expr,
1222            else_expr,
1223        } => {
1224            work.push(*cond);
1225            work.push(*then_expr);
1226            work.push(*else_expr);
1227        }
1228        SLTNode::Concat(parts) => work.extend(parts.iter().map(|(part, _)| *part)),
1229        SLTNode::Slice { expr, .. } => work.push(*expr),
1230        SLTNode::ForFold {
1231            start,
1232            end,
1233            result,
1234            initials,
1235            updates,
1236            effects,
1237            continue_cond,
1238            ..
1239        } => {
1240            if let crate::SLTLoopBound::Expr(node) = start {
1241                work.push(*node);
1242            }
1243            if let crate::SLTLoopBound::Expr(node) = end {
1244                work.push(*node);
1245            }
1246            if let crate::SLTForFoldResult::Transient { initial, update } = result {
1247                work.push(*initial);
1248                work.push(*update);
1249            }
1250            work.extend(initials.iter().map(|state| state.expr));
1251            work.extend(updates.iter().map(|state| state.expr));
1252            for effect in effects {
1253                match effect {
1254                    crate::SLTForEffect::Event { guard, args, .. } => {
1255                        work.extend(*guard);
1256                        work.extend(args.iter().copied());
1257                    }
1258                    crate::SLTForEffect::Runner(runner) => work.push(*runner),
1259                }
1260            }
1261            work.push(*continue_cond);
1262        }
1263        SLTNode::ForFoldGroup {
1264            entry_guard,
1265            states,
1266            ..
1267        } => {
1268            work.push(*entry_guard);
1269            for state in states {
1270                work.push(state.initial);
1271                work.push(state.update);
1272            }
1273        }
1274    }
1275}
1276
1277#[derive(Clone, PartialEq, Eq, Hash)]
1278enum NormalizedIndexExpr<Addr: Clone + Eq + Hash> {
1279    LoopValue {
1280        signed: bool,
1281        access: BitAccess,
1282    },
1283    Input {
1284        variable: Addr,
1285        signed: bool,
1286        access: BitAccess,
1287        index: Vec<(NormalizedIndexExpr<Addr>, usize)>,
1288    },
1289    Constant(num_bigint::BigUint, num_bigint::BigUint, usize, bool),
1290    Binary(
1291        Box<NormalizedIndexExpr<Addr>>,
1292        BinaryOp,
1293        Box<NormalizedIndexExpr<Addr>>,
1294    ),
1295    Unary(UnaryOp, Box<NormalizedIndexExpr<Addr>>),
1296    Mux {
1297        cond: Box<NormalizedIndexExpr<Addr>>,
1298        then_expr: Box<NormalizedIndexExpr<Addr>>,
1299        else_expr: Box<NormalizedIndexExpr<Addr>>,
1300    },
1301    Concat(Vec<(NormalizedIndexExpr<Addr>, usize)>),
1302    Slice(Box<NormalizedIndexExpr<Addr>>, BitAccess),
1303}
1304
1305impl<Addr: Clone + Eq + Hash> NormalizedIndexExpr<Addr> {
1306    fn contains_loop_value(&self) -> bool {
1307        match self {
1308            Self::LoopValue { .. } => true,
1309            Self::Input { index, .. } | Self::Concat(index) => {
1310                index.iter().any(|(expr, _)| expr.contains_loop_value())
1311            }
1312            Self::Constant(..) => false,
1313            Self::Binary(lhs, _, rhs) => lhs.contains_loop_value() || rhs.contains_loop_value(),
1314            Self::Unary(_, inner) | Self::Slice(inner, _) => inner.contains_loop_value(),
1315            Self::Mux {
1316                cond,
1317                then_expr,
1318                else_expr,
1319            } => {
1320                cond.contains_loop_value()
1321                    || then_expr.contains_loop_value()
1322                    || else_expr.contains_loop_value()
1323            }
1324        }
1325    }
1326
1327    fn operation_cost(&self) -> u128 {
1328        match self {
1329            Self::LoopValue { .. } | Self::Constant(..) => 0,
1330            Self::Input { index, .. } => 4u128.saturating_add(
1331                index
1332                    .iter()
1333                    .map(|(expr, _)| 2u128.saturating_add(expr.operation_cost()))
1334                    .sum(),
1335            ),
1336            Self::Binary(lhs, _, rhs) => 1u128
1337                .saturating_add(lhs.operation_cost())
1338                .saturating_add(rhs.operation_cost()),
1339            Self::Unary(_, inner) | Self::Slice(inner, _) => {
1340                1u128.saturating_add(inner.operation_cost())
1341            }
1342            Self::Mux {
1343                cond,
1344                then_expr,
1345                else_expr,
1346            } => 1u128
1347                .saturating_add(cond.operation_cost())
1348                .saturating_add(then_expr.operation_cost())
1349                .saturating_add(else_expr.operation_cost()),
1350            Self::Concat(parts) => parts.iter().fold(1u128, |cost, (part, _)| {
1351                cost.saturating_add(part.operation_cost())
1352            }),
1353        }
1354    }
1355}
1356
1357fn normalize_index_expr<Addr: Clone + Eq + Hash + Copy>(
1358    node: NodeId,
1359    loop_var: Addr,
1360    arena: &SLTNodeArena<Addr>,
1361    memo: &mut HashMap<NodeId, Option<NormalizedIndexExpr<Addr>>>,
1362) -> Option<NormalizedIndexExpr<Addr>> {
1363    if let Some(found) = memo.get(&node) {
1364        return found.clone();
1365    }
1366    let normalized = match arena.get(node) {
1367        SLTNode::Input {
1368            variable,
1369            signed,
1370            index,
1371            access,
1372        } if *variable == loop_var => {
1373            if !index.is_empty() {
1374                None
1375            } else {
1376                Some(NormalizedIndexExpr::LoopValue {
1377                    signed: *signed,
1378                    access: *access,
1379                })
1380            }
1381        }
1382        SLTNode::Input {
1383            variable,
1384            signed,
1385            index,
1386            access,
1387        } => Some(NormalizedIndexExpr::Input {
1388            variable: *variable,
1389            signed: *signed,
1390            access: *access,
1391            index: index
1392                .iter()
1393                .map(|entry| {
1394                    normalize_index_expr(entry.node, loop_var, arena, memo)
1395                        .map(|node| (node, entry.stride))
1396                })
1397                .collect::<Option<Vec<_>>>()?,
1398        }),
1399        SLTNode::Constant(payload, mask, width, signed) => Some(NormalizedIndexExpr::Constant(
1400            payload.clone(),
1401            mask.clone(),
1402            *width,
1403            *signed,
1404        )),
1405        SLTNode::Binary(lhs, op, rhs) => Some(NormalizedIndexExpr::Binary(
1406            Box::new(normalize_index_expr(*lhs, loop_var, arena, memo)?),
1407            *op,
1408            Box::new(normalize_index_expr(*rhs, loop_var, arena, memo)?),
1409        )),
1410        SLTNode::Unary(op, inner) => Some(NormalizedIndexExpr::Unary(
1411            *op,
1412            Box::new(normalize_index_expr(*inner, loop_var, arena, memo)?),
1413        )),
1414        SLTNode::Capture { expr, .. } => normalize_index_expr(*expr, loop_var, arena, memo),
1415        SLTNode::Mux {
1416            cond,
1417            then_expr,
1418            else_expr,
1419        } => Some(NormalizedIndexExpr::Mux {
1420            cond: Box::new(normalize_index_expr(*cond, loop_var, arena, memo)?),
1421            then_expr: Box::new(normalize_index_expr(*then_expr, loop_var, arena, memo)?),
1422            else_expr: Box::new(normalize_index_expr(*else_expr, loop_var, arena, memo)?),
1423        }),
1424        SLTNode::Concat(parts) => Some(NormalizedIndexExpr::Concat(
1425            parts
1426                .iter()
1427                .map(|(part, width)| {
1428                    normalize_index_expr(*part, loop_var, arena, memo).map(|part| (part, *width))
1429                })
1430                .collect::<Option<Vec<_>>>()?,
1431        )),
1432        SLTNode::Slice { expr, access } => Some(NormalizedIndexExpr::Slice(
1433            Box::new(normalize_index_expr(*expr, loop_var, arena, memo)?),
1434            *access,
1435        )),
1436        SLTNode::ForFold { .. } | SLTNode::ForFoldGroup { .. } => None,
1437    };
1438    memo.insert(node, normalized.clone());
1439    normalized
1440}
1441
1442#[derive(Clone, PartialEq, Eq, Hash)]
1443struct ExactIndexedLoadKey<Addr: Clone + Eq + Hash> {
1444    base: Addr,
1445    access: BitAccess,
1446    index: Vec<(NormalizedIndexExpr<Addr>, usize)>,
1447}
1448
1449impl<Addr: Clone + Eq + Hash> ExactIndexedLoadKey<Addr> {
1450    fn saved_runtime_cost(&self) -> u128 {
1451        let width = self.access.msb - self.access.lsb + 1;
1452        let chunks = width.div_ceil(64) as u128;
1453        let address_cost = self.index.iter().fold(1u128, |cost, (expr, _)| {
1454            cost.saturating_add(2).saturating_add(expr.operation_cost())
1455        });
1456        6u128.saturating_mul(chunks).saturating_add(address_cost)
1457    }
1458}
1459
1460#[derive(Clone)]
1461struct FoldGroupReadFacts<Addr: Clone + Eq + Hash> {
1462    loop_var: Addr,
1463    state_targets: Vec<VarAtomBase<Addr>>,
1464    guard_reads: Vec<SchedulerInputRead<Addr>>,
1465    initial_reads: Vec<SchedulerInputRead<Addr>>,
1466    update_reads: Vec<SchedulerInputRead<Addr>>,
1467    indexed_loads: HashSet<ExactIndexedLoadKey<Addr>>,
1468    carried_chunks: u128,
1469}
1470
1471#[derive(Clone)]
1472struct ExactFoldGroup<Addr: Clone + Eq + Hash> {
1473    root: NodeId,
1474    facts: FoldGroupReadFacts<Addr>,
1475}
1476
1477struct FoldGroupScheduleInfo<Addr: Clone + Eq + Hash> {
1478    projection_paths: Vec<usize>,
1479    read_facts: Option<FoldGroupReadFacts<Addr>>,
1480    exact_and_exclusive: bool,
1481}
1482
1483struct FoldGroupScheduleIndex<Addr: Clone + Eq + Hash> {
1484    groups: BTreeMap<NodeId, FoldGroupScheduleInfo<Addr>>,
1485    direct_group_by_path: Vec<Option<NodeId>>,
1486}
1487
1488fn collect_reachable_scheduled_groups<Addr: Clone + Eq + Hash>(
1489    root: NodeId,
1490    arena: &SLTNodeArena<Addr>,
1491    reaches_scheduled_group: &[bool],
1492    scheduled_groups: &HashSet<NodeId>,
1493    result: &mut HashSet<NodeId>,
1494) {
1495    if !reaches_scheduled_group[root.0] {
1496        return;
1497    }
1498    let mut visited = HashSet::default();
1499    let mut work = vec![root];
1500    let mut children = Vec::new();
1501    while let Some(node) = work.pop() {
1502        if !visited.insert(node) || !reaches_scheduled_group[node.0] {
1503            continue;
1504        }
1505        if scheduled_groups.contains(&node) {
1506            result.insert(node);
1507        }
1508        children.clear();
1509        push_scheduler_node_children(node, arena, &mut children);
1510        work.extend(
1511            children
1512                .iter()
1513                .copied()
1514                .filter(|child| reaches_scheduled_group[child.0]),
1515        );
1516    }
1517}
1518
1519fn exact_fold_group_paths<Addr: Clone + Eq + Hash + Copy>(
1520    group: NodeId,
1521    paths: &[usize],
1522    input: &[LogicPath<Addr>],
1523    arena: &SLTNodeArena<Addr>,
1524) -> bool {
1525    let SLTNode::ForFoldGroup { states, .. } = arena.get(group) else {
1526        return false;
1527    };
1528    let Some(packed_accesses) = packed_fold_group_state_accesses(states) else {
1529        return false;
1530    };
1531    let mut covered = vec![Vec::<BitAccess>::new(); states.len()];
1532
1533    for &path_index in paths {
1534        let path = &input[path_index];
1535        if !path.local_inputs.is_empty() || !path.pre_lower_nodes.is_empty() {
1536            return false;
1537        }
1538        let Some(target) = path.target.var() else {
1539            return false;
1540        };
1541        let Some(projection) = fold_group_projection_access(path.expr, group, arena) else {
1542            return false;
1543        };
1544        let matches = states
1545            .iter()
1546            .zip(&packed_accesses)
1547            .enumerate()
1548            .filter_map(|(state_index, (state, packed))| {
1549                if target.id != state.target.id
1550                    || target.access.lsb < state.target.access.lsb
1551                    || target.access.msb > state.target.access.msb
1552                {
1553                    return None;
1554                }
1555                let relative_lsb = target.access.lsb - state.target.access.lsb;
1556                let relative_msb = target.access.msb - state.target.access.lsb;
1557                let expected = BitAccess::new(
1558                    packed.lsb.checked_add(relative_lsb)?,
1559                    packed.lsb.checked_add(relative_msb)?,
1560                );
1561                (projection == expected).then_some(state_index)
1562            })
1563            .collect::<Vec<_>>();
1564        if matches.len() != 1 {
1565            return false;
1566        }
1567        covered[matches[0]].push(target.access);
1568    }
1569
1570    states.iter().zip(&mut covered).all(|(state, ranges)| {
1571        ranges.sort_unstable_by_key(|range| (range.lsb, range.msb));
1572        let mut next = state.target.access.lsb;
1573        for range in ranges.iter() {
1574            if range.lsb != next || range.msb > state.target.access.msb {
1575                return false;
1576            }
1577            let Some(after) = range.msb.checked_add(1) else {
1578                return range.msb == state.target.access.msb;
1579            };
1580            next = after;
1581        }
1582        next == state.target.access.msb.saturating_add(1)
1583    })
1584}
1585
1586fn build_fold_group_schedule_index<Addr: Clone + Eq + Ord + Hash + Copy>(
1587    input: &[LogicPath<Addr>],
1588    arena: &SLTNodeArena<Addr>,
1589) -> FoldGroupScheduleIndex<Addr> {
1590    let direct_group_by_path = input
1591        .iter()
1592        .map(|path| projected_for_fold_group(path.expr, arena))
1593        .collect::<Vec<_>>();
1594    let direct_roots = direct_group_by_path
1595        .iter()
1596        .flatten()
1597        .copied()
1598        .collect::<HashSet<_>>();
1599    if direct_roots.is_empty() {
1600        return FoldGroupScheduleIndex {
1601            groups: BTreeMap::new(),
1602            direct_group_by_path,
1603        };
1604    }
1605
1606    // SLT children always precede their owners. Compute one boolean per node,
1607    // then traverse only semantic roots that can actually reach a scheduled
1608    // fold. This avoids carrying a root set through the entire arena.
1609    let mut reaches_scheduled_group = Vec::<bool>::with_capacity(arena.len());
1610    let mut children = Vec::new();
1611    for raw in 0..arena.len() {
1612        let node = NodeId(raw);
1613        children.clear();
1614        push_scheduler_node_children(node, arena, &mut children);
1615        reaches_scheduled_group.push(
1616            direct_roots.contains(&node)
1617                || children
1618                    .iter()
1619                    .any(|child| reaches_scheduled_group[child.0]),
1620        );
1621    }
1622
1623    let mut groups = BTreeMap::<NodeId, FoldGroupScheduleInfo<Addr>>::new();
1624    for (path_index, path) in input.iter().enumerate() {
1625        let direct = direct_group_by_path[path_index];
1626        if let Some(root) = direct {
1627            groups
1628                .entry(root)
1629                .or_insert_with(|| FoldGroupScheduleInfo {
1630                    projection_paths: Vec::new(),
1631                    read_facts: None,
1632                    exact_and_exclusive: true,
1633                })
1634                .projection_paths
1635                .push(path_index);
1636        }
1637
1638        let mut reached = HashSet::default();
1639        collect_reachable_scheduled_groups(
1640            path.expr,
1641            arena,
1642            &reaches_scheduled_group,
1643            &direct_roots,
1644            &mut reached,
1645        );
1646        for root in reached.drain() {
1647            let info = groups.entry(root).or_insert_with(|| FoldGroupScheduleInfo {
1648                projection_paths: Vec::new(),
1649                read_facts: None,
1650                exact_and_exclusive: true,
1651            });
1652            if direct != Some(root) {
1653                info.exact_and_exclusive = false;
1654            }
1655        }
1656        let auxiliary_roots = path
1657            .local_inputs
1658            .iter()
1659            .map(|(_, node)| *node)
1660            .chain(path.pre_lower_nodes.iter().copied())
1661            .chain(match &path.target {
1662                LogicPathTarget::Var(_) => Vec::new(),
1663                LogicPathTarget::CombCaptureEvent {
1664                    guard,
1665                    args,
1666                    loop_runner,
1667                    ..
1668                } => guard
1669                    .iter()
1670                    .chain(args)
1671                    .chain(loop_runner)
1672                    .copied()
1673                    .collect(),
1674            });
1675        for node in auxiliary_roots {
1676            reached.clear();
1677            collect_reachable_scheduled_groups(
1678                node,
1679                arena,
1680                &reaches_scheduled_group,
1681                &direct_roots,
1682                &mut reached,
1683            );
1684            for root in reached.drain() {
1685                groups
1686                    .entry(root)
1687                    .or_insert_with(|| FoldGroupScheduleInfo {
1688                        projection_paths: Vec::new(),
1689                        read_facts: None,
1690                        exact_and_exclusive: true,
1691                    })
1692                    .exact_and_exclusive = false;
1693            }
1694        }
1695    }
1696
1697    for (&root, info) in &mut groups {
1698        info.exact_and_exclusive &= !info.projection_paths.is_empty()
1699            && exact_fold_group_paths(root, &info.projection_paths, input, arena);
1700        if info.exact_and_exclusive {
1701            info.read_facts = collect_fold_group_read_facts(root, arena);
1702            info.exact_and_exclusive &= info.read_facts.is_some();
1703        }
1704    }
1705
1706    FoldGroupScheduleIndex {
1707        groups,
1708        direct_group_by_path,
1709    }
1710}
1711
1712fn discover_exact_fold_groups<Addr: Clone + Eq + Ord + Hash + Copy>(
1713    indices: &[usize],
1714    schedule_index: &FoldGroupScheduleIndex<Addr>,
1715) -> Vec<ExactFoldGroup<Addr>> {
1716    let scheduled_indices = indices.iter().copied().collect::<HashSet<_>>();
1717    let mut roots = indices
1718        .iter()
1719        .filter_map(|&index| schedule_index.direct_group_by_path[index])
1720        .collect::<Vec<_>>();
1721    roots.sort_unstable();
1722    roots.dedup();
1723
1724    roots
1725        .into_iter()
1726        .filter_map(|root| {
1727            let info = schedule_index.groups.get(&root)?;
1728            if !info.exact_and_exclusive
1729                || info
1730                    .projection_paths
1731                    .iter()
1732                    .any(|index| !scheduled_indices.contains(index))
1733            {
1734                return None;
1735            }
1736            Some(ExactFoldGroup {
1737                root,
1738                facts: info.read_facts.clone()?,
1739            })
1740        })
1741        .collect()
1742}
1743
1744fn same_fold_group_domain<Addr: Clone + Eq + Hash>(
1745    lhs: NodeId,
1746    rhs: NodeId,
1747    arena: &SLTNodeArena<Addr>,
1748) -> bool {
1749    let SLTNode::ForFoldGroup {
1750        loop_width: lhs_width,
1751        loop_signed: lhs_signed,
1752        start: lhs_start,
1753        step: lhs_step,
1754        trip_count: lhs_count,
1755        entry_guard: lhs_guard,
1756        ..
1757    } = arena.get(lhs)
1758    else {
1759        return false;
1760    };
1761    let SLTNode::ForFoldGroup {
1762        loop_width: rhs_width,
1763        loop_signed: rhs_signed,
1764        start: rhs_start,
1765        step: rhs_step,
1766        trip_count: rhs_count,
1767        entry_guard: rhs_guard,
1768        ..
1769    } = arena.get(rhs)
1770    else {
1771        return false;
1772    };
1773    lhs_width == rhs_width
1774        && lhs_signed == rhs_signed
1775        && lhs_start == rhs_start
1776        && lhs_step == rhs_step
1777        && lhs_count == rhs_count
1778        && lhs_guard == rhs_guard
1779}
1780
1781#[derive(Clone, Copy)]
1782struct SchedulerInputRead<Addr> {
1783    id: Addr,
1784    access: BitAccess,
1785    indexed: bool,
1786}
1787
1788fn collect_scheduler_plain_reads<Addr: Clone + Eq + Hash + Copy>(
1789    root: NodeId,
1790    arena: &SLTNodeArena<Addr>,
1791    reads: &mut Vec<SchedulerInputRead<Addr>>,
1792) -> bool {
1793    let mut visited = HashSet::default();
1794    let mut work = vec![root];
1795    while let Some(node) = work.pop() {
1796        if !visited.insert(node) {
1797            continue;
1798        }
1799        match arena.get(node) {
1800            SLTNode::Input {
1801                variable,
1802                index,
1803                access,
1804                ..
1805            } => {
1806                reads.push(SchedulerInputRead {
1807                    id: *variable,
1808                    access: *access,
1809                    indexed: !index.is_empty(),
1810                });
1811                work.extend(index.iter().map(|entry| entry.node));
1812            }
1813            SLTNode::ForFold { .. } | SLTNode::ForFoldGroup { .. } => return false,
1814            _ => push_scheduler_node_children(node, arena, &mut work),
1815        }
1816    }
1817    true
1818}
1819
1820fn collect_fold_group_read_facts<Addr: Clone + Eq + Hash + Copy>(
1821    root: NodeId,
1822    arena: &SLTNodeArena<Addr>,
1823) -> Option<FoldGroupReadFacts<Addr>> {
1824    let SLTNode::ForFoldGroup {
1825        loop_var,
1826        entry_guard,
1827        states,
1828        ..
1829    } = arena.get(root)
1830    else {
1831        return None;
1832    };
1833
1834    let mut guard_reads = Vec::new();
1835    collect_scheduler_plain_reads(*entry_guard, arena, &mut guard_reads).then_some(())?;
1836    let mut initial_reads = Vec::new();
1837    let mut update_reads = Vec::new();
1838    for state in states {
1839        collect_scheduler_plain_reads(state.initial, arena, &mut initial_reads).then_some(())?;
1840        collect_scheduler_plain_reads(state.update, arena, &mut update_reads).then_some(())?;
1841    }
1842
1843    let state_ids = states
1844        .iter()
1845        .map(|state| state.target.id)
1846        .collect::<HashSet<_>>();
1847    let mut indexed_loads = HashSet::default();
1848    let mut normalize_memo = HashMap::default();
1849    let mut visited = HashSet::default();
1850    let mut work = states.iter().map(|state| state.update).collect::<Vec<_>>();
1851    while let Some(node) = work.pop() {
1852        if !visited.insert(node) {
1853            continue;
1854        }
1855        match arena.get(node) {
1856            SLTNode::Input {
1857                variable,
1858                index,
1859                access,
1860                ..
1861            } => {
1862                if *variable != *loop_var && !state_ids.contains(variable) && !index.is_empty() {
1863                    let normalized = index
1864                        .iter()
1865                        .map(|entry| {
1866                            normalize_index_expr(entry.node, *loop_var, arena, &mut normalize_memo)
1867                                .map(|node| (node, entry.stride))
1868                        })
1869                        .collect::<Option<Vec<_>>>();
1870                    if let Some(normalized) = normalized
1871                        && normalized
1872                            .iter()
1873                            .any(|(expr, _)| expr.contains_loop_value())
1874                    {
1875                        indexed_loads.insert(ExactIndexedLoadKey {
1876                            base: *variable,
1877                            access: *access,
1878                            index: normalized,
1879                        });
1880                    }
1881                }
1882                work.extend(index.iter().map(|entry| entry.node));
1883            }
1884            _ => push_scheduler_node_children(node, arena, &mut work),
1885        }
1886    }
1887
1888    let carried_chunks = states.iter().fold(0u128, |chunks, state| {
1889        let width = state.target.access.msb - state.target.access.lsb + 1;
1890        chunks.saturating_add(width.div_ceil(64) as u128)
1891    });
1892    let mut state_ranges: HashMap<Addr, Vec<BitAccess>> = HashMap::default();
1893    for state in states {
1894        if state.target.id == *loop_var
1895            || state_ranges
1896                .entry(state.target.id)
1897                .or_default()
1898                .iter()
1899                .any(|range| range.overlaps(&state.target.access))
1900        {
1901            return None;
1902        }
1903        state_ranges
1904            .entry(state.target.id)
1905            .or_default()
1906            .push(state.target.access);
1907    }
1908    if guard_reads.iter().any(|read| read.id == *loop_var)
1909        || initial_reads.iter().any(|read| read.id == *loop_var)
1910        || update_reads
1911            .iter()
1912            .any(|read| read.id == *loop_var && read.indexed)
1913    {
1914        return None;
1915    }
1916    Some(FoldGroupReadFacts {
1917        loop_var: *loop_var,
1918        state_targets: states.iter().map(|state| state.target).collect(),
1919        guard_reads,
1920        initial_reads,
1921        update_reads,
1922        indexed_loads,
1923        carried_chunks,
1924    })
1925}
1926
1927fn scheduler_read_overlaps_targets<Addr: Clone + Eq + Hash>(
1928    read: &SchedulerInputRead<Addr>,
1929    targets: &[VarAtomBase<Addr>],
1930) -> bool {
1931    targets.iter().any(|target| {
1932        target.id == read.id && (read.indexed || target.access.overlaps(&read.access))
1933    })
1934}
1935
1936fn fold_groups_are_pairwise_independent<Addr: Clone + Eq + Hash + Copy>(
1937    lhs: &FoldGroupReadFacts<Addr>,
1938    rhs: &FoldGroupReadFacts<Addr>,
1939) -> bool {
1940    if lhs.loop_var == rhs.loop_var
1941        || lhs
1942            .state_targets
1943            .iter()
1944            .any(|target| target.id == rhs.loop_var)
1945        || rhs
1946            .state_targets
1947            .iter()
1948            .any(|target| target.id == lhs.loop_var)
1949        || lhs.state_targets.iter().any(|left| {
1950            rhs.state_targets
1951                .iter()
1952                .any(|right| left.id == right.id && left.access.overlaps(&right.access))
1953        })
1954    {
1955        return false;
1956    }
1957
1958    let all_targets = lhs
1959        .state_targets
1960        .iter()
1961        .chain(&rhs.state_targets)
1962        .copied()
1963        .collect::<Vec<_>>();
1964    let guard_is_independent = lhs.guard_reads.iter().chain(&rhs.guard_reads).all(|read| {
1965        read.id != lhs.loop_var
1966            && read.id != rhs.loop_var
1967            && !scheduler_read_overlaps_targets(read, &all_targets)
1968    });
1969    let initials_are_independent = lhs
1970        .initial_reads
1971        .iter()
1972        .chain(&rhs.initial_reads)
1973        .all(|read| read.id != lhs.loop_var && read.id != rhs.loop_var);
1974    let lhs_update_is_independent = lhs.update_reads.iter().all(|read| {
1975        read.id != rhs.loop_var && !scheduler_read_overlaps_targets(read, &rhs.state_targets)
1976    });
1977    let rhs_update_is_independent = rhs.update_reads.iter().all(|read| {
1978        read.id != lhs.loop_var && !scheduler_read_overlaps_targets(read, &lhs.state_targets)
1979    });
1980    guard_is_independent
1981        && initials_are_independent
1982        && lhs_update_is_independent
1983        && rhs_update_is_independent
1984}
1985
1986fn fold_groups_share_exact_load<Addr: Clone + Eq + Hash>(
1987    lhs: &FoldGroupReadFacts<Addr>,
1988    rhs: &FoldGroupReadFacts<Addr>,
1989) -> bool {
1990    let (small, large) = if lhs.indexed_loads.len() <= rhs.indexed_loads.len() {
1991        (&lhs.indexed_loads, &rhs.indexed_loads)
1992    } else {
1993        (&rhs.indexed_loads, &lhs.indexed_loads)
1994    };
1995    small.iter().any(|key| large.contains(key))
1996}
1997
1998#[derive(Clone)]
1999struct WeightedFoldFamily {
2000    members: Vec<usize>,
2001    benefit: u128,
2002    pressure: u128,
2003}
2004
2005impl WeightedFoldFamily {
2006    fn is_positive(&self) -> bool {
2007        self.members.len() >= 2 && self.benefit > self.pressure
2008    }
2009
2010    fn cmp_net(&self, other: &Self) -> std::cmp::Ordering {
2011        self.benefit
2012            .saturating_add(other.pressure)
2013            .cmp(&other.benefit.saturating_add(self.pressure))
2014    }
2015}
2016
2017fn weighted_fold_family<Addr: Clone + Eq + Hash>(
2018    members: &[usize],
2019    candidates: &[ExactFoldGroup<Addr>],
2020    four_state: bool,
2021) -> WeightedFoldFamily {
2022    const SAVED_LOOP_CONTROL_COST: u128 = 6;
2023    const CARRIED_CHUNK_PRESSURE_COST: u128 = 4;
2024
2025    let mut users_by_load = HashMap::<ExactIndexedLoadKey<Addr>, usize>::default();
2026    let mut total_chunks = 0u128;
2027    let mut largest_separate_group = 0u128;
2028    for &member in members {
2029        let facts = &candidates[member].facts;
2030        total_chunks = total_chunks.saturating_add(facts.carried_chunks);
2031        largest_separate_group = largest_separate_group.max(facts.carried_chunks);
2032        for key in &facts.indexed_loads {
2033            *users_by_load.entry(key.clone()).or_insert(0) += 1;
2034        }
2035    }
2036    let load_benefit = users_by_load
2037        .into_iter()
2038        .filter(|(_, users)| *users >= 2)
2039        .fold(0u128, |benefit, (key, users)| {
2040            benefit.saturating_add(key.saved_runtime_cost().saturating_mul((users - 1) as u128))
2041        });
2042    let control_benefit =
2043        SAVED_LOOP_CONTROL_COST.saturating_mul(members.len().saturating_sub(1) as u128);
2044    let state_multiplier = if four_state { 2 } else { 1 };
2045    let pressure = total_chunks
2046        .saturating_sub(largest_separate_group)
2047        .saturating_mul(state_multiplier)
2048        .saturating_mul(CARRIED_CHUNK_PRESSURE_COST);
2049    WeightedFoldFamily {
2050        members: members.to_vec(),
2051        benefit: load_benefit.saturating_add(control_benefit),
2052        pressure,
2053    }
2054}
2055
2056fn family_signature<Addr: Clone + Eq + Hash>(
2057    family: &WeightedFoldFamily,
2058    candidates: &[ExactFoldGroup<Addr>],
2059) -> Vec<NodeId> {
2060    let mut roots = family
2061        .members
2062        .iter()
2063        .map(|member| candidates[*member].root)
2064        .collect::<Vec<_>>();
2065    roots.sort_unstable();
2066    roots
2067}
2068
2069fn better_weighted_family<Addr: Clone + Eq + Hash>(
2070    candidate: &WeightedFoldFamily,
2071    current: &WeightedFoldFamily,
2072    groups: &[ExactFoldGroup<Addr>],
2073) -> bool {
2074    candidate.cmp_net(current).is_gt()
2075        || candidate.cmp_net(current).is_eq()
2076            && family_signature(candidate, groups) < family_signature(current, groups)
2077}
2078
2079fn grow_weighted_family<Addr: Clone + Eq + Hash>(
2080    seed: usize,
2081    first: usize,
2082    candidates: &[ExactFoldGroup<Addr>],
2083    available: &[bool],
2084    compatible: &[Vec<bool>],
2085    shared_load: &[Vec<bool>],
2086    rejected: &HashSet<Vec<NodeId>>,
2087    four_state: bool,
2088) -> Option<WeightedFoldFamily> {
2089    let mut members = vec![seed, first];
2090    let mut best = None;
2091    loop {
2092        let weighted = weighted_fold_family(&members, candidates, four_state);
2093        let signature = family_signature(&weighted, candidates);
2094        if weighted.is_positive()
2095            && !rejected.contains(&signature)
2096            && best
2097                .as_ref()
2098                .is_none_or(|current| better_weighted_family(&weighted, current, candidates))
2099        {
2100            best = Some(weighted);
2101        }
2102
2103        let mut best_expansion: Option<(usize, WeightedFoldFamily)> = None;
2104        for candidate in 0..candidates.len() {
2105            if !available[candidate]
2106                || members.contains(&candidate)
2107                || !members.iter().all(|member| compatible[*member][candidate])
2108                || !members.iter().any(|member| shared_load[*member][candidate])
2109            {
2110                continue;
2111            }
2112            let mut expanded = members.clone();
2113            expanded.push(candidate);
2114            let weighted = weighted_fold_family(&expanded, candidates, four_state);
2115            if best_expansion
2116                .as_ref()
2117                .is_none_or(|(current_index, current)| {
2118                    better_weighted_family(&weighted, current, candidates)
2119                        || weighted.cmp_net(current).is_eq()
2120                            && candidates[candidate].root < candidates[*current_index].root
2121                })
2122            {
2123                best_expansion = Some((candidate, weighted));
2124            }
2125        }
2126        let Some((next, _)) = best_expansion else {
2127            break;
2128        };
2129        members.push(next);
2130    }
2131    best
2132}
2133
2134fn best_weighted_fold_family<Addr: Clone + Eq + Hash>(
2135    candidates: &[ExactFoldGroup<Addr>],
2136    available: &[bool],
2137    compatible: &[Vec<bool>],
2138    shared_load: &[Vec<bool>],
2139    rejected: &HashSet<Vec<NodeId>>,
2140    four_state: bool,
2141) -> Option<WeightedFoldFamily> {
2142    let mut best = None;
2143    for seed in 0..candidates.len() {
2144        if !available[seed] {
2145            continue;
2146        }
2147        // Force every compatible first edge from every seed. The subsequent
2148        // growth is weighted, so a low-root conflicting candidate cannot hide
2149        // a more profitable compatible partition.
2150        for first in 0..candidates.len() {
2151            if seed == first
2152                || !available[first]
2153                || !compatible[seed][first]
2154                || !shared_load[seed][first]
2155            {
2156                continue;
2157            }
2158            let Some(family) = grow_weighted_family(
2159                seed,
2160                first,
2161                candidates,
2162                available,
2163                compatible,
2164                shared_load,
2165                rejected,
2166                four_state,
2167            ) else {
2168                continue;
2169            };
2170            if best
2171                .as_ref()
2172                .is_none_or(|current| better_weighted_family(&family, current, candidates))
2173            {
2174                best = Some(family);
2175            }
2176        }
2177    }
2178    best
2179}
2180
2181fn jointly_lower_fold_group_families<Addr: Clone + Eq + Ord + Hash + Debug + Copy + Display>(
2182    indices: &[usize],
2183    schedule_index: &FoldGroupScheduleIndex<Addr>,
2184    lowerer: &crate::SLTToSIRLowerer,
2185    builder: &mut SIRBuilder<Addr>,
2186    arena: &SLTNodeArena<Addr>,
2187    lower_cache: &mut HashMap<NodeId, RegisterId>,
2188    four_state: bool,
2189) -> HashSet<NodeId> {
2190    // `indices` is one bounded run of exact-fold roots. Direct dependencies
2191    // split runs before this point, and the fixed root window bounds the
2192    // pairwise compatibility matrix independently of design size.
2193    let candidates = discover_exact_fold_groups(indices, schedule_index);
2194    let mut compatible = vec![vec![false; candidates.len()]; candidates.len()];
2195    let mut shared_load = vec![vec![false; candidates.len()]; candidates.len()];
2196    for lhs in 0..candidates.len() {
2197        for rhs in lhs + 1..candidates.len() {
2198            let domains_match =
2199                same_fold_group_domain(candidates[lhs].root, candidates[rhs].root, arena);
2200            let independent = fold_groups_are_pairwise_independent(
2201                &candidates[lhs].facts,
2202                &candidates[rhs].facts,
2203            );
2204            let shares =
2205                fold_groups_share_exact_load(&candidates[lhs].facts, &candidates[rhs].facts);
2206            compatible[lhs][rhs] = domains_match && independent;
2207            compatible[rhs][lhs] = compatible[lhs][rhs];
2208            shared_load[lhs][rhs] = shares;
2209            shared_load[rhs][lhs] = shares;
2210        }
2211    }
2212
2213    let mut available = vec![true; candidates.len()];
2214    let mut rejected = HashSet::<Vec<NodeId>>::default();
2215    let mut lowered = HashSet::default();
2216    while let Some(family) = best_weighted_fold_family(
2217        &candidates,
2218        &available,
2219        &compatible,
2220        &shared_load,
2221        &rejected,
2222        four_state,
2223    ) {
2224        let roots = family
2225            .members
2226            .iter()
2227            .map(|member| candidates[*member].root)
2228            .collect::<Vec<_>>();
2229        if lowerer.lower_fold_groups_jointly(builder, &roots, arena, lower_cache) {
2230            for member in family.members {
2231                available[member] = false;
2232            }
2233            lowered.extend(roots);
2234        } else {
2235            rejected.insert(family_signature(&family, &candidates));
2236        }
2237    }
2238    lowered
2239}
2240
2241#[derive(Clone, Copy)]
2242struct PreparedFoldProjection {
2243    packed_result: RegisterId,
2244    access: BitAccess,
2245}
2246
2247/// Materialize each shared grouped fold once, but leave its projections
2248/// deferred.  A target Store may invalidate the ordinary lowering cache, so
2249/// the packed result is retained explicitly until every projection has been
2250/// consumed.  Each narrow projection is created immediately before its Store
2251/// instead of keeping all projected registers live simultaneously.
2252fn prepare_atomic_fold_group_results<Addr: Clone + Eq + Ord + Hash + Debug + Copy + Display>(
2253    indices: &[usize],
2254    input: &[LogicPath<Addr>],
2255    fold_group_schedule_index: &FoldGroupScheduleIndex<Addr>,
2256    lowerer: &crate::SLTToSIRLowerer,
2257    builder: &mut SIRBuilder<Addr>,
2258    arena: &SLTNodeArena<Addr>,
2259    lower_cache: &mut HashMap<NodeId, RegisterId>,
2260    dep_memo: &mut HashMap<NodeId, HashSet<Addr>>,
2261    inverse_dep_memo: &mut HashMap<Addr, HashSet<NodeId>>,
2262    four_state: bool,
2263) -> HashMap<usize, PreparedFoldProjection> {
2264    let jointly_lowered = jointly_lower_fold_group_families(
2265        indices,
2266        fold_group_schedule_index,
2267        lowerer,
2268        builder,
2269        arena,
2270        lower_cache,
2271        four_state,
2272    );
2273    let mut counts = HashMap::default();
2274    for &idx in indices {
2275        let path = &input[idx];
2276        if path.target.var().is_none()
2277            || !path.local_inputs.is_empty()
2278            || !path.pre_lower_nodes.is_empty()
2279        {
2280            continue;
2281        }
2282        if let Some(group) = projected_for_fold_group(path.expr, arena) {
2283            *counts.entry(group).or_insert(0usize) += 1;
2284        }
2285    }
2286
2287    let mut prepared = HashMap::default();
2288    for &idx in indices {
2289        let path = &input[idx];
2290        if !path.local_inputs.is_empty() || !path.pre_lower_nodes.is_empty() {
2291            continue;
2292        }
2293        let Some(group) = projected_for_fold_group(path.expr, arena) else {
2294            continue;
2295        };
2296        if counts.get(&group).copied().unwrap_or(0) < 2 && !jointly_lowered.contains(&group) {
2297            continue;
2298        }
2299        let Some(access) = fold_group_projection_access(path.expr, group, arena) else {
2300            continue;
2301        };
2302        collect_logic_path_input_deps(path, arena, dep_memo, inverse_dep_memo);
2303        let packed_result = lower_cache
2304            .get(&group)
2305            .copied()
2306            .unwrap_or_else(|| lowerer.lower(builder, group, arena, lower_cache));
2307        prepared.insert(
2308            idx,
2309            PreparedFoldProjection {
2310                packed_result,
2311                access,
2312            },
2313        );
2314    }
2315    prepared
2316}
2317
2318#[derive(Error, Debug, PartialEq, Eq)]
2319pub enum SchedulerError<A: Display + Debug + Eq + Hash + Clone> {
2320    #[error("Combinational loop detected: {}", .blocks.iter().map(|v| format!("{}", v)).collect::<Vec<_>>().join(" -> "))]
2321    CombinationalLoop { blocks: Vec<LogicPath<A>> },
2322    #[error("Multiple driver detected: {}", .blocks.iter().map(|v| format!("{}", v)).collect::<Vec<_>>().join(","))]
2323    MultipleDriver { blocks: Vec<LogicPath<A>> },
2324    #[error("internal logic-path SCC condensation graph is invalid")]
2325    InvalidDependencyGraph,
2326}
2327
2328impl<A: Display + Debug + Eq + Hash + Clone> SchedulerError<A> {
2329    pub fn map_addr<B: Display + Debug + Eq + Hash + Clone, F>(
2330        self,
2331        arena: &SLTNodeArena<A>,
2332        target_arena: &mut SLTNodeArena<B>,
2333        f: &F,
2334    ) -> Result<SchedulerError<B>, SLTNodeFactsError>
2335    where
2336        F: Fn(&A) -> B,
2337    {
2338        let mut cache = HashMap::default();
2339        Ok(match self {
2340            SchedulerError::CombinationalLoop { blocks } => SchedulerError::CombinationalLoop {
2341                blocks: blocks
2342                    .into_iter()
2343                    .map(|b| b.map_addr(arena, target_arena, &mut cache, f))
2344                    .collect::<Result<Vec<_>, _>>()?,
2345            },
2346            SchedulerError::MultipleDriver { blocks } => SchedulerError::MultipleDriver {
2347                blocks: blocks
2348                    .into_iter()
2349                    .map(|b| b.map_addr(arena, target_arena, &mut cache, f))
2350                    .collect::<Result<Vec<_>, _>>()?,
2351            },
2352            SchedulerError::InvalidDependencyGraph => SchedulerError::InvalidDependencyGraph,
2353        })
2354    }
2355}
2356
2357pub struct ScheduleResult<Addr> {
2358    pub execution_units: Vec<ExecutionUnit<Addr>>,
2359    pub runtime_errors: HashMap<i64, RuntimeErrorInfo<Addr>>,
2360    /// Persistent-state ranges published directly by FF actions in the shared
2361    /// comb/FF schedule. These Stores are semantic state updates rather than
2362    /// disposable comb publications.
2363    pub direct_ff_writes: Vec<VarAtomBase<Addr>>,
2364}
2365
2366pub trait ClockFfLowering<Addr> {
2367    type Error;
2368
2369    fn summaries(&self) -> &[FfAccessSummary<Addr>];
2370    fn begin(
2371        &mut self,
2372        builder: &mut SIRBuilder<Addr>,
2373        direct_writes: &[Vec<VarAtomBase<Addr>>],
2374    ) -> Result<(), Self::Error>;
2375    fn lower(
2376        &mut self,
2377        index: usize,
2378        direct_writes: &[VarAtomBase<Addr>],
2379        builder: &mut SIRBuilder<Addr>,
2380    ) -> Result<(), Self::Error>;
2381    fn finish(
2382        &mut self,
2383        builder: &mut SIRBuilder<Addr>,
2384        direct_writes: &[Vec<VarAtomBase<Addr>>],
2385    ) -> Result<(), Self::Error>;
2386}
2387
2388pub enum ClockSortError<Addr: Display + Debug + Eq + Hash + Clone, E> {
2389    Scheduler(SchedulerError<Addr>),
2390    Lowering(E),
2391}
2392
2393enum ScheduledWork {
2394    CombPath(usize),
2395    CombScc(Vec<usize>),
2396    /// A dependency-ordered run of state publications selected by one SLT
2397    /// condition.  The path scheduler has already fixed the order; lowering
2398    /// only preserves the exclusivity which would otherwise become several
2399    /// independent Muxes and later need to be rediscovered from SIR.
2400    GuardedComb {
2401        condition: NodeId,
2402        paths: Vec<usize>,
2403    },
2404    Ff(usize),
2405}
2406
2407fn scheduled_root_mux<A: Clone + Eq + Hash>(
2408    path: &LogicPath<A>,
2409    arena: &SLTNodeArena<A>,
2410) -> Option<(NodeId, NodeId, NodeId)> {
2411    if path.target.var().is_none()
2412        || !path.local_inputs.is_empty()
2413        || !path.pre_lower_nodes.is_empty()
2414    {
2415        return None;
2416    }
2417    let SLTNode::Mux {
2418        cond,
2419        then_expr,
2420        else_expr,
2421    } = arena.get(path.expr)
2422    else {
2423        return None;
2424    };
2425    Some((*cond, *then_expr, *else_expr))
2426}
2427
2428fn collect_pure_scheduled_nodes<A: Clone + Eq + Hash>(
2429    root: NodeId,
2430    arena: &SLTNodeArena<A>,
2431    nodes: &mut HashSet<NodeId>,
2432) -> bool {
2433    if !nodes.insert(root) {
2434        return true;
2435    }
2436    match arena.get(root) {
2437        SLTNode::Input { index, .. } => index
2438            .iter()
2439            .all(|index| collect_pure_scheduled_nodes(index.node, arena, nodes)),
2440        SLTNode::Constant(..) => true,
2441        SLTNode::Binary(lhs, _, rhs) => {
2442            collect_pure_scheduled_nodes(*lhs, arena, nodes)
2443                && collect_pure_scheduled_nodes(*rhs, arena, nodes)
2444        }
2445        SLTNode::Unary(_, inner)
2446        | SLTNode::Capture { expr: inner, .. }
2447        | SLTNode::Slice { expr: inner, .. } => collect_pure_scheduled_nodes(*inner, arena, nodes),
2448        SLTNode::Mux {
2449            cond,
2450            then_expr,
2451            else_expr,
2452        } => {
2453            collect_pure_scheduled_nodes(*cond, arena, nodes)
2454                && collect_pure_scheduled_nodes(*then_expr, arena, nodes)
2455                && collect_pure_scheduled_nodes(*else_expr, arena, nodes)
2456        }
2457        SLTNode::Concat(parts) => parts
2458            .iter()
2459            .all(|(part, _)| collect_pure_scheduled_nodes(*part, arena, nodes)),
2460        SLTNode::ForFold { .. } | SLTNode::ForFoldGroup { .. } => false,
2461    }
2462}
2463
2464fn scheduled_node_cost<A: Clone + Eq + Hash>(node: NodeId, arena: &SLTNodeArena<A>) -> u128 {
2465    match arena.get(node) {
2466        SLTNode::Input { .. } | SLTNode::Constant(..) => 0,
2467        _ => crate::get_width(node, arena).div_ceil(64).max(1) as u128,
2468    }
2469}
2470
2471fn scheduled_condition_probability<A: Clone + Eq + Hash>(
2472    mut condition: NodeId,
2473    arena: &SLTNodeArena<A>,
2474) -> (u128, u128) {
2475    let mut inverted = false;
2476    loop {
2477        match arena.get(condition) {
2478            SLTNode::Unary(UnaryOp::Ident | UnaryOp::ToTwoState, inner) => condition = *inner,
2479            SLTNode::Unary(UnaryOp::LogicNot, inner) => {
2480                inverted = !inverted;
2481                condition = *inner;
2482            }
2483            _ => break,
2484        }
2485    }
2486    let constant = |node| matches!(arena.get(node), SLTNode::Constant(..));
2487    let equality = matches!(
2488        arena.get(condition),
2489        SLTNode::Binary(lhs, BinaryOp::Eq | BinaryOp::EqWildcard, rhs)
2490            if constant(*lhs) || constant(*rhs)
2491    );
2492    let inequality = matches!(
2493        arena.get(condition),
2494        SLTNode::Binary(lhs, BinaryOp::Ne | BinaryOp::NeWildcard, rhs)
2495            if constant(*lhs) || constant(*rhs)
2496    );
2497    let true_weight = if equality {
2498        1
2499    } else if inequality {
2500        4
2501    } else {
2502        return (1, 2);
2503    };
2504    if inverted {
2505        (5 - true_weight, 5)
2506    } else {
2507        (true_weight, 5)
2508    }
2509}
2510
2511fn scheduled_guard_region_is_profitable<A: Clone + Eq + Hash>(
2512    condition: NodeId,
2513    paths: &[usize],
2514    input: &[LogicPath<A>],
2515    arena: &SLTNodeArena<A>,
2516) -> bool {
2517    const CONTROL_COST: u128 = 2;
2518    const MISPREDICT_COST: u128 = 16;
2519
2520    let mut true_nodes = HashSet::default();
2521    let mut false_nodes = HashSet::default();
2522    for &path in paths {
2523        let Some((actual, then_expr, else_expr)) = scheduled_root_mux(&input[path], arena) else {
2524            return false;
2525        };
2526        if actual != condition
2527            || !collect_pure_scheduled_nodes(then_expr, arena, &mut true_nodes)
2528            || !collect_pure_scheduled_nodes(else_expr, arena, &mut false_nodes)
2529        {
2530            return false;
2531        }
2532    }
2533    let true_owned = true_nodes
2534        .difference(&false_nodes)
2535        .map(|node| scheduled_node_cost(*node, arena))
2536        .sum::<u128>();
2537    let false_owned = false_nodes
2538        .difference(&true_nodes)
2539        .map(|node| scheduled_node_cost(*node, arena))
2540        .sum::<u128>();
2541    let (true_weight, total_weight) = scheduled_condition_probability(condition, arena);
2542    let false_weight = total_weight - true_weight;
2543    let removed_mux_cost = total_weight.saturating_mul(paths.len() as u128);
2544    let saved = false_weight
2545        .saturating_mul(true_owned)
2546        .saturating_add(true_weight.saturating_mul(false_owned))
2547        .saturating_add(removed_mux_cost);
2548    let introduced = total_weight.saturating_mul(CONTROL_COST).saturating_add(
2549        true_weight
2550            .min(false_weight)
2551            .saturating_mul(MISPREDICT_COST),
2552    );
2553    saved > introduced
2554}
2555
2556fn condition_reads_target<A: Clone + Eq + Hash>(
2557    condition: NodeId,
2558    targets: &HashSet<A>,
2559    arena: &SLTNodeArena<A>,
2560) -> bool {
2561    let mut visited = HashSet::default();
2562    let mut work = vec![condition];
2563    while let Some(node) = work.pop() {
2564        if !visited.insert(node) {
2565            continue;
2566        }
2567        match arena.get(node) {
2568            SLTNode::Input {
2569                variable, index, ..
2570            } => {
2571                if targets.contains(variable) {
2572                    return true;
2573                }
2574                work.extend(index.iter().map(|index| index.node));
2575            }
2576            SLTNode::Constant(..) => {}
2577            SLTNode::Binary(lhs, _, rhs) => work.extend([*lhs, *rhs]),
2578            SLTNode::Unary(_, inner)
2579            | SLTNode::Capture { expr: inner, .. }
2580            | SLTNode::Slice { expr: inner, .. } => work.push(*inner),
2581            SLTNode::Mux {
2582                cond,
2583                then_expr,
2584                else_expr,
2585            } => work.extend([*cond, *then_expr, *else_expr]),
2586            SLTNode::Concat(parts) => work.extend(parts.iter().map(|(part, _)| *part)),
2587            SLTNode::ForFold { .. } | SLTNode::ForFoldGroup { .. } => return true,
2588        }
2589    }
2590    false
2591}
2592
2593fn form_scheduled_guard_regions<A: Clone + Eq + Hash>(
2594    work: Vec<ScheduledWork>,
2595    input: &[LogicPath<A>],
2596    arena: &SLTNodeArena<A>,
2597    four_state: bool,
2598) -> Vec<ScheduledWork> {
2599    if four_state {
2600        return work;
2601    }
2602    let mut result = Vec::with_capacity(work.len());
2603    let mut pending = work.into_iter().peekable();
2604    while let Some(item) = pending.next() {
2605        let ScheduledWork::CombPath(first_path) = item else {
2606            result.push(item);
2607            continue;
2608        };
2609        let Some((condition, _, _)) = scheduled_root_mux(&input[first_path], arena) else {
2610            result.push(ScheduledWork::CombPath(first_path));
2611            continue;
2612        };
2613
2614        let mut paths = vec![first_path];
2615        while let Some(ScheduledWork::CombPath(next_path)) = pending.peek() {
2616            if scheduled_root_mux(&input[*next_path], arena)
2617                .is_none_or(|(next_condition, _, _)| next_condition != condition)
2618            {
2619                break;
2620            }
2621            paths.push(*next_path);
2622            pending.next();
2623        }
2624
2625        let targets = paths
2626            .iter()
2627            .filter_map(|path| input[*path].target.var().map(|target| target.id.clone()))
2628            .collect::<HashSet<_>>();
2629        if paths.len() >= 2
2630            && !condition_reads_target(condition, &targets, arena)
2631            && scheduled_guard_region_is_profitable(condition, &paths, input, arena)
2632        {
2633            result.push(ScheduledWork::GuardedComb { condition, paths });
2634        } else {
2635            result.extend(paths.into_iter().map(ScheduledWork::CombPath));
2636        }
2637    }
2638    result
2639}
2640
2641/// Lower a consecutive set of exact grouped-fold paths.  The packed fold
2642/// results are computed atomically, then each projection is created and stored
2643/// in topological order.  Ordinary paths bypass this buffer entirely.
2644fn flush_pending_fold_paths<Addr: Clone + Eq + Ord + Hash + Debug + Copy + Display>(
2645    pending: &mut Vec<usize>,
2646    input: &[LogicPath<Addr>],
2647    fold_group_schedule_index: &FoldGroupScheduleIndex<Addr>,
2648    lowerer: &crate::SLTToSIRLowerer,
2649    builder: &mut SIRBuilder<Addr>,
2650    arena: &SLTNodeArena<Addr>,
2651    lower_cache: &mut HashMap<NodeId, RegisterId>,
2652    dep_memo: &mut HashMap<NodeId, HashSet<Addr>>,
2653    inverse_dep_memo: &mut HashMap<Addr, HashSet<NodeId>>,
2654    unpacked_element_widths: &HashMap<Addr, usize>,
2655    four_state: bool,
2656) {
2657    if pending.is_empty() {
2658        return;
2659    }
2660
2661    let prepared_results = prepare_atomic_fold_group_results(
2662        pending,
2663        input,
2664        fold_group_schedule_index,
2665        lowerer,
2666        builder,
2667        arena,
2668        lower_cache,
2669        dep_memo,
2670        inverse_dep_memo,
2671        four_state,
2672    );
2673    for idx in pending.drain(..) {
2674        let path = &input[idx];
2675        collect_logic_path_input_deps(path, arena, dep_memo, inverse_dep_memo);
2676        let prepared_result = prepared_results.get(&idx).map(|projection| {
2677            lowerer.project_materialized(builder, projection.packed_result, projection.access)
2678        });
2679        emit_logic_path_store_with_result(
2680            lowerer,
2681            builder,
2682            path,
2683            arena,
2684            lower_cache,
2685            unpacked_element_widths,
2686            prepared_result,
2687        );
2688        invalidate_logic_path_target(path, inverse_dep_memo, lower_cache);
2689    }
2690}
2691
2692fn is_exact_fold_path<Addr: Clone + Eq + Hash + Copy>(
2693    path: usize,
2694    fold_groups: &FoldGroupScheduleIndex<Addr>,
2695) -> bool {
2696    if let Some(root) = fold_groups.direct_group_by_path[path]
2697        && fold_groups
2698            .groups
2699            .get(&root)
2700            .is_some_and(|info| info.exact_and_exclusive)
2701    {
2702        return true;
2703    }
2704    false
2705}
2706
2707/// Assign a scheduling domain to paths whose memory effects or packed fold
2708/// result benefit from staying adjacent.  A domain is only a ready-queue
2709/// preference: it never contracts paths into one graph node and therefore
2710/// cannot make a not-yet-ready path execute early.
2711fn logic_path_scheduling_domains<Addr: Clone + Eq + Ord + Hash + Copy>(
2712    input: &[LogicPath<Addr>],
2713    fold_groups: &FoldGroupScheduleIndex<Addr>,
2714) -> Vec<Option<usize>> {
2715    let mut next_domain = 0usize;
2716    let mut fold_domains = BTreeMap::<NodeId, usize>::new();
2717    let mut target_domains = BTreeMap::<Addr, usize>::new();
2718    let mut domains = Vec::with_capacity(input.len());
2719
2720    for (path, logic_path) in input.iter().enumerate() {
2721        let exact_fold_root = fold_groups.direct_group_by_path[path].filter(|root| {
2722            fold_groups
2723                .groups
2724                .get(root)
2725                .is_some_and(|info| info.exact_and_exclusive)
2726        });
2727        let domain = if let Some(root) = exact_fold_root {
2728            Some(*fold_domains.entry(root).or_insert_with(|| {
2729                let domain = next_domain;
2730                next_domain += 1;
2731                domain
2732            }))
2733        } else {
2734            logic_path.target.var().map(|target| {
2735                *target_domains.entry(target.id).or_insert_with(|| {
2736                    let domain = next_domain;
2737                    next_domain += 1;
2738                    domain
2739                })
2740            })
2741        };
2742        domains.push(domain);
2743    }
2744    domains
2745}
2746
2747/// Produce a deterministic topological order of the SCC condensation graph.
2748/// A ready effect domain is drained while possible, but each component remains
2749/// an independent queue item.  This preserves target/fold locality without
2750/// treating a layer or a complete ready frontier as a lowering batch.  Actual
2751/// register-pressure scheduling is performed later on MIR instructions and
2752/// their real VRegs.
2753fn stable_topological_sccs(
2754    sccs: Vec<Vec<usize>>,
2755    adj: &[Vec<usize>],
2756    path_domains: &[Option<usize>],
2757) -> Option<(Vec<Vec<usize>>, Vec<usize>)> {
2758    if path_domains.len() != adj.len() {
2759        return None;
2760    }
2761    let component_count = sccs.len();
2762    let mut component_by_path = vec![usize::MAX; adj.len()];
2763    let mut keys = Vec::with_capacity(component_count);
2764    let mut component_domains = Vec::with_capacity(component_count);
2765    for (component, scc) in sccs.iter().enumerate() {
2766        keys.push(*scc.iter().min()?);
2767        for &path in scc {
2768            if path >= adj.len() || component_by_path[path] != usize::MAX {
2769                return None;
2770            }
2771            component_by_path[path] = component;
2772        }
2773        let singleton = (scc.len() == 1).then_some(scc[0]);
2774        let acyclic = singleton.is_some_and(|path| !adj[path].contains(&path));
2775        component_domains.push(
2776            acyclic
2777                .then(|| path_domains[singleton.expect("acyclic SCC is a singleton")])
2778                .flatten(),
2779        );
2780    }
2781    if component_by_path.contains(&usize::MAX) {
2782        return None;
2783    }
2784
2785    let mut outgoing = vec![Vec::<usize>::new(); component_count];
2786    for (definition, users) in adj.iter().enumerate() {
2787        let source = component_by_path[definition];
2788        for &user in users {
2789            let target = *component_by_path.get(user)?;
2790            if source != target {
2791                outgoing[source].push(target);
2792            }
2793        }
2794    }
2795    let mut indegree = vec![0usize; component_count];
2796    for edges in &mut outgoing {
2797        edges.sort_unstable();
2798        edges.dedup();
2799        for &target in edges.iter() {
2800            indegree[target] = indegree[target].checked_add(1)?;
2801        }
2802    }
2803
2804    let mut ready = BTreeSet::new();
2805    let mut ready_by_domain = HashMap::<usize, BTreeSet<(usize, usize)>>::default();
2806    for (component, degree) in indegree.iter().enumerate() {
2807        if *degree == 0 {
2808            let entry = (keys[component], component);
2809            ready.insert(entry);
2810            if let Some(domain) = component_domains[component] {
2811                ready_by_domain.entry(domain).or_default().insert(entry);
2812            }
2813        }
2814    }
2815    let mut components = sccs.into_iter().map(Some).collect::<Vec<_>>();
2816    let mut ordered = Vec::with_capacity(component_count);
2817    let mut active_domain = None;
2818    while !ready.is_empty() {
2819        let selected = active_domain
2820            .and_then(|domain| ready_by_domain.get(&domain)?.iter().next().copied())
2821            .or_else(|| ready.iter().next().copied())?;
2822        let (_, component) = selected;
2823        ready.remove(&selected);
2824        if let Some(domain) = component_domains[component] {
2825            ready_by_domain.get_mut(&domain)?.remove(&selected);
2826        }
2827        active_domain = component_domains[component];
2828        ordered.push(components[component].take()?);
2829        for &target in &outgoing[component] {
2830            indegree[target] = indegree[target].checked_sub(1)?;
2831            if indegree[target] == 0 {
2832                let entry = (keys[target], target);
2833                ready.insert(entry);
2834                if let Some(domain) = component_domains[target] {
2835                    ready_by_domain.entry(domain).or_default().insert(entry);
2836                }
2837            }
2838        }
2839    }
2840    (ordered.len() == component_count).then_some((ordered, component_by_path))
2841}
2842
2843fn logic_path_is_scheduling_barrier<Addr: Clone + Eq + Hash>(path: &LogicPath<Addr>) -> bool {
2844    matches!(path.target, LogicPathTarget::CombCaptureEvent { .. })
2845        || !path.comb_capture_enable_sites.is_empty()
2846}
2847
2848fn cached_logic_path_roots<Addr: Clone + Eq + Hash>(path: &LogicPath<Addr>) -> Vec<NodeId> {
2849    let mut roots = path
2850        .local_inputs
2851        .iter()
2852        .map(|(_, node)| *node)
2853        .collect::<Vec<_>>();
2854    if path.local_inputs.is_empty() && matches!(path.target, LogicPathTarget::Var(_)) {
2855        roots.extend(path.pre_lower_nodes.iter().copied());
2856        roots.push(path.expr);
2857    }
2858    roots
2859}
2860
2861/// Build the sparse incidence relation between LogicPaths and SLT values
2862/// which can actually survive in `lower_cache` across path boundaries.
2863///
2864/// A tree-only node cannot be independently reused: a cached ancestor hides
2865/// it from later lowering.  Only DAG joins (including roots named by multiple
2866/// paths) therefore become materialization tokens.  Constants are deliberately
2867/// omitted because target lowering can rematerialize them without retaining a
2868/// source register.
2869fn logic_path_materialization_tokens<Addr: Clone + Eq + Hash>(
2870    input: &[LogicPath<Addr>],
2871    arena: &SLTNodeArena<Addr>,
2872) -> (Vec<Vec<usize>>, Vec<usize>) {
2873    let mut references = vec![0usize; arena.len()];
2874    let mut children = Vec::new();
2875    for raw in 0..arena.len() {
2876        children.clear();
2877        push_scheduler_node_children(NodeId(raw), arena, &mut children);
2878        children.sort_unstable();
2879        children.dedup();
2880        for &child in &children {
2881            references[child.0] = references[child.0].saturating_add(1);
2882        }
2883    }
2884    for path in input {
2885        let mut roots = cached_logic_path_roots(path);
2886        roots.sort_unstable();
2887        roots.dedup();
2888        for root in roots {
2889            references[root.0] = references[root.0].saturating_add(1);
2890        }
2891    }
2892
2893    let candidates = references
2894        .iter()
2895        .enumerate()
2896        .map(|(raw, references)| {
2897            *references > 1 && !matches!(arena.get(NodeId(raw)), SLTNode::Constant(..))
2898        })
2899        .collect::<Vec<_>>();
2900    let mut raw_tokens = vec![Vec::<usize>::new(); input.len()];
2901    let mut token_users = vec![0usize; arena.len()];
2902    let mut visited = vec![0usize; arena.len()];
2903    let mut epoch = 0usize;
2904    let mut work = Vec::new();
2905    for (path_index, path) in input.iter().enumerate() {
2906        epoch = epoch.wrapping_add(1);
2907        if epoch == 0 {
2908            visited.fill(0);
2909            epoch = 1;
2910        }
2911        work.extend(cached_logic_path_roots(path));
2912        while let Some(node) = work.pop() {
2913            if visited[node.0] == epoch {
2914                continue;
2915            }
2916            visited[node.0] = epoch;
2917            if candidates[node.0] {
2918                raw_tokens[path_index].push(node.0);
2919            }
2920            push_scheduler_node_children(node, arena, &mut work);
2921        }
2922        raw_tokens[path_index].sort_unstable();
2923        raw_tokens[path_index].dedup();
2924        for &token in &raw_tokens[path_index] {
2925            token_users[token] = token_users[token].saturating_add(1);
2926        }
2927    }
2928
2929    let mut dense_by_raw = vec![usize::MAX; arena.len()];
2930    let mut weights = Vec::new();
2931    for (raw, users) in token_users.into_iter().enumerate() {
2932        if users < 2 {
2933            continue;
2934        }
2935        dense_by_raw[raw] = weights.len();
2936        let width = crate::get_width(NodeId(raw), arena);
2937        weights.push(width.div_ceil(64).max(1));
2938    }
2939    for row in &mut raw_tokens {
2940        row.retain(|raw| dense_by_raw[*raw] != usize::MAX);
2941        for token in row.iter_mut() {
2942            *token = dense_by_raw[*token];
2943        }
2944    }
2945    (raw_tokens, weights)
2946}
2947
2948fn exact_native_storage_store(width: usize) -> bool {
2949    width != 0 && width <= 64 && matches!(width.div_ceil(8), 1 | 2 | 4 | 8)
2950}
2951
2952fn static_store_avoids_native_rmw<Addr: Copy + Eq + Hash>(
2953    write: &VarAtomBase<Addr>,
2954    var_widths: &HashMap<Addr, usize>,
2955    unpacked_element_widths: &HashMap<Addr, usize>,
2956) -> bool {
2957    let Some(width) = write
2958        .access
2959        .msb
2960        .checked_sub(write.access.lsb)
2961        .and_then(|width| width.checked_add(1))
2962    else {
2963        return false;
2964    };
2965
2966    if let Some(&element_width) = unpacked_element_widths.get(&write.id) {
2967        // Unpacked layout selection happens after scheduling. Only a complete
2968        // byte-aligned native-width element avoids RMW in both packed and
2969        // independently strided layouts.
2970        return element_width != 0
2971            && width == element_width
2972            && write.access.lsb.is_multiple_of(element_width)
2973            && matches!(width, 8 | 16 | 32 | 64);
2974    }
2975
2976    let whole_object = write.access.lsb == 0
2977        && var_widths.get(&write.id).copied() == Some(width)
2978        && exact_native_storage_store(width);
2979    let aligned_native_range =
2980        write.access.lsb.is_multiple_of(8) && matches!(width, 8 | 16 | 32 | 64);
2981    whole_object || aligned_native_range
2982}
2983
2984fn direct_ff_write_ranges<Addr: Copy + Eq + Hash>(
2985    input: &[LogicPath<Addr>],
2986    summaries: &[FfAccessSummary<Addr>],
2987    plan: &FfCombSchedulePlan,
2988    retained: &HashSet<(usize, usize)>,
2989    ff_node_base: usize,
2990    var_widths: &HashMap<Addr, usize>,
2991    unpacked_element_widths: &HashMap<Addr, usize>,
2992) -> Vec<Vec<VarAtomBase<Addr>>> {
2993    summaries
2994        .iter()
2995        .enumerate()
2996        .map(|(writer, summary)| {
2997            let mut direct = summary
2998                .writes
2999                .iter()
3000                .map(|write| {
3001                    // Direct publication is only a throughput optimization.
3002                    // Dynamic destinations and static bitfields requiring a
3003                    // physical read-modify-write stay staged.
3004                    if summary.dynamic_writes.contains(&write.id)
3005                        || !static_store_avoids_native_rmw(
3006                            write,
3007                            var_widths,
3008                            unpacked_element_widths,
3009                        )
3010                    {
3011                        return false;
3012                    }
3013
3014                    // A read in the same always_ff action cannot be ordered
3015                    // before only one of its own Stores. Keep just the
3016                    // overlapping range staged so lowering reads pre-edge
3017                    // state without penalizing disjoint writes.
3018                    if summary
3019                        .reads
3020                        .iter()
3021                        .any(|read| read.id == write.id && read.access.overlaps(&write.access))
3022                    {
3023                        return false;
3024                    }
3025
3026                    let writer_node = ff_node_base + writer;
3027                    let comb_readers_proven = plan.comb_before_direct_write[writer]
3028                        .iter()
3029                        .filter(|&&reader| {
3030                            input[reader]
3031                                .sources
3032                                .iter()
3033                                .chain(&input[reader].previous_sources)
3034                                .any(|read| {
3035                                    read.id == write.id && read.access.overlaps(&write.access)
3036                                })
3037                        })
3038                        .all(|&reader| retained.contains(&(reader, writer_node)));
3039                    let ff_readers_proven = plan.ff_before_direct_write[writer]
3040                        .iter()
3041                        .filter(|&&reader| {
3042                            summaries[reader].reads.iter().any(|read| {
3043                                read.id == write.id && read.access.overlaps(&write.access)
3044                            })
3045                        })
3046                        .all(|&reader| retained.contains(&(ff_node_base + reader, writer_node)));
3047                    comb_readers_proven && ff_readers_proven
3048                })
3049                .collect::<Vec<_>>();
3050
3051            // Direct Stores publish immediately, while staged Stores publish
3052            // together in finish(). Mixing those destinations inside one
3053            // overlapping write component can reverse procedural last-write-
3054            // wins order. Propagate staging through the component while
3055            // leaving disjoint ranges eligible for direct publication.
3056            let mut staged = direct
3057                .iter()
3058                .enumerate()
3059                .filter_map(|(index, &direct)| (!direct).then_some(index))
3060                .collect::<Vec<_>>();
3061            while let Some(staged_index) = staged.pop() {
3062                let staged_write = summary.writes[staged_index];
3063                for (candidate, candidate_write) in summary.writes.iter().enumerate() {
3064                    if direct[candidate]
3065                        && staged_write.id == candidate_write.id
3066                        && staged_write.access.overlaps(&candidate_write.access)
3067                    {
3068                        direct[candidate] = false;
3069                        staged.push(candidate);
3070                    }
3071                }
3072            }
3073
3074            summary
3075                .writes
3076                .iter()
3077                .zip(direct)
3078                .filter_map(|(write, direct)| direct.then_some(write))
3079                .copied()
3080                .collect()
3081        })
3082        .collect()
3083}
3084
3085fn schedule_acyclic_path_region(
3086    paths: &[usize],
3087    dependencies: &LogicPathEdges,
3088    values: &LogicPathEdges,
3089    materialization_tokens: &[Vec<usize>],
3090    token_weights: &[usize],
3091    local_by_path: &mut [usize],
3092) -> Option<Vec<usize>> {
3093    for (local, &path) in paths.iter().enumerate() {
3094        if path >= local_by_path.len() || local_by_path[path] != usize::MAX {
3095            for &mapped in &paths[..local] {
3096                local_by_path[mapped] = usize::MAX;
3097            }
3098            return None;
3099        }
3100        local_by_path[path] = local;
3101    }
3102
3103    let result = (|| {
3104        let dependency_predecessors =
3105            MappedGraphRows::new(&dependencies.predecessors, paths, local_by_path).ok()?;
3106        let value_predecessors =
3107            MappedGraphRows::new(&values.predecessors, paths, local_by_path).ok()?;
3108        let tokens = MappedNodeRows::new(materialization_tokens, paths).ok()?;
3109        let successors = MappedGraphRows::new(&dependencies.users, paths, local_by_path).ok()?;
3110        let value_users = MappedGraphRows::new(&values.users, paths, local_by_path).ok()?;
3111        let order = schedule_min_live_values_and_tokens_with_mapped_rows(
3112            dependency_predecessors,
3113            value_predecessors,
3114            tokens,
3115            token_weights,
3116            successors,
3117            value_users,
3118        )
3119        .ok()?;
3120        Some(order.into_iter().map(|local| paths[local]).collect())
3121    })();
3122
3123    for &path in paths {
3124        local_by_path[path] = usize::MAX;
3125    }
3126    result
3127}
3128
3129/// Schedule maximal acyclic runs by MemoryDef/MemoryUse liveness. Loop SCCs
3130/// and observable events are synchronization boundaries, so no path is moved
3131/// across an iteration or externally visible effect.
3132fn schedule_logic_path_regions<Addr: Clone + Eq + Hash>(
3133    topological_sccs: Vec<Vec<usize>>,
3134    dependencies: &LogicPathEdges,
3135    values: &LogicPathEdges,
3136    materialization_tokens: &[Vec<usize>],
3137    token_weights: &[usize],
3138    input: &[LogicPath<Addr>],
3139) -> Option<Vec<ScheduledWork>> {
3140    if dependencies.users.len() != values.users.len()
3141        || dependencies.users.len() != materialization_tokens.len()
3142        || dependencies.users.len() < input.len()
3143    {
3144        return None;
3145    }
3146    let mut local_by_path = vec![usize::MAX; dependencies.users.len()];
3147    let mut result = Vec::with_capacity(topological_sccs.len());
3148    let mut pending = Vec::new();
3149    let flush = |pending: &mut Vec<usize>,
3150                 result: &mut Vec<ScheduledWork>,
3151                 local_by_path: &mut [usize]|
3152     -> Option<()> {
3153        if pending.is_empty() {
3154            return Some(());
3155        }
3156        let ordered = schedule_acyclic_path_region(
3157            pending,
3158            dependencies,
3159            values,
3160            materialization_tokens,
3161            token_weights,
3162            local_by_path,
3163        )?;
3164        result.extend(ordered.into_iter().map(|path| {
3165            if path < input.len() {
3166                ScheduledWork::CombPath(path)
3167            } else {
3168                ScheduledWork::Ff(path - input.len())
3169            }
3170        }));
3171        pending.clear();
3172        Some(())
3173    };
3174
3175    for scc in topological_sccs {
3176        if let [path] = scc.as_slice()
3177            && !dependencies.users[*path].contains(path)
3178            && (*path >= input.len() || !logic_path_is_scheduling_barrier(&input[*path]))
3179        {
3180            pending.push(*path);
3181        } else {
3182            flush(&mut pending, &mut result, &mut local_by_path)?;
3183            if let [path] = scc.as_slice() {
3184                if *path < input.len() {
3185                    result.push(ScheduledWork::CombPath(*path));
3186                } else {
3187                    result.push(ScheduledWork::Ff(*path - input.len()));
3188                }
3189            } else if scc.iter().all(|path| *path < input.len()) {
3190                result.push(ScheduledWork::CombScc(scc));
3191            } else {
3192                return None;
3193            }
3194        }
3195    }
3196    flush(&mut pending, &mut result, &mut local_by_path)?;
3197    Some(result)
3198}
3199
3200/// Schedules and transforms LogicPaths into Simulation Intermediate Representation (SIR).
3201///
3202/// This process performs:
3203/// 1. Dependency analysis to detect multiple drivers and combinational loops.
3204/// 2. SCC detection via Tarjan's algorithm.
3205/// 3. Scheduling based on two primary strategies:
3206///    - **Strategy A (Static Unrolling)**: For DAG parts or loops with small, predictable convergence bounds.
3207///    - **Strategy B (Dynamic Convergence)**: For complex SCCs or potential "True Loops", implementing
3208///      runtime oscillation detection and convergence-based repetition.
3209fn sort_impl<Addr: Clone + Eq + Ord + Hash + Debug + Copy + Display, E>(
3210    input: Vec<LogicPath<Addr>>,
3211    arena: &SLTNodeArena<Addr>,
3212    ignored_loops: &HashSet<(Addr, Addr)>,
3213    true_loops: &HashMap<(Addr, Addr), usize>,
3214    four_state: bool,
3215    var_widths: &HashMap<Addr, usize>,
3216    unpacked_element_widths: &HashMap<Addr, usize>,
3217    first_runtime_error_code: i64,
3218    mut ff: Option<&mut dyn ClockFfLowering<Addr, Error = E>>,
3219) -> Result<ScheduleResult<Addr>, ClockSortError<Addr, E>> {
3220    let (input, ff_plan) = if let Some(ff_lowering) = ff.as_deref() {
3221        let mut plan = plan_ff_comb_schedule(&input, ff_lowering.summaries())
3222            .map_err(ClockSortError::Scheduler)?;
3223        let mut old_to_new = vec![usize::MAX; input.len()];
3224        let mut filtered = Vec::with_capacity(
3225            plan.required_comb
3226                .iter()
3227                .filter(|required| **required)
3228                .count(),
3229        );
3230        for (old, path) in input.into_iter().enumerate() {
3231            if plan.required_comb[old] {
3232                old_to_new[old] = filtered.len();
3233                filtered.push(path);
3234            }
3235        }
3236        let remap_paths = |paths: &mut Vec<usize>| {
3237            for path in paths.iter_mut() {
3238                *path = old_to_new[*path];
3239            }
3240            paths.retain(|path| *path != usize::MAX);
3241            paths.sort_unstable();
3242            paths.dedup();
3243        };
3244        for paths in plan
3245            .comb_value_predecessors
3246            .iter_mut()
3247            .chain(&mut plan.comb_before_direct_write)
3248        {
3249            remap_paths(paths);
3250        }
3251        plan.required_comb = vec![true; filtered.len()];
3252        (filtered, Some(plan))
3253    } else {
3254        (input, None)
3255    };
3256
3257    // 1. Build bit-range MemoryDefs/MemoryUses and their semantic graph.
3258    // This is the source dataflow adapter; interval indexing and DAG
3259    // scheduling remain IR-independent in celox-analysis.
3260    let n = input.len();
3261    let (mut materialization_tokens, token_weights) =
3262        logic_path_materialization_tokens(&input, arena);
3263    let LogicPathMemorySsa {
3264        mut dependencies,
3265        mut values,
3266    } = build_logic_path_memory_ssa(&input).map_err(ClockSortError::Scheduler)?;
3267
3268    // FF actions are ordinary sinks in the existing comb dependency graph.
3269    // The added nodes go through the same SCC, topological, and ready-queue
3270    // scheduling as LogicPaths; there is no second clock-specific scheduler.
3271    let ff_count = ff
3272        .as_deref()
3273        .map_or(0, |lowering| lowering.summaries().len());
3274    let mut direct_ff_writes_by_action = vec![Vec::new(); ff_count];
3275    if let Some(plan) = &ff_plan {
3276        dependencies.resize(n + ff_count);
3277        values.resize(n + ff_count);
3278        materialization_tokens.resize_with(n + ff_count, Vec::new);
3279        for (ff_index, predecessors) in plan.comb_value_predecessors.iter().enumerate() {
3280            let ff_node = n + ff_index;
3281            for &definition in predecessors {
3282                dependencies.push(definition, ff_node);
3283                values.push(definition, ff_node);
3284            }
3285        }
3286        let write_order_edges = plan
3287            .comb_before_direct_write
3288            .iter()
3289            .zip(&plan.ff_before_direct_write)
3290            .enumerate()
3291            .map(|(writer, (comb_readers, ff_readers))| {
3292                comb_readers
3293                    .iter()
3294                    .copied()
3295                    .map(|reader| (reader, n + writer))
3296                    .chain(
3297                        ff_readers
3298                            .iter()
3299                            .copied()
3300                            .map(|reader| (n + reader, n + writer)),
3301                    )
3302                    .collect::<Vec<_>>()
3303            })
3304            .collect::<Vec<_>>();
3305        let retained = add_acyclic_ff_write_order_edges(
3306            &mut dependencies.users,
3307            write_order_edges.iter().flatten().copied(),
3308        );
3309        for &(predecessor, user) in &retained {
3310            dependencies.predecessors[user].push(predecessor);
3311        }
3312        direct_ff_writes_by_action = direct_ff_write_ranges(
3313            &input,
3314            ff.as_deref()
3315                .expect("an FF schedule plan requires an FF lowering callback")
3316                .summaries(),
3317            plan,
3318            &retained,
3319            n,
3320            var_widths,
3321            unpacked_element_widths,
3322        );
3323        dependencies.canonicalize();
3324        values.canonicalize();
3325    }
3326    let direct_ff_writes = direct_ff_writes_by_action
3327        .iter()
3328        .flatten()
3329        .copied()
3330        .collect();
3331
3332    // 2. SCC extraction identifies the synchronization boundaries at which
3333    // the dataflow equations must be iterated rather than freely reordered.
3334    let mut ctx = TarjanContext {
3335        index: 0,
3336        stack: Vec::new(),
3337        on_stack: HashSet::default(),
3338        indices: vec![None; n + ff_count],
3339        lowlink: vec![None; n + ff_count],
3340        sccs: Vec::new(),
3341    };
3342    for i in 0..(n + ff_count) {
3343        if ctx.indices[i].is_none() {
3344            strong_connect(i, &dependencies.users, &mut ctx);
3345        }
3346    }
3347    let fold_group_schedule_index = build_fold_group_schedule_index(&input, arena);
3348    let mut path_domains = logic_path_scheduling_domains(&input, &fold_group_schedule_index);
3349    path_domains.resize(n + ff_count, None);
3350    let (topological_sccs, component_by_path) =
3351        stable_topological_sccs(ctx.sccs, &dependencies.users, &path_domains).ok_or(
3352            ClockSortError::Scheduler(SchedulerError::InvalidDependencyGraph),
3353        )?;
3354    let scheduled_work = schedule_logic_path_regions(
3355        topological_sccs,
3356        &dependencies,
3357        &values,
3358        &materialization_tokens,
3359        &token_weights,
3360        &input,
3361    )
3362    .ok_or(ClockSortError::Scheduler(
3363        SchedulerError::InvalidDependencyGraph,
3364    ))?;
3365    let adj = dependencies.users;
3366    drop(dependencies.predecessors);
3367    drop(values);
3368    let scheduled_work = form_scheduled_guard_regions(scheduled_work, &input, arena, four_state);
3369
3370    let mut builder = SIRBuilder::new();
3371    if let Some(ff_lowering) = ff.as_deref_mut() {
3372        ff_lowering
3373            .begin(&mut builder, &direct_ff_writes_by_action)
3374            .map_err(ClockSortError::Lowering)?;
3375    }
3376    let lowerer = crate::SLTToSIRLowerer::new(four_state)
3377        .with_unpacked_input_types(arena, unpacked_element_widths);
3378
3379    let mut lower_cache = HashMap::default();
3380    let mut dep_memo = HashMap::default();
3381    let mut inverse_dep_memo = HashMap::default();
3382
3383    const UNROLL_THRESHOLD: usize = 32;
3384
3385    // Helper: Emits SIR for a logic path and manages the lowering cache.
3386    // lowerer.lower allocates registers and emits instructions for sub-expressions.
3387    let emit_node = |builder: &mut SIRBuilder<Addr>,
3388                     idx: usize,
3389                     lower_cache: &mut HashMap<NodeId, RegisterId>,
3390                     dep_memo: &mut HashMap<NodeId, HashSet<Addr>>,
3391                     inverse_dep_memo: &mut HashMap<Addr, HashSet<NodeId>>| {
3392        let path = &input[idx];
3393
3394        collect_logic_path_input_deps(path, arena, dep_memo, inverse_dep_memo);
3395        emit_logic_path_store(
3396            &lowerer,
3397            builder,
3398            path,
3399            arena,
3400            lower_cache,
3401            unpacked_element_widths,
3402        );
3403        invalidate_logic_path_target(path, inverse_dep_memo, lower_cache);
3404    };
3405    // Maximum blocks in a single EU before flushing to a new one.
3406    // This prevents Cranelift from choking on massive functions.
3407    const EU_BLOCK_LIMIT: usize = 20_000;
3408
3409    let mut result_eus: Vec<ExecutionUnit<Addr>> = Vec::new();
3410    let mut runtime_errors: HashMap<i64, RuntimeErrorInfo<Addr>> = HashMap::default();
3411    let mut next_runtime_error_code = first_runtime_error_code;
3412
3413    // Pairwise joint-fold profitability is intentionally local. Keep all
3414    // projections of one packed root together, but never build an unbounded
3415    // compatibility matrix for a long run of independent roots.
3416    const MAX_JOINT_FOLD_ROOTS: usize = 16;
3417    let mut pending_fold_indices: Vec<usize> = Vec::new();
3418    let mut pending_fold_roots = HashSet::default();
3419
3420    // 4. Lower each scheduled component, selecting static unrolling or
3421    // dynamic convergence for cyclic SCCs.
3422    for work in scheduled_work {
3423        let singleton;
3424        let scc = match &work {
3425            ScheduledWork::CombPath(path) => {
3426                singleton = [*path];
3427                singleton.as_slice()
3428            }
3429            ScheduledWork::CombScc(scc) => scc.as_slice(),
3430            ScheduledWork::GuardedComb { condition, paths } => {
3431                flush_pending_fold_paths(
3432                    &mut pending_fold_indices,
3433                    &input,
3434                    &fold_group_schedule_index,
3435                    &lowerer,
3436                    &mut builder,
3437                    arena,
3438                    &mut lower_cache,
3439                    &mut dep_memo,
3440                    &mut inverse_dep_memo,
3441                    unpacked_element_widths,
3442                    four_state,
3443                );
3444                pending_fold_roots.clear();
3445                emit_scheduled_guard_region(
3446                    &lowerer,
3447                    &mut builder,
3448                    *condition,
3449                    paths,
3450                    &input,
3451                    arena,
3452                    &mut lower_cache,
3453                    &mut dep_memo,
3454                    &mut inverse_dep_memo,
3455                    unpacked_element_widths,
3456                );
3457                continue;
3458            }
3459            ScheduledWork::Ff(index) => {
3460                flush_pending_fold_paths(
3461                    &mut pending_fold_indices,
3462                    &input,
3463                    &fold_group_schedule_index,
3464                    &lowerer,
3465                    &mut builder,
3466                    arena,
3467                    &mut lower_cache,
3468                    &mut dep_memo,
3469                    &mut inverse_dep_memo,
3470                    unpacked_element_widths,
3471                    four_state,
3472                );
3473                pending_fold_roots.clear();
3474                ff.as_deref_mut()
3475                    .ok_or(ClockSortError::Scheduler(
3476                        SchedulerError::InvalidDependencyGraph,
3477                    ))?
3478                    .lower(*index, &direct_ff_writes_by_action[*index], &mut builder)
3479                    .map_err(ClockSortError::Lowering)?;
3480                // Any directly published range is ordered after all of its
3481                // old-state readers. Other FF state stays invisible in
3482                // WORKING/SPARSE_WORKING until the final publish, so values
3483                // computed before the FF CFG keep the ordinary lowering cache
3484                // valid.
3485                continue;
3486            }
3487        };
3488        let component = component_by_path[scc[0]];
3489        let mut user_safety_limit = None;
3490        for &v_idx in scc {
3491            for &u_idx in &adj[v_idx] {
3492                if component_by_path[u_idx] == component {
3493                    if let (Some(v_target), Some(u_target)) =
3494                        (input[v_idx].target.var(), input[u_idx].target.var())
3495                    {
3496                        let edge = (v_target.id, u_target.id);
3497                        if let Some(&limit) = true_loops.get(&edge) {
3498                            user_safety_limit =
3499                                Some(user_safety_limit.map_or(limit, |l: usize| l.max(limit)));
3500                        }
3501                    }
3502                }
3503            }
3504        }
3505        let is_loop = scc.len() > 1 || (scc.len() == 1 && adj[scc[0]].contains(&scc[0]));
3506
3507        if is_loop {
3508            // Exact grouped folds are atomic with respect to a loop SCC.
3509            flush_pending_fold_paths(
3510                &mut pending_fold_indices,
3511                &input,
3512                &fold_group_schedule_index,
3513                &lowerer,
3514                &mut builder,
3515                arena,
3516                &mut lower_cache,
3517                &mut dep_memo,
3518                &mut inverse_dep_memo,
3519                unpacked_element_widths,
3520                four_state,
3521            );
3522            pending_fold_roots.clear();
3523            let mut authorized = user_safety_limit.is_some();
3524            'check_scc: for &v_idx in scc {
3525                for &u_idx in &adj[v_idx] {
3526                    if component_by_path[u_idx] == component
3527                        && input[v_idx]
3528                            .target
3529                            .var()
3530                            .zip(input[u_idx].target.var())
3531                            .is_some_and(|(v, u)| ignored_loops.contains(&(v.id, u.id)))
3532                    {
3533                        // Some loops are explicitly allowed by the user (e.g., false loops).
3534                        authorized = true;
3535                        break 'check_scc;
3536                    }
3537                }
3538            }
3539
3540            if !authorized {
3541                return Err(ClockSortError::Scheduler(
3542                    SchedulerError::CombinationalLoop {
3543                        blocks: scc.iter().map(|idx| input[*idx].clone()).collect(),
3544                    },
3545                ));
3546            }
3547
3548            // FAS Sort
3549            let optimized_scc_order = greedy_fas_sort(scc, &adj);
3550            let force_strategy_b = user_safety_limit.is_some();
3551            let iterations = calculate_required_iterations(&adj, &optimized_scc_order);
3552            let total_ops_estimate = optimized_scc_order.len().saturating_mul(iterations);
3553            if !force_strategy_b && total_ops_estimate <= UNROLL_THRESHOLD {
3554                // Strategy A: Static Unrolling
3555                // The loop is unrolled a fixed number of times based on structural dependency depth (iterations).
3556                for _ in 0..iterations {
3557                    for &idx in &optimized_scc_order {
3558                        emit_node(
3559                            &mut builder,
3560                            idx,
3561                            &mut lower_cache,
3562                            &mut dep_memo,
3563                            &mut inverse_dep_memo,
3564                        );
3565                    }
3566                }
3567            } else {
3568                let runtime_error_code = next_runtime_error_code;
3569                next_runtime_error_code += 1;
3570                let mut seen = HashSet::default();
3571                let sources = scc
3572                    .iter()
3573                    .filter_map(|idx| {
3574                        let addr = input[*idx].target.var()?.id;
3575                        seen.insert(addr).then_some(addr)
3576                    })
3577                    .collect::<Vec<_>>();
3578                runtime_errors.insert(
3579                    runtime_error_code,
3580                    RuntimeErrorInfo {
3581                        message: "Detected True Loop".to_string(),
3582                        signals: sources,
3583                    },
3584                );
3585
3586                // Strategy B: Dynamic Convergence
3587                // Implements a runtime loop that continues executing the SCC until all signals converge (dirty flag is false).
3588                // Includes a safety limit to detect non-converging "True Loops" and avoid infinite hang.
3589
3590                // 1. Determine the runtime repetition limit.
3591                let safety_limit = user_safety_limit.unwrap_or(iterations + 1);
3592
3593                // 2. Prepare Constants and Counters
3594                let zero_reg = builder.alloc_bit(64, false);
3595                builder.emit(SIRInstruction::Imm(zero_reg, SIRValue::new(0u64)));
3596
3597                let limit_reg = builder.alloc_bit(64, false);
3598                builder.emit(SIRInstruction::Imm(
3599                    limit_reg,
3600                    SIRValue::new(safety_limit as u64),
3601                ));
3602
3603                // 3. Blocks
3604                let current_counter = builder.alloc_bit(64, false);
3605                let header_block = builder.new_block_with(vec![current_counter]); // [counter]
3606                let body_block = builder.new_block();
3607                let exit_block = builder.new_block();
3608                let error_block = builder.new_block(); // For True Loop detection
3609
3610                // Start: Jump to header with counter = 0
3611                builder.seal_block(SIRTerminator::Jump(header_block, vec![zero_reg]));
3612
3613                // --- Header Block ---
3614                builder.switch_to_block(header_block);
3615
3616                // Check: counter < safety_limit
3617                let can_continue_reg = builder.alloc_bit(1, false);
3618                builder.emit(SIRInstruction::Binary(
3619                    can_continue_reg,
3620                    current_counter,
3621                    BinaryOp::LtU,
3622                    limit_reg,
3623                ));
3624
3625                // If counter exceeded limit, we might have an oscillating True Loop
3626                builder.seal_block(SIRTerminator::Branch {
3627                    cond: can_continue_reg,
3628                    true_block: (body_block, vec![]),
3629                    false_block: (error_block, vec![]),
3630                });
3631                builder.switch_to_block(body_block);
3632                let mut current_dirty_reg = builder.alloc_bit(1, false);
3633                builder.emit(SIRInstruction::Imm(current_dirty_reg, SIRValue::new(0u32)));
3634                for &idx in &optimized_scc_order {
3635                    let path = &input[idx];
3636                    let Some(target) = path.target.var() else {
3637                        emit_logic_path_store(
3638                            &lowerer,
3639                            &mut builder,
3640                            path,
3641                            arena,
3642                            &mut lower_cache,
3643                            unpacked_element_widths,
3644                        );
3645                        continue;
3646                    };
3647                    let width = 1 + target.access.msb - target.access.lsb;
3648                    let addr = target.id;
3649                    let offset =
3650                        static_access_offset(&target.id, target.access, unpacked_element_widths);
3651
3652                    // --- Dynamic Convergence Check Logic ---
3653                    // For each node in the SCC, we verify if its value changed after this iteration.
3654                    //
3655                    // a. Load the current value (pre-update benchmark)
3656                    let old_val_reg = builder.alloc_bit(width, false);
3657                    builder.emit(SIRInstruction::Load(
3658                        old_val_reg,
3659                        addr,
3660                        offset.clone(),
3661                        width,
3662                    ));
3663                    collect_logic_path_input_deps(
3664                        path,
3665                        arena,
3666                        &mut dep_memo,
3667                        &mut inverse_dep_memo,
3668                    );
3669                    // b. Compute the new value
3670                    let new_val_reg = lower_logic_path_expr(
3671                        &lowerer,
3672                        &mut builder,
3673                        path,
3674                        arena,
3675                        &mut lower_cache,
3676                    );
3677
3678                    // c. Compare: changed = (old != new)
3679                    let is_changed_reg = builder.alloc_bit(1, false);
3680                    builder.emit(SIRInstruction::Binary(
3681                        is_changed_reg,
3682                        old_val_reg,
3683                        BinaryOp::Ne, // Not Equal
3684                        new_val_reg,
3685                    ));
3686                    let new_dirty_reg = builder.alloc_bit(1, false);
3687
3688                    // d. Accumulate dirty flag: dirty = dirty | is_changed
3689                    // If any signal in the SCC changes, the entire SCC requires another iteration.
3690                    builder.emit(SIRInstruction::Binary(
3691                        new_dirty_reg,
3692                        current_dirty_reg,
3693                        BinaryOp::Or,
3694                        is_changed_reg,
3695                    ));
3696                    current_dirty_reg = new_dirty_reg;
3697                    // e. Store the new value
3698                    builder.emit(SIRInstruction::Store(
3699                        addr,
3700                        offset,
3701                        width,
3702                        new_val_reg,
3703                        Vec::new(),
3704                        Vec::new(),
3705                    ));
3706                    if !path.comb_capture_enable_sites.is_empty() {
3707                        let (old, new) = if path.comb_capture_enable_always {
3708                            let old = builder.alloc_bit(1, false);
3709                            let new = builder.alloc_bit(1, false);
3710                            builder.emit(SIRInstruction::Imm(old, SIRValue::new(0u8)));
3711                            builder.emit(SIRInstruction::Imm(new, SIRValue::new(1u8)));
3712                            (old, new)
3713                        } else {
3714                            (old_val_reg, new_val_reg)
3715                        };
3716                        builder.emit(SIRInstruction::CombCaptureEnableIfChanged {
3717                            old,
3718                            new,
3719                            sites: path.comb_capture_enable_sites.clone(),
3720                        });
3721                    }
3722                    if let Some(to_remove) = inverse_dep_memo.get(&addr) {
3723                        for node in to_remove {
3724                            lower_cache.remove(node);
3725                        }
3726                    }
3727                    // -------------------------------
3728                }
3729
3730                // 4. Branch: Loop if dirty
3731                let one_reg = builder.alloc_bit(64, false);
3732                builder.emit(SIRInstruction::Imm(one_reg, SIRValue::new(1u64)));
3733                let next_counter = builder.alloc_bit(64, false);
3734                builder.emit(SIRInstruction::Binary(
3735                    next_counter,
3736                    current_counter,
3737                    BinaryOp::Add,
3738                    one_reg,
3739                ));
3740
3741                // Increment the iteration counter and branch.
3742                // If 'dirty' is true, return to the header block; otherwise, exit the loop.
3743                builder.seal_block(SIRTerminator::Branch {
3744                    cond: current_dirty_reg,
3745                    true_block: (header_block, vec![next_counter]),
3746                    false_block: (exit_block, vec![]),
3747                });
3748
3749                // --- Error/Exit Blocks ---
3750                builder.switch_to_block(error_block);
3751                // Emit a trap or special instruction to indicate "Combinational Loop Oscillation"
3752                // builder.emit(SIRInstruction::Trap(1));
3753                builder.seal_block(SIRTerminator::Error(runtime_error_code));
3754
3755                // 5. Exit Block
3756                builder.switch_to_block(exit_block);
3757            }
3758        } else {
3759            // DAG Part — flush before emitting if the EU has grown too large
3760            if ff.is_none() && builder.block_count() >= EU_BLOCK_LIMIT {
3761                flush_pending_fold_paths(
3762                    &mut pending_fold_indices,
3763                    &input,
3764                    &fold_group_schedule_index,
3765                    &lowerer,
3766                    &mut builder,
3767                    arena,
3768                    &mut lower_cache,
3769                    &mut dep_memo,
3770                    &mut inverse_dep_memo,
3771                    unpacked_element_widths,
3772                    four_state,
3773                );
3774                pending_fold_roots.clear();
3775                if let Some(eu) = builder.flush_eu() {
3776                    result_eus.push(eu);
3777                    // Clear the lowering cache — register IDs are EU-scoped
3778                    lower_cache.clear();
3779                }
3780            }
3781
3782            let idx = scc[0];
3783            let exact_fold = is_exact_fold_path(idx, &fold_group_schedule_index);
3784            let fold_root = exact_fold
3785                .then_some(fold_group_schedule_index.direct_group_by_path[idx])
3786                .flatten();
3787            let depends_on_pending = pending_fold_indices
3788                .iter()
3789                .any(|pending| adj[*pending].binary_search(&idx).is_ok());
3790            let starts_new_root = fold_root.is_some_and(|root| !pending_fold_roots.contains(&root));
3791            let window_full = starts_new_root && pending_fold_roots.len() >= MAX_JOINT_FOLD_ROOTS;
3792            if exact_fold && !depends_on_pending && !window_full {
3793                pending_fold_indices.push(idx);
3794                pending_fold_roots.extend(fold_root);
3795            } else {
3796                flush_pending_fold_paths(
3797                    &mut pending_fold_indices,
3798                    &input,
3799                    &fold_group_schedule_index,
3800                    &lowerer,
3801                    &mut builder,
3802                    arena,
3803                    &mut lower_cache,
3804                    &mut dep_memo,
3805                    &mut inverse_dep_memo,
3806                    unpacked_element_widths,
3807                    four_state,
3808                );
3809                pending_fold_roots.clear();
3810                if exact_fold {
3811                    pending_fold_indices.push(idx);
3812                    pending_fold_roots.extend(fold_root);
3813                } else {
3814                    emit_node(
3815                        &mut builder,
3816                        idx,
3817                        &mut lower_cache,
3818                        &mut dep_memo,
3819                        &mut inverse_dep_memo,
3820                    );
3821                }
3822            }
3823        }
3824    }
3825
3826    // Flush the final exact grouped-fold run after the SCC loop.
3827    flush_pending_fold_paths(
3828        &mut pending_fold_indices,
3829        &input,
3830        &fold_group_schedule_index,
3831        &lowerer,
3832        &mut builder,
3833        arena,
3834        &mut lower_cache,
3835        &mut dep_memo,
3836        &mut inverse_dep_memo,
3837        unpacked_element_widths,
3838        four_state,
3839    );
3840    pending_fold_roots.clear();
3841
3842    if let Some(ff_lowering) = ff {
3843        ff_lowering
3844            .finish(&mut builder, &direct_ff_writes_by_action)
3845            .map_err(ClockSortError::Lowering)?;
3846    }
3847    builder.seal_block(SIRTerminator::Return);
3848    let (blocks, reg_map, _) = builder.drain();
3849    result_eus.push(ExecutionUnit {
3850        entry_block_id: BlockId(0),
3851        blocks,
3852        register_map: reg_map,
3853    });
3854    Ok(ScheduleResult {
3855        execution_units: result_eus,
3856        runtime_errors,
3857        direct_ff_writes,
3858    })
3859}
3860
3861#[cfg(test)]
3862pub fn sort<Addr: Clone + Eq + Ord + Hash + Debug + Copy + Display>(
3863    input: Vec<LogicPath<Addr>>,
3864    arena: &SLTNodeArena<Addr>,
3865    ignored_loops: &HashSet<(Addr, Addr)>,
3866    true_loops: &HashMap<(Addr, Addr), usize>,
3867    four_state: bool,
3868    var_widths: &HashMap<Addr, usize>,
3869    first_runtime_error_code: i64,
3870) -> Result<ScheduleResult<Addr>, SchedulerError<Addr>> {
3871    sort_with_unpacked_element_widths(
3872        input,
3873        arena,
3874        ignored_loops,
3875        true_loops,
3876        four_state,
3877        var_widths,
3878        &HashMap::default(),
3879        first_runtime_error_code,
3880    )
3881}
3882
3883pub fn sort_with_unpacked_element_widths<Addr: Clone + Eq + Ord + Hash + Debug + Copy + Display>(
3884    input: Vec<LogicPath<Addr>>,
3885    arena: &SLTNodeArena<Addr>,
3886    ignored_loops: &HashSet<(Addr, Addr)>,
3887    true_loops: &HashMap<(Addr, Addr), usize>,
3888    four_state: bool,
3889    var_widths: &HashMap<Addr, usize>,
3890    unpacked_element_widths: &HashMap<Addr, usize>,
3891    first_runtime_error_code: i64,
3892) -> Result<ScheduleResult<Addr>, SchedulerError<Addr>> {
3893    match sort_impl::<Addr, std::convert::Infallible>(
3894        input,
3895        arena,
3896        ignored_loops,
3897        true_loops,
3898        four_state,
3899        var_widths,
3900        unpacked_element_widths,
3901        first_runtime_error_code,
3902        None,
3903    ) {
3904        Ok(result) => Ok(result),
3905        Err(ClockSortError::Scheduler(error)) => Err(error),
3906        Err(ClockSortError::Lowering(_)) => {
3907            unreachable!("ordinary comb scheduling has no FF lowering callback")
3908        }
3909    }
3910}
3911
3912pub fn sort_clock<Addr: Clone + Eq + Ord + Hash + Debug + Copy + Display, E>(
3913    input: Vec<LogicPath<Addr>>,
3914    arena: &SLTNodeArena<Addr>,
3915    ignored_loops: &HashSet<(Addr, Addr)>,
3916    true_loops: &HashMap<(Addr, Addr), usize>,
3917    four_state: bool,
3918    var_widths: &HashMap<Addr, usize>,
3919    unpacked_element_widths: &HashMap<Addr, usize>,
3920    first_runtime_error_code: i64,
3921    ff: &mut dyn ClockFfLowering<Addr, Error = E>,
3922) -> Result<ScheduleResult<Addr>, ClockSortError<Addr, E>> {
3923    sort_impl(
3924        input,
3925        arena,
3926        ignored_loops,
3927        true_loops,
3928        four_state,
3929        var_widths,
3930        unpacked_element_widths,
3931        first_runtime_error_code,
3932        Some(ff),
3933    )
3934}
3935
3936#[cfg(test)]
3937mod tests {
3938    use num_bigint::{BigInt, BigUint};
3939
3940    use super::{
3941        ExactFoldGroup, ExactIndexedLoadKey, FfAccessSummary, FoldGroupReadFacts,
3942        NormalizedIndexExpr, add_acyclic_ff_write_order_edges, best_weighted_fold_family,
3943        build_fold_group_schedule_index, build_logic_path_memory_ssa, collect_node_input_deps,
3944        direct_ff_write_ranges, plan_ff_comb_schedule, prepare_atomic_fold_group_results, sort,
3945        stable_topological_sccs,
3946    };
3947    use crate::{HashMap, HashSet};
3948    use crate::{
3949        LogicPath, LogicPathTarget, SLTForFoldGroupState, SLTNode, SLTNodeArena, SLTToSIRLowerer,
3950    };
3951    use celox_design::{BinaryOp, BitAccess, VarAtomBase};
3952    use celox_sir::{SIRBuilder, SIRInstruction, SIRTerminator};
3953
3954    #[test]
3955    fn stable_scc_order_preserves_source_order_when_dependencies_allow_it() {
3956        // 0 and 1 are initially ready, both feed 2, and 2 feeds 3. Tarjan is
3957        // free to return components in any order; lowering remains stable.
3958        let adj = vec![vec![2], vec![2], vec![3], vec![]];
3959        let unordered = vec![vec![3], vec![1], vec![2], vec![0]];
3960
3961        let (ordered, component_by_path) =
3962            stable_topological_sccs(unordered, &adj, &[None; 4]).unwrap();
3963
3964        assert_eq!(ordered, vec![vec![0], vec![1], vec![2], vec![3]]);
3965        assert_eq!(component_by_path.len(), adj.len());
3966    }
3967
3968    #[test]
3969    fn stable_scc_order_drains_only_ready_members_of_an_effect_domain() {
3970        let adj = vec![vec![], vec![], vec![], vec![]];
3971        let unordered = vec![vec![3], vec![1], vec![2], vec![0]];
3972
3973        let (ordered, _) =
3974            stable_topological_sccs(unordered, &adj, &[Some(0), Some(1), Some(0), Some(1)])
3975                .unwrap();
3976
3977        assert_eq!(ordered, vec![vec![0], vec![2], vec![1], vec![3]]);
3978    }
3979
3980    #[test]
3981    fn effect_domain_preference_never_pulls_a_path_across_its_dependency() {
3982        // Path 1 shares path 0's destination domain, but it is not ready until
3983        // the independent path 2 has executed.
3984        let adj = vec![vec![], vec![], vec![1]];
3985        let unordered = vec![vec![1], vec![2], vec![0]];
3986
3987        let (ordered, _) =
3988            stable_topological_sccs(unordered, &adj, &[Some(0), Some(0), Some(1)]).unwrap();
3989
3990        assert_eq!(ordered, vec![vec![0], vec![2], vec![1]]);
3991    }
3992
3993    fn simple_path(
3994        arena: &mut SLTNodeArena<u32>,
3995        target: u32,
3996        source: Option<u32>,
3997    ) -> LogicPath<u32> {
3998        let expr = if let Some(source) = source {
3999            arena
4000                .alloc(SLTNode::Input {
4001                    variable: source,
4002                    signed: false,
4003                    index: Vec::new(),
4004                    access: BitAccess::new(0, 7),
4005                })
4006                .unwrap()
4007        } else {
4008            arena
4009                .alloc(SLTNode::Constant(
4010                    BigUint::from(target),
4011                    BigUint::from(0u8),
4012                    8,
4013                    false,
4014                ))
4015                .unwrap()
4016        };
4017        LogicPath {
4018            target: LogicPathTarget::Var(VarAtomBase::new(target, 0, 7)),
4019            sources: source
4020                .map(|source| [VarAtomBase::new(source, 0, 7)].into_iter().collect())
4021                .unwrap_or_default(),
4022            previous_sources: crate::HashSet::default(),
4023            address_sources: crate::HashSet::default(),
4024            local_inputs: Vec::new(),
4025            order_before: crate::HashSet::default(),
4026            comb_capture_enable_sites: Vec::new(),
4027            comb_capture_enable_always: false,
4028            pre_lower_nodes: Vec::new(),
4029            expr,
4030        }
4031    }
4032
4033    #[test]
4034    fn scheduled_paths_share_one_high_level_guard_without_changing_store_order() {
4035        let mut arena = SLTNodeArena::new();
4036        let condition = arena
4037            .alloc(SLTNode::Input {
4038                variable: 100,
4039                signed: false,
4040                index: Vec::new(),
4041                access: BitAccess::new(0, 0),
4042            })
4043            .unwrap();
4044        let mut paths = Vec::new();
4045        for target in 0u32..12 {
4046            let then_expr = arena
4047                .alloc(SLTNode::Constant(
4048                    BigUint::from(target),
4049                    BigUint::from(0u8),
4050                    8,
4051                    false,
4052                ))
4053                .unwrap();
4054            let else_expr = arena
4055                .alloc(SLTNode::Constant(
4056                    BigUint::from(target + 32),
4057                    BigUint::from(0u8),
4058                    8,
4059                    false,
4060                ))
4061                .unwrap();
4062            let expr = arena
4063                .alloc(SLTNode::Mux {
4064                    cond: condition,
4065                    then_expr,
4066                    else_expr,
4067                })
4068                .unwrap();
4069            paths.push(LogicPath {
4070                target: LogicPathTarget::Var(VarAtomBase::new(target, 0, 7)),
4071                sources: [VarAtomBase::new(100, 0, 0)].into_iter().collect(),
4072                previous_sources: crate::HashSet::default(),
4073                address_sources: crate::HashSet::default(),
4074                local_inputs: Vec::new(),
4075                order_before: crate::HashSet::default(),
4076                comb_capture_enable_sites: Vec::new(),
4077                comb_capture_enable_always: false,
4078                pre_lower_nodes: Vec::new(),
4079                expr,
4080            });
4081        }
4082
4083        let result = sort(
4084            paths,
4085            &arena,
4086            &crate::HashSet::default(),
4087            &crate::HashMap::default(),
4088            false,
4089            &crate::HashMap::default(),
4090            0,
4091        )
4092        .unwrap();
4093        let eu = &result.execution_units[0];
4094        let branches = eu
4095            .blocks
4096            .values()
4097            .filter(|block| matches!(block.terminator, SIRTerminator::Branch { .. }))
4098            .collect::<Vec<_>>();
4099        assert_eq!(branches.len(), 1);
4100        assert_eq!(
4101            eu.blocks
4102                .values()
4103                .flat_map(|block| &block.instructions)
4104                .filter(|instruction| matches!(instruction, SIRInstruction::Mux(..)))
4105                .count(),
4106            0
4107        );
4108
4109        let SIRTerminator::Branch {
4110            true_block,
4111            false_block,
4112            ..
4113        } = &branches[0].terminator
4114        else {
4115            unreachable!()
4116        };
4117        for block in [true_block.0, false_block.0] {
4118            let stores = eu.blocks[&block]
4119                .instructions
4120                .iter()
4121                .filter_map(|instruction| match instruction {
4122                    SIRInstruction::Store(target, ..) => Some(*target),
4123                    _ => None,
4124                })
4125                .collect::<Vec<_>>();
4126            assert_eq!(stores, (0..12).collect::<Vec<_>>());
4127        }
4128    }
4129
4130    #[test]
4131    fn scheduled_outer_guard_keeps_nested_mux_lowering_decisions() {
4132        let mut arena = SLTNodeArena::new();
4133        let outer_condition = arena
4134            .alloc(SLTNode::Input {
4135                variable: 100,
4136                signed: false,
4137                index: Vec::new(),
4138                access: BitAccess::new(0, 0),
4139            })
4140            .unwrap();
4141        let inner_condition = arena
4142            .alloc(SLTNode::Input {
4143                variable: 101,
4144                signed: false,
4145                index: Vec::new(),
4146                access: BitAccess::new(0, 0),
4147            })
4148            .unwrap();
4149        let input_value = arena
4150            .alloc(SLTNode::Input {
4151                variable: 102,
4152                signed: false,
4153                index: Vec::new(),
4154                access: BitAccess::new(0, 7),
4155            })
4156            .unwrap();
4157        let one = arena
4158            .alloc(SLTNode::Constant(
4159                BigUint::from(1u8),
4160                BigUint::from(0u8),
4161                8,
4162                false,
4163            ))
4164            .unwrap();
4165        let mut inner_true = input_value;
4166        let mut inner_false = input_value;
4167        for _ in 0..16 {
4168            inner_true = arena
4169                .alloc(SLTNode::Binary(inner_true, BinaryOp::Add, one))
4170                .unwrap();
4171            inner_false = arena
4172                .alloc(SLTNode::Binary(inner_false, BinaryOp::Xor, one))
4173                .unwrap();
4174        }
4175        let nested = arena
4176            .alloc(SLTNode::Mux {
4177                cond: inner_condition,
4178                then_expr: inner_true,
4179                else_expr: inner_false,
4180            })
4181            .unwrap();
4182
4183        let mut paths = Vec::new();
4184        for target in 0u32..12 {
4185            let then_expr = if target == 0 {
4186                nested
4187            } else {
4188                arena
4189                    .alloc(SLTNode::Constant(
4190                        BigUint::from(target),
4191                        BigUint::from(0u8),
4192                        8,
4193                        false,
4194                    ))
4195                    .unwrap()
4196            };
4197            let else_expr = arena
4198                .alloc(SLTNode::Constant(
4199                    BigUint::from(target + 32),
4200                    BigUint::from(0u8),
4201                    8,
4202                    false,
4203                ))
4204                .unwrap();
4205            let expr = arena
4206                .alloc(SLTNode::Mux {
4207                    cond: outer_condition,
4208                    then_expr,
4209                    else_expr,
4210                })
4211                .unwrap();
4212            paths.push(LogicPath {
4213                target: LogicPathTarget::Var(VarAtomBase::new(target, 0, 7)),
4214                sources: [
4215                    VarAtomBase::new(100, 0, 0),
4216                    VarAtomBase::new(101, 0, 0),
4217                    VarAtomBase::new(102, 0, 7),
4218                ]
4219                .into_iter()
4220                .collect(),
4221                previous_sources: crate::HashSet::default(),
4222                address_sources: crate::HashSet::default(),
4223                local_inputs: Vec::new(),
4224                order_before: crate::HashSet::default(),
4225                comb_capture_enable_sites: Vec::new(),
4226                comb_capture_enable_always: false,
4227                pre_lower_nodes: Vec::new(),
4228                expr,
4229            });
4230        }
4231
4232        let result = sort(
4233            paths,
4234            &arena,
4235            &crate::HashSet::default(),
4236            &crate::HashMap::default(),
4237            false,
4238            &crate::HashMap::default(),
4239            0,
4240        )
4241        .unwrap();
4242        let branch_count = result.execution_units[0]
4243            .blocks
4244            .values()
4245            .filter(|block| matches!(block.terminator, SIRTerminator::Branch { .. }))
4246            .count();
4247        assert_eq!(
4248            branch_count, 2,
4249            "outer region and nested Mux must both branch"
4250        );
4251    }
4252
4253    #[test]
4254    fn ff_plan_models_comb_and_ff_old_state_readers_before_writer() {
4255        let mut arena = SLTNodeArena::new();
4256        let comb_value = simple_path(&mut arena, 10, None);
4257        let comb_old_state_reader = simple_path(&mut arena, 11, Some(20));
4258        let ff = vec![
4259            FfAccessSummary {
4260                reads: vec![VarAtomBase::new(10, 0, 7), VarAtomBase::new(11, 0, 7)],
4261                writes: vec![VarAtomBase::new(20, 0, 7)],
4262                dynamic_writes: HashSet::default(),
4263            },
4264            FfAccessSummary {
4265                reads: vec![VarAtomBase::new(20, 0, 7)],
4266                writes: vec![VarAtomBase::new(30, 0, 7)],
4267                dynamic_writes: HashSet::default(),
4268            },
4269        ];
4270
4271        let plan = plan_ff_comb_schedule(&[comb_value, comb_old_state_reader], &ff).unwrap();
4272
4273        assert_eq!(plan.required_comb, vec![true, true]);
4274        assert_eq!(plan.comb_value_predecessors, vec![vec![0, 1], vec![]]);
4275        assert_eq!(plan.comb_before_direct_write, vec![vec![1], vec![]]);
4276        assert_eq!(plan.ff_before_direct_write, vec![vec![1], vec![]]);
4277    }
4278
4279    #[test]
4280    fn ff_plan_keeps_same_recipe_old_state_reads_local_to_lowering() {
4281        let summary = FfAccessSummary {
4282            reads: vec![VarAtomBase::new(20, 0, 7)],
4283            writes: vec![VarAtomBase::new(20, 0, 7)],
4284            dynamic_writes: HashSet::default(),
4285        };
4286
4287        let plan = plan_ff_comb_schedule(&[], &[summary]).unwrap();
4288
4289        assert_eq!(plan.ff_before_direct_write, vec![Vec::<usize>::new()]);
4290    }
4291
4292    #[test]
4293    fn ff_direct_write_proof_is_kept_per_bit_range() {
4294        let summaries = vec![
4295            FfAccessSummary {
4296                reads: vec![VarAtomBase::new(30, 0, 7)],
4297                writes: vec![VarAtomBase::new(20, 0, 7), VarAtomBase::new(20, 8, 15)],
4298                dynamic_writes: HashSet::default(),
4299            },
4300            FfAccessSummary {
4301                reads: vec![VarAtomBase::new(20, 0, 7)],
4302                writes: vec![VarAtomBase::new(30, 0, 7)],
4303                dynamic_writes: HashSet::default(),
4304            },
4305        ];
4306        let plan = plan_ff_comb_schedule(&[], &summaries).unwrap();
4307        // The FF0 -> FF1 edge is retained, while the inverse edge which only
4308        // protects FF0's low byte is dropped to break the cycle.
4309        let retained = [(0, 1)].into_iter().collect();
4310        let var_widths = [(20, 16), (30, 8)].into_iter().collect();
4311
4312        let direct = direct_ff_write_ranges(
4313            &[],
4314            &summaries,
4315            &plan,
4316            &retained,
4317            0,
4318            &var_widths,
4319            &HashMap::default(),
4320        );
4321
4322        assert_eq!(direct[0], vec![VarAtomBase::new(20, 8, 15)]);
4323        assert_eq!(direct[1], vec![VarAtomBase::new(30, 0, 7)]);
4324
4325        let local_old_read = vec![FfAccessSummary {
4326            reads: vec![VarAtomBase::new(20, 0, 7)],
4327            writes: vec![VarAtomBase::new(20, 0, 7), VarAtomBase::new(20, 8, 15)],
4328            dynamic_writes: HashSet::default(),
4329        }];
4330        let local_plan = plan_ff_comb_schedule(&[], &local_old_read).unwrap();
4331        let local_direct = direct_ff_write_ranges(
4332            &[],
4333            &local_old_read,
4334            &local_plan,
4335            &HashSet::default(),
4336            0,
4337            &var_widths,
4338            &HashMap::default(),
4339        );
4340
4341        assert_eq!(local_direct[0], vec![VarAtomBase::new(20, 8, 15)]);
4342    }
4343
4344    #[test]
4345    fn ff_direct_write_rejects_ranges_requiring_rmw() {
4346        let summary = FfAccessSummary {
4347            reads: Vec::new(),
4348            writes: vec![
4349                VarAtomBase::new(20, 0, 3),
4350                VarAtomBase::new(21, 0, 0),
4351                VarAtomBase::new(22, 0, 7),
4352                VarAtomBase::new(23, 6, 11),
4353                VarAtomBase::new(24, 8, 15),
4354            ],
4355            dynamic_writes: [22].into_iter().collect(),
4356        };
4357        let plan = plan_ff_comb_schedule(&[], std::slice::from_ref(&summary)).unwrap();
4358        let var_widths = [(20, 8), (21, 1), (22, 8), (23, 24), (24, 24)]
4359            .into_iter()
4360            .collect();
4361        let unpacked_element_widths = [(23, 6), (24, 8)].into_iter().collect();
4362
4363        let direct = direct_ff_write_ranges(
4364            &[],
4365            &[summary],
4366            &plan,
4367            &HashSet::default(),
4368            0,
4369            &var_widths,
4370            &unpacked_element_widths,
4371        );
4372
4373        assert_eq!(
4374            direct[0],
4375            vec![VarAtomBase::new(21, 0, 0), VarAtomBase::new(24, 8, 15)]
4376        );
4377    }
4378
4379    #[test]
4380    fn ff_direct_write_stages_complete_overlapping_components() {
4381        let summary = FfAccessSummary {
4382            reads: Vec::new(),
4383            writes: vec![
4384                VarAtomBase::new(20, 4, 19),
4385                VarAtomBase::new(20, 16, 31),
4386                VarAtomBase::new(20, 24, 31),
4387                VarAtomBase::new(20, 40, 47),
4388            ],
4389            dynamic_writes: HashSet::default(),
4390        };
4391        let plan = plan_ff_comb_schedule(&[], std::slice::from_ref(&summary)).unwrap();
4392        let var_widths = [(20, 64)].into_iter().collect();
4393
4394        let direct = direct_ff_write_ranges(
4395            &[],
4396            &[summary],
4397            &plan,
4398            &HashSet::default(),
4399            0,
4400            &var_widths,
4401            &HashMap::default(),
4402        );
4403
4404        // The unaligned first write is staged. Staging propagates through the
4405        // two transitively overlapping native ranges, but not the disjoint byte.
4406        assert_eq!(direct[0], vec![VarAtomBase::new(20, 40, 47)]);
4407    }
4408
4409    #[test]
4410    fn cyclic_direct_write_preference_falls_back_to_staging() {
4411        let mut adj = vec![vec![1], Vec::new(), Vec::new()];
4412
4413        add_acyclic_ff_write_order_edges(&mut adj, [(1, 0), (1, 2)]);
4414
4415        assert_eq!(adj, vec![vec![1], vec![2], Vec::new()]);
4416    }
4417
4418    #[test]
4419    fn memory_ssa_scheduler_lowers_independent_single_use_chains_contiguously() {
4420        let mut arena = SLTNodeArena::new();
4421        let paths = vec![
4422            simple_path(&mut arena, 10, None),
4423            simple_path(&mut arena, 20, None),
4424            simple_path(&mut arena, 11, Some(10)),
4425            simple_path(&mut arena, 21, Some(20)),
4426            simple_path(&mut arena, 12, Some(11)),
4427            simple_path(&mut arena, 22, Some(21)),
4428        ];
4429        let result = sort(
4430            paths,
4431            &arena,
4432            &crate::HashSet::default(),
4433            &crate::HashMap::default(),
4434            false,
4435            &crate::HashMap::default(),
4436            1,
4437        )
4438        .unwrap();
4439        let unit = &result.execution_units[0];
4440        let stores = unit.blocks[&unit.entry_block_id]
4441            .instructions
4442            .iter()
4443            .filter_map(|instruction| match instruction {
4444                SIRInstruction::Store(address, ..) => Some(*address),
4445                _ => None,
4446            })
4447            .collect::<Vec<_>>();
4448
4449        assert_eq!(stores, vec![10, 11, 12, 20, 21, 22]);
4450    }
4451
4452    #[test]
4453    fn previous_value_use_is_an_order_edge_not_a_forwarded_value() {
4454        let mut arena = SLTNodeArena::new();
4455        let mut previous_user = simple_path(&mut arena, 20, Some(10));
4456        previous_user.sources.clear();
4457        previous_user.previous_sources = [VarAtomBase::new(10, 0, 7)].into_iter().collect();
4458        let writer = simple_path(&mut arena, 10, None);
4459        let paths = vec![previous_user, writer];
4460
4461        let memory_ssa = build_logic_path_memory_ssa(&paths).unwrap();
4462        assert_eq!(memory_ssa.dependencies.users, vec![vec![1], vec![]]);
4463        assert_eq!(memory_ssa.dependencies.predecessors, vec![vec![], vec![0]]);
4464        assert_eq!(
4465            memory_ssa.values.users,
4466            vec![Vec::<usize>::new(), Vec::new()]
4467        );
4468        assert_eq!(
4469            memory_ssa.values.predecessors,
4470            vec![Vec::<usize>::new(), Vec::new()]
4471        );
4472
4473        let result = sort(
4474            paths,
4475            &arena,
4476            &crate::HashSet::default(),
4477            &crate::HashMap::default(),
4478            false,
4479            &crate::HashMap::default(),
4480            1,
4481        )
4482        .unwrap();
4483        let unit = &result.execution_units[0];
4484        let stores = unit.blocks[&unit.entry_block_id]
4485            .instructions
4486            .iter()
4487            .filter_map(|instruction| match instruction {
4488                SIRInstruction::Store(address, ..) => Some(*address),
4489                _ => None,
4490            })
4491            .collect::<Vec<_>>();
4492        assert_eq!(stores, vec![20, 10]);
4493    }
4494
4495    #[test]
4496    fn memory_use_depends_on_each_overlapping_bit_range_definition() {
4497        let mut arena = SLTNodeArena::new();
4498        let mut low = simple_path(&mut arena, 10, None);
4499        low.target = LogicPathTarget::Var(VarAtomBase::new(10, 0, 3));
4500        let mut high = simple_path(&mut arena, 10, None);
4501        high.target = LogicPathTarget::Var(VarAtomBase::new(10, 4, 7));
4502        let mut user = simple_path(&mut arena, 20, Some(10));
4503        user.sources = [VarAtomBase::new(10, 2, 5)].into_iter().collect();
4504
4505        let memory_ssa = build_logic_path_memory_ssa(&[low, high, user]).unwrap();
4506        assert_eq!(
4507            memory_ssa.dependencies.users,
4508            vec![vec![2], vec![2], vec![]]
4509        );
4510        assert_eq!(
4511            memory_ssa.dependencies.predecessors,
4512            vec![vec![], vec![], vec![0, 1]]
4513        );
4514        assert_eq!(memory_ssa.values.users, vec![vec![2], vec![2], vec![]]);
4515        assert_eq!(
4516            memory_ssa.values.predecessors,
4517            vec![vec![], vec![], vec![0, 1]]
4518        );
4519    }
4520
4521    fn fixed_group_path(
4522        arena: &mut SLTNodeArena<u32>,
4523        guard: crate::NodeId,
4524        loop_var: u32,
4525        target: u32,
4526        external: u32,
4527        trip_count: usize,
4528    ) -> LogicPath<u32> {
4529        fixed_group_path_with_index(
4530            arena,
4531            guard,
4532            loop_var,
4533            target,
4534            external,
4535            trip_count,
4536            1,
4537            BitAccess::new(0, 7),
4538        )
4539    }
4540
4541    #[allow(clippy::too_many_arguments)]
4542    fn fixed_group_path_with_index(
4543        arena: &mut SLTNodeArena<u32>,
4544        guard: crate::NodeId,
4545        loop_var: u32,
4546        target: u32,
4547        external: u32,
4548        trip_count: usize,
4549        index_scale: u8,
4550        load_access: BitAccess,
4551    ) -> LogicPath<u32> {
4552        let target = VarAtomBase::new(target, 0, 7);
4553        let initial = arena
4554            .alloc(SLTNode::Input {
4555                variable: target.id,
4556                signed: false,
4557                index: Vec::new(),
4558                access: target.access,
4559            })
4560            .unwrap();
4561        let loop_index = arena
4562            .alloc(SLTNode::Input {
4563                variable: loop_var,
4564                signed: false,
4565                index: Vec::new(),
4566                access: BitAccess::new(0, 7),
4567            })
4568            .unwrap();
4569        let loop_index = if index_scale == 1 {
4570            loop_index
4571        } else {
4572            let scale = arena
4573                .alloc(SLTNode::Constant(
4574                    BigUint::from(index_scale),
4575                    BigUint::from(0u8),
4576                    8,
4577                    false,
4578                ))
4579                .unwrap();
4580            arena
4581                .alloc(SLTNode::Binary(loop_index, BinaryOp::Mul, scale))
4582                .unwrap()
4583        };
4584        let update = arena
4585            .alloc(SLTNode::Input {
4586                variable: external,
4587                signed: false,
4588                index: vec![
4589                    serde_json::from_value(serde_json::json!({
4590                        "node": loop_index,
4591                        "stride": 8,
4592                        "kind": "Packed",
4593                    }))
4594                    .unwrap(),
4595                ],
4596                access: load_access,
4597            })
4598            .unwrap();
4599        let group = arena
4600            .alloc(SLTNode::ForFoldGroup {
4601                loop_var,
4602                loop_width: 8,
4603                loop_signed: false,
4604                start: BigInt::from(0),
4605                step: BigInt::from(1),
4606                trip_count,
4607                entry_guard: guard,
4608                states: vec![SLTForFoldGroupState {
4609                    target,
4610                    initial,
4611                    update,
4612                }],
4613            })
4614            .unwrap();
4615        LogicPath {
4616            target: LogicPathTarget::Var(target),
4617            sources: [VarAtomBase::new(external, 0, 63)].into_iter().collect(),
4618            previous_sources: crate::HashSet::default(),
4619            address_sources: crate::HashSet::default(),
4620            local_inputs: Vec::new(),
4621            order_before: crate::HashSet::default(),
4622            comb_capture_enable_sites: Vec::new(),
4623            comb_capture_enable_always: false,
4624            pre_lower_nodes: Vec::new(),
4625            expr: group,
4626        }
4627    }
4628
4629    fn fixed_group_fixture(
4630        left_trip_count: usize,
4631        right_trip_count: usize,
4632    ) -> (
4633        SLTNodeArena<u32>,
4634        Vec<LogicPath<u32>>,
4635        crate::HashMap<u32, usize>,
4636    ) {
4637        let mut arena = SLTNodeArena::new();
4638        let guard = arena
4639            .alloc(SLTNode::Constant(
4640                BigUint::from(1u8),
4641                BigUint::from(0u8),
4642                1,
4643                false,
4644            ))
4645            .unwrap();
4646        let paths = vec![
4647            fixed_group_path(&mut arena, guard, 100, 10, 50, left_trip_count),
4648            fixed_group_path(&mut arena, guard, 101, 11, 50, right_trip_count),
4649        ];
4650        let widths = [(10, 8), (11, 8)].into_iter().collect();
4651        (arena, paths, widths)
4652    }
4653
4654    fn schedule_branch_count(result: &super::ScheduleResult<u32>) -> usize {
4655        result
4656            .execution_units
4657            .iter()
4658            .flat_map(|unit| unit.blocks.values())
4659            .filter(|block| matches!(block.terminator, SIRTerminator::Branch { .. }))
4660            .count()
4661    }
4662
4663    fn synthetic_exact_load(base: u32) -> ExactIndexedLoadKey<u32> {
4664        ExactIndexedLoadKey {
4665            base,
4666            access: BitAccess::new(0, 7),
4667            index: vec![(
4668                NormalizedIndexExpr::LoopValue {
4669                    signed: false,
4670                    access: BitAccess::new(0, 7),
4671                },
4672                8,
4673            )],
4674        }
4675    }
4676
4677    fn synthetic_fold_group(
4678        root: usize,
4679        target: u32,
4680        loop_var: u32,
4681        indexed_loads: impl IntoIterator<Item = ExactIndexedLoadKey<u32>>,
4682        carried_chunks: u128,
4683    ) -> ExactFoldGroup<u32> {
4684        ExactFoldGroup {
4685            root: crate::NodeId(root),
4686            facts: FoldGroupReadFacts {
4687                loop_var,
4688                state_targets: vec![VarAtomBase::new(target, 0, 7)],
4689                guard_reads: Vec::new(),
4690                initial_reads: Vec::new(),
4691                update_reads: Vec::new(),
4692                indexed_loads: indexed_loads.into_iter().collect(),
4693                carried_chunks,
4694            },
4695        }
4696    }
4697
4698    #[test]
4699    fn independent_exact_fold_groups_lower_jointly_and_keep_store_order() {
4700        let (arena, paths, widths) = fixed_group_fixture(4, 4);
4701        let result = sort(
4702            paths,
4703            &arena,
4704            &crate::HashSet::default(),
4705            &crate::HashMap::default(),
4706            false,
4707            &widths,
4708            1,
4709        )
4710        .unwrap();
4711
4712        assert_eq!(schedule_branch_count(&result), 2);
4713        let store_block = result
4714            .execution_units
4715            .iter()
4716            .flat_map(|unit| unit.blocks.values())
4717            .find(|block| {
4718                block
4719                    .instructions
4720                    .iter()
4721                    .filter(|instruction| matches!(instruction, SIRInstruction::Store(..)))
4722                    .count()
4723                    == 2
4724            })
4725            .expect("joint results must be materialized before the ordered stores");
4726        let stores = store_block
4727            .instructions
4728            .iter()
4729            .filter_map(|instruction| match instruction {
4730                SIRInstruction::Store(address, _, _, _, _, _) => Some(*address),
4731                _ => None,
4732            })
4733            .collect::<Vec<_>>();
4734        assert_eq!(stores, vec![10, 11]);
4735    }
4736
4737    #[test]
4738    fn different_index_expression_or_load_slice_prevents_joint_lowering() {
4739        let mut arena = SLTNodeArena::new();
4740        let guard = arena
4741            .alloc(SLTNode::Constant(
4742                BigUint::from(1u8),
4743                BigUint::from(0u8),
4744                1,
4745                false,
4746            ))
4747            .unwrap();
4748        let paths = vec![
4749            fixed_group_path_with_index(&mut arena, guard, 100, 10, 50, 4, 1, BitAccess::new(0, 7)),
4750            fixed_group_path_with_index(&mut arena, guard, 101, 11, 50, 4, 2, BitAccess::new(0, 7)),
4751            fixed_group_path_with_index(
4752                &mut arena,
4753                guard,
4754                102,
4755                12,
4756                50,
4757                4,
4758                1,
4759                BitAccess::new(8, 15),
4760            ),
4761        ];
4762        let widths = [(10, 8), (11, 8), (12, 8)].into_iter().collect();
4763        let result = sort(
4764            paths,
4765            &arena,
4766            &crate::HashSet::default(),
4767            &crate::HashMap::default(),
4768            false,
4769            &widths,
4770            1,
4771        )
4772        .unwrap();
4773
4774        assert_eq!(schedule_branch_count(&result), 6);
4775    }
4776
4777    #[test]
4778    fn weighted_family_selection_ignores_conflicting_first_root() {
4779        let shared = synthetic_exact_load(50);
4780        let candidates = vec![
4781            synthetic_fold_group(0, 10, 100, [shared.clone()], 1),
4782            synthetic_fold_group(1, 11, 101, [shared.clone()], 1),
4783            synthetic_fold_group(2, 12, 102, [shared], 1),
4784        ];
4785        let available = vec![true; 3];
4786        let compatible = vec![
4787            vec![false, false, false],
4788            vec![false, false, true],
4789            vec![false, true, false],
4790        ];
4791        let shared_load = vec![
4792            vec![false, true, true],
4793            vec![true, false, true],
4794            vec![true, true, false],
4795        ];
4796
4797        let family = best_weighted_fold_family(
4798            &candidates,
4799            &available,
4800            &compatible,
4801            &shared_load,
4802            &crate::HashSet::default(),
4803            false,
4804        )
4805        .expect("the compatible B+C family has positive net benefit");
4806        let mut roots = family
4807            .members
4808            .iter()
4809            .map(|member| candidates[*member].root.0)
4810            .collect::<Vec<_>>();
4811        roots.sort_unstable();
4812        assert_eq!(roots, vec![1, 2]);
4813    }
4814
4815    #[test]
4816    fn weighted_family_selection_rejects_benefit_not_exceeding_pressure() {
4817        let shared = synthetic_exact_load(50);
4818        let candidates = vec![
4819            synthetic_fold_group(0, 10, 100, [shared.clone()], 8),
4820            synthetic_fold_group(1, 11, 101, [shared], 8),
4821        ];
4822        let available = vec![true; 2];
4823        let compatible = vec![vec![false, true], vec![true, false]];
4824        let shared_load = compatible.clone();
4825
4826        assert!(
4827            best_weighted_fold_family(
4828                &candidates,
4829                &available,
4830                &compatible,
4831                &shared_load,
4832                &crate::HashSet::default(),
4833                false,
4834            )
4835            .is_none()
4836        );
4837    }
4838
4839    #[test]
4840    fn dependency_edge_separates_otherwise_joint_fold_groups() {
4841        let (arena, mut paths, widths) = fixed_group_fixture(4, 4);
4842        paths[1].sources.insert(VarAtomBase::new(10, 0, 7));
4843        let result = sort(
4844            paths,
4845            &arena,
4846            &crate::HashSet::default(),
4847            &crate::HashMap::default(),
4848            false,
4849            &widths,
4850            1,
4851        )
4852        .unwrap();
4853
4854        assert_eq!(schedule_branch_count(&result), 4);
4855    }
4856
4857    #[test]
4858    fn different_fold_group_domains_do_not_lower_jointly() {
4859        let (arena, paths, widths) = fixed_group_fixture(4, 5);
4860        let result = sort(
4861            paths,
4862            &arena,
4863            &crate::HashSet::default(),
4864            &crate::HashMap::default(),
4865            false,
4866            &widths,
4867            1,
4868        )
4869        .unwrap();
4870
4871        assert_eq!(schedule_branch_count(&result), 4);
4872    }
4873
4874    #[test]
4875    fn joint_fold_preparation_does_not_mutate_logic_path_metadata() {
4876        let (arena, mut paths, _) = fixed_group_fixture(4, 4);
4877        paths[0].previous_sources = [VarAtomBase::new(60, 0, 7)].into_iter().collect();
4878        paths[0].sources.insert(VarAtomBase::new(61, 0, 7));
4879        paths[0].address_sources = [VarAtomBase::new(61, 0, 7)].into_iter().collect();
4880        paths[0].comb_capture_enable_sites = vec![3, 7];
4881        let snapshot = paths.clone();
4882        let mut builder = SIRBuilder::new();
4883        let mut cache = crate::HashMap::default();
4884        let mut dependencies = crate::HashMap::default();
4885        let mut inverse_dependencies = crate::HashMap::default();
4886        let schedule_index = build_fold_group_schedule_index(&paths, &arena);
4887
4888        let prepared = prepare_atomic_fold_group_results(
4889            &[0, 1],
4890            &paths,
4891            &schedule_index,
4892            &SLTToSIRLowerer::new(false),
4893            &mut builder,
4894            &arena,
4895            &mut cache,
4896            &mut dependencies,
4897            &mut inverse_dependencies,
4898            false,
4899        );
4900
4901        assert_eq!(prepared.len(), 2);
4902        assert_eq!(paths, snapshot);
4903    }
4904
4905    #[test]
4906    fn for_fold_group_dependencies_keep_initial_but_hide_loop_scoped_updates() {
4907        let mut arena = SLTNodeArena::<u32>::new();
4908        let input = |arena: &mut SLTNodeArena<u32>, variable| {
4909            arena
4910                .alloc(SLTNode::Input {
4911                    variable,
4912                    signed: false,
4913                    index: Vec::new(),
4914                    access: BitAccess::new(0, 7),
4915                })
4916                .unwrap()
4917        };
4918        let guard = arena
4919            .alloc(SLTNode::Constant(
4920                BigUint::from(1u8),
4921                BigUint::from(0u8),
4922                1,
4923                false,
4924            ))
4925            .unwrap();
4926        let initial = input(&mut arena, 2);
4927        let state_input = input(&mut arena, 2);
4928        let loop_input = input(&mut arena, 1);
4929        let external_input = input(&mut arena, 3);
4930        let scoped_sum = arena
4931            .alloc(SLTNode::Binary(state_input, BinaryOp::Add, loop_input))
4932            .unwrap();
4933        let update = arena
4934            .alloc(SLTNode::Binary(scoped_sum, BinaryOp::Add, external_input))
4935            .unwrap();
4936        let group = arena
4937            .alloc(SLTNode::ForFoldGroup {
4938                loop_var: 1,
4939                loop_width: 8,
4940                loop_signed: false,
4941                start: BigInt::from(0),
4942                step: BigInt::from(1),
4943                trip_count: 2,
4944                entry_guard: guard,
4945                states: vec![SLTForFoldGroupState {
4946                    target: VarAtomBase::new(2, 0, 7),
4947                    initial,
4948                    update,
4949                }],
4950            })
4951            .unwrap();
4952        let mut memo = crate::HashMap::default();
4953        let mut inverse_memo = crate::HashMap::default();
4954
4955        let dependencies = collect_node_input_deps(group, &arena, &mut memo, &mut inverse_memo);
4956
4957        assert!(
4958            dependencies.contains(&2),
4959            "initial state is an external dependency"
4960        );
4961        assert!(
4962            dependencies.contains(&3),
4963            "ordinary update input remains external"
4964        );
4965        assert!(
4966            !dependencies.contains(&1),
4967            "loop variable is supplied by the fold"
4968        );
4969    }
4970
4971    #[test]
4972    fn partial_for_fold_group_state_keeps_uncovered_variable_dependency() {
4973        let mut arena = SLTNodeArena::<u32>::new();
4974        let guard = arena
4975            .alloc(SLTNode::Constant(
4976                BigUint::from(1u8),
4977                BigUint::from(0u8),
4978                1,
4979                false,
4980            ))
4981            .unwrap();
4982        let initial = arena
4983            .alloc(SLTNode::Constant(
4984                BigUint::from(0u8),
4985                BigUint::from(0u8),
4986                8,
4987                false,
4988            ))
4989            .unwrap();
4990        let carried = arena
4991            .alloc(SLTNode::Input {
4992                variable: 2,
4993                signed: false,
4994                index: Vec::new(),
4995                access: BitAccess::new(0, 7),
4996            })
4997            .unwrap();
4998        let uncovered = arena
4999            .alloc(SLTNode::Input {
5000                variable: 2,
5001                signed: false,
5002                index: Vec::new(),
5003                access: BitAccess::new(8, 15),
5004            })
5005            .unwrap();
5006        let update = arena
5007            .alloc(SLTNode::Binary(carried, BinaryOp::Add, uncovered))
5008            .unwrap();
5009        let group = arena
5010            .alloc(SLTNode::ForFoldGroup {
5011                loop_var: 1,
5012                loop_width: 8,
5013                loop_signed: false,
5014                start: BigInt::from(0),
5015                step: BigInt::from(1),
5016                trip_count: 2,
5017                entry_guard: guard,
5018                states: vec![SLTForFoldGroupState {
5019                    target: VarAtomBase::new(2, 0, 7),
5020                    initial,
5021                    update,
5022                }],
5023            })
5024            .unwrap();
5025        let mut memo = crate::HashMap::default();
5026        let mut inverse_memo = crate::HashMap::default();
5027
5028        let dependencies = collect_node_input_deps(group, &arena, &mut memo, &mut inverse_memo);
5029        assert!(
5030            dependencies.contains(&2),
5031            "the uncovered high byte remains an external dependency"
5032        );
5033    }
5034
5035    #[test]
5036    fn shared_for_fold_group_projections_materialize_once_and_store_sequentially() {
5037        let mut arena = SLTNodeArena::<u32>::new();
5038        let guard = arena
5039            .alloc(SLTNode::Constant(
5040                BigUint::from(1u8),
5041                BigUint::from(0u8),
5042                1,
5043                false,
5044            ))
5045            .unwrap();
5046        let input = |arena: &mut SLTNodeArena<u32>, variable| {
5047            arena
5048                .alloc(SLTNode::Input {
5049                    variable,
5050                    signed: false,
5051                    index: Vec::new(),
5052                    access: BitAccess::new(0, 7),
5053                })
5054                .unwrap()
5055        };
5056        let initial_a = input(&mut arena, 1);
5057        let initial_b = input(&mut arena, 2);
5058        let previous_a = input(&mut arena, 1);
5059        let previous_b = input(&mut arena, 2);
5060        let group = arena
5061            .alloc(SLTNode::ForFoldGroup {
5062                loop_var: 3,
5063                loop_width: 8,
5064                loop_signed: false,
5065                start: BigInt::from(0),
5066                step: BigInt::from(1),
5067                trip_count: 3,
5068                entry_guard: guard,
5069                states: vec![
5070                    SLTForFoldGroupState {
5071                        target: VarAtomBase::new(1, 0, 7),
5072                        initial: initial_a,
5073                        update: previous_b,
5074                    },
5075                    SLTForFoldGroupState {
5076                        target: VarAtomBase::new(2, 0, 7),
5077                        initial: initial_b,
5078                        update: previous_a,
5079                    },
5080                ],
5081            })
5082            .unwrap();
5083        let high = arena
5084            .alloc(SLTNode::Slice {
5085                expr: group,
5086                access: BitAccess::new(8, 15),
5087            })
5088            .unwrap();
5089        let low = arena
5090            .alloc(SLTNode::Slice {
5091                expr: group,
5092                access: BitAccess::new(0, 7),
5093            })
5094            .unwrap();
5095        let path = |target, expr| LogicPath {
5096            target: LogicPathTarget::Var(VarAtomBase::new(target, 0, 7)),
5097            // Keep the scheduling graph acyclic in this focused cache test;
5098            // dependency memoization still sees both initial reads.
5099            sources: crate::HashSet::default(),
5100            previous_sources: crate::HashSet::default(),
5101            address_sources: crate::HashSet::default(),
5102            local_inputs: Vec::new(),
5103            order_before: crate::HashSet::default(),
5104            comb_capture_enable_sites: Vec::new(),
5105            comb_capture_enable_always: false,
5106            pre_lower_nodes: Vec::new(),
5107            expr,
5108        };
5109        let mut widths = crate::HashMap::default();
5110        widths.insert(1, 8);
5111        widths.insert(2, 8);
5112
5113        let result = sort(
5114            vec![path(1, high), path(2, low)],
5115            &arena,
5116            &crate::HashSet::default(),
5117            &crate::HashMap::default(),
5118            false,
5119            &widths,
5120            1,
5121        )
5122        .unwrap();
5123        assert_eq!(result.execution_units.len(), 1);
5124        let eu = &result.execution_units[0];
5125        assert_eq!(
5126            eu.blocks
5127                .values()
5128                .filter(|block| matches!(block.terminator, SIRTerminator::Branch { .. }))
5129                .count(),
5130            2,
5131            "one grouped fold has only its entry and counted-loop branches"
5132        );
5133        assert_eq!(
5134            eu.blocks
5135                .values()
5136                .flat_map(|block| &block.instructions)
5137                .filter(|instruction| matches!(instruction, SIRInstruction::Load(..)))
5138                .count(),
5139            2,
5140            "both initial values are loaded once, not once per projection"
5141        );
5142
5143        let store_block = eu
5144            .blocks
5145            .values()
5146            .find(|block| {
5147                block
5148                    .instructions
5149                    .iter()
5150                    .filter(|instruction| matches!(instruction, SIRInstruction::Store(..)))
5151                    .count()
5152                    == 2
5153            })
5154            .expect("both atomic projection stores share the materialization exit");
5155        let store_positions = store_block
5156            .instructions
5157            .iter()
5158            .enumerate()
5159            .filter_map(|(position, instruction)| {
5160                matches!(instruction, SIRInstruction::Store(..)).then_some(position)
5161            })
5162            .collect::<Vec<_>>();
5163        let stored_values = store_block
5164            .instructions
5165            .iter()
5166            .filter_map(|instruction| match instruction {
5167                SIRInstruction::Store(_, _, _, value, _, _) => Some(*value),
5168                _ => None,
5169            })
5170            .collect::<Vec<_>>();
5171        assert_eq!(store_positions.len(), stored_values.len());
5172        for (projection, value) in stored_values.into_iter().enumerate() {
5173            let definition = store_block
5174                .instructions
5175                .iter()
5176                .position(|instruction| match instruction {
5177                    SIRInstruction::Binary(dst, ..)
5178                    | SIRInstruction::Unary(dst, ..)
5179                    | SIRInstruction::Slice(dst, ..)
5180                    | SIRInstruction::Concat(dst, ..)
5181                    | SIRInstruction::Mux(dst, ..) => *dst == value,
5182                    _ => false,
5183                })
5184                .expect("stored projection has a local definition");
5185            assert!(definition < store_positions[projection]);
5186            if projection != 0 {
5187                assert!(
5188                    definition > store_positions[projection - 1],
5189                    "a later projection must not remain live across an earlier Store"
5190                );
5191            }
5192        }
5193    }
5194}