corevm_host/

range_set.rs

use alloc::vec::Vec;
use codec::{Compact, CompactLen, Decode, Encode, Input, MaxEncodedLen};

// TODO @ivan Proper B-Tree-based implementation would be useful for both CoreVM and PolkaVM.
#[derive(Clone, PartialEq, Eq, Encode, Default)]
pub struct RangeSet {
	ranges: Vec<Range>,
}

impl RangeSet {
	pub const fn new() -> Self {
		Self { ranges: Vec::new() }
	}

	pub fn with_capacity(capacity: usize) -> Self {
		Self { ranges: Vec::with_capacity(capacity) }
	}

	pub fn insert(&mut self, mut new_range: Range) {
		if new_range.is_empty() {
			return;
		}
		if self.ranges.is_empty() {
			self.ranges.push(new_range);
			return;
		}
		let (at, from) =
			match self.ranges.binary_search_by(|range| range.start.cmp(&new_range.start)) {
				Ok(i) => (i, i),
				Err(i) => (i, i.saturating_sub(1)),
			};
		let mut splice_range = at..at;
		for i in from..self.ranges.len() {
			let r = &mut self.ranges[i];
			let contains_start = (r.start..=r.end).contains(&new_range.start);
			let contains_end = (r.start..=r.end).contains(&new_range.end);
			if contains_start && contains_end {
				return;
			}
			if contains_start {
				new_range.start = r.start;
				splice_range = i..i + 1;
				continue;
			}
			if contains_end {
				new_range.end = r.end;
				splice_range.end = i + 1;
				break;
			}
		}
		self.ranges.splice(splice_range, [new_range]);
	}

	pub fn remove(&mut self, range: &Range) {
		if range.is_empty() {
			return;
		}
		let from = match self.ranges.binary_search_by(|r| r.start.cmp(&range.start)) {
			Ok(i) => i,
			Err(i) => i.saturating_sub(1),
		};
		#[allow(clippy::reversed_empty_ranges)]
		let mut drain_range = usize::MAX..0;
		for i in from..self.ranges.len() {
			let r = &mut self.ranges[i];
			let contains_start = (r.start..r.end).contains(&range.start);
			let contains_end = (r.start..=r.end).contains(&range.end);
			if contains_start && contains_end {
				let old_end = r.end;
				r.end = range.start;
				let new_range = Range::new(range.end, old_end);
				if r.is_empty() && new_range.is_empty() {
					drain_range = i..i + 1;
					break;
				}
				if r.is_empty() {
					self.ranges[i] = new_range;
					break;
				}
				if !new_range.is_empty() {
					self.ranges.insert(i + 1, new_range);
				}
				break;
			}
			if range.contains_range(r) {
				if i < drain_range.start {
					drain_range.start = i;
				}
				if i + 1 > drain_range.end {
					drain_range.end = i + 1;
				}
			} else if contains_start {
				r.end = range.start;
				drain_range.start = if r.is_empty() { i } else { i + 1 };
			}
			if contains_end {
				r.start = range.end;
				drain_range.end = if r.is_empty() { i + 1 } else { i };
				break;
			}
		}
		if !drain_range.is_empty() {
			self.ranges.drain(drain_range);
		}
	}

	pub fn overlap(&self, range: &Range) -> bool {
		if range.is_empty() {
			return false;
		}
		for r in &self.ranges {
			if r.start < range.end && range.start < r.end {
				return true;
			}
		}
		false
	}

	pub fn clear(&mut self) {
		self.ranges.clear();
	}

	pub fn enclosing_range(&self) -> Option<Range> {
		if self.ranges.is_empty() {
			return None;
		}
		let start = self.ranges[0].start;
		let end = self.ranges[self.ranges.len() - 1].end;
		Some(Range::new(start, end))
	}

	pub fn count(&self) -> u64 {
		self.ranges.iter().map(|range| range.end - range.start).sum()
	}

	pub fn contains_index(&self, index: u64) -> bool {
		let i = match self.ranges.binary_search_by(|range| range.start.cmp(&index)) {
			Ok(i) => i,
			Err(i) => i.saturating_sub(1),
		};
		self.ranges.get(i).map(|range| range.contains(index)).unwrap_or(false)
	}

	pub fn as_slice(&self) -> &[Range] {
		self.ranges.as_slice()
	}

	pub fn encoded_len(&self) -> usize {
		let len = self.ranges.len();
		Compact::<u64>::compact_len(&(len as u64)) +
			self.ranges.iter().map(|range| range.encoded_len()).sum::<usize>()
	}
}

impl AsRef<[Range]> for RangeSet {
	fn as_ref(&self) -> &[Range] {
		&self.ranges[..]
	}
}

