weirflow 0.1.0

use super::{NodeAttr, PathError, PrepareWitnessGraphError};

/// Prepared witness graph for extracting many paths over the same CSR.
///
/// Build this once per program graph with [`prepare_witness_graph`] and
/// call [`extract_path_prepared`] for each source/sink seed. This avoids
/// rebuilding the reverse edge CSR for every finding.
#[derive(Clone, Debug, PartialEq)]
pub struct PreparedWitnessGraph<'a> {
    pub(super) reverse_edges: ReverseEdges,
    pub(super) pg_node_attrs: &'a [NodeAttr],
    pub(super) file_table: &'a [String],
    pub(super) descriptions: &'a [String],
    pub(super) adapter: &'a str,
}

impl PreparedWitnessGraph<'_> {
    /// Number of statement nodes covered by this prepared witness graph.
    #[must_use]
    pub fn node_count(&self) -> usize {
        self.pg_node_attrs.len()
    }

    /// Number of reverse CSR edges prepared once for witness extraction.
    #[must_use]
    pub fn reverse_edge_count(&self) -> usize {
        self.reverse_edges.preds.len()
    }

    /// Number of reverse CSR offsets retained by this prepared graph.
    #[must_use]
    pub fn reverse_offset_count(&self) -> usize {
        self.reverse_edges.offsets.len()
    }
}

/// Prepare reusable witness graph state for repeated path extraction.
///
/// The expensive part is constructing the incoming-edge CSR from the
/// forward program graph. Release proof generation may extract many paths
/// over one graph, so callers should prepare once and reuse.
#[must_use]
#[cfg(any(test, feature = "legacy-infallible"))]
pub fn prepare_witness_graph<'a>(
    edge_offsets: &[u32],
    edge_targets: &[u32],
    edge_kind_mask: &[u32],
    pg_node_attrs: &'a [NodeAttr],
    file_table: &'a [String],
    descriptions: &'a [String],
    adapter: &'a str,
) -> PreparedWitnessGraph<'a> {
    try_prepare_witness_graph(
        edge_offsets,
        edge_targets,
        edge_kind_mask,
        pg_node_attrs,
        file_table,
        descriptions,
        adapter,
    )
    .unwrap_or_else(|error| {
        panic!("malformed Weir witness graph. {error:?}");
    })
}

/// Checked version of [`prepare_witness_graph`].
///
/// Prefer this in production proof assembly when malformed CSR data should be
/// reported as a structured finding/build error instead of a panic.
pub fn try_prepare_witness_graph<'a>(
    edge_offsets: &[u32],
    edge_targets: &[u32],
    edge_kind_mask: &[u32],
    pg_node_attrs: &'a [NodeAttr],
    file_table: &'a [String],
    descriptions: &'a [String],
    adapter: &'a str,
) -> Result<PreparedWitnessGraph<'a>, PathError> {
    let node_count = u32::try_from(pg_node_attrs.len()).map_err(|error| {
        PathError::MalformedGraph {
            reason: PrepareWitnessGraphError::TargetOutOfBounds {
                edge_index: 0,
                target: u32::MAX,
                node_count: u32::MAX,
            },
            fix: format!(
                "Fix: witness graph has {} node attribute rows, which does not fit u32: {error}. Shard the ProgramGraph before witness extraction.",
                pg_node_attrs.len()
            ),
        }
    })?;
    Ok(PreparedWitnessGraph {
        reverse_edges: build_reverse_edges(edge_offsets, edge_targets, edge_kind_mask, node_count)?,
        pg_node_attrs,
        file_table,
        descriptions,
        adapter,
    })
}

/// Compact incoming-edge CSR. `offsets[n]..offsets[n+1]` indexes
/// predecessors and edge-kind masks for node `n`.
#[derive(Clone, Debug, PartialEq)]
pub(super) struct ReverseEdges {
    pub(super) offsets: Vec<u32>,
    pub(super) preds: Vec<u32>,
    pub(super) kinds: Vec<u32>,
}

fn witness_node_count_to_usize(node_count: u32) -> Result<usize, PathError> {
    usize::try_from(node_count).map_err(|source| PathError::MalformedGraph {
        reason: PrepareWitnessGraphError::OffsetCountOverflow { nodes: usize::MAX },
        fix: format!(
            "Fix: witness graph node_count cannot fit usize: {source}; shard the ProgramGraph before witness extraction."
        ),
    })
}

