byte_mutator/
undo_buffer.rs

1//! `UndoBuffer` is a structure that is used to expose an 'undo' interface to a fixed size buffer.
2//! It internally maintains two fixed size `ArrayVec` structures. These arrays are identical on
3//! construction, but only one can be written to. `UndoBuffer` keeps track of which areas of the
4//! buffer have been mutably borrowed, and can then reset the buffer back to its original state.
5//!
6//! todo: Variably sized buffers.
7
8use arrayvec::ArrayVec;
9use std::cmp::{max, min};
10use std::io::Write;
11
12// todo fix -- env var?
13const DEFAULT_BUFFER_SIZE: usize = 1024;
14
15#[derive(Debug, Clone)]
16/// Fixed size buffer with an original state and a writeable buffer. Tracks which region of the
17/// buffer has been exposed for changes and enables an undo of those changes.
18pub struct UndoBuffer {
19    buffer: ArrayVec<[u8; DEFAULT_BUFFER_SIZE]>,
20    original: ArrayVec<[u8; DEFAULT_BUFFER_SIZE]>,
21    dirty: Option<(usize, usize)>,
22}
23
24impl UndoBuffer {
25    pub fn new(buf: &[u8]) -> Self {
26        let mut original = ArrayVec::<[u8; DEFAULT_BUFFER_SIZE]>::new();
27        let mut buffer = ArrayVec::<[u8; DEFAULT_BUFFER_SIZE]>::new();
28
29        // todo: will panic if buf.len() > original.len()
30        (&mut original)
31            .write_all(buf)
32            .expect("Failed to copy into UndoBuffer");
33        (&mut buffer)
34            .write_all(buf)
35            .expect("Failed to copy into UndoBuffer");
36
37        Self {
38            original,
39            buffer,
40            dirty: None,
41        }
42    }
43
44    /// Used length of the writable buffer
45    pub fn len(&self) -> usize {
46        self.buffer.len()
47    }
48
49    /// Whether the writable buffer is empty
50    pub fn is_empty(&self) -> bool {
51        self.buffer.len() == 0
52    }
53
54    /// Returns a full mutable slice of the write buffer, and marks the entire thing as dirty.
55    pub fn get_mut(&mut self) -> &mut [u8] {
56        self.dirty = Some((0, self.buffer.len()));
57        &mut self.buffer[..]
58    }
59
60    /// Returns a mutable subslice of the buffer.
61    /// Marks that region as dirty for future undo operations.
62    pub fn get_mut_range(&mut self, start: usize, end: usize) -> &mut [u8] {
63        // protect against running off the end of the buffer
64        let end = min(self.buffer.len(), end);
65        self.dirty = match self.dirty {
66            Some(range) => {
67                // expand to cover range
68                Some((min(range.0, start), max(range.1, end)))
69            }
70            None => Some((start, end)),
71        };
72
73        &mut self.buffer[start..end]
74    }
75
76    /// Returns an immutable reference to the buffer in its current state
77    pub fn read(&self) -> &[u8] {
78        &self.buffer[..]
79    }
80
81    /// Undo all changes and set the write/readable buffer back to the original state
82    pub fn undo(&mut self) {
83        let (start, end) = match self.dirty {
84            None => {
85                return; // no-op
86            }
87            Some(range) => range,
88        };
89
90        (&mut self.buffer[start..end])
91            .write_all(&self.original[start..end])
92            .expect("Failed to write");
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99    use crate::mutators::bitflipper::BitFlipper;
100
101    #[test]
102    fn mutate_and_reset() {
103        let mut buffer = UndoBuffer::new(b"foo");
104
105        // first bit should flip resulting in 'goo'
106        // 0b1100110 -> 0b1100111, 103 -> 102, f -> g
107        BitFlipper::mutate(buffer.get_mut(), 0, 1);
108        assert_eq!(buffer.read(), b"goo");
109
110        // should be back to 'foo'
111        buffer.undo();
112        assert_eq!(buffer.read(), b"foo");
113    }
114
115    #[test]
116    fn mutate_reset_range() {
117        // clamp changes to the last byte
118        let (min, max) = (2, 3);
119        let mut buffer = UndoBuffer::new(b"foo");
120        let range = buffer.get_mut_range(min, max);
121
122        // flip a bit
123        BitFlipper::mutate(range, 0, 1);
124
125        // assert that something changed
126        assert_ne!(buffer.read()[0..3], b"foo"[..]);
127
128        // set it back
129        buffer.undo();
130
131        // make sure we match
132        assert_eq!(buffer.read()[0..3], b"foo"[..]);
133    }
134}