shadow_rs/utility/
counter.rs

Help
1/*
2 * The Shadow Simulator
3 * See LICENSE for licensing information
4 */
5
6/*!
7A counter that can be used to count frequencies of a set of objects. The counter starts
8with no keys. When an unknown key is incremented, the counter adds a new key to an
9internal map and sets the count for that key to 1. When a known key is incremented, the
10count for that key increases. The state of the counter can be extracted by converting it
11to a string, which lists the counts for all keys sorted with the heaviest hitters first.
12Currently, only String types are supported, but we may eventually support counting
13generic types.
14*/
15
16use std::collections::HashMap;
17use std::fmt::{Display, Formatter};
18use std::ops::{Add, Sub};
19
20use serde::ser::SerializeMap;
21
22/// The main counter object that maps individual keys to count values.
23#[derive(Debug, Clone, PartialEq, Eq)]
24pub struct Counter {
25    // TODO: convert this so we could count generic types instead of Strings
26    items: HashMap<String, i64>,
27}
28
29/// The supported operations on the values stored in this counter.
30enum CounterOperation {
31    Add,
32    Set,
33    Sub,
34}
35
36/// A collection of counters that can store and modify values for a set of keys.
37impl Counter {
38    /// Initializes a new counter map that starts with no keys.
39    pub fn new() -> Counter {
40        Counter {
41            items: HashMap::new(),
42        }
43    }
44
45    /// Increment the counter value by one for the key given by id.
46    /// Returns the value of the counter after it was incremented.
47    pub fn add_one(&mut self, id: &str) -> i64 {
48        self.operate(id, CounterOperation::Add, 1)
49    }
50
51    /// Decrement the counter value by one for the key given by id.
52    /// Returns the value of the counter after it was decremented.
53    pub fn sub_one(&mut self, id: &str) -> i64 {
54        self.operate(id, CounterOperation::Sub, 1)
55    }
56
57    /// Increment the counter value by the given value for the key given by id.
58    /// Returns the value of the counter after it was incremented.
59    pub fn add_value(&mut self, id: &str, value: i64) -> i64 {
60        self.operate(id, CounterOperation::Add, value)
61    }
62
63    /// Decrement the counter value by the given value for the key given by id.
64    /// Returns the value of the counter after it was decremented.
65    pub fn sub_value(&mut self, id: &str, value: i64) -> i64 {
66        self.operate(id, CounterOperation::Sub, value)
67    }
68
69    /// Sets the counter value to the given value for the key given by id.
70    /// Returns the value of the counter after it was set.
71    pub fn set_value(&mut self, id: &str, value: i64) -> i64 {
72        self.operate(id, CounterOperation::Set, value)
73    }
74
75    /// Returns the counter value for the key given by id, or 0 if no operations have
76    /// been performed on the key.
77    pub fn get_value(&self, id: &str) -> i64 {
78        match self.items.get(id) {
79            Some(val) => *val,
80            None => 0,
81        }
82    }
83
84    /// Add all values for all keys in `other` to this counter.
85    pub fn add_counter(&mut self, other: &Counter) {
86        for (key, val) in other.items.iter() {
87            self.add_value(key, *val);
88        }
89    }
90
91    /// Subtract all values for all keys in `other` from this counter.
92    pub fn sub_counter(&mut self, other: &Counter) {
93        for (key, val) in other.items.iter() {
94            self.sub_value(key, *val);
95        }
96    }
97
98    /// Perform a supported operation on the counter value.
99    fn operate(&mut self, id: &str, op: CounterOperation, value: i64) -> i64 {
100        match self.items.get_mut(id) {
101            Some(val) => {
102                // Update and return the existing value without allocating new key.
103                match op {
104                    CounterOperation::Add => *val += value,
105                    CounterOperation::Sub => *val -= value,
106                    CounterOperation::Set => *val = value,
107                }
108                // Remove the key if the value reached 0.
109                if *val == 0 {
110                    assert_eq!(self.items.remove(id), Some(0));
111                    0
112                } else {
113                    *val
114                }
115            }
116            None => {
117                // Allocate new key and insert it with initial value of 0.
118                assert_eq!(self.items.insert(id.to_string(), 0), None);
119                // Now that we inserted it, we can do the operation.
120                self.operate(id, op, value)
121            }
122        }
123    }
124
125    /// Get an iterator that returns elements in the order best suited for human-readable output
126    /// (currently sorted by value with the largest value first).
127    fn sorted_for_display(
128        &self,
129    ) -> impl IntoIterator<
130        IntoIter = impl ExactSizeIterator<Item = (&String, &i64)>,
131        Item = (&String, &i64),
132    > {
133        // Get the items in a vector so we can sort them.
134        let mut item_vec = Vec::from_iter(&self.items);
135
136        // Sort the counts so our string is consistent.
137        // Use reverse on vals to get the heaviest hitters first, but sort keys normally.
138        item_vec.sort_by(|&(key_a, val_a), &(key_b, val_b)| {
139            val_a.cmp(val_b).reverse().then(key_a.cmp(key_b))
140        });
141
142        item_vec
143    }
144}
145
146impl Default for Counter {
147    fn default() -> Self {
148        Self::new()
149    }
150}
151
152impl Display for Counter {
153    /// Returns a string representation of the counter in the form
154    ///   `{key1:value1, key2:value2, ..., keyN:valueN}`
155    /// for known keys and values.
156    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
157        let items = self.sorted_for_display().into_iter();
158        let items_len = items.len();
159
160        // Create a string representation of the counts by iterating over the items.
161        write!(f, "{{")?;
162        for (i, item) in items.enumerate() {
163            write!(f, "{}:{}", item.0, item.1)?;
164            if i < items_len - 1 {
165                write!(f, ", ")?;
166            }
167        }
168        write!(f, "}}")
169    }
170}
171
172impl serde::Serialize for Counter {
173    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
174    where
175        S: serde::Serializer,
176    {
177        let items = self.sorted_for_display().into_iter();
178        let mut map = serializer.serialize_map(Some(items.len()))?;
179        for (k, v) in items {
180            map.serialize_entry(k, v)?;
181        }
182        map.end()
183    }
184}
185
186impl Add for Counter {
187    type Output = Self;
188    /// Combines two counters by adding all values for all keys of `other` to `self`.
189    fn add(mut self, other: Self) -> Self {
190        self.add_counter(&other);
191        self
192    }
193}
194
195impl Sub for Counter {
196    type Output = Self;
197    /// Combines two counters by subtracting all values for all keys of `other` from `self`.
198    fn sub(mut self, other: Self) -> Self {
199        self.sub_counter(&other);
200        self
201    }
202}
203
204mod export {
205    use std::ffi::CStr;
206    use std::ffi::CString;
207    use std::os::raw::c_char;
208
209    use super::*;
210
211    #[unsafe(no_mangle)]
212    pub extern "C-unwind" fn counter_new() -> *mut Counter {
213        Box::into_raw(Box::new(Counter::new()))
214    }
215
216    #[unsafe(no_mangle)]
217    pub extern "C-unwind" fn counter_free(counter_ptr: *mut Counter) {
218        if counter_ptr.is_null() {
219            return;
220        }
221        drop(unsafe { Box::from_raw(counter_ptr) });
222    }
223
224    #[unsafe(no_mangle)]
225    pub extern "C-unwind" fn counter_add_value(
226        counter: *mut Counter,
227        id: *const c_char,
228        value: i64,
229    ) -> i64 {
230        assert!(!counter.is_null());
231        assert!(!id.is_null());
232
233        let counter = unsafe { &mut *counter };
234        let id = unsafe { CStr::from_ptr(id) };
235
236        counter.add_value(id.to_str().unwrap(), value)
237    }
238
239    #[unsafe(no_mangle)]
240    pub extern "C-unwind" fn counter_sub_value(
241        counter: *mut Counter,
242        id: *const c_char,
243        value: i64,
244    ) -> i64 {
245        assert!(!counter.is_null());
246        assert!(!id.is_null());
247
248        let counter = unsafe { &mut *counter };
249        let id = unsafe { CStr::from_ptr(id) };
250
251        counter.sub_value(id.to_str().unwrap(), value)
252    }
253
254    #[unsafe(no_mangle)]
255    pub extern "C-unwind" fn counter_add_counter(counter: *mut Counter, other: *const Counter) {
256        assert!(!counter.is_null());
257        assert!(!other.is_null());
258
259        let counter = unsafe { &mut *counter };
260        let other = unsafe { &*other };
261
262        counter.add_counter(other)
263    }
264
265    #[unsafe(no_mangle)]
266    pub extern "C-unwind" fn counter_sub_counter(counter: *mut Counter, other: *const Counter) {
267        assert!(!counter.is_null());
268        assert!(!other.is_null());
269
270        let counter = unsafe { &mut *counter };
271        let other = unsafe { &*other };
272
273        counter.sub_counter(other)
274    }
275
276    #[unsafe(no_mangle)]
277    pub extern "C-unwind" fn counter_equals_counter(
278        counter: *const Counter,
279        other: *const Counter,
280    ) -> bool {
281        assert!(!counter.is_null());
282        assert!(!other.is_null());
283
284        let counter = unsafe { &*counter };
285        let other = unsafe { &*other };
286
287        counter == other
288    }
289
290    /// Creates a new string representation of the counter, e.g., for logging.
291    /// The returned string must be free'd by passing it to counter_free_string.
292    #[unsafe(no_mangle)]
293    pub extern "C-unwind" fn counter_alloc_string(counter: *mut Counter) -> *mut c_char {
294        assert!(!counter.is_null());
295
296        let counter = unsafe { &mut *counter };
297        let string = counter.to_string();
298
299        // Transfer ownership back to caller
300        CString::new(string).unwrap().into_raw()
301    }
302
303    /// Frees a string previously returned from counter_alloc_string.
304    #[unsafe(no_mangle)]
305    pub extern "C-unwind" fn counter_free_string(counter: *mut Counter, ptr: *mut c_char) {
306        assert!(!counter.is_null());
307        assert!(!ptr.is_null());
308        // Free the previously alloc'd string
309        drop(unsafe { CString::from_raw(ptr) });
310    }
311}
312
313#[cfg(test)]
314mod tests {
315    use super::*;
316
317    #[test]
318    fn test_set_value() {
319        let mut counter = Counter::new();
320        assert_eq!(counter.set_value("read", 100), 100);
321        assert_eq!(counter.set_value("read", 10), 10);
322        assert_eq!(counter.set_value("read", 0), 0);
323        assert_eq!(counter.set_value("read", 10), 10);
324    }
325
326    #[test]
327    fn test_get_value() {
328        let mut counter = Counter::new();
329        assert_eq!(counter.get_value("read"), 0);
330        assert_eq!(counter.get_value("write"), 0);
331        assert_eq!(counter.get_value("close"), 0);
332        counter.add_one("write");
333        counter.add_one("write");
334        counter.add_one("read");
335        counter.add_one("write");
336        assert_eq!(counter.get_value("read"), 1);
337        assert_eq!(counter.get_value("write"), 3);
338        assert_eq!(counter.get_value("close"), 0);
339    }
340
341    #[test]
342    fn test_add_one() {
343        let mut counter = Counter::new();
344        assert_eq!(counter.add_one("read"), 1);
345        assert_eq!(counter.add_one("read"), 2);
346        assert_eq!(counter.add_one("write"), 1);
347        assert_eq!(counter.add_one("read"), 3);
348    }
349
350    #[test]
351    fn test_sub_one() {
352        let mut counter = Counter::new();
353        counter.set_value("read", 2);
354        assert_eq!(counter.sub_one("read"), 1);
355        assert_eq!(counter.sub_one("read"), 0);
356        assert_eq!(counter.sub_one("read"), -1);
357        assert_eq!(counter.get_value("read"), -1);
358        counter.set_value("read", 100);
359        counter.set_value("write", 100);
360        assert_eq!(counter.sub_one("read"), 99);
361    }
362
363    #[test]
364    fn test_add_value() {
365        let mut counter = Counter::new();
366        assert_eq!(counter.add_value("read", 10), 10);
367        assert_eq!(counter.add_value("read", 10), 20);
368        assert_eq!(counter.add_value("write", 10), 10);
369        assert_eq!(counter.add_value("read", 10), 30);
370    }
371
372    #[test]
373    fn test_sub_value() {
374        let mut counter = Counter::new();
375        counter.set_value("read", 100);
376        assert_eq!(counter.sub_value("read", 10), 90);
377        assert_eq!(counter.sub_value("read", 10), 80);
378        assert_eq!(counter.sub_value("write", 10), -10);
379        assert_eq!(counter.sub_value("read", 10), 70);
380    }
381
382    #[test]
383    fn test_operator_add() {
384        let mut counter_a = Counter::new();
385        counter_a.set_value("read", 100);
386
387        let mut counter_b = Counter::new();
388        counter_b.set_value("read", 50);
389
390        let mut counter_sum = Counter::new();
391        counter_sum.set_value("read", 150);
392
393        // This transfers ownership of a and b to counter_added
394        let counter_added = counter_a + counter_b;
395        assert_eq!(counter_added.get_value("read"), 150);
396        assert_eq!(counter_added, counter_sum);
397    }
398
399    #[test]
400    fn test_operator_sub() {
401        let mut counter_a = Counter::new();
402        counter_a.set_value("read", 100);
403
404        let mut counter_b = Counter::new();
405        counter_b.set_value("read", 25);
406
407        let mut counter_sum = Counter::new();
408        counter_sum.set_value("read", 75);
409
410        // This transfers ownership of a and b to counter_subbed
411        let counter_subbed = counter_a - counter_b;
412        assert_eq!(counter_subbed.get_value("read"), 75);
413        assert_eq!(counter_subbed, counter_sum);
414    }
415
416    #[test]
417    fn test_operator_sub_multi() {
418        let mut counter_a = Counter::new();
419        counter_a.set_value("read", 100);
420
421        let mut counter_b = Counter::new();
422        counter_b.set_value("read", 25);
423
424        let mut counter_c = Counter::new();
425        counter_c.set_value("read", 75);
426
427        // Subtract to get negative first, then add back to zero.
428        assert_eq!(counter_b - counter_a + counter_c, Counter::new());
429    }
430
431    #[test]
432    fn test_add_counter() {
433        let mut counter_a = Counter::new();
434        counter_a.set_value("read", 100);
435        counter_a.set_value("write", 1);
436
437        let mut counter_b = Counter::new();
438        counter_b.set_value("read", 50);
439        counter_b.set_value("write", 2);
440
441        let mut counter_sum = Counter::new();
442        counter_sum.set_value("read", 150);
443        counter_sum.set_value("write", 3);
444
445        counter_a.add_counter(&counter_b);
446
447        assert_eq!(counter_a.get_value("read"), 150);
448        assert_eq!(counter_a.get_value("write"), 3);
449        assert_eq!(counter_b.get_value("read"), 50);
450        assert_eq!(counter_b.get_value("write"), 2);
451        assert_eq!(counter_a, counter_sum);
452    }
453
454    #[test]
455    fn test_sub_counter() {
456        let mut counter_a = Counter::new();
457        counter_a.set_value("read", 100);
458        counter_a.set_value("write", 3);
459
460        let mut counter_b = Counter::new();
461        counter_b.set_value("read", 25);
462        counter_b.set_value("write", 1);
463
464        let mut counter_sum = Counter::new();
465        counter_sum.set_value("read", 75);
466        counter_sum.set_value("write", 2);
467
468        counter_a.sub_counter(&counter_b);
469
470        assert_eq!(counter_a.get_value("read"), 75);
471        assert_eq!(counter_a.get_value("write"), 2);
472        assert_eq!(counter_b.get_value("read"), 25);
473        assert_eq!(counter_b.get_value("write"), 1);
474        assert_eq!(counter_a, counter_sum);
475    }
476
477    #[test]
478    fn test_counter_equality_nonzero() {
479        let mut counter_a = Counter::new();
480        counter_a.set_value("read", 1);
481        counter_a.set_value("write", 2);
482
483        let mut counter_b = Counter::new();
484        counter_b.set_value("read", 1);
485        counter_b.set_value("write", 2);
486
487        let mut counter_c = Counter::new();
488        counter_c.set_value("read", 10);
489        counter_c.set_value("write", 20);
490
491        assert_eq!(counter_a, counter_b);
492        assert_ne!(counter_a, counter_c);
493        assert_ne!(counter_b, counter_c);
494
495        let mut counter_d = Counter::new();
496        counter_d.set_value("read", 1);
497        counter_d.set_value("write", 2);
498        counter_d.set_value("close", 1);
499        counter_d.sub_value("close", 1);
500
501        assert_eq!(counter_a, counter_d);
502    }
503
504    #[test]
505    fn test_counter_equality_zero() {
506        let mut counter_a = Counter::new();
507        counter_a.set_value("read", 1);
508        counter_a.set_value("write", 2);
509
510        let mut counter_d = Counter::new();
511        counter_d.set_value("read", 1);
512        counter_d.set_value("write", 2);
513        counter_d.set_value("close", 1);
514        counter_d.sub_value("close", 1);
515
516        assert_eq!(counter_a, counter_d);
517    }
518
519    #[test]
520    fn test_to_string() {
521        let mut counter = Counter::new();
522
523        counter.add_one("read");
524        counter.add_one("read");
525        counter.add_one("close");
526        counter.add_one("write");
527        counter.add_one("write");
528        counter.add_one("write");
529
530        // Make sure the keys are sorted with the largest count first
531        assert_eq!(
532            counter.to_string(),
533            String::from("{write:3, read:2, close:1}")
534        );
535
536        counter.add_one("read");
537        counter.add_one("read");
538
539        // The order should have changed with read first now
540        assert_eq!(
541            counter.to_string(),
542            String::from("{read:4, write:3, close:1}")
543        );
544    }
545
546    #[test]
547    fn test_to_string_order() {
548        let mut counter_a = Counter::new();
549        counter_a.add_one("write");
550        counter_a.add_one("close");
551        counter_a.add_one("read");
552        let mut counter_b = Counter::new();
553        counter_b.add_one("read");
554        counter_b.add_one("write");
555        counter_b.add_one("close");
556
557        // Make sure the counters of equal value are sorted based on key.
558        assert_eq!(counter_a.to_string(), counter_b.to_string());
559        assert_eq!(
560            counter_a.to_string(),
561            String::from("{close:1, read:1, write:1}")
562        );
563    }
564}
shadow_rs/utility/counter.rs

shadow_rs/utility/
counter.rs
