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