weirflow 0.1.0

use super::bitmask::try_bit_is_set;
use super::graph::{try_prepare_witness_graph, PreparedWitnessGraph, ReverseEdges};
use super::{
    ExtractedPath, ExtractedStatement, NodeAttr, PathError, PathSeed, PrepareWitnessGraphError,
};
use vyre_primitives::predicate::edge_kind;

const IFDS_PATH_EDGE_MASK: u32 = edge_kind::ASSIGNMENT
    | edge_kind::CALL_ARG
    | edge_kind::RETURN
    | edge_kind::PHI
    | edge_kind::ALIAS
    | edge_kind::MEM_STORE
    | edge_kind::MEM_LOAD
    | edge_kind::MUT_REF;

pub(super) const MAX_PATH_DEPTH: usize = 1024;

/// One prepared witness extraction request.
///
/// `source_reach` is intentionally per request because witness soundness
/// requires a per-source reachability mask. Aggregating masks across sources
/// can produce a path that starts at the wrong source.
#[derive(Clone, Copy, Debug)]
pub struct PathExtractionRequest<'a> {
    /// Source/sink pair to explain.
    pub seed: &'a PathSeed,
    /// Per-source reachability mask for `seed.source_node`.
    pub source_reach: &'a [u32],
    /// Per-rule sanitizer mask. Empty means the rule has no sanitizers.
    pub sanitizer_mask: &'a [u32],
}

/// Batch witness extraction counters for one prepared graph.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct PreparedWitnessBatchStats {
    /// Request rows consumed by the batch call.
    pub requests: u64,
    /// Paths reconstructed successfully.
    pub successes: u64,
    /// Request rows that returned a structured witness error.
    pub errors: u64,
    /// Prepared reverse-CSR graph reuses. This equals `requests` for the
    /// prepared batch API and is exposed so release evidence can prove callers
    /// did not rebuild the reverse graph per finding.
    pub prepared_graph_reuses: u64,
    /// Statement nodes in the prepared witness graph.
    pub prepared_node_count: u64,
    /// Reverse CSR edge entries retained by the prepared witness graph.
    pub prepared_reverse_edges: u64,
}

fn path_node_index(node: u32, node_count: u32) -> Result<usize, PathError> {
    if node >= node_count {
        return Err(PathError::NodeOutOfBounds { node, node_count });
    }
    usize::try_from(node).map_err(|error| PathError::MalformedGraph {
        reason: PrepareWitnessGraphError::TargetOutOfBounds {
            edge_index: 0,
            target: node,
            node_count,
        },
        fix: format!(
            "Fix: witness path node id {node} cannot fit usize: {error}; shard the ProgramGraph before witness extraction."
        ),
    })
}

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

fn witness_bit_is_set(bitmask: &[u32], node: u32, field: &'static str) -> Result<bool, PathError> {
    try_bit_is_set(bitmask, node).map_err(|error| PathError::MalformedGraph {
        reason: PrepareWitnessGraphError::TargetOutOfBounds {
            edge_index: 0,
            target: node,
            node_count: u32::MAX,
        },
        fix: format!("Fix: witness {field} bitmask cannot address node {node}: {error}"),
    })
}

/// Reconstruct a source→sink path by walking the CSR backward from
/// `seed.sink_node`, greedily choosing a *source-reached*,
/// *non-sanitized* predecessor at each step until `seed.source_node`
/// is hit.
///
/// `source_reach` MUST be a per-source reachability mask: bit `n` is
/// set iff the *specific source named in the `PathSeed`* reaches
/// node `n` along an IFDS-eligible path. Aggregate-over-all-sources
/// masks must NOT be passed here  -  they admit witnesses that
/// originate from a different source than the seed (audit 2026-04-27
/// finding 3).
///
/// `sanitizer_mask` is a per-node bitmask whose bit `n` is set iff
/// node `n` is a sanitizer for this rule. The walker rejects any
/// sanitizer-tagged predecessor (audit finding 2). Pass an empty
/// slice (`&[]`) iff the rule has no sanitizers.
///
/// `pg_node_attrs` is a slice of `(byte_start, byte_end, file_idx)`
/// tuples  -  one per node  -  that a consumer pipeline emits alongside
/// the CSR. `file_table` maps `file_idx` to a repository-relative
/// path. `adapter` is the language-adapter id (e.g. `"c-c11"`).
/// `descriptions` is per-node short text.
///
/// Returns:
/// - `Ok(path)` when source is reached from sink along an
///   IFDS-eligible chain that does not cross any sanitizer.
/// - `Err(PathError::NoPath)` when no such path exists.
/// - `Err(PathError::DepthExceeded { partial_chain })` when the
///   chain exceeded `MAX_PATH_DEPTH`.
/// - `Err(PathError::NodeOutOfBounds { .. })` when the seed names a
///   node id that is not present in `pg_node_attrs`.
#[allow(clippy::too_many_arguments)]
pub fn extract_path(
    seed: &PathSeed,
    source_reach: &[u32],
    sanitizer_mask: &[u32],
    edge_offsets: &[u32],
    edge_targets: &[u32],
    edge_kind_mask: &[u32],
    pg_node_attrs: &[NodeAttr],
    file_table: &[String],
    descriptions: &[String],
    adapter: &str,
) -> Result<ExtractedPath, PathError> {
    let graph = try_prepare_witness_graph(
        edge_offsets,
        edge_targets,
        edge_kind_mask,
        pg_node_attrs,
        file_table,
        descriptions,
        adapter,
    )?;
    extract_path_prepared(seed, source_reach, sanitizer_mask, &graph)
}