impl core::fmt::Debug for RangeSet {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		f.debug_list().entries(self.ranges.iter()).finish()
	}
}

impl FromIterator<Range> for RangeSet {
	fn from_iter<I: IntoIterator<Item = Range>>(items: I) -> Self {
		let iter = items.into_iter();
		let (min_size, max_size) = iter.size_hint();
		let mut set = RangeSet::with_capacity(max_size.unwrap_or(min_size));
		set.extend(iter);
		set
	}
}

impl Extend<Range> for RangeSet {
	fn extend<I: IntoIterator<Item = Range>>(&mut self, iter: I) {
		for range in iter.into_iter() {
			self.insert(range);
		}
	}
}

impl Decode for RangeSet {
	fn decode<I: Input>(input: &mut I) -> Result<Self, codec::Error> {
		let ranges = Vec::<Range>::decode(input)?;
		if !validate_ranges(&ranges) {
			return Err("RangeSet: out-of-order/overlapping/empty ranges".into());
		}
		Ok(Self { ranges })
	}
}

fn validate_ranges(ranges: &[Range]) -> bool {
	if ranges.iter().any(|r| r.is_empty()) {
		return false;
	}
	for window in ranges.windows(2) {
		let a = &window[0];
		let b = &window[1];
		if b.start <= a.start || b.start <= a.end {
			return false;
		}
	}
	true
}

#[derive(Clone, PartialEq, Eq, Encode, Decode, MaxEncodedLen)]
pub struct Range {
	/// The start of the range (inclusive).
	#[codec(compact)]
	pub start: u64,
	/// The end of the range (exclusive).
	#[codec(compact)]
	pub end: u64,
}

impl Range {
	pub const fn new(start: u64, end: u64) -> Self {
		Self { start, end }
	}

	pub const fn is_empty(&self) -> bool {
		self.start >= self.end
	}

	/// Returns `true` if `self` is fully contains `other`.
	///
	/// If `other` range is empty, returns `false`.
	pub fn contains_range(&self, other: &Range) -> bool {
		!other.is_empty() &&
			(self.start..self.end).contains(&other.start) &&
			(self.start..=self.end).contains(&other.end)
	}

	pub const fn contains(&self, i: u64) -> bool {
		self.start <= i && i < self.end
	}

	pub fn encoded_len(&self) -> usize {
		Compact::<u64>::compact_len(&self.start) + Compact::<u64>::compact_len(&self.end)
	}
}

impl core::fmt::Display for Range {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		if self.start.saturating_add(1) == self.end {
			write!(f, "{}", self.start)
		} else {
			write!(f, "{}..{}", self.start, self.end)
		}
	}
}

impl core::fmt::Debug for Range {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		core::fmt::Display::fmt(self, f)
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use rand::{seq::SliceRandom, Rng};

	#[test]
	fn insert_works() {
		let mut ranges = RangeSet::new();
		assert_eq!(0, ranges.as_ref().len());
		ranges.insert(Range::new(0, 0));
		assert_eq!(0, ranges.as_ref().len());
		ranges.insert(Range::new(0, 1));
		assert_eq!(1, ranges.as_ref().len());
		ranges.insert(Range::new(0, 1));
		assert_eq!(1, ranges.as_ref().len());
		ranges.insert(Range::new(2, 3));
		assert_eq!([Range::new(0, 1), Range::new(2, 3)].as_slice(), ranges.as_ref());
		ranges.insert(Range::new(1, 2));
		assert_eq!([Range::new(0, 3)].as_slice(), ranges.as_ref());
		ranges.insert(Range::new(3, 10));
		assert_eq!([Range::new(0, 10)].as_slice(), ranges.as_ref());
	}

	#[test]
	fn remove_works() {
		let mut ranges = RangeSet::new();
		assert_eq!(0, ranges.as_ref().len());
		// Remove exact.
		ranges.insert(Range::new(0, 1));
		assert_eq!(1, ranges.as_ref().len());
		ranges.remove(&Range::new(0, 1));
		assert_eq!(0, ranges.as_ref().len());
		// Remove in the middle.
		ranges.insert(Range::new(0, 4));
		ranges.remove(&Range::new(1, 3));
		assert_eq!([Range::new(0, 1), Range::new(3, 4)].as_slice(), ranges.as_ref());
		// Remove at the start.
		ranges.clear();
		ranges.insert(Range::new(0, 4));
		ranges.remove(&Range::new(0, 1));
		assert_eq!([Range::new(1, 4)].as_slice(), ranges.as_ref());
		// Remove at the end.
		ranges.clear();
		ranges.insert(Range::new(0, 4));
		ranges.remove(&Range::new(3, 4));
		assert_eq!([Range::new(0, 3)].as_slice(), ranges.as_ref());
		// Remove enclosing range.
		ranges.clear();
		ranges.insert(Range::new(1, 4));
		ranges.remove(&Range::new(0, 5));
		assert_eq!(([] as [Range; 0]).as_slice(), ranges.as_ref());
		// Remove overlaping-at-the-start range.
		ranges.clear();
		ranges.insert(Range::new(1, 4));
		ranges.remove(&Range::new(0, 2));
		assert_eq!([Range::new(2, 4)].as_slice(), ranges.as_ref());
		// Remove overlaping-at-the-end range.
		ranges.clear();
		ranges.insert(Range::new(1, 4));
		ranges.remove(&Range::new(3, 5));
		assert_eq!([Range::new(1, 3)].as_slice(), ranges.as_ref());
	}

