weir/

def_use.rs

//! DF-11  -  Def-use chain query over a constructed SSA form.
//!
//! For a given SSA-renamed definition site (say `%tmp.3 = …`),
//! return the bitset of nodes that USE that definition. This is the
//! canonical "what reads this def" query that CodeQL exposes as
//! `DataFlow::Node::getUses()` and that every taint rule reaches for
//! when answering "did the value defined here reach a sink without
//! being overwritten."
//!
//! Composes:
//!   * [`super::ssa::SsaForm::def_use_chains`]  -  the per-SSA-version
//!     use list built by Cytron + variable renaming during SSA
//!     construction. We index it by the SSA version of the supplied
//!     definition site and emit a bitset over node ids.
//!
//! Soundness: [`Exact`](super::soundness::Soundness::Exact). The
//! def-use chain is computed by SSA construction, which is sound by
//! construction; this primitive is a pure query against that
//! pre-built table.

use super::ssa::SsaForm;
use vyre::ir::Program;
use vyre_primitives::bitset::and::bitset_and;
use vyre_primitives::graph::csr_forward_traverse::bitset_words;

pub(crate) const OP_ID: &str = "weir::def_use";

/// Look up uses of an SSA definition version. Returns a sorted dense
/// vector of `node_id` values for every use of `def_version`.
///
/// Missing definition versions are rejected because treating an
/// absent SSA definition as "known but unused" creates silent
/// false-negatives in taint and def-use queries.
pub fn try_def_use_chain(form: &SsaForm, def_version: u32) -> Result<Vec<u32>, String> {
    let Some(recorded_uses) = form.def_use_chains.get(&def_version) else {
        return Err(format!(
            "unknown SSA definition version {def_version}. Fix: query only versions emitted by SsaForm::def_use_chains; an absent definition is not an empty use-set."
        ));
    };
    let mut uses = recorded_uses.clone();
    uses.sort_unstable();
    uses.dedup();
    Ok(uses)
}

/// Look up uses of an SSA definition version.
///
/// Panics with an actionable diagnostic when `def_version` is absent.
/// Use [`try_def_use_chain`] when caller-controlled versions must be
/// reported as recoverable input errors.
#[must_use]
#[cfg(any(test, feature = "legacy-infallible"))]
pub fn def_use_chain(form: &SsaForm, def_version: u32) -> Vec<u32> {
    try_def_use_chain(form, def_version)
        .expect("weir::def_use_chain received an unknown SSA definition version")
}

/// Pack the use-set for `def_version` into a bitset over `node_count`.
/// Each bit `i` is `1` iff node `i` is recorded as a use of the
/// supplied definition version. Bits beyond `node_count` are zero.
pub fn try_def_use_chain_bitset(
    form: &SsaForm,
    def_version: u32,
    node_count: u32,
) -> Result<Vec<u32>, String> {
    let uses = try_def_use_chain(form, def_version)?;
    let words = crate::dispatch_decode::bitset_word_capacity("def_use", node_count)?;
    let mut bits = vec![0u32; words];
    for use_id in uses {
        if use_id < node_count {
            let w = crate::dispatch_decode::u32_to_usize(use_id / 32, "def_use bitset word")?;
            let b = use_id % 32;
            bits[w] |= 1u32 << b;
        }
    }
    Ok(bits)
}

/// Pack the use-set for `def_version` into a bitset over `node_count`.
///
/// Panics with an actionable diagnostic when `def_version` is absent.
/// Use [`try_def_use_chain_bitset`] for caller-controlled query ids.
#[must_use]
#[cfg(any(test, feature = "legacy-infallible"))]
pub fn def_use_chain_bitset(form: &SsaForm, def_version: u32, node_count: u32) -> Vec<u32> {
    try_def_use_chain_bitset(form, def_version, node_count)
        .expect("weir::def_use_chain_bitset received an unknown SSA definition version")
}

/// Build a def-use query Program over a precomputed use bitset.
///
/// `def_use_buf` is the bitset produced by [`def_use_chain_bitset`]
/// or an equivalent frontend/importer fact buffer. `query_set` is the
/// set of nodes being queried. The output bitset contains nodes that
/// both read the selected SSA definition and are in the query set.
#[must_use]
pub fn def_use_query(node_count: u32, def_use_buf: &str, query_set: &str, out: &str) -> Program {
    let words = bitset_words(node_count);
    vyre_harness::region::tag_program(OP_ID, bitset_and(def_use_buf, query_set, out, words))
}