/// Reconstruct a source→sink path using a reusable prepared graph.
pub fn extract_path_prepared(
    seed: &PathSeed,
    source_reach: &[u32],
    sanitizer_mask: &[u32],
    graph: &PreparedWitnessGraph<'_>,
) -> Result<ExtractedPath, PathError> {
    let node_count = u32::try_from(graph.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: prepared witness graph has {} node attribute rows, which does not fit u32: {error}. Shard the ProgramGraph before witness extraction.",
                graph.pg_node_attrs.len()
            ),
        }
    })?;
    // Bound-check the seed before indexing into visited/attrs.
    if seed.source_node >= node_count {
        return Err(PathError::NodeOutOfBounds {
            node: seed.source_node,
            node_count,
        });
    }
    if seed.sink_node >= node_count {
        return Err(PathError::NodeOutOfBounds {
            node: seed.sink_node,
            node_count,
        });
    }
    let _source_index = path_node_index(seed.source_node, node_count)?;
    let sink_index = path_node_index(seed.sink_node, node_count)?;

    if seed.source_node == seed.sink_node {
        // Trivial path of length 1. Sanitizer check still applies:
        // if the source itself is tagged a sanitizer the witness is
        // unsound.
        if witness_bit_is_set(sanitizer_mask, seed.source_node, "sanitizer")? {
            return Err(PathError::NoPath);
        }
        let stmt = build_statement(
            seed.source_node,
            node_count,
            graph.pg_node_attrs,
            graph.file_table,
            graph.descriptions,
            graph.adapter,
        )?;
        return Ok(ExtractedPath {
            statements: vec![stmt],
        });
    }

    if !witness_bit_is_set(source_reach, seed.sink_node, "source reachability")?
        || !witness_bit_is_set(source_reach, seed.source_node, "source reachability")?
    {
        // The seeded source did not reach the seeded sink.
        return Err(PathError::NoPath);
    }
    // The endpoints themselves cannot be sanitizers  -  that would
    // mean the proof crosses a sanitizer at iteration 0.
    if witness_bit_is_set(sanitizer_mask, seed.sink_node, "sanitizer")?
        || witness_bit_is_set(sanitizer_mask, seed.source_node, "sanitizer")?
    {
        return Err(PathError::NoPath);
    }

    let mut visited = vec![false; graph.pg_node_attrs.len()];
    let mut chain: Vec<u32> = Vec::new();
    let mut current = seed.sink_node;
    chain.push(current);
    visited[sink_index] = true;

    for _ in 0..MAX_PATH_DEPTH {
        if current == seed.source_node {
            break;
        }
        let Some(pred) = find_reached_predecessor(
            current,
            node_count,
            source_reach,
            sanitizer_mask,
            &graph.reverse_edges,
            &visited,
        )?
        else {
            return Err(PathError::NoPath);
        };
        chain.push(pred);
        let pred_index = path_node_index(pred, node_count)?;
        visited[pred_index] = true;
        current = pred;
    }

    if current != seed.source_node {
        // Hit the depth limit before reaching the source. Hand the
        // partial chain back as evidence for an analysis-capacity
        // failure. Callers must increase the witness budget rather
        // than downgrade an incomplete proof.
        return Err(PathError::DepthExceeded {
            partial_chain: chain,
        });
    }

    // chain is sink-first; flip to source-first.
    chain.reverse();
    let mut statements =
        crate::staging_reserve::reserved_vec(chain.len(), "reachability witness statements")
            .map_err(|error| PathError::ResourceExhausted {
                field: "reachability witness statements".to_string(),
                fix: format!(
                    "Fix: {error}; reduce the witness batch size or shard the ProgramGraph before extracting paths."
                ),
            })?;
    for node in chain {
        statements.push(build_statement(
            node,
            node_count,
            graph.pg_node_attrs,
            graph.file_table,
            graph.descriptions,
            graph.adapter,
        )?);
    }
    Ok(ExtractedPath { statements })
}

