Skip to main content

netgauze_analytics/
flow.rs

1// Copyright (C) 2023-present The NetGauze Authors.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//    http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
12// implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16use serde::{Deserialize, Serialize};
17use std::error::Error;
18use strum_macros::Display;
19
20#[derive(Clone, Eq, PartialEq, Hash, Debug, Serialize, Deserialize)]
21pub enum AggrOp {
22    Key(Option<Vec<usize>>),
23    Add,
24    Min,
25    Max,
26    BoolMapOr,
27}
28
29#[derive(Debug, Clone, Display)]
30pub enum AggregationError {
31    FlatSetIsNotData,
32    FlatFlowInfoNFv9NotSupported,
33    OperationNotSupported,
34}
35
36impl Error for AggregationError {}
37
38// Flow Aggregation Generic Helper Functions
39#[inline]
40pub fn reduce_num<T: Copy + std::ops::AddAssign + Ord>(
41    lhs: &mut Option<Vec<T>>,
42    rhs: &Option<Vec<T>>,
43    op: &AggrOp,
44) -> Result<(), AggregationError> {
45    match op {
46        AggrOp::Key(indices) => set_field(lhs, rhs, indices),
47        AggrOp::Add => aggr_add(lhs, rhs),
48        AggrOp::Min => aggr_min(lhs, rhs),
49        AggrOp::Max => aggr_max(lhs, rhs),
50        AggrOp::BoolMapOr => Err(AggregationError::OperationNotSupported),
51    }
52}
53
54#[inline]
55pub fn reduce_boolmap<T: Copy + std::ops::BitOrAssign>(
56    lhs: &mut Option<Vec<T>>,
57    rhs: &Option<Vec<T>>,
58    op: &AggrOp,
59) -> Result<(), AggregationError> {
60    match op {
61        AggrOp::Key(indices) => set_field(lhs, rhs, indices),
62        AggrOp::Add => Err(AggregationError::OperationNotSupported),
63        AggrOp::Min => Err(AggregationError::OperationNotSupported),
64        AggrOp::Max => Err(AggregationError::OperationNotSupported),
65        AggrOp::BoolMapOr => aggr_bitor(lhs, rhs),
66    }
67}
68
69#[inline]
70pub fn reduce_misc<T: Copy>(
71    lhs: &mut Option<Vec<T>>,
72    rhs: &Option<Vec<T>>,
73    op: &AggrOp,
74) -> Result<(), AggregationError> {
75    match op {
76        AggrOp::Key(indices) => set_field(lhs, rhs, indices),
77        AggrOp::Add => Err(AggregationError::OperationNotSupported),
78        AggrOp::Min => Err(AggregationError::OperationNotSupported),
79        AggrOp::Max => Err(AggregationError::OperationNotSupported),
80        AggrOp::BoolMapOr => Err(AggregationError::OperationNotSupported),
81    }
82}
83
84#[inline]
85pub fn reduce_misc_clone<T: Clone>(
86    lhs: &mut Option<Vec<T>>,
87    rhs: &Option<Vec<T>>,
88    op: &AggrOp,
89) -> Result<(), AggregationError> {
90    match op {
91        AggrOp::Key(indices) => set_field_clone(lhs, rhs, indices),
92        AggrOp::Add => Err(AggregationError::OperationNotSupported),
93        AggrOp::Min => Err(AggregationError::OperationNotSupported),
94        AggrOp::Max => Err(AggregationError::OperationNotSupported),
95        AggrOp::BoolMapOr => Err(AggregationError::OperationNotSupported),
96    }
97}
98
99#[inline]
100fn set_field<T: Copy>(
101    lhs: &mut Option<Vec<T>>,
102    rhs: &Option<Vec<T>>,
103    indices: &Option<Vec<usize>>,
104) -> Result<(), AggregationError> {
105    if let Some(idxs) = indices {
106        if let Some(rhs_vec) = rhs {
107            let res: Vec<T> = idxs
108                .iter()
109                .filter_map(|&idx| rhs_vec.get(idx).copied())
110                .collect();
111            *lhs = if res.is_empty() { None } else { Some(res) };
112        }
113    } else if let Some(rhs_vec) = rhs {
114        *lhs = Some(rhs_vec.to_vec());
115    } else {
116        *lhs = None;
117    }
118    Ok(())
119}
120
121#[inline]
122fn set_field_clone<T: Clone>(
123    lhs: &mut Option<Vec<T>>,
124    rhs: &Option<Vec<T>>,
125    indices: &Option<Vec<usize>>,
126) -> Result<(), AggregationError> {
127    if let Some(idxs) = indices {
128        if let Some(rhs_vec) = rhs {
129            let res: Vec<T> = idxs
130                .iter()
131                .filter_map(|&idx| rhs_vec.get(idx).cloned())
132                .collect();
133            *lhs = if res.is_empty() { None } else { Some(res) };
134        }
135    } else {
136        *lhs = rhs.clone();
137    }
138    Ok(())
139}
140
141#[inline]
142fn aggr_add<T: Copy + std::ops::AddAssign>(
143    lhs: &mut Option<Vec<T>>,
144    rhs: &Option<Vec<T>>,
145) -> Result<(), AggregationError> {
146    match (lhs.as_mut(), rhs) {
147        (None, _) => *lhs = rhs.clone(),
148        (Some(_), None) => {
149            // Incoming doesn't have any value
150        }
151        (Some(lhs), Some(rhs)) => {
152            lhs.iter_mut().zip(rhs.iter()).for_each(|(a, b)| *a += *b);
153
154            // If the other value is longer, just append to the vector
155            if rhs.len() > lhs.len() {
156                for i in &rhs[lhs.len()..] {
157                    lhs.push(*i);
158                }
159            }
160        }
161    }
162    Ok(())
163}
164
165#[inline]
166fn aggr_min<T: Copy + std::cmp::Ord>(
167    lhs: &mut Option<Vec<T>>,
168    rhs: &Option<Vec<T>>,
169) -> Result<(), AggregationError> {
170    match (lhs.as_mut(), rhs) {
171        (None, _) => *lhs = rhs.clone(),
172        (Some(_), None) => {
173            // Incoming doesn't have any value
174        }
175        (Some(lhs), Some(rhs)) => {
176            lhs.iter_mut()
177                .zip(rhs.iter())
178                .for_each(|(a, b)| *a = std::cmp::min(*a, *b));
179
180            // If the other value is longer, just append to the vector
181            if rhs.len() > lhs.len() {
182                for i in &rhs[lhs.len()..] {
183                    lhs.push(*i);
184                }
185            }
186        }
187    }
188    Ok(())
189}
190
191#[inline]
192fn aggr_max<T: Copy + std::cmp::Ord>(
193    lhs: &mut Option<Vec<T>>,
194    rhs: &Option<Vec<T>>,
195) -> Result<(), AggregationError> {
196    match (lhs.as_mut(), rhs) {
197        (None, _) => *lhs = rhs.clone(),
198        (Some(_), None) => {
199            // Incoming doesn't have any value
200        }
201        (Some(lhs), Some(rhs)) => {
202            lhs.iter_mut()
203                .zip(rhs.iter())
204                .for_each(|(a, b)| *a = std::cmp::max(*a, *b));
205
206            // If the other value is longer, just append to the vector
207            if rhs.len() > lhs.len() {
208                for i in &rhs[lhs.len()..] {
209                    lhs.push(*i);
210                }
211            }
212        }
213    }
214    Ok(())
215}
216
217#[inline]
218fn aggr_bitor<T: Copy + std::ops::BitOrAssign>(
219    lhs: &mut Option<Vec<T>>,
220    rhs: &Option<Vec<T>>,
221) -> Result<(), AggregationError> {
222    match (lhs.as_mut(), rhs) {
223        (None, _) => *lhs = rhs.clone(),
224        (Some(_), None) => {
225            // Incoming doesn't have any value
226        }
227        (Some(lhs), Some(rhs)) => {
228            lhs.iter_mut().zip(rhs.iter()).for_each(|(a, b)| *a |= *b);
229
230            // If the other value is longer, just append to the vector
231            if rhs.len() > lhs.len() {
232                for i in &rhs[lhs.len()..] {
233                    lhs.push(*i);
234                }
235            }
236        }
237    }
238    Ok(())
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244
245    #[test]
246    fn test_reduce_num_add_lhs_longer() {
247        let mut lhs = Some(vec![1, 1, 1]);
248        let rhs = Some(vec![1, 1]);
249        let op = AggrOp::Add;
250        reduce_num(&mut lhs, &rhs, &op).unwrap();
251        assert_eq!(lhs, Some(vec![2, 2, 1]));
252    }
253
254    #[test]
255    fn test_reduce_num_add_rhs_longer() {
256        let mut lhs = Some(vec![1, 1]);
257        let rhs = Some(vec![1, 1, 1]);
258        let op = AggrOp::Add;
259        reduce_num(&mut lhs, &rhs, &op).unwrap();
260        assert_eq!(lhs, Some(vec![2, 2, 1]));
261    }
262
263    #[test]
264    fn test_reduce_num_min() {
265        let mut lhs = Some(vec![5, 6, 7]);
266        let rhs = Some(vec![4, 7, 6]);
267        let op = AggrOp::Min;
268        reduce_num(&mut lhs, &rhs, &op).unwrap();
269        assert_eq!(lhs, Some(vec![4, 6, 6]));
270    }
271
272    #[test]
273    fn test_reduce_num_max() {
274        let mut lhs = Some(vec![1, 2, 3]);
275        let rhs = Some(vec![4, 1, 5]);
276        let op = AggrOp::Max;
277        reduce_num(&mut lhs, &rhs, &op).unwrap();
278        assert_eq!(lhs, Some(vec![4, 2, 5]));
279    }
280
281    #[test]
282    fn test_reduce_boolmap_bitor() {
283        let mut lhs = Some(vec![0b001, 0b001, 0b010]);
284        let rhs = Some(vec![0b001, 0b0100, 0b100]);
285        let op = AggrOp::BoolMapOr;
286        reduce_boolmap(&mut lhs, &rhs, &op).unwrap();
287        assert_eq!(lhs, Some(vec![0b001, 0b101, 0b110]));
288    }
289
290    #[test]
291    fn test_reduce_misc() {
292        let mut lhs = None;
293        let rhs = Some(vec![1, 2, 3, 4, 5]);
294        let op = AggrOp::Key(Some(vec![0, 2, 4]));
295        reduce_misc(&mut lhs, &rhs, &op).unwrap();
296        assert_eq!(lhs, Some(vec![1, 3, 5]));
297    }
298
299    #[test]
300    fn test_reduce_misc_clone() {
301        let mut lhs = None;
302        let rhs = Some(vec![
303            String::from("a"),
304            String::from("b"),
305            String::from("c"),
306        ]);
307        let op = AggrOp::Key(Some(vec![0, 2]));
308        reduce_misc_clone(&mut lhs, &rhs, &op).unwrap();
309        assert_eq!(lhs, Some(vec![String::from("a"), String::from("c")]));
310    }
311}