iroh_blobs/api/proto/
bitfield.rs

1use std::{cmp::Ordering, num::NonZeroU64};
2
3use bao_tree::{ChunkNum, ChunkRanges};
4use range_collections::range_set::RangeSetRange;
5use serde::{Deserialize, Deserializer, Serialize};
6use smallvec::SmallVec;
7
8use crate::store::util::{
9    observer::{Combine, CombineInPlace},
10    RangeSetExt,
11};
12
13pub(crate) fn is_validated(size: NonZeroU64, ranges: &ChunkRanges) -> bool {
14    let size = size.get();
15    // ChunkNum::chunks will be at least 1, so this is safe.
16    let last_chunk = ChunkNum::chunks(size) - 1;
17    ranges.contains(&last_chunk)
18}
19
20pub fn is_complete(size: NonZeroU64, ranges: &ChunkRanges) -> bool {
21    let complete = ChunkRanges::from(..ChunkNum::chunks(size.get()));
22    // is_subset is a bit weirdly named. This means that complete is a subset of ranges.
23    complete.is_subset(ranges)
24}
25
26/// The state of a bitfield, or an update to a bitfield
27#[derive(Debug, PartialEq, Eq, Clone, Default)]
28pub struct Bitfield {
29    /// Possible update to the size information. can this be just a u64?
30    pub(crate) size: u64,
31    /// The ranges that were added
32    pub ranges: ChunkRanges,
33}
34
35#[derive(Debug, PartialEq, Eq, Clone, Copy, Default)]
36pub struct BitfieldState {
37    pub complete: bool,
38    pub validated_size: Option<u64>,
39}
40
41impl Serialize for Bitfield {
42    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
43        let mut numbers = SmallVec::<[_; 4]>::new();
44        numbers.push(self.size);
45        numbers.extend(self.ranges.boundaries().iter().map(|x| x.0));
46        numbers.serialize(serializer)
47    }
48}
49
50impl<'de> Deserialize<'de> for Bitfield {
51    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
52        let mut numbers: SmallVec<[u64; 4]> = SmallVec::deserialize(deserializer)?;
53
54        // Need at least 1 u64 for size
55        if numbers.is_empty() {
56            return Err(serde::de::Error::custom("Bitfield needs at least size"));
57        }
58
59        // Split: first is size, rest are ranges
60        let size = numbers.remove(0);
61        let mut ranges = SmallVec::<[_; 2]>::new();
62        for num in numbers {
63            ranges.push(ChunkNum(num));
64        }
65        let Some(ranges) = ChunkRanges::new(ranges) else {
66            return Err(serde::de::Error::custom("Boundaries are not sorted"));
67        };
68        Ok(Bitfield::new(ranges, size))
69    }
70}
71
72impl Bitfield {
73    #[cfg(feature = "fs-store")]
74    pub(crate) fn new_unchecked(ranges: ChunkRanges, size: u64) -> Self {
75        Self { ranges, size }
76    }
77
78    pub fn new(mut ranges: ChunkRanges, size: u64) -> Self {
79        // for zero size, we have to trust the caller
80        if let Some(size) = NonZeroU64::new(size) {
81            let end = ChunkNum::chunks(size.get());
82            if ChunkRanges::from(..end).is_subset(&ranges) {
83                // complete bitfield, canonicalize to all
84                ranges = ChunkRanges::all();
85            } else if ranges.contains(&(end - 1)) {
86                // validated bitfield, canonicalize to open end
87                ranges |= ChunkRanges::from(end..);
88            }
89        }
90        Self { ranges, size }
91    }
92
93    pub fn size(&self) -> u64 {
94        self.size
95    }
96
97    /// An empty bitfield. This is the neutral element for the combine operation.
98    pub fn empty() -> Self {
99        Self {
100            ranges: ChunkRanges::empty(),
101            size: 0,
102        }
103    }
104
105    /// Create a complete bitfield for the given size
106    pub fn complete(size: u64) -> Self {
107        Self {
108            ranges: ChunkRanges::all(),
109            size,
110        }
111    }
112
113    /// True if the chunk corresponding to the size is included in the ranges
114    pub fn is_validated(&self) -> bool {
115        if let Some(size) = NonZeroU64::new(self.size) {
116            is_validated(size, &self.ranges)
117        } else {
118            self.ranges.is_all()
119        }
120    }
121
122    /// True if all chunks up to size are included in the ranges
123    pub fn is_complete(&self) -> bool {
124        if let Some(size) = NonZeroU64::new(self.size) {
125            is_complete(size, &self.ranges)
126        } else {
127            self.ranges.is_all()
128        }
129    }
130
131    /// Get the validated size if the bitfield is validated
132    pub fn validated_size(&self) -> Option<u64> {
133        if self.is_validated() {
134            Some(self.size)
135        } else {
136            None
137        }
138    }
139
140    /// Total valid bytes in this bitfield.
141    pub fn total_bytes(&self) -> u64 {
142        let mut total = 0;
143        for range in self.ranges.iter() {
144            let (start, end) = match range {
145                RangeSetRange::Range(range) => {
146                    (range.start.to_bytes(), range.end.to_bytes().min(self.size))
147                }
148                RangeSetRange::RangeFrom(range) => (range.start.to_bytes(), self.size),
149            };
150            total += end - start;
151        }
152        total
153    }
154
155    pub fn state(&self) -> BitfieldState {
156        BitfieldState {
157            complete: self.is_complete(),
158            validated_size: self.validated_size(),
159        }
160    }
161
162    /// Update the bitfield with a new value, and gives detailed information about the change.
163    ///
164    /// returns a tuple of (changed, Some((old, new))). If the bitfield changed at all, the flag
165    /// is true. If there was a significant change, the old and new states are returned.
166    pub fn update(&mut self, update: &Bitfield) -> UpdateResult {
167        let s0 = self.state();
168        // todo: this is very inefficient because of all the clones
169        let bitfield1 = self.clone().combine(update.clone());
170        if bitfield1 != *self {
171            let s1 = bitfield1.state();
172            *self = bitfield1;
173            if s0 != s1 {
174                UpdateResult::MajorChange(s0, s1)
175            } else {
176                UpdateResult::MinorChange(s1)
177            }
178        } else {
179            UpdateResult::NoChange(s0)
180        }
181    }
182
183    pub fn is_empty(&self) -> bool {
184        self.ranges.is_empty() && self.size == 0
185    }
186
187    /// A diff between two bitfields. This is an inverse of the combine operation.
188    pub fn diff(&self, that: &Bitfield) -> Self {
189        let size = choose_size(self, that);
190        let ranges_diff = &self.ranges ^ &that.ranges;
191        Self {
192            ranges: ranges_diff,
193            size,
194        }
195    }
196}
197
198pub fn choose_size(a: &Bitfield, b: &Bitfield) -> u64 {
199    match (a.ranges.upper_bound(), b.ranges.upper_bound()) {
200        (Some(ac), Some(bc)) => match ac.cmp(&bc) {
201            Ordering::Less => b.size,
202            Ordering::Greater => a.size,
203            Ordering::Equal => a.size.max(b.size),
204        },
205        (Some(_), None) => b.size,
206        (None, Some(_)) => a.size,
207        (None, None) => a.size.max(b.size),
208    }
209}
210
211impl Combine for Bitfield {
212    fn combine(self, that: Self) -> Self {
213        // the size of the chunk with the larger last chunk wins
214        let size = choose_size(&self, &that);
215        let ranges = self.ranges | that.ranges;
216        Self::new(ranges, size)
217    }
218}
219
220impl CombineInPlace for Bitfield {
221    fn combine_with(&mut self, other: Self) -> Self {
222        let new = &other.ranges - &self.ranges;
223        if new.is_empty() {
224            return Bitfield::empty();
225        }
226        self.ranges.union_with(&new);
227        self.size = choose_size(self, &other);
228        Bitfield {
229            ranges: new,
230            size: self.size,
231        }
232    }
233
234    fn is_neutral(&self) -> bool {
235        self.ranges.is_empty() && self.size == 0
236    }
237}
238
239#[allow(clippy::enum_variant_names)]
240#[derive(Debug, Clone, Copy)]
241pub enum UpdateResult {
242    NoChange(BitfieldState),
243    MinorChange(BitfieldState),
244    MajorChange(BitfieldState, BitfieldState),
245}
246
247impl UpdateResult {
248    pub fn new_state(&self) -> &BitfieldState {
249        match self {
250            UpdateResult::NoChange(new) => new,
251            UpdateResult::MinorChange(new) => new,
252            UpdateResult::MajorChange(_, new) => new,
253        }
254    }
255
256    /// True if this change went from non-validated to validated
257    pub fn was_validated(&self) -> bool {
258        match self {
259            UpdateResult::NoChange(_) => false,
260            UpdateResult::MinorChange(_) => false,
261            UpdateResult::MajorChange(old, new) => {
262                new.validated_size.is_some() && old.validated_size.is_none()
263            }
264        }
265    }
266
267    pub fn changed(&self) -> bool {
268        match self {
269            UpdateResult::NoChange(_) => false,
270            UpdateResult::MinorChange(_) => true,
271            UpdateResult::MajorChange(_, _) => true,
272        }
273    }
274}
275
276#[cfg(test)]
277mod tests {
278    use bao_tree::{ChunkNum, ChunkRanges};
279    use proptest::prelude::{prop, Strategy};
280    use smallvec::SmallVec;
281    use test_strategy::proptest;
282
283    use super::Bitfield;
284    use crate::store::util::observer::{Combine, CombineInPlace};
285
286    fn gen_chunk_ranges(max: ChunkNum, k: usize) -> impl Strategy<Value = ChunkRanges> {
287        prop::collection::btree_set(0..=max.0, 0..=k).prop_map(|vec| {
288            let bounds = vec.into_iter().map(ChunkNum).collect::<SmallVec<[_; 2]>>();
289            ChunkRanges::new(bounds).unwrap()
290        })
291    }
292
293    fn gen_bitfields(size: u64, k: usize) -> impl Strategy<Value = Bitfield> {
294        (0..size).prop_flat_map(move |size| {
295            let chunks = ChunkNum::full_chunks(size);
296            gen_chunk_ranges(chunks, k).prop_map(move |ranges| Bitfield::new(ranges, size))
297        })
298    }
299
300    fn gen_non_empty_bitfields(size: u64, k: usize) -> impl Strategy<Value = Bitfield> {
301        gen_bitfields(size, k).prop_filter("non-empty", |x| !x.is_neutral())
302    }
303
304    #[proptest]
305    fn test_combine_empty(#[strategy(gen_non_empty_bitfields(32768, 4))] a: Bitfield) {
306        assert_eq!(a.clone().combine(Bitfield::empty()), a);
307        assert_eq!(Bitfield::empty().combine(a.clone()), a);
308    }
309
310    #[proptest]
311    fn test_combine_order(
312        #[strategy(gen_non_empty_bitfields(32768, 4))] a: Bitfield,
313        #[strategy(gen_non_empty_bitfields(32768, 4))] b: Bitfield,
314    ) {
315        let ab = a.clone().combine(b.clone());
316        let ba = b.combine(a);
317        assert_eq!(ab, ba);
318    }
319
320    #[proptest]
321    fn test_complete_normalized(#[strategy(gen_non_empty_bitfields(32768, 4))] a: Bitfield) {
322        if a.is_complete() {
323            assert_eq!(a.ranges, ChunkRanges::all());
324        }
325    }
326}
iroh_blobs/api/proto/bitfield.rs

iroh_blobs/api/proto/
bitfield.rs