/// Reconstruct many source→sink paths over one prepared witness graph.
///
/// Each request produces its own `Result`, so an out-of-bounds seed, sanitizer
/// block, or depth limit on one finding does not discard paths for unrelated
/// findings in the same program graph.
pub fn extract_paths_prepared(
    requests: &[PathExtractionRequest<'_>],
    graph: &PreparedWitnessGraph<'_>,
) -> Result<(Vec<Result<ExtractedPath, PathError>>, PreparedWitnessBatchStats), PathError> {
    let mut results = crate::staging_reserve::reserved_vec(
        requests.len(),
        "reachability prepared witness batch results",
    )
    .map_err(|error| PathError::ResourceExhausted {
        field: "reachability prepared witness batch results".to_string(),
        fix: format!(
            "Fix: {error}; reduce the witness batch size or shard the ProgramGraph before extracting paths."
        ),
    })?;
    let stats = extract_paths_prepared_into(requests, graph, &mut results)?;
    Ok((results, stats))
}

/// Reconstruct many source→sink paths into caller-owned result storage over
/// one prepared witness graph.
pub fn extract_paths_prepared_into(
    requests: &[PathExtractionRequest<'_>],
    graph: &PreparedWitnessGraph<'_>,
    results: &mut Vec<Result<ExtractedPath, PathError>>,
) -> Result<PreparedWitnessBatchStats, PathError> {
    results.clear();
    if results.capacity() < requests.len() {
        results
            .try_reserve(requests.len() - results.capacity())
            .map_err(|error| PathError::ResourceExhausted {
                field: "reachability prepared witness batch results".to_string(),
                fix: format!(
                    "Fix: could not reserve {} prepared witness result slot(s): {error}; reduce the witness batch size or shard the ProgramGraph before extracting paths.",
                    requests.len()
                ),
            })?;
    }

    let mut successes = 0u64;
    let mut errors = 0u64;
    for request in requests {
        let result = extract_path_prepared(
            request.seed,
            request.source_reach,
            request.sanitizer_mask,
            graph,
        );
        if result.is_ok() {
            successes = successes.checked_add(1).ok_or_else(|| {
                PathError::ResourceExhausted {
                    field: "reachability prepared witness batch success counter".to_string(),
                    fix: "Fix: prepared witness success counter overflowed u64; shard the witness request batch.".to_string(),
                }
            })?;
        } else {
            errors = errors.checked_add(1).ok_or_else(|| {
                PathError::ResourceExhausted {
                    field: "reachability prepared witness batch error counter".to_string(),
                    fix: "Fix: prepared witness error counter overflowed u64; shard the witness request batch.".to_string(),
                }
            })?;
        }
        results.push(result);
    }

    let requests_u64 = u64::try_from(requests.len()).map_err(|error| {
        PathError::ResourceExhausted {
            field: "reachability prepared witness batch request counter".to_string(),
            fix: format!(
                "Fix: witness request count does not fit u64: {error}; shard the witness request batch."
            ),
        }
    })?;
    let prepared_node_count = u64::try_from(graph.node_count()).map_err(|error| {
        PathError::ResourceExhausted {
            field: "reachability prepared witness graph node counter".to_string(),
            fix: format!(
                "Fix: prepared witness node count does not fit u64: {error}; shard the ProgramGraph before extracting paths."
            ),
        }
    })?;
    let prepared_reverse_edges = u64::try_from(graph.reverse_edge_count()).map_err(|error| {
        PathError::ResourceExhausted {
            field: "reachability prepared witness reverse edge counter".to_string(),
            fix: format!(
                "Fix: prepared witness reverse edge count does not fit u64: {error}; shard the ProgramGraph before extracting paths."
            ),
        }
    })?;
    Ok(PreparedWitnessBatchStats {
        requests: requests_u64,
        successes,
        errors,
        prepared_graph_reuses: requests_u64,
        prepared_node_count,
        prepared_reverse_edges,
    })
}

/// Find any unvisited node `pred` such that there's an incoming
/// IFDS-eligible edge `pred -> current` AND `pred` is reached by the
/// seeded source's reachability mask AND `pred` is NOT a sanitizer.
/// Returns `None` when no such pred exists.
fn find_reached_predecessor(
    current: u32,
    node_count: u32,
    source_reach: &[u32],
    sanitizer_mask: &[u32],
    reverse_edges: &ReverseEdges,
    visited: &[bool],
) -> Result<Option<u32>, PathError> {
    let current_index = path_node_index(current, node_count)?;
    let next_index = current_index.checked_add(1).ok_or_else(|| PathError::MalformedGraph {
        reason: PrepareWitnessGraphError::InvalidOffsetCount {
            offsets: reverse_edges.offsets.len(),
            expected: usize::MAX,
        },
        fix: format!(
            "Fix: witness reverse-edge offset index overflowed after node {current}; rebuild the prepared witness graph."
        ),
    })?;
    let start_value =
        *reverse_edges
            .offsets
            .get(current_index)
            .ok_or_else(|| PathError::MalformedGraph {
                reason: PrepareWitnessGraphError::InvalidOffsetCount {
                    offsets: reverse_edges.offsets.len(),
                    expected: next_index,
                },
                fix: format!(
                    "Fix: witness reverse-edge offsets are missing row start for node {current}."
                ),
            })?;
    let end_value =
        *reverse_edges
            .offsets
            .get(next_index)
            .ok_or_else(|| PathError::MalformedGraph {
                reason: PrepareWitnessGraphError::InvalidOffsetCount {
                    offsets: reverse_edges.offsets.len(),
                    expected: match next_index.checked_add(1) {
                        Some(expected) => expected,
                        None => usize::MAX,
                    },
                },
                fix: format!(
                    "Fix: witness reverse-edge offsets are missing row end for node {current}."
                ),
            })?;
    let start = reverse_offset_to_usize(start_value, current_index, "reverse row start")?;
    let end = reverse_offset_to_usize(end_value, current_index, "reverse row end")?;
    for edge_index in start..end {
        let pred =
            *reverse_edges
                .preds
                .get(edge_index)
                .ok_or_else(|| PathError::MalformedGraph {
                    reason: PrepareWitnessGraphError::TerminalOffsetMismatch {
                        terminal: end,
                        edges: reverse_edges.preds.len(),
                    },
                    fix: format!(
                    "Fix: witness reverse predecessor array is shorter than reverse offset {end}."
                ),
                })?;
        let mask =
            *reverse_edges
                .kinds
                .get(edge_index)
                .ok_or_else(|| PathError::MalformedGraph {
                    reason: PrepareWitnessGraphError::TerminalOffsetMismatch {
                        terminal: end,
                        edges: reverse_edges.kinds.len(),
                    },
                    fix: format!(
                    "Fix: witness reverse edge-kind array is shorter than reverse offset {end}."
                ),
                })?;
        let pred_index = path_node_index(pred, node_count)?;
        if visited[pred_index] {
            continue;
        }
        if !witness_bit_is_set(source_reach, pred, "source reachability")? {
            continue;
        }
        if witness_bit_is_set(sanitizer_mask, pred, "sanitizer")? {
            // A sanitizer-tagged predecessor would make the witness
            // cross a sanitizer; reject (audit finding 2).
            continue;
        }
        if mask & IFDS_PATH_EDGE_MASK != 0 {
            return Ok(Some(pred));
        }
    }
    Ok(None)
}

fn build_statement(
    node_id: u32,
    node_count: u32,
    pg_node_attrs: &[NodeAttr],
    file_table: &[String],
    descriptions: &[String],
    adapter: &str,
) -> Result<ExtractedStatement, PathError> {
    let node_index = path_node_index(node_id, node_count)?;
    let attr = pg_node_attrs
        .get(node_index)
        .ok_or(PathError::NodeOutOfBounds {
            node: node_id,
            node_count,
        })?;
    let file = file_table
        .get(usize::try_from(attr.file_idx).ok().unwrap_or(usize::MAX))
        .cloned()
        .unwrap_or_else(|| String::from("<unknown>"));
    let description = descriptions
        .get(node_index)
        .cloned()
        .unwrap_or_else(|| format!("node {node_id}"));
    Ok(ExtractedStatement {
        adapter: adapter.to_string(),
        description,
        file,
        node_id,
        byte_start: attr.byte_start,
        byte_end: attr.byte_end,
    })
}