	#[test]
	fn insert_remove_random() {
		let mut rng = rand::rng();
		for _ in 0..1000 {
			let mut ranges = Vec::new();
			let mut offset = 0;
			for _ in 0..10 {
				let start = offset;
				let end = rng.random_range(start..=start + 20);
				offset = end;
				ranges.push(Range::new(start, end));
			}
			// Insert all ranges in random order.
			ranges.shuffle(&mut rng);
			let mut set = RangeSet::new();
			for range in ranges.iter() {
				set.insert(range.clone());
			}
			assert_eq!(1, set.as_ref().len());
			assert_eq!(&[Range::new(0, offset)], set.as_ref());
			// Remove all ranges in random order.
			ranges.shuffle(&mut rng);
			for range in ranges.iter() {
				set.remove(range);
			}
			assert_eq!(0, set.as_ref().len());
		}
	}

	#[test]
	fn remove_reverts_insert() {
		let mut rng = rand::rng();
		for _ in 0..1000 {
			let mut ranges = RangeSet::new();
			{
				let mut offset = 0;
				for _ in 0..10 {
					let start = rng.random_range(offset..=20);
					let end = rng.random_range(start..=20);
					offset = end;
					ranges.insert(Range::new(start, end));
				}
			}
			let range = {
				let start = rng.random_range(0..=20);
				Range::new(start, rng.random_range(start..=20))
			};
			if ranges.overlap(&range) {
				// `remove` reverts `insert` only if there is no overlap
				continue;
			}
			let expected = ranges.clone();
			ranges.insert(range.clone());
			let middle = ranges.clone();
			ranges.remove(&range);
			assert_eq!(expected, ranges,
                "ranges = {expected:?}, insert/remove {range:?}, after insert = {middle:?}, after remove = {ranges:?}");
		}
	}

	#[test]
	fn insert_reverts_remove() {
		let mut rng = rand::rng();
		for _ in 0..1000 {
			let mut ranges = RangeSet::new();
			{
				let mut offset = 0;
				for _ in 0..10 {
					let start = rng.random_range(offset..=20);
					let end = rng.random_range(start..=20);
					offset = end;
					ranges.insert(Range::new(start, end));
				}
			}
			let range = {
				let start = rng.random_range(0..=20);
				Range::new(start, rng.random_range(start..=20))
			};
			if !ranges.as_ref().iter().any(|r| r.contains_range(&range)) {
				// `insert` reverts `remove` only if the new range is fully contained in one of the
				// existing ranges
				continue;
			}
			let expected = ranges.clone();
			ranges.remove(&range);
			let middle = ranges.clone();
			ranges.insert(range.clone());
			assert_eq!(expected, ranges,
                "ranges = {expected:?}, remove/insert {range:?}, after remove = {middle:?}, after insert = {ranges:?}");
		}
	}

	#[test]
	fn validate_works() {
		assert!(!validate_ranges(&[Range::new(0, 0)]));
		assert!(!validate_ranges(&[Range::new(0, 1), Range::new(0, 2)]));
		assert!(!validate_ranges(&[Range::new(0, 2), Range::new(1, 3)]));
		assert!(!validate_ranges(&[Range::new(0, 1), Range::new(1, 2)]));
		assert!(!validate_ranges(&[Range::new(2, 3), Range::new(0, 1)]));
		assert!(validate_ranges(&[Range::new(0, 1), Range::new(2, 3)]));
	}

	#[test]
	fn encoded_len_works() {
		let mut rng = rand::rng();
		for _ in 0..1000 {
			let mut ranges = RangeSet::new();
			{
				let mut offset = 0;
				for _ in 0..10 {
					let start = rng.random_range(offset..=20);
					let end = rng.random_range(start..=20);
					offset = end;
					ranges.insert(Range::new(start, end));
				}
			}
			let actual = ranges.encoded_len();
			let expected = ranges.encode().len();
			assert_eq!(expected, actual);
		}
	}
}