# map.rs - source
Source of the Rust file `crates/tlsn/src/map.rs`.

Skip to main content

#
tlsn/
map.rs

    1use std::ops::Range;
    2
    3use mpz_memory_core::{Vector, binary::U8};
    4use rangeset::set::RangeSet;
    5
    6#[derive(Debug, Clone, PartialEq)]
    7pub(crate) struct RangeMap<T> {
    8    map: Vec<(usize, T)>,
    9}
    10
    11impl<T> Default for RangeMap<T>
    12where
    13    T: Item,
    14{
    15    fn default() -> Self {
    16        Self { map: Vec::new() }
    17    }
    18}
    19
    20impl<T> RangeMap<T>
    21where
    22    T: Item,
    23{
    24    /// Returns `true` if the map is empty.
    25    pub(crate) fn is_empty(&self) -> bool {
    26        self.map.is_empty()
    27    }
    28
    29    /// Returns the keys of the map.
    30    pub(crate) fn keys(&self) -> impl Iterator<Item = Range<usize>> {
    31        self.map
    32            .iter()
    33            .map(|(idx, item)| *idx..*idx + item.length())
    34    }
    35
    36    pub(crate) fn iter(&self) -> impl Iterator<Item = (Range<usize>, &T)> {
    37        self.map
    38            .iter()
    39            .map(|(idx, item)| (*idx..*idx + item.length(), item))
    40    }
    41
    42    pub(crate) fn get(&self, range: Range<usize>) -> Option<T::Slice<'_>> {
    43        if range.start >= range.end {
    44            return None;
    45        }
    46
    47        // Find the item with the greatest start index <= range.start
    48        let pos = match self.map.binary_search_by(|(idx, _)| idx.cmp(&range.start)) {
    49            Ok(i) => i,
    50            Err(0) => return None,
    51            Err(i) => i - 1,
    52        };
    53
    54        let (base, item) = &self.map[pos];
    55
    56        item.slice(range.start - *base..range.end - *base)
    57    }
    58
    59    pub(crate) fn index(&self, idx: &RangeSet<usize>) -> Option<Self> {
    60        let mut map = Vec::new();
    61        for idx in idx.iter() {
    62            let pos = match self.map.binary_search_by(|(base, _)| base.cmp(&idx.start)) {
    63                Ok(i) => i,
    64                Err(0) => return None,
    65                Err(i) => i - 1,
    66            };
    67
    68            let (base, item) = self.map.get(pos)?;
    69            if idx.start < *base || idx.end > *base + item.length() {
    70                return None;
    71            }
    72
    73            let start = idx.start - *base;
    74            let end = start + idx.len();
    75
    76            map.push((
    77                idx.start,
    78                item.slice(start..end)
    79                    .expect("slice length is checked")
    80                    .into(),
    81            ));
    82        }
    83
    84        Some(Self { map })
    85    }
    86}
    87
    88impl<T> FromIterator<(usize, T)> for RangeMap<T>
    89where
    90    T: Item,
    91{
    92    fn from_iter<I: IntoIterator<Item = (usize, T)>>(items: I) -> Self {
    93        let mut pos = 0;
    94        let mut map = Vec::new();
    95        for (idx, item) in items {
    96            assert!(
    97                idx >= pos,
    98                "items must be sorted by index and non-overlapping"
    99            );
    100
    101            pos = idx + item.length();
    102            map.push((idx, item));
    103        }
    104
    105        Self { map }
    106    }
    107}
    108
    109pub(crate) trait Item: Sized {
    110    type Slice<'a>: Into<Self>
    111    where
    112        Self: 'a;
    113
    114    fn length(&self) -> usize;
    115
    116    fn slice<'a>(&'a self, range: Range<usize>) -> Option<Self::Slice<'a>>;
    117}
    118
    119impl Item for Vector<U8> {
    120    type Slice<'a> = Vector<U8>;
    121
    122    fn length(&self) -> usize {
    123        self.len()
    124    }
    125
    126    fn slice<'a>(&'a self, range: Range<usize>) -> Option<Self::Slice<'a>> {
    127        self.get(range)
    128    }
    129}
    130
    131#[cfg(test)]
    132mod tests {
    133    use super::*;
    134
    135    impl Item for Range<usize> {
    136        type Slice<'a> = Range<usize>;
    137
    138        fn length(&self) -> usize {
    139            self.end - self.start
    140        }
    141
    142        fn slice(&self, range: Range<usize>) -> Option<Self> {
    143            if range.end > self.end - self.start {
    144                return None;
    145            }
    146
    147            Some(range.start + self.start..range.end + self.start)
    148        }
    149    }
    150
    151    #[test]
    152    fn test_range_map() {
    153        let map = RangeMap::from_iter([(0, 10..14), (10, 20..24), (20, 30..32)]);
    154
    155        assert_eq!(map.get(0..4), Some(10..14));
    156        assert_eq!(map.get(10..14), Some(20..24));
    157        assert_eq!(map.get(20..22), Some(30..32));
    158        assert_eq!(map.get(0..2), Some(10..12));
    159        assert_eq!(map.get(11..13), Some(21..23));
    160        assert_eq!(map.get(0..10), None);
    161        assert_eq!(map.get(10..20), None);
    162        assert_eq!(map.get(20..30), None);
    163    }
    164
    165    #[test]
    166    fn test_range_map_index() {
    167        let map = RangeMap::from_iter([(0, 10..14), (10, 20..24), (20, 30..32)]);
    168
    169        let idx = RangeSet::from([0..4, 10..14, 20..22]);
    170        assert_eq!(map.index(&idx), Some(map.clone()));
    171
    172        let idx = RangeSet::from(25..30);
    173        assert_eq!(map.index(&idx), None);
    174
    175        let idx = RangeSet::from(15..20);
    176        assert_eq!(map.index(&idx), None);
    177
    178        let idx = RangeSet::from([1..3, 11..12, 13..14, 21..22]);
    179        assert_eq!(
    180            map.index(&idx),
    181            Some(RangeMap::from_iter([
    182                (1, 11..13),
    183                (11, 21..22),
    184                (13, 23..24),
    185                (21, 31..32)
    186            ]))
    187        );
    188    }
    189}