rstared/
lib.rs

1// SPDX-FileCopyrightText: 2025 rstared contributors
2//
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5#![doc(html_root_url = "https://docs.rs/rstared")]
6#![doc = include_str!("../README.md")]
7//#![deny(missing_docs)]
8#![forbid(unsafe_code)]
9#![no_std]
10
11#[cfg(feature = "std")]
12extern crate std;
13
14// No feature for `alloc` because it would be always enabled anyway.
15extern crate alloc;
16
17use maplike::IntoIter;
18use rstar::{RTree, RTreeObject, primitives::GeomWithData};
19
20pub use maplike::{Get, Insert, Keyed, Map, Push, Remove, StableRemove};
21
22#[derive(Clone, Debug)]
23pub struct RTreed<K, V: RTreeObject, C> {
24    collection: C,
25    rtree: RTree<GeomWithData<V, K>>,
26}
27
28impl<K, V: RTreeObject, C> RTreed<K, V, C> {
29    pub fn new(collection: C) -> Self {
30        Self {
31            collection,
32            rtree: RTree::new(),
33        }
34    }
35
36    pub fn collection(&self) -> &C {
37        &self.collection
38    }
39
40    pub fn rtree(&self) -> &RTree<GeomWithData<V, K>> {
41        &self.rtree
42    }
43}
44
45impl<K, V: RTreeObject, C: Default> Default for RTreed<K, V, C> {
46    fn default() -> Self {
47        RTreed::new(C::default())
48    }
49}
50
51impl<K, V: RTreeObject, C> Map for RTreed<K, V, C> {
52    type Item = V;
53}
54
55impl<K, V: RTreeObject, C> Keyed for RTreed<K, V, C> {
56    type Key = K;
57}
58
59impl<K, V: RTreeObject, C: Get<K, Item = V>> Get<K> for RTreed<K, V, C> {
60    #[inline(always)]
61    fn get(&self, key: &K) -> Option<&V> {
62        self.collection.get(key)
63    }
64}
65
66impl<K: Clone, V: Clone + RTreeObject, C: Get<K, Item = V> + Insert<K>> Insert<K>
67    for RTreed<K, V, C>
68{
69    #[inline(always)]
70    fn insert(&mut self, key: K, value: V) {
71        self.rtree
72            .insert(GeomWithData::new(value.clone(), key.clone()));
73        self.collection.insert(key, value);
74    }
75}
76
77impl<K: Clone + PartialEq, V: Clone + PartialEq + RTreeObject, C: StableRemove<K, Item = V>>
78    Remove<K> for RTreed<K, V, C>
79{
80    #[inline(always)]
81    fn remove(&mut self, key: &K) -> Option<V> {
82        let value = self.collection.remove(&key.clone())?;
83        self.rtree
84            .remove(&GeomWithData::new(value.clone(), key.clone()));
85
86        Some(value)
87    }
88}
89
90impl<K: Clone + PartialEq, V: Clone + PartialEq + RTreeObject, C: StableRemove<K, Item = V>>
91    StableRemove<K> for RTreed<K, V, C>
92{
93}
94
95impl<K: Clone, V: Clone + RTreeObject, C: Push<K, Item = V>> Push<K> for RTreed<K, V, C> {
96    #[inline(always)]
97    fn push(&mut self, value: V) -> K {
98        let key = self.collection.push(value.clone());
99        self.rtree.insert(GeomWithData::new(value, key.clone()));
100
101        key
102    }
103}
104
105impl<K, V: RTreeObject, C: IntoIter<K, Item = V>> IntoIter<K> for RTreed<K, V, C> {
106    type IntoIter = C::IntoIter;
107
108    #[inline(always)]
109    fn into_iter(self) -> C::IntoIter {
110        self.collection.into_iter()
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use rand::Rng;
117    use rstar::{AABB, primitives::Rectangle};
118
119    use super::*;
120
121    #[cfg(feature = "stable-vec")]
122    #[test]
123    fn test_push_and_remove_random_aars_in_stable_vec() {
124        use stable_vec::StableVec;
125        test_push_and_remove_random_aars::<usize, StableVec<Rectangle<(i32, i32)>>>(StableVec::<
126            Rectangle<(i32, i32)>,
127        >::new(
128        ));
129    }
130
131    #[cfg(feature = "thunderdome")]
132    #[test]
133    fn test_push_and_remove_random_aars_in_thunderdome() {
134        use thunderdome::{Arena, Index};
135        test_push_and_remove_random_aars::<Index, Arena<Rectangle<(i32, i32)>>>(Arena::<
136            Rectangle<(i32, i32)>,
137        >::new());
138    }
139
140    /// "AAR" stands for "axis-aligned rectangle".
141    fn test_push_and_remove_random_aars<
142        K: Clone + PartialEq,
143        C: Get<K, Item = Rectangle<(i32, i32)>> + Push<K> + StableRemove<K>,
144    >(
145        collection: C,
146    ) {
147        let mut rtreed: RTreed<K, Rectangle<(i32, i32)>, C> = RTreed::new(collection);
148        let mut rng = rand::rng();
149        let mut keys = alloc::vec![];
150
151        for _ in 0..100 {
152            let x = rng.random_range(0..=90);
153            let y = rng.random_range(0..=90);
154            let width = rng.random_range(0..=10);
155            let height = rng.random_range(0..=10);
156
157            keys.push(rtreed.push(Rectangle::from_corners((x, y), (x + width, y + height))));
158        }
159
160        assert_eq!(
161            rtreed
162                .rtree()
163                .locate_in_envelope(&AABB::from_corners((0, 0), (100, 100)))
164                .count(),
165            100
166        );
167
168        for i in 0..10 {
169            rtreed.remove(&keys[i]);
170        }
171
172        assert_eq!(
173            rtreed
174                .rtree()
175                .locate_in_envelope(&AABB::from_corners((0, 0), (100, 100)))
176                .count(),
177            90
178        );
179    }
180}
rstared/lib.rs

rstared/
lib.rs
