Shared via Rust Playground

playground.rs

#![feature(collections, collections_bound, btree_range)]

use std::collections::Bound;
use std::borrow::Borrow;
use std::collections::btree_map::{self, BTreeMap};
use std::marker::PhantomData;

pub struct Unbounded<'a, Q: ?Sized + 'a>(PhantomData<&'a Q>);
pub struct Inclusive<'a, Q: ?Sized + 'a>(&'a Q);
pub struct Exclusive<'a, Q: ?Sized + 'a>(&'a Q);

trait IntoBound<'a> {
    type Bound: ?Sized + 'a;
    fn into_bound(self) -> Bound<&'a Self::Bound>;
}

impl<'a, Q: ?Sized + 'a> IntoBound<'a> for Inclusive<'a, Q> {
    type Bound = Q;
    fn into_bound(self) -> Bound<&'a Self::Bound> { Bound::Included(self.0) }
}
impl<'a, Q: ?Sized + 'a> IntoBound<'a> for Exclusive<'a, Q> {
    type Bound = Q;
    fn into_bound(self) -> Bound<&'a Self::Bound> { Bound::Excluded(self.0) }
}
impl<'a, Q: ?Sized + 'a> IntoBound<'a> for Unbounded<'a, Q> {
    type Bound = Q;
    fn into_bound(self) -> Bound<&'a Self::Bound> { Bound::Unbounded }
}






// Range

trait Ranged<K, V> {
    fn ranged(&self) -> Range<K, V, Unbounded<K>, Unbounded<K>>;
}

impl<K, V> Ranged<K, V> for BTreeMap<K, V> {
    fn ranged(&self) -> Range<K, V, Unbounded<K>, Unbounded<K>> {
        Range(self, Unbounded(PhantomData), Unbounded(PhantomData))
    }
}

pub struct Range<'a, K: 'a, V: 'a, L, R>(&'a BTreeMap<K, V>, L, R);

impl<'a, K, V, R> Range<'a, K, V, Unbounded<'a, K>, R> { 
    pub fn gt<Q:?Sized>(self, min: &Q) -> Range<'a, K, V, Exclusive<Q>, R> { Range(self.0, Exclusive(min), self.2) }
    pub fn ge<Q:?Sized>(self, min: &Q) -> Range<'a, K, V, Inclusive<Q>, R> { Range(self.0, Inclusive(min), self.2) }
}
impl<'a, K, V, L> Range<'a, K, V, L, Unbounded<'a, K>> { 
    pub fn lt<Q:?Sized>(self, max: &Q) -> Range<'a, K, V, L, Exclusive<Q>> { Range(self.0, self.1, Exclusive(max)) }
    pub fn le<Q:?Sized>(self, max: &Q) -> Range<'a, K, V, L, Inclusive<Q>> { Range(self.0, self.1, Inclusive(max)) }
}

impl<'map, K, V, L, R> IntoIterator for Range<'map, K, V, L, R>
    where L: IntoBound<'map> + 'map,
          R: IntoBound<'map> + 'map,
          L::Bound: Ord + 'map,
          R::Bound: Ord + 'map,
          K: Ord + Borrow<L::Bound> + Borrow<R::Bound>,
{
    type IntoIter = btree_map::Range<'map, K, V>;
    type Item = (&'map K, &'map V);
    
    fn into_iter(self) -> Self::IntoIter {
        let Range(map, l, r) = self;
        map.range(l.into_bound(), r.into_bound())
    }
}




// RangeMut

trait RangedMut<K, V> {
    fn ranged_mut(&mut self) -> RangeMut<K, V, Unbounded<K>, Unbounded<K>>;
}

impl<K, V> RangedMut<K, V> for BTreeMap<K, V> {
    fn ranged_mut(&mut self) -> RangeMut<K, V, Unbounded<K>, Unbounded<K>> {
        RangeMut(self, Unbounded(PhantomData), Unbounded(PhantomData))
    }
}

pub struct RangeMut<'a, K: 'a, V: 'a, L, R>(&'a mut BTreeMap<K, V>, L, R);

impl<'a, K, V, R> RangeMut<'a, K, V, Unbounded<'a, K>, R> { 
    pub fn gt<Q:?Sized>(self, min: &Q) -> RangeMut<'a, K, V, Exclusive<Q>, R> { RangeMut(self.0, Exclusive(min), self.2) }
    pub fn ge<Q:?Sized>(self, min: &Q) -> RangeMut<'a, K, V, Inclusive<Q>, R> { RangeMut(self.0, Inclusive(min), self.2) }
}
impl<'a, K, V, L> RangeMut<'a, K, V, L, Unbounded<'a, K>> { 
    pub fn lt<Q:?Sized>(self, max: &Q) -> RangeMut<'a, K, V, L, Exclusive<Q>> { RangeMut(self.0, self.1, Exclusive(max)) }
    pub fn le<Q:?Sized>(self, max: &Q) -> RangeMut<'a, K, V, L, Inclusive<Q>> { RangeMut(self.0, self.1, Inclusive(max)) }
}

impl<'map, K, V, L, R> IntoIterator for RangeMut<'map, K, V, L, R>
    where L: IntoBound<'map> + 'map,
          R: IntoBound<'map> + 'map,
          L::Bound: Ord + 'map,
          R::Bound: Ord + 'map,
          K: Ord + Borrow<L::Bound> + Borrow<R::Bound>,
{
    type IntoIter = btree_map::RangeMut<'map, K, V>;
    type Item = (&'map K, &'map mut V);
    
    fn into_iter(self) -> Self::IntoIter {
        let RangeMut(map, l, r) = self;
        map.range_mut(l.into_bound(), r.into_bound())
    }
}




fn main() {
    let mut map = BTreeMap::new();
    map.insert(0, 0);
    map.insert(10, 10);
    map.insert(20, 20);
    
    for (k, v) in map.ranged().lt(&20) {
        println!("{} {}", k, v); 
    }
    
    for (k, v) in map.ranged().ge(&10).lt(&30) {
        println!("{} {}", k, v); 
    }
}