Merge pull request #409 from filecoin-project/feat/bitfield-refactor

bitfield improvements

Author: Stebalien

ipld/bitfield/Cargo.toml

@@ -4,14 +4,15 @@ description = "Bitfield logic for use in Filecoin actors"
 version = "0.3.0"
 license = "MIT OR Apache-2.0"
 authors = ["ChainSafe Systems <[email protected]>", "Protocol Labs", "Filecoin Core Devs"]
-edition = "2018"
+edition = "2021"
 repository = "https://github.com/filecoin-project/ref-fvm"
 
 [dependencies]
 unsigned-varint = "0.7.1"
 serde = { version = "1.0", features = ["derive"] }
 serde_bytes = { package = "cs_serde_bytes", version = "0.12" }
 fvm_shared = { version = "0.3.1", path = "../../shared" }
+thiserror = "1.0.30"
 
 [dev-dependencies]
 rand_xorshift = "0.2.0"

ipld/bitfield/src/iter/mod.rs

@@ -133,6 +133,9 @@ where
     I: Iterator<Item = Range<u64>>,
 {
     /// Creates a new `Ranges` instance.
+    ///
+    /// WARNING: This is asserting that the underlying iterator obeys the `RangeIterator`
+    /// constraints. Using this incorrectly could lead to panics, etc.
     pub fn new<II>(iter: II) -> Self
     where
         II: IntoIterator<IntoIter = I, Item = Range<u64>>,
@@ -155,16 +158,21 @@ where
 impl<I> RangeIterator for Ranges<I> where I: Iterator<Item = Range<u64>> {}
 
 /// Returns a `RangeIterator` which ranges contain the values from the provided iterator.
-/// The values need to be in ascending order — if not, the returned iterator may not satisfy
-/// all `RangeIterator` requirements.
-pub fn ranges_from_bits(bits: impl IntoIterator<Item = u64>) -> impl RangeIterator {
+/// The values need to be in ascending order and may not include u64::MAX. Otherwise, the iterator
+/// will panic.
+pub(crate) fn ranges_from_bits(bits: impl IntoIterator<Item = u64>) -> impl RangeIterator {
     let mut iter = bits.into_iter().peekable();
 
     Ranges::new(iter::from_fn(move || {
         let start = iter.next()?;
-        let mut end = start + 1;
-        while iter.peek() == Some(&end) {
-            end += 1;
+        let mut end = start.checked_add(1).expect("bitfield overflow");
+        while let Some(&next) = iter.peek() {
+            if next < end {
+                panic!("out of order bitfield")
+            } else if next > end {
+                break;
+            }
+            end = end.checked_add(1).expect("bitfield overflow");
             iter.next();
         }
         Some(start..end)

ipld/bitfield/src/lib.rs

@@ -1,22 +1,35 @@
 // Copyright 2019-2022 ChainSafe Systems
 // SPDX-License-Identifier: Apache-2.0, MIT
 
+// disable this lint because it can actually cause performance regressions, and usually leads to
+// hard to read code.
+#![allow(clippy::comparison_chain)]
+
 pub mod iter;
 mod range;
 mod rleplus;
 mod unvalidated;
 
 use std::collections::BTreeSet;
-use std::iter::FromIterator;
 use std::ops::{
     BitAnd, BitAndAssign, BitOr, BitOrAssign, BitXor, BitXorAssign, Range, Sub, SubAssign,
 };
 
 use iter::{ranges_from_bits, RangeIterator};
 pub(crate) use range::RangeSize;
+pub use rleplus::Error;
+use thiserror::Error;
 pub use unvalidated::{UnvalidatedBitField, Validate};
 
-type Result<T> = std::result::Result<T, &'static str>;
+#[derive(Clone, Error, Debug)]
+#[error("bitfields may not include u64::MAX")]
+pub struct OutOfRangeError;
+
+impl From<OutOfRangeError> for Error {
+    fn from(_: OutOfRangeError) -> Self {
+        Error::RLEOverflow
+    }
+}
 
 /// A bit field with buffered insertion/removal that serializes to/from RLE+. Similar to
 /// `HashSet<u64>`, but more memory-efficient when long runs of 1s and 0s are present.
@@ -36,22 +49,79 @@ impl PartialEq for BitField {
     }
 }
 
-impl FromIterator<u64> for BitField {
-    fn from_iter<I: IntoIterator<Item = u64>>(iter: I) -> Self {
-        let mut vec: Vec<_> = iter.into_iter().collect();
-        if vec.is_empty() {
-            Self::new()
+/// Possibly a valid bitfield, or an out of bounds error. Ideally we'd just use a result, but we
+/// can't implement [`FromIterator`] on a [`Result`] due to coherence.
+///
+/// You probably want to call [`BitField::try_from_bits`] instead of using this directly.
+#[doc(hidden)]
+pub enum MaybeBitField {
+    /// A valid bitfield.
+    Ok(BitField),
+    /// Out of bounds.
+    OutOfBounds,
+}
+
+impl MaybeBitField {
+    pub fn unwrap(self) -> BitField {
+        use MaybeBitField::*;
+        match self {
+            Ok(bf) => bf,
+            OutOfBounds => panic!("bitfield bit out of bounds"),
+        }
+    }
+
+    pub fn expect(self, message: &str) -> BitField {
+        use MaybeBitField::*;
+        match self {
+            Ok(bf) => bf,
+            OutOfBounds => panic!("{}", message),
+        }
+    }
+}
+
+impl TryFrom<MaybeBitField> for BitField {
+    type Error = OutOfRangeError;
+
+    fn try_from(value: MaybeBitField) -> Result<Self, Self::Error> {
+        match value {
+            MaybeBitField::Ok(bf) => Ok(bf),
+            MaybeBitField::OutOfBounds => Err(OutOfRangeError),
+        }
+    }
+}
+
+impl FromIterator<bool> for MaybeBitField {
+    fn from_iter<T: IntoIterator<Item = bool>>(iter: T) -> MaybeBitField {
+        let mut iter = iter.into_iter().fuse();
+        let bits = (0u64..u64::MAX)
+            .zip(&mut iter)
+            .filter(|&(_, b)| b)
+            .map(|(i, _)| i);
+        let bf = BitField::from_ranges(ranges_from_bits(bits));
+
+        // Now, if we have remaining bits, raise an error. Otherwise, we're good.
+        if iter.next().is_some() {
+            MaybeBitField::OutOfBounds
         } else {
-            vec.sort_unstable();
-            Self::from_ranges(ranges_from_bits(vec))
+            MaybeBitField::Ok(bf)
         }
     }
 }
 
-impl FromIterator<bool> for BitField {
-    fn from_iter<I: IntoIterator<Item = bool>>(iter: I) -> Self {
-        let bits = (0u64..).zip(iter).filter(|&(_, b)| b).map(|(i, _)| i);
-        Self::from_ranges(ranges_from_bits(bits))
+impl FromIterator<u64> for MaybeBitField {
+    fn from_iter<T: IntoIterator<Item = u64>>(iter: T) -> MaybeBitField {
+        let mut vec: Vec<_> = iter.into_iter().collect();
+        if vec.is_empty() {
+            MaybeBitField::Ok(BitField::new())
+        } else {
+            vec.sort_unstable();
+            vec.dedup();
+            if vec.last() == Some(&u64::MAX) {
+                MaybeBitField::OutOfBounds
+            } else {
+                MaybeBitField::Ok(BitField::from_ranges(ranges_from_bits(vec)))
+            }
+        }
     }
 }
 
@@ -69,10 +139,33 @@ impl BitField {
         }
     }
 
-    /// Adds the bit at a given index to the bit field.
+    /// Tries to create a new bitfield from a bit iterator. It fails if the resulting bitfield would
+    /// contain values not in the range `0..u64::MAX` (non-inclusive).
+    pub fn try_from_bits<I>(iter: I) -> Result<Self, OutOfRangeError>
+    where
+        I: IntoIterator,
+        MaybeBitField: FromIterator<I::Item>,
+    {
+        iter.into_iter().collect::<MaybeBitField>().try_into()
+    }
+
+    /// Adds the bit at a given index to the bit field, panicing if it's out of range.
+    ///
+    /// # Panics
+    ///
+    /// Panics if `bit` is `u64::MAX`.
     pub fn set(&mut self, bit: u64) {
+        self.try_set(bit).unwrap()
+    }
+
+    /// Adds the bit at a given index to the bit field, returning an error if it's out of range.
+    pub fn try_set(&mut self, bit: u64) -> Result<(), OutOfRangeError> {
+        if bit == u64::MAX {
+            return Err(OutOfRangeError);
+        }
         self.unset.remove(&bit);
         self.set.insert(bit);
+        Ok(())
     }
 
     /// Removes the bit at a given index from the bit field.
@@ -114,8 +207,7 @@ impl BitField {
             self.set.iter().min().copied(),
             self.ranges
                 .iter()
-                .filter_map(|r| r.clone().find(|i| !self.unset.contains(i)))
-                .next(),
+                .find_map(|r| r.clone().find(|i| !self.unset.contains(i))),
         ) {
             (None, None) => None,
             (Some(v), None) | (None, Some(v)) => Some(v),
@@ -165,12 +257,12 @@ impl BitField {
     }
 
     /// Returns an iterator over the indices of the bit field's set bits if the number
-    /// of set bits in the bit field does not exceed `max`. Returns an error otherwise.
-    pub fn bounded_iter(&self, max: u64) -> Result<impl Iterator<Item = u64> + '_> {
+    /// of set bits in the bit field does not exceed `max`. Returns `None` otherwise.
+    pub fn bounded_iter(&self, max: u64) -> Option<impl Iterator<Item = u64> + '_> {
         if self.len() <= max {
-            Ok(self.iter())
+            Some(self.iter())
         } else {
-            Err("Bits set exceeds max in retrieval")
+            None
         }
     }
 
@@ -197,15 +289,15 @@ impl BitField {
     }
 
     /// Returns a slice of the bit field with the start index of set bits
-    /// and number of bits to include in the slice. Returns an error if the
-    /// bit field contains fewer than `start + len` set bits.
-    pub fn slice(&self, start: u64, len: u64) -> Result<Self> {
+    /// and number of bits to include in the slice. Returns `None` if the bit
+    /// field contains fewer than `start + len` set bits.
+    pub fn slice(&self, start: u64, len: u64) -> Option<Self> {
         let slice = BitField::from_ranges(self.ranges().skip_bits(start).take_bits(len));
 
         if slice.len() == len {
-            Ok(slice)
+            Some(slice)
         } else {
-            Err("Not enough bits")
+            None
         }
     }
 
@@ -330,7 +422,7 @@ macro_rules! bitfield {
         std::iter::once($head != 0_u32).chain(bitfield!(@iter $($tail),*))
     };
     ($($val:literal),* $(,)?) => {
-        bitfield!(@iter $($val),*).collect::<$crate::BitField>()
+        bitfield!(@iter $($val),*).collect::<$crate::MaybeBitField>().unwrap()
     };
 }
 

ipld/bitfield/src/rleplus/error.rs

@@ -0,0 +1,13 @@
+use thiserror::Error;
+
+#[derive(PartialEq, Eq, Clone, Debug, Error)]
+pub enum Error {
+    #[error("bitfield not minimally encoded")]
+    NotMinimal,
+    #[error("bitfield specifies an unsupported version")]
+    UnsupportedVersion,
+    #[error("bitfield overflows 2^63-1")]
+    RLEOverflow,
+    #[error("invalid varint")]
+    InvalidVarint,
+}