fn witness_edge_count_to_usize(value: u32, label: &'static str) -> Result<usize, PathError> {
    usize::try_from(value).map_err(|source| PathError::MalformedGraph {
        reason: PrepareWitnessGraphError::EdgeCountOverflow {
            edges: usize::MAX,
        },
        fix: format!(
            "Fix: witness graph {label} cannot fit usize: {source}; shard the ProgramGraph before witness extraction."
        ),
    })
}

fn witness_node_id_to_usize(value: u32, label: &'static str) -> Result<usize, PathError> {
    usize::try_from(value).map_err(|source| PathError::MalformedGraph {
        reason: PrepareWitnessGraphError::ReverseEdgeCountOverflow {
            node: usize::MAX,
        },
        fix: format!(
            "Fix: witness graph {label} cannot fit usize: {source}; shard the ProgramGraph before witness extraction."
        ),
    })
}

fn witness_reverse_count_to_usize(
    value: u32,
    node: usize,
    label: &'static str,
) -> Result<usize, PathError> {
    usize::try_from(value).map_err(|source| PathError::MalformedGraph {
        reason: PrepareWitnessGraphError::ReverseEdgeCountOverflow { node },
        fix: format!(
            "Fix: witness graph {label} cannot fit usize: {source}; shard the ProgramGraph before witness extraction."
        ),
    })
}

