weirflow 0.1.0

#![allow(clippy::too_many_arguments)]

use super::FixedPointAnalysisKind;

/// Packed forward CSR graph buffers ready for Vyre dispatch.
#[derive(Clone, Debug, PartialEq)]
pub struct FixedPointForwardGraph {
    pub(crate) kind: FixedPointAnalysisKind,
    pub(crate) node_count: u32,
    pub(crate) edge_count: u32,
    pub(crate) pg_nodes_bytes: Vec<u8>,
    pub(crate) edge_offsets_bytes: Vec<u8>,
    pub(crate) edge_targets_bytes: Vec<u8>,
    pub(crate) edge_kind_mask_bytes: Vec<u8>,
    pub(crate) pg_node_tags_bytes: Vec<u8>,
    pub(crate) stable_layout_hash: u64,
    pub(crate) normalized_graph: crate::graph_layout::NormalizedCsrGraph,
    pub(crate) normalization_scratch: crate::graph_layout::CsrGraphNormalizationScratch,
}

impl FixedPointForwardGraph {
    /// Validate and pack invariant forward-CSR buffers once.
    pub fn new(
        stage: &str,
        node_count: u32,
        edge_offsets: &[u32],
        edge_targets: &[u32],
        edge_kind_mask: &[u32],
    ) -> Result<Self, String> {
        Self::new_for_kind(
            FixedPointAnalysisKind::Generic,
            stage,
            node_count,
            edge_offsets,
            edge_targets,
            edge_kind_mask,
        )
    }

    /// Validate and pack invariant forward-CSR buffers for one analysis family.
    pub fn new_for_kind(
        kind: FixedPointAnalysisKind,
        stage: &str,
        node_count: u32,
        edge_offsets: &[u32],
        edge_targets: &[u32],
        edge_kind_mask: &[u32],
    ) -> Result<Self, String> {
        let mut normalized_graph = crate::graph_layout::NormalizedCsrGraph::empty();
        let mut normalization_scratch = crate::graph_layout::CsrGraphNormalizationScratch::new();
        crate::graph_layout::CsrGraph::new(node_count, edge_offsets, edge_targets, edge_kind_mask)
            .normalize_into(stage, &mut normalized_graph, &mut normalization_scratch)?;
        Self::from_normalized(
            kind,
            stage,
            node_count,
            normalized_graph,
            normalization_scratch,
        )
    }

    /// Validate and pack invariant CSR buffers once after masking edge kinds.
    pub fn new_with_edge_kind_mask(
        stage: &str,
        node_count: u32,
        edge_offsets: &[u32],
        edge_targets: &[u32],
        edge_kind_mask: &[u32],
        allowed_mask: u32,
    ) -> Result<Self, String> {
        Self::new_with_edge_kind_mask_for_kind(
            FixedPointAnalysisKind::Generic,
            stage,
            node_count,
            edge_offsets,
            edge_targets,
            edge_kind_mask,
            allowed_mask,
        )
    }

    /// Validate and pack masked CSR buffers for one analysis family.
    pub fn new_with_edge_kind_mask_for_kind(
        kind: FixedPointAnalysisKind,
        stage: &str,
        node_count: u32,
        edge_offsets: &[u32],
        edge_targets: &[u32],
        edge_kind_mask: &[u32],
        allowed_mask: u32,
    ) -> Result<Self, String> {
        let mut normalized_graph = crate::graph_layout::NormalizedCsrGraph::empty();
        let mut normalization_scratch = crate::graph_layout::CsrGraphNormalizationScratch::new();
        crate::graph_layout::CsrGraph::new(node_count, edge_offsets, edge_targets, edge_kind_mask)
            .normalize_with_edge_kind_mask_into(
                stage,
                allowed_mask,
                &mut normalized_graph,
                &mut normalization_scratch,
            )?;
        Self::from_normalized(
            kind,
            stage,
            node_count,
            normalized_graph,
            normalization_scratch,
        )
    }

    fn from_normalized(
        kind: FixedPointAnalysisKind,
        stage: &str,
        node_count: u32,
        normalized_graph: crate::graph_layout::NormalizedCsrGraph,
        normalization_scratch: crate::graph_layout::CsrGraphNormalizationScratch,
    ) -> Result<Self, String> {
        let mut graph = Self {
            kind,
            node_count: 0,
            edge_count: 0,
            pg_nodes_bytes: Vec::new(),
            edge_offsets_bytes: Vec::new(),
            edge_targets_bytes: Vec::new(),
            edge_kind_mask_bytes: Vec::new(),
            pg_node_tags_bytes: Vec::new(),
            stable_layout_hash: 0,
            normalized_graph,
            normalization_scratch,
        };
        graph.replace_from_current_normalized(kind, stage, node_count)?;
        Ok(graph)
    }

    /// Validate and repack this graph in place while retaining buffer capacity.
    pub fn replace(
        &mut self,
        stage: &str,
        node_count: u32,
        edge_offsets: &[u32],
        edge_targets: &[u32],
        edge_kind_mask: &[u32],
    ) -> Result<(), String> {
        self.replace_for_kind(
            self.kind,
            stage,
            node_count,
            edge_offsets,
            edge_targets,
            edge_kind_mask,
        )
    }

