Skip to main content

vyre_libs/security/
sink_intersection.rs

1//! `sink_intersection`  -  count how many of a query set are also in
2//! a sink-family bitset. Used by rules that want a fractional
3//! confidence ("X% of nodes reachable from source landed in sinks").
4
5use std::sync::Arc;
6
7use vyre::ir::model::expr::Ident;
8use vyre::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program, UnOp};
9use vyre_primitives::graph::csr_forward_traverse::bitset_words;
10
11pub(crate) const OP_ID: &str = "vyre-libs::security::sink_intersection";
12
13/// Build a sink-intersection-count Program. AND query with sink_set,
14/// popcount the result, write to out_scalar.
15#[must_use]
16pub fn sink_intersection(
17    node_count: u32,
18    query_set: &str,
19    sink_set: &str,
20    intersect_buf: &str,
21    out_scalar: &str,
22) -> Program {
23    let words = bitset_words(node_count);
24    let t = Expr::InvocationId { axis: 0 };
25    let body = vec![
26        Node::let_bind(
27            "intersect",
28            Expr::bitand(
29                Expr::load(query_set, t.clone()),
30                Expr::load(sink_set, t.clone()),
31            ),
32        ),
33        Node::store(intersect_buf, t.clone(), Expr::var("intersect")),
34        Node::let_bind(
35            "count",
36            Expr::UnOp {
37                op: UnOp::Popcount,
38                operand: Box::new(Expr::var("intersect")),
39            },
40        ),
41        Node::let_bind(
42            "_",
43            Expr::atomic_add(out_scalar, Expr::u32(0), Expr::var("count")),
44        ),
45    ];
46    Program::wrapped(
47        vec![
48            BufferDecl::storage(query_set, 0, BufferAccess::ReadOnly, DataType::U32)
49                .with_count(words),
50            BufferDecl::storage(sink_set, 1, BufferAccess::ReadOnly, DataType::U32)
51                .with_count(words),
52            BufferDecl::storage(intersect_buf, 2, BufferAccess::ReadWrite, DataType::U32)
53                .with_count(words),
54            BufferDecl::output(out_scalar, 3, DataType::U32).with_count(1),
55        ],
56        [256, 1, 1],
57        vec![Node::Region {
58            generator: Ident::from(OP_ID),
59            source_region: None,
60            body: Arc::new(vec![Node::if_then(
61                Expr::lt(t.clone(), Expr::u32(words)),
62                body,
63            )]),
64        }],
65    )
66}
67
68/// CPU oracle: count of bits set in `query AND sink`.
69#[must_use]
70#[cfg(test)]
71pub(crate) fn cpu_ref(query_set: &[u32], sink_set: &[u32]) -> u32 {
72    let inter = vyre_primitives::bitset::and::cpu_ref(query_set, sink_set);
73    inter.iter().map(|w| w.count_ones()).sum()
74}
75
76/// Soundness marker for [`sink_intersection`].
77pub struct SinkIntersection;
78impl vyre::soundness::SoundnessTagged for SinkIntersection {
79    fn soundness(&self) -> vyre::soundness::Soundness {
80        vyre::soundness::Soundness::Exact
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87
88    #[test]
89    fn full_overlap_counts_all_set_bits() {
90        assert_eq!(cpu_ref(&[0b1111], &[0b1111]), 4);
91    }
92
93    #[test]
94    fn no_overlap_returns_zero() {
95        assert_eq!(cpu_ref(&[0b1010], &[0b0101]), 0);
96    }
97
98    #[test]
99    fn partial_overlap_counts_intersection() {
100        assert_eq!(cpu_ref(&[0b1110], &[0b0111]), 2);
101    }
102
103    #[test]
104    fn distributes_across_words() {
105        assert_eq!(cpu_ref(&[0xFF00, 0x00FF], &[0xFFFF, 0xFFFF]), 16);
106    }
107}