/// reference oracle for [`def_use_query`].
#[must_use]
#[cfg(any(test, feature = "cpu-parity"))]
#[deprecated(
    note = "reference oracle only; production code must dispatch the Weir Program on a concrete GPU backend or use weir::oracle for parity evidence"
)]
pub(crate) fn cpu_ref(def_use_buf: &[u32], query_set: &[u32]) -> Vec<u32> {
    vyre_primitives::bitset::and::cpu_ref(def_use_buf, query_set)
}

/// Marker type for the def-use dataflow primitive.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct DefUse;

impl DefUse {
    /// Stable operation id for the def-use primitive.
    #[must_use]
    pub const fn op_id(&self) -> &'static str {
        OP_ID
    }
}

impl super::soundness::SoundnessTagged for DefUse {
    fn soundness(&self) -> super::soundness::Soundness {
        super::soundness::Soundness::Exact
    }
}

inventory::submit! {
    vyre_harness::OpEntry::new(
        OP_ID,
        || def_use_query(4, "def_use", "query", "out"),
        Some(|| {
            let u32s = crate::dispatch_decode::pack_u32;
            vec![vec![u32s(&[0b1110]), u32s(&[0b1010]), u32s(&[0])]]
        }),
        Some(|| {
            let u32s = crate::dispatch_decode::pack_u32;
            vec![vec![u32s(&[0b1010])]]
        }),
    )
}

#[cfg(test)]
#[allow(deprecated)]
mod tests {
    use super::*;
    use std::collections::HashMap;

    fn form_with(def_use: &[(u32, &[u32])]) -> SsaForm {
        let mut chains = HashMap::new();
        for (def, uses) in def_use {
            chains.insert(*def, uses.to_vec());
        }
        SsaForm {
            phi_nodes: HashMap::new(),
            renamed_usages: HashMap::new(),
            def_use_chains: chains,
        }
    }

    #[test]
    fn known_definition_returns_sorted_dedup_uses() {
        let form = form_with(&[(7, &[5, 3, 5, 9, 3])]);
        assert_eq!(def_use_chain(&form, 7), vec![3, 5, 9]);
    }

    #[test]
    fn known_definition_with_no_uses_remains_empty() {
        let form = form_with(&[(7, &[])]);
        assert_eq!(
            try_def_use_chain(&form, 7).expect("known empty definition"),
            Vec::<u32>::new()
        );
    }

    #[test]
    fn unknown_definition_is_rejected_not_treated_as_empty() {
        let form = form_with(&[(7, &[3, 5])]);
        let error = try_def_use_chain(&form, 99).expect_err("unknown definition must fail");
        assert!(
            error.contains("unknown SSA definition version 99"),
            "{error}"
        );
    }

    #[test]
    fn bitset_packs_use_ids() {
        let form = form_with(&[(7, &[1, 33, 65])]);
        let bits = def_use_chain_bitset(&form, 7, 96);
        // Bit 1 in word 0; bit 1 in word 1 (offset 33-32); bit 1 in word 2 (65-64).
        assert_eq!(bits, vec![1u32 << 1, 1u32 << 1, 1u32 << 1]);
    }

    #[test]
    fn bitset_skips_uses_beyond_node_count() {
        let form = form_with(&[(7, &[3, 200])]);
        let bits = def_use_chain_bitset(&form, 7, 64);
        // Only bit 3 in word 0 is set; the 200 use exceeds node_count.
        assert_eq!(bits, vec![1u32 << 3, 0]);
    }

    #[test]
    fn empty_form_rejects_unknown_definition() {
        let form = form_with(&[]);
        let err = try_def_use_chain_bitset(&form, 7, 64)
            .expect_err("unknown definition in empty form must fail loudly");
        assert!(
            err.contains("unknown SSA definition version 7"),
            "unexpected diagnostic: {err}"
        );
    }

    #[test]
    fn def_use_query_cpu_ref_intersects_query() {
        assert_eq!(cpu_ref(&[0b1110], &[0b1010]), vec![0b1010]);
    }
}