vyre_libs/security/
sink_intersection.rs1use 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#[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#[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
76pub 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}