1#[cfg(all(not(target_arch = "wasm32"), target_arch = "x86_64"))]
4use crate::bundle_simd::{finalize_simd_avx2, update_counts_simd_avx2};
5#[cfg(all(not(target_arch = "wasm32"), target_arch = "aarch64"))]
6use crate::bundle_simd::{finalize_simd_neon, update_counts_simd_neon};
7use crate::error::{MemoryError, Result};
8use crate::hyperdim::HVec10240;
9
10#[derive(Debug, Clone)]
15pub struct BundleAccumulator {
16 counts: Box<[i32; HVec10240::DIMENSION]>,
17 n: u32,
18}
19
20impl Default for BundleAccumulator {
21 fn default() -> Self {
22 Self {
23 counts: Box::new([0i32; HVec10240::DIMENSION]),
24 n: 0,
25 }
26 }
27}
28
29impl BundleAccumulator {
30 pub fn new() -> Self {
32 Self {
33 counts: Box::new([0i32; HVec10240::DIMENSION]),
34 n: 0,
35 }
36 }
37
38 pub fn add(&mut self, hv: &HVec10240) {
40 #[cfg(all(not(target_arch = "wasm32"), target_arch = "x86_64"))]
41 {
42 if is_x86_feature_detected!("avx2") {
43 unsafe { update_counts_simd_avx2(&mut self.counts, &hv.data, 1) };
45 self.n += 1;
46 return;
47 }
48 }
49
50 #[cfg(all(not(target_arch = "wasm32"), target_arch = "aarch64"))]
51 {
52 unsafe { update_counts_simd_neon(&mut self.counts, &hv.data, 1) };
54 self.n += 1;
55 }
56
57 #[cfg(not(all(not(target_arch = "wasm32"), target_arch = "aarch64")))]
58 {
59 for i in 0..80 {
60 let mut val = hv.data[i];
61 while val != 0 {
62 let j = val.trailing_zeros() as usize;
63 self.counts[i * 128 + j] += 1;
64 val &= val - 1;
65 }
66 }
67 self.n += 1;
68 }
69 }
70
71 pub fn remove(&mut self, hv: &HVec10240) {
76 if self.n == 0 {
77 return;
78 }
79
80 #[cfg(all(not(target_arch = "wasm32"), target_arch = "x86_64"))]
81 {
82 if is_x86_feature_detected!("avx2") {
83 unsafe { update_counts_simd_avx2(&mut self.counts, &hv.data, -1) };
85 self.n -= 1;
86 return;
87 }
88 }
89
90 #[cfg(all(not(target_arch = "wasm32"), target_arch = "aarch64"))]
91 {
92 unsafe { update_counts_simd_neon(&mut self.counts, &hv.data, -1) };
94 self.n -= 1;
95 }
96
97 #[cfg(not(all(not(target_arch = "wasm32"), target_arch = "aarch64")))]
98 {
99 for i in 0..80 {
100 let mut val = hv.data[i];
101 while val != 0 {
102 let j = val.trailing_zeros() as usize;
103 self.counts[i * 128 + j] -= 1;
104 val &= val - 1;
105 }
106 }
107 self.n -= 1;
108 }
109 }
110
111 pub fn try_remove(&mut self, hv: &HVec10240) -> Result<()> {
115 if self.n == 0 {
116 return Err(MemoryError::InvalidInput {
117 field: "accumulator".to_string(),
118 reason: "cannot remove from empty BundleAccumulator".to_string(),
119 });
120 }
121
122 #[cfg(all(not(target_arch = "wasm32"), target_arch = "x86_64"))]
123 {
124 if is_x86_feature_detected!("avx2") {
125 unsafe { update_counts_simd_avx2(&mut self.counts, &hv.data, -1) };
127 self.n -= 1;
128 return Ok(());
129 }
130 }
131
132 #[cfg(all(not(target_arch = "wasm32"), target_arch = "aarch64"))]
133 {
134 unsafe { update_counts_simd_neon(&mut self.counts, &hv.data, -1) };
136 self.n -= 1;
137 return Ok(());
138 }
139
140 #[cfg(not(all(not(target_arch = "wasm32"), target_arch = "aarch64")))]
141 {
142 for i in 0..80 {
143 let mut val = hv.data[i];
144 while val != 0 {
145 let j = val.trailing_zeros() as usize;
146 self.counts[i * 128 + j] -= 1;
147 val &= val - 1;
148 }
149 }
150 self.n -= 1;
151 Ok(())
152 }
153 }
154
155 pub fn finalize(&self) -> HVec10240 {
160 if self.n == 0 {
161 return HVec10240::zero();
162 }
163
164 let threshold = (self.n / 2) as i32;
165
166 #[cfg(all(not(target_arch = "wasm32"), target_arch = "x86_64"))]
167 {
168 if is_x86_feature_detected!("avx2") {
169 return HVec10240 {
171 data: unsafe { finalize_simd_avx2(&self.counts, threshold) },
172 };
173 }
174 }
175
176 #[cfg(all(not(target_arch = "wasm32"), target_arch = "aarch64"))]
177 {
178 return HVec10240 {
180 data: unsafe { finalize_simd_neon(&self.counts, threshold) },
181 };
182 }
183
184 #[cfg(not(all(not(target_arch = "wasm32"), target_arch = "aarch64")))]
185 {
186 let mut data = [0u128; 80];
187
188 for (i, word) in data.iter_mut().enumerate() {
189 let offset = i * 128;
190 for j in 0..128 {
191 let condition = self.counts[offset + j] > threshold;
193 *word |= (condition as u128) << j;
194 }
195 }
196
197 HVec10240 { data }
198 }
199 }
200
201 pub const fn len(&self) -> u32 {
203 self.n
204 }
205
206 pub const fn is_empty(&self) -> bool {
208 self.n == 0
209 }
210
211 pub fn clear(&mut self) {
213 *self.counts = [0i32; HVec10240::DIMENSION];
214 self.n = 0;
215 }
216}
217
218#[cfg(test)]
219mod tests {
220 #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
221 use super::*;
222
223 #[test]
224 fn test_bundle_accumulator_add_finalize() {
225 let v1 = HVec10240::random();
226 let v2 = HVec10240::random();
227 let v3 = HVec10240::random();
228
229 let mut acc = BundleAccumulator::new();
230 acc.add(&v1);
231 acc.add(&v2);
232 acc.add(&v3);
233
234 let bundled = acc.finalize();
235 assert_ne!(bundled, HVec10240::zero());
237 assert_eq!(acc.len(), 3);
239 }
240
241 #[test]
242 fn test_bundle_accumulator_remove() {
243 let v1 = HVec10240::random();
244 let v2 = HVec10240::random();
245
246 let mut acc = BundleAccumulator::new();
247 acc.add(&v1);
248 acc.add(&v2);
249 acc.remove(&v2);
250
251 assert_eq!(acc.len(), 1);
252 let bundled = acc.finalize();
253 assert!(bundled.cosine_similarity(&v1) > 0.9);
255 }
256
257 #[test]
258 fn test_bundle_accumulator_empty() {
259 let acc = BundleAccumulator::new();
260 assert!(acc.is_empty());
261 assert_eq!(acc.finalize(), HVec10240::zero());
262 }
263
264 #[test]
265 fn test_bundle_accumulator_clear() {
266 let mut acc = BundleAccumulator::new();
267 acc.add(&HVec10240::random());
268 acc.clear();
269 assert!(acc.is_empty());
270 }
271}