plotnik_compiler/emit/

layout.rs

//! Cache-aligned instruction layout.
//!
//! Extracts linear chains from the control flow graph and places them
//! contiguously. Pads instructions to prevent cache line straddling.

use std::collections::{BTreeMap, HashSet};

use crate::bytecode::{InstructionIR, Label, LayoutResult};

const CACHE_LINE: usize = 64;
const STEP_SIZE: usize = 8;

/// Successor graph for layout analysis.
struct Graph {
    /// label -> list of successor labels
    successors: BTreeMap<Label, Vec<Label>>,
    /// label -> list of predecessor labels
    predecessors: BTreeMap<Label, Vec<Label>>,
}

impl Graph {
    fn build(instructions: &[InstructionIR]) -> Self {
        let mut successors: BTreeMap<Label, Vec<Label>> = BTreeMap::new();
        let mut predecessors: BTreeMap<Label, Vec<Label>> = BTreeMap::new();

        for instr in instructions {
            let label = instr.label();
            successors.entry(label).or_default();

            for succ in instr.successors() {
                successors.entry(label).or_default().push(succ);
                predecessors.entry(succ).or_default().push(label);
            }
        }

        Self {
            successors,
            predecessors,
        }
    }

    fn successors(&self, label: Label) -> &[Label] {
        self.successors
            .get(&label)
            .map(|v| v.as_slice())
            .unwrap_or(&[])
    }

    fn predecessor_count(&self, label: Label) -> usize {
        self.predecessors.get(&label).map(|v| v.len()).unwrap_or(0)
    }
}

/// Cache-aligned layout strategy.
pub struct CacheAligned;

impl CacheAligned {
    /// Compute layout for instructions with given entry points.
    ///
    /// Returns mapping from labels to step IDs and total step count.
    pub fn layout(instructions: &[InstructionIR], entries: &[Label]) -> LayoutResult {
        if instructions.is_empty() {
            return LayoutResult::empty();
        }

        let graph = Graph::build(instructions);
        let label_to_instr: BTreeMap<Label, &InstructionIR> =
            instructions.iter().map(|i| (i.label(), i)).collect();

        let chains = extract_chains(&graph, instructions, entries);
        let ordered = order_chains(chains, entries);

        assign_step_ids(ordered, &label_to_instr)
    }
}

/// Extract linear chains from the control flow graph.
fn extract_chains(
    graph: &Graph,
    instructions: &[InstructionIR],
    entries: &[Label],
) -> Vec<Vec<Label>> {
    let mut visited = HashSet::new();
    let mut chains = Vec::new();

    // Start with entry points (hot paths)
    for &entry in entries {
        if visited.contains(&entry) {
            continue;
        }
        chains.push(build_chain(entry, graph, &mut visited));
    }

    // Then remaining unvisited instructions
    for instr in instructions {
        let label = instr.label();
        if visited.contains(&label) {
            continue;
        }
        chains.push(build_chain(label, graph, &mut visited));
    }

    chains
}

/// Build a single chain starting from a label.
///
/// Extends the chain while there's a single unvisited successor with a single predecessor.
fn build_chain(start: Label, graph: &Graph, visited: &mut HashSet<Label>) -> Vec<Label> {
    let mut chain = vec![start];
    visited.insert(start);

    let mut current = start;
    while let [next] = graph.successors(current)
        && !visited.contains(next)
        && graph.predecessor_count(*next) == 1
    {
        chain.push(*next);
        visited.insert(*next);
        current = *next;
    }

    chain
}

/// Order chains: entries first, then by size (larger = hotter assumption).
fn order_chains(mut chains: Vec<Vec<Label>>, entries: &[Label]) -> Vec<Vec<Label>> {
    let entry_set: HashSet<Label> = entries.iter().copied().collect();

    // Partition into entry chains and non-entry chains
    let (mut entry_chains, mut other_chains): (Vec<_>, Vec<_>) =
        chains.drain(..).partition(|chain| {
            chain
                .first()
                .map(|l| entry_set.contains(l))
                .unwrap_or(false)
        });

    // Sort other chains by size (descending) for better locality
    other_chains.sort_by_key(|chain| std::cmp::Reverse(chain.len()));

    // Entry chains first, then others
    entry_chains.extend(other_chains);
    entry_chains
}

/// Assign step IDs with cache line awareness.
fn assign_step_ids(
    chains: Vec<Vec<Label>>,
    label_to_instr: &BTreeMap<Label, &InstructionIR>,
) -> LayoutResult {
    let mut mapping = BTreeMap::new();

    let mut current_step = 0u16;
    let mut current_offset = 0usize; // Byte offset for cache alignment

    for chain in chains {
        for label in chain {
            let Some(instr) = label_to_instr.get(&label) else {
                continue;
            };
            let size = instr.size();

            // Pad if instruction would straddle cache line boundary
            let line_offset = current_offset % CACHE_LINE;
            if line_offset + size > CACHE_LINE {
                let padding_bytes = CACHE_LINE - line_offset;
                let padding_steps = (padding_bytes / STEP_SIZE) as u16;
                current_step += padding_steps;
                current_offset += padding_bytes;
            }

            // Invariant: instruction must not straddle cache line
            assert!(
                current_offset % CACHE_LINE + size <= CACHE_LINE,
                "instruction at offset {} with size {} straddles 64-byte cache line",
                current_offset,
                size
            );

            mapping.insert(label, current_step);
            let step_count = (size / STEP_SIZE) as u16;
            current_step += step_count;
            current_offset += size;
        }
    }

    LayoutResult::new(mapping, current_step)
}