/// Build incoming-edge CSR once so path reconstruction does not rescan
/// every node for every predecessor lookup or allocate one vector per node.
fn build_reverse_edges(
    edge_offsets: &[u32],
    edge_targets: &[u32],
    edge_kind_mask: &[u32],
    node_count: u32,
) -> Result<ReverseEdges, PathError> {
    let node_count_usize = witness_node_count_to_usize(node_count)?;
    let expected_offsets =
        node_count_usize
            .checked_add(1)
            .ok_or_else(|| {
                PathError::MalformedGraph {
            reason: PrepareWitnessGraphError::OffsetCountOverflow {
                nodes: node_count_usize,
            },
            fix:
                "Fix: witness graph node_count overflows usize when adding the CSR sentinel offset."
                    .to_string(),
        }
            })?;
    if edge_offsets.len() != expected_offsets {
        return Err(PathError::MalformedGraph {
            reason: PrepareWitnessGraphError::InvalidOffsetCount {
                offsets: edge_offsets.len(),
                expected: expected_offsets,
            },
            fix: format!(
                "Fix: witness graph edge_offsets has {} entries, expected exactly node_count + 1 = {expected_offsets}.",
                edge_offsets.len()
            ),
        });
    }
    let edge_count =
        witness_edge_count_to_usize(edge_offsets[node_count_usize], "final CSR offset")?;
    if edge_targets.len() != edge_count || edge_kind_mask.len() != edge_count {
        let reason = if edge_targets.len() != edge_kind_mask.len() {
            PrepareWitnessGraphError::EdgeArrayLengthMismatch {
                targets: edge_targets.len(),
                kinds: edge_kind_mask.len(),
            }
        } else {
            PrepareWitnessGraphError::TerminalOffsetMismatch {
                terminal: edge_count,
                edges: edge_targets.len(),
            }
        };
        return Err(PathError::MalformedGraph {
            reason,
            fix: format!(
                "Fix: witness graph declared {edge_count} edges but edge_targets has {} and edge_kind_mask has {}; CSR edge arrays must match the final offset exactly.",
                edge_targets.len(),
                edge_kind_mask.len()
            ),
        });
    }
    if edge_offsets[0] != 0 {
        return Err(PathError::MalformedGraph {
            reason: PrepareWitnessGraphError::NonZeroInitialOffset {
                first: edge_offsets[0],
            },
            fix: format!(
                "Fix: witness graph edge_offsets[0] is {}, expected 0; rebuild CSR offsets from zero so no orphan prefix edges are ignored.",
                edge_offsets[0]
            ),
        });
    }
    for src in 0..node_count_usize {
        let start = witness_edge_count_to_usize(edge_offsets[src], "row start offset")?;
        let declared_end = witness_edge_count_to_usize(edge_offsets[src + 1], "row end offset")?;
        if start > declared_end {
            return Err(PathError::MalformedGraph {
                reason: PrepareWitnessGraphError::NonMonotonicOffsets {
                    source: src,
                    start,
                    end: declared_end,
                },
                fix: format!(
                    "Fix: witness graph edge_offsets is not monotonic at node {src}: start={start}, end={declared_end}."
                ),
            });
        }
    }
    let mut counts = vec![0u32; node_count_usize];
    for src in 0..node_count_usize {
        let start = witness_edge_count_to_usize(edge_offsets[src], "row start offset")?;
        let declared_end = witness_edge_count_to_usize(edge_offsets[src + 1], "row end offset")?;
        if declared_end > edge_count {
            return Err(PathError::MalformedGraph {
                reason: PrepareWitnessGraphError::OffsetExceedsEdgeCount {
                    offset_index: src + 1,
                    offset: declared_end,
                    edges: edge_count,
                },
                fix: format!(
                    "Fix: witness graph edge_offsets[{next}]={declared_end} exceeds declared edge count {edge_count}.",
                    next = src + 1
                ),
            });
        }
        for (edge_index, &dst) in edge_targets
            .iter()
            .enumerate()
            .take(declared_end)
            .skip(start)
        {
            if dst >= node_count {
                return Err(PathError::MalformedGraph {
                    reason: PrepareWitnessGraphError::TargetOutOfBounds {
                        edge_index,
                        target: dst,
                        node_count,
                    },
                    fix: format!(
                        "Fix: witness graph edge {edge_index} targets node {dst}, but node_count is {node_count}."
                    ),
                });
            }
            let dst_index = witness_node_id_to_usize(dst, "edge target")?;
            counts[dst_index] = counts[dst_index].checked_add(1).ok_or_else(|| {
                PathError::MalformedGraph {
                    reason: PrepareWitnessGraphError::ReverseEdgeCountOverflow { node: dst_index },
                    fix: format!(
                        "Fix: witness graph node {dst_index} has more incoming edges than u32 can encode; shard the ProgramGraph before witness extraction."
                    ),
                }
            })?;
        }
    }

    let mut offsets = vec![0u32; node_count_usize + 1];
    for index in 0..node_count_usize {
        offsets[index + 1] = offsets[index].checked_add(counts[index]).ok_or_else(|| {
            PathError::MalformedGraph {
                reason: PrepareWitnessGraphError::ReverseEdgeCountOverflow { node: index },
                fix: format!(
                    "Fix: witness graph reverse edge offsets overflowed u32 at node {index}; shard the ProgramGraph before witness extraction."
                ),
            }
        })?;
    }
    let reverse_edge_count = witness_reverse_count_to_usize(
        offsets[node_count_usize],
        node_count_usize,
        "reverse edge count",
    )?;
    counts.copy_from_slice(&offsets[..node_count_usize]);
    let mut preds = vec![0u32; reverse_edge_count];
    let mut kinds = vec![0u32; reverse_edge_count];

    for src_u32 in 0..node_count {
        let src = witness_node_id_to_usize(src_u32, "source row")?;
        let start = witness_edge_count_to_usize(edge_offsets[src], "CSR start offset")?;
        let end = witness_edge_count_to_usize(edge_offsets[src + 1], "CSR end offset")?;
        for edge_index in start..end {
            let dst = edge_targets[edge_index];
            let dst_index = witness_node_id_to_usize(dst, "destination node")?;
            let write =
                witness_reverse_count_to_usize(counts[dst_index], dst_index, "reverse cursor")?;
            counts[dst_index] = counts[dst_index].checked_add(1).ok_or_else(|| {
                PathError::MalformedGraph {
                    reason: PrepareWitnessGraphError::ReverseEdgeCountOverflow {
                        node: dst_index,
                    },
                    fix: format!(
                        "Fix: witness graph reverse cursor overflowed u32 for node {dst}; shard the ProgramGraph before witness extraction."
                    ),
                }
            })?;
            preds[write] = src_u32;
            kinds[write] = edge_kind_mask[edge_index];
        }
    }

    Ok(ReverseEdges {
        offsets,
        preds,
        kinds,
    })
}