harmos_storage/
batch.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10/* === Modules === */
11
12/* === Re-Exports === */
13
14/* === Imports === */
15
16// Standard Library Imports
17use std::ops::Range;
18use std::sync::Arc;
19// Third Party Imports
20// Local Imports
21
22/* === Constants === */
23
24/* === Macros === */
25
26/* === Type Aliases === */
27
28pub type FileChunk = StorageBatch<u8>;
29
30/* === Structs === */
31
32#[derive(Debug)]
33pub struct StorageBatch<T> {
34    values: Arc<Vec<T>>,
35    range: Range<usize>,
36}
37
38/* === Enums === */
39
40/* === Traits === */
41
42/* === Implementations === */
43
44impl<T> StorageBatch<T> {
45    pub fn new(values: Vec<T>) -> Self {
46        let range = 0..values.len();
47        Self { values: Arc::new(values), range }
48    }
49
50    pub(crate) fn merge(batches: &mut [&mut Self]) -> Option<Self> {
51        let first = batches.first()?;
52        let mut range = first.range.clone();
53        let adjoining = batches.iter().skip(1).all(|batch| {
54            let joined = Arc::ptr_eq(&first.values, &batch.values) && range.end == batch.range.start;
55            range.end = batch.range.end;
56            joined
57        });
58        if adjoining {
59            return Some(Self { values: Arc::clone(&first.values), range });
60        }
61        if batches.iter().any(|batch| Arc::strong_count(&batch.values) != 1 || batch.range != (0..batch.values.len())) {
62            return None;
63        }
64        let count = batches.iter().map(|batch| batch.len()).sum();
65        let lengths: Vec<_> = batches.iter().map(|batch| batch.len()).collect();
66        let mut values = Vec::with_capacity(count);
67        for batch in batches.iter_mut() {
68            let source = std::mem::take(Arc::get_mut(&mut batch.values).expect("unique queued backing"));
69            values.extend(source);
70        }
71        let merged = Self {
72            values: Arc::new(values),
73            range: 0..count,
74        };
75        let mut start = 0;
76        for (batch, length) in batches.iter_mut().zip(lengths) {
77            **batch = merged.slice(start..start + length);
78            start += length;
79        }
80        Some(merged)
81    }
82
83    pub fn as_slice(&self) -> &[T] {
84        &self.values[self.range.clone()]
85    }
86
87    pub fn slice(&self, range: Range<usize>) -> Self {
88        assert!(range.start <= range.end && range.end <= self.len(), "batch slice is outside its span");
89        let range = self.range.start + range.start..self.range.start + range.end;
90        Self { values: Arc::clone(&self.values), range }
91    }
92
93    pub fn len(&self) -> usize {
94        self.range.len()
95    }
96
97    pub fn is_empty(&self) -> bool {
98        self.range.is_empty()
99    }
100
101    pub fn capacity(&self) -> usize {
102        self.values.capacity()
103    }
104
105    pub fn allocation_bytes(&self) -> usize {
106        self.values.capacity().saturating_mul(std::mem::size_of::<T>())
107    }
108}
109
110impl From<Vec<u8>> for StorageBatch<u8> {
111    fn from(values: Vec<u8>) -> Self {
112        Self::new(values)
113    }
114}
115
116impl<T> Clone for StorageBatch<T> {
117    fn clone(&self) -> Self {
118        Self {
119            values: Arc::clone(&self.values),
120            range: self.range.clone(),
121        }
122    }
123}
124
125/* === Public Functions === */
126
127/* === Private Functions === */
128
129/* === Tests === */
130
131#[cfg(test)]
132mod tests {
133    use super::{FileChunk, StorageBatch};
134
135    #[test]
136    fn chunk_allocation() {
137        let mut bytes = Vec::with_capacity(1024);
138        bytes.extend_from_slice(&[1, 2, 3, 4]);
139        let pointer = bytes.as_ptr();
140        let capacity = bytes.capacity();
141        let batch = FileChunk::from(bytes);
142        assert_eq!(batch.as_slice().as_ptr(), pointer);
143        let retained = batch.slice(1..3);
144        let empty = batch.slice(4..4);
145        assert_eq!(retained.allocation_bytes(), capacity);
146        assert!(empty.is_empty());
147        drop(batch);
148        assert_eq!(retained.as_slice(), &[2, 3]);
149        assert_eq!(retained.as_slice().as_ptr(), pointer.wrapping_add(1));
150    }
151
152    #[test]
153    fn payload_merge() {
154        let left = vec![7u8; 32];
155        let right = vec![8u8; 32];
156        let pointers = [left.as_ptr(), right.as_ptr()];
157        let mut first = StorageBatch::new(vec![left]);
158        let mut second = StorageBatch::new(vec![right]);
159        let merged = StorageBatch::merge(&mut [&mut first, &mut second]).unwrap();
160        assert_eq!(merged.as_slice()[0].as_ptr(), pointers[0]);
161        assert_eq!(merged.as_slice()[1].as_ptr(), pointers[1]);
162        let detached = merged.clone();
163        let retry = StorageBatch::merge(&mut [&mut first, &mut second]).unwrap();
164        assert_eq!(retry.len(), 2);
165        assert_eq!(retry.as_slice()[1].as_ptr(), pointers[1]);
166        drop(detached);
167    }
168
169    #[test]
170    fn shared_merge_refused() {
171        let mut first = StorageBatch::new(vec![vec![7u8; 32]]);
172        let mut second = StorageBatch::new(vec![vec![8u8; 32]]);
173        let detached = first.clone();
174        assert!(StorageBatch::merge(&mut [&mut first, &mut second]).is_none());
175        assert_eq!(first.as_slice(), detached.as_slice());
176        assert_eq!(second.as_slice()[0], vec![8; 32]);
177    }
178}
harmos_storage/batch.rs

harmos_storage/
batch.rs