    /// Validate and repack this graph for the same analysis family.
    pub fn replace_for_kind(
        &mut self,
        kind: FixedPointAnalysisKind,
        stage: &str,
        node_count: u32,
        edge_offsets: &[u32],
        edge_targets: &[u32],
        edge_kind_mask: &[u32],
    ) -> Result<(), String> {
        self.require_kind(kind, stage)?;
        crate::graph_layout::CsrGraph::new(node_count, edge_offsets, edge_targets, edge_kind_mask)
            .normalize_into(
                stage,
                &mut self.normalized_graph,
                &mut self.normalization_scratch,
            )?;
        self.replace_from_current_normalized(kind, stage, node_count)
    }

    /// Validate and repack this graph after masking edge kinds in place while
    /// retaining buffer capacity.
    pub fn replace_with_edge_kind_mask(
        &mut self,
        stage: &str,
        node_count: u32,
        edge_offsets: &[u32],
        edge_targets: &[u32],
        edge_kind_mask: &[u32],
        allowed_mask: u32,
    ) -> Result<(), String> {
        self.replace_with_edge_kind_mask_for_kind(
            self.kind,
            stage,
            node_count,
            edge_offsets,
            edge_targets,
            edge_kind_mask,
            allowed_mask,
        )
    }

    /// Validate and repack masked CSR buffers for the same analysis family.
    pub fn replace_with_edge_kind_mask_for_kind(
        &mut self,
        kind: FixedPointAnalysisKind,
        stage: &str,
        node_count: u32,
        edge_offsets: &[u32],
        edge_targets: &[u32],
        edge_kind_mask: &[u32],
        allowed_mask: u32,
    ) -> Result<(), String> {
        self.require_kind(kind, stage)?;
        crate::graph_layout::CsrGraph::new(node_count, edge_offsets, edge_targets, edge_kind_mask)
            .normalize_with_edge_kind_mask_into(
                stage,
                allowed_mask,
                &mut self.normalized_graph,
                &mut self.normalization_scratch,
            )?;
        self.replace_from_current_normalized(kind, stage, node_count)
    }

    fn replace_from_current_normalized(
        &mut self,
        kind: FixedPointAnalysisKind,
        stage: &str,
        node_count: u32,
    ) -> Result<(), String> {
        let normalized = &self.normalized_graph;
        let nodes = usize::try_from(normalized.node_count()).map_err(|_| {
            format!("{stage} node_count={node_count} cannot fit usize. Fix: shard the graph before dispatch.")
        })?;
        self.kind = kind;
        self.node_count = normalized.node_count();
        self.edge_count = normalized.edge_count();
        self.stable_layout_hash = normalized.stable_layout_hash();
        crate::dispatch_decode::pack_repeated_u32_into(0, nodes, &mut self.pg_nodes_bytes)?;
        crate::dispatch_decode::try_pack_u32_into(
            normalized.edge_offsets(),
            &mut self.edge_offsets_bytes,
            "fixed-point graph edge offsets bytes",
        )?;
        crate::dispatch_decode::try_pack_u32_into(
            normalized.edge_targets(),
            &mut self.edge_targets_bytes,
            "fixed-point graph edge targets bytes",
        )?;
        crate::dispatch_decode::try_pack_u32_into(
            normalized.edge_kind_mask(),
            &mut self.edge_kind_mask_bytes,
            "fixed-point graph edge kind mask bytes",
        )?;
        crate::dispatch_decode::pack_repeated_u32_into(0, nodes, &mut self.pg_node_tags_bytes)?;
        Ok(())
    }

    /// Analysis family this prepared graph layout belongs to.
    #[must_use]
    pub fn kind(&self) -> FixedPointAnalysisKind {
        self.kind
    }

    /// Verify this graph layout is consumed by the analysis family that
    /// prepared it.
    pub fn require_kind(
        &self,
        expected: FixedPointAnalysisKind,
        consumer: &str,
    ) -> Result<(), String> {
        if self.kind == expected {
            return Ok(());
        }
        Err(prepared_kind_mismatch(consumer, self.kind, expected))
    }

    /// Number of graph nodes in the prepared dispatch domain.
    #[must_use]
    pub fn node_count(&self) -> u32 {
        self.node_count
    }

    /// Number of graph edges in the prepared dispatch domain.
    #[must_use]
    pub fn edge_count(&self) -> u32 {
        self.edge_count
    }

    /// Total retained bytes across invariant packed graph buffers.
    #[must_use]
    pub fn retained_bytes(&self) -> usize {
        self.pg_nodes_bytes.capacity()
            + self.edge_offsets_bytes.capacity()
            + self.edge_targets_bytes.capacity()
            + self.edge_kind_mask_bytes.capacity()
            + self.pg_node_tags_bytes.capacity()
    }

    /// Deterministic hash of the normalized graph layout.
    ///
    /// This is intentionally not `HashMap`'s randomized hasher. Plan and
    /// resident-resource caches need a stable key for equivalent CSR layouts
    /// across runs, devices, and processes.
    #[must_use]
    pub fn stable_layout_hash(&self) -> u64 {
        self.stable_layout_hash
    }
}

#[cold]
fn prepared_kind_mismatch(
    consumer: &str,
    got: FixedPointAnalysisKind,
    expected: FixedPointAnalysisKind,
) -> String {
    let mut scratch = crate::error_format::ErrorFormatScratch::default();
    let _ = std::fmt::Write::write_fmt(
        &mut scratch.buf,
        format_args!(
            "{consumer} received a {got:?} fixed-point graph, expected {expected:?}. Fix: prepare the graph with the matching Weir analysis constructor."
        ),
    );
    scratch.finish()
}