Skip to main content

catscope_edge_generator/primitive/
tree.rs

1use std::collections::HashMap;
2
3use once_cell::sync::Lazy;
4use solana_sdk::{clock::Slot, pubkey::Pubkey};
5
6use super::err::CatscopeWasmError;
7
8pub type Weight = u32;
9pub const WEIGHT_IS_OUTGOING: Weight = 1 << 0;
10pub const WEIGHT_SLOT: Weight = 1 << 1;
11// this is for the host to map accounts to a program for when clients sync.
12pub const WEIGHT_CLIENT: Weight = 1 << 2;
13/// upload the destination node to a client.
14pub const WEIGHT_UPLOAD: Weight = 1 << 3;
15
16pub const WEIGHT_NON_ACCOUNT: Weight =
17    WEIGHT_SLOT | WEIGHT_CLIENT | WEIGHT_UPLOAD | WEIGHT_IS_OUTGOING;
18pub const WEIGHT_ACCOUNT: Weight = !WEIGHT_NON_ACCOUNT;
19
20pub const MAX_WEIGHT_NONACCOUNT_EXPONENT: u8 = 3;
21
22pub const WEIGHT_PROGRAM: Weight = 1 << 4;
23pub const WEIGHT_SPLTOKEN_OWNER: Weight = 1 << 5;
24pub const WEIGHT_SPLTOKEN_MINT: Weight = 1 << 6;
25pub const WEIGHT_DIRECT: Weight = 1 << 7;
26pub const WEIGHT_SYMLINK: Weight = 1 << 8;
27
28// the memory requirements grow exponentionally (doubles) for every integer increment of this
29// variable.
30pub const MAX_WEIGHT_ACCOUNT_EXPONENT: u8 = 9;
31pub const MAX_WEIGHT: Weight = 1 << MAX_WEIGHT_ACCOUNT_EXPONENT;
32
33/// Index weights for use when updating subscriptions.
34static WEIGHT_HASH_MAP: Lazy<HashMap<Weight, Vec<Weight>>> = Lazy::new(|| {
35    let mut lookup: HashMap<Weight, Vec<Weight>> = HashMap::new();
36    // 0..32
37    for mask in 0..(1 << (MAX_WEIGHT_ACCOUNT_EXPONENT - MAX_WEIGHT_NONACCOUNT_EXPONENT)) {
38        let mut subset = mask;
39        // Generate all subsets (masks) of 'b'
40        loop {
41            lookup
42                .entry(mask)
43                .or_default()
44                .push(subset << MAX_WEIGHT_NONACCOUNT_EXPONENT);
45            if subset == 0 {
46                break;
47            }
48            subset = (subset - 1) & mask; // Generate next subset
49        }
50    }
51    lookup
52});
53
54#[repr(C, align(8))]
55#[derive(Debug, Clone, Default)]
56pub struct ProgramList {
57    pub count: u16,
58    pub list: [Pubkey; 32], // have a max length
59}
60
61/// The edge goes in the graph determined by the `from` `program_id`.
62/// `from` is the account_id.
63#[repr(C, align(8))]
64#[derive(Debug, Clone, Default)]
65pub struct FilterEdge {
66    pub slot: Slot,
67    pub to: Pubkey,
68    pub from: Pubkey,
69    pub weight: Weight, // weight of zero is not allowed
70}
71impl FilterEdge {
72    pub fn from_raw_parts<'a, 'b: 'a>(data: &'b [u8]) -> Result<&'a Self, CatscopeWasmError> {
73        if data.len() < std::mem::size_of::<Self>() {
74            return Err(CatscopeWasmError::InsufficientBuffer);
75        }
76        let filter: &Self = unsafe { &*(data.as_ptr() as *const Self) };
77        Ok(filter)
78    }
79    pub fn set_outgoing(&mut self, id: &Pubkey) {
80        if self.from.eq(id) {
81            self.weight |= WEIGHT_IS_OUTGOING;
82        }
83    }
84}
85
86#[inline(always)]
87pub fn edge_is_outgoing(weight: &Weight) -> bool {
88    0 < *weight & WEIGHT_IS_OUTGOING
89}
90
91#[inline]
92fn zero_out_bits_above_n(value: u32, n: u32) -> u32 {
93    // Create a mask with 1s for bits <= n and 0s for bits > n.
94    // We use wrapping_shl to handle potential overflow if n is close to 32.
95    let mask = if n >= 31 {
96        u32::MAX // If n is 31 or more, keep all bits.
97    } else {
98        (1u32 << (n + 1)).wrapping_sub(1) // Create the mask (e.g., n=2 -> 0b111).
99    };
100
101    value & mask // Apply the mask to zero out the unwanted bits.
102}
103
104// Return all bit map combinations covered by this weight.
105pub fn weight_list(weight: &Weight) -> &'static [Weight] {
106    let reduced = zero_out_bits_above_n(*weight, MAX_WEIGHT_ACCOUNT_EXPONENT as u32 - 1);
107    let k = reduced >> MAX_WEIGHT_NONACCOUNT_EXPONENT;
108    let result = match WEIGHT_HASH_MAP.get(&k) {
109        Some(x) => x,
110        None => panic!(
111            "failed to get k {}; reduced {}; weight {};",
112            k, reduced, weight
113        ),
114    };
115    result
116}
117pub fn parse_program_list(input: &[u8]) -> Result<Vec<Pubkey>, CatscopeWasmError> {
118    let input_str = match std::str::from_utf8(input) {
119        Ok(x) => x,
120        Err(e) => return Err(CatscopeWasmError::Unknown(e.to_string())),
121    };
122    let pre_list: Vec<&str> = input_str.split(',').collect();
123    let mut list = Vec::with_capacity(pre_list.len());
124    for i in 0..pre_list.len() {
125        let y = pre_list[i].trim();
126        let x: Pubkey = match y.try_into() {
127            Ok(z) => z,
128            Err(e) => return Err(CatscopeWasmError::Unknown(e.to_string())),
129        };
130        list.push(x);
131    }
132    Ok(list)
133}