1use super::interval_tree::IntervalTree;
2use super::packed_coord::PackedCoord;
3use super::vertex::VertexId;
4use std::collections::HashSet;
5
6#[derive(Debug, Default)]
38pub struct SheetIndex {
39 row_tree: IntervalTree<VertexId>,
42
43 col_tree: IntervalTree<VertexId>,
46}
47
48impl SheetIndex {
49 pub fn new() -> Self {
51 Self {
52 row_tree: IntervalTree::new(),
53 col_tree: IntervalTree::new(),
54 }
55 }
56
57 pub fn build_from_sorted(&mut self, items: &[(PackedCoord, VertexId)]) {
59 self.add_vertices_batch(items);
60 }
61
62 pub fn add_vertex(&mut self, coord: PackedCoord, vertex_id: VertexId) {
67 let row = coord.row();
68 let col = coord.col();
69
70 self.row_tree
72 .entry(row, row)
73 .or_insert_with(HashSet::new)
74 .insert(vertex_id);
75
76 self.col_tree
78 .entry(col, col)
79 .or_insert_with(HashSet::new)
80 .insert(vertex_id);
81 }
82
83 pub fn add_vertices_batch(&mut self, items: &[(PackedCoord, VertexId)]) {
85 if items.is_empty() {
86 return;
87 }
88 if self.row_tree.is_empty() && self.col_tree.is_empty() {
90 let mut row_items: Vec<(u32, HashSet<VertexId>)> = Vec::with_capacity(items.len());
92 let mut col_items: Vec<(u32, HashSet<VertexId>)> = Vec::with_capacity(items.len());
93 use rustc_hash::FxHashMap;
95 let mut row_map: FxHashMap<u32, HashSet<VertexId>> = FxHashMap::default();
96 let mut col_map: FxHashMap<u32, HashSet<VertexId>> = FxHashMap::default();
97 for (coord, vid) in items {
98 row_map.entry(coord.row()).or_default().insert(*vid);
99 col_map.entry(coord.col()).or_default().insert(*vid);
100 }
101 row_items.reserve(row_map.len());
102 for (k, v) in row_map.into_iter() {
103 row_items.push((k, v));
104 }
105 col_items.reserve(col_map.len());
106 for (k, v) in col_map.into_iter() {
107 col_items.push((k, v));
108 }
109 self.row_tree.bulk_build_points(row_items);
110 self.col_tree.bulk_build_points(col_items);
111 return;
112 }
113 for (coord, vid) in items {
115 let row = coord.row();
116 let col = coord.col();
117 self.row_tree
118 .entry(row, row)
119 .or_insert_with(HashSet::new)
120 .insert(*vid);
121 self.col_tree
122 .entry(col, col)
123 .or_insert_with(HashSet::new)
124 .insert(*vid);
125 }
126 }
127
128 pub fn remove_vertex(&mut self, coord: PackedCoord, vertex_id: VertexId) {
133 let row = coord.row();
134 let col = coord.col();
135
136 if let Some(vertices) = self.row_tree.get_mut(row, row) {
138 vertices.remove(&vertex_id);
139 }
142
143 if let Some(vertices) = self.col_tree.get_mut(col, col) {
145 vertices.remove(&vertex_id);
146 }
147 }
148
149 pub fn update_vertex(
154 &mut self,
155 old_coord: PackedCoord,
156 new_coord: PackedCoord,
157 vertex_id: VertexId,
158 ) {
159 self.remove_vertex(old_coord, vertex_id);
160 self.add_vertex(new_coord, vertex_id);
161 }
162
163 pub fn vertices_in_row_range(&self, start: u32, end: u32) -> Vec<VertexId> {
168 let mut result = HashSet::new();
169
170 for (_low, _high, values) in self.row_tree.query(start, end) {
171 result.extend(values.iter());
172 }
173
174 result.into_iter().collect()
175 }
176
177 pub fn vertices_in_col_range(&self, start: u32, end: u32) -> Vec<VertexId> {
182 let mut result = HashSet::new();
183
184 for (_low, _high, values) in self.col_tree.query(start, end) {
185 result.extend(values.iter());
186 }
187
188 result.into_iter().collect()
189 }
190
191 pub fn vertices_in_rect(
196 &self,
197 start_row: u32,
198 end_row: u32,
199 start_col: u32,
200 end_col: u32,
201 ) -> Vec<VertexId> {
202 let row_vertices: HashSet<_> = self
204 .vertices_in_row_range(start_row, end_row)
205 .into_iter()
206 .collect();
207
208 let col_vertices: HashSet<_> = self
210 .vertices_in_col_range(start_col, end_col)
211 .into_iter()
212 .collect();
213
214 row_vertices.intersection(&col_vertices).copied().collect()
216 }
217
218 pub fn len(&self) -> usize {
220 let mut unique: HashSet<VertexId> = HashSet::new();
221
222 for (_low, _high, values) in self.row_tree.query(0, u32::MAX) {
224 unique.extend(values.iter());
225 }
226
227 unique.len()
228 }
229
230 pub fn is_empty(&self) -> bool {
232 self.row_tree.is_empty()
233 }
234
235 pub fn clear(&mut self) {
237 self.row_tree = IntervalTree::new();
238 self.col_tree = IntervalTree::new();
239 }
240}
241
242#[cfg(test)]
243mod tests {
244 use super::*;
245
246 #[test]
247 fn test_add_and_query_single_vertex() {
248 let mut index = SheetIndex::new();
249 let coord = PackedCoord::new(5, 10);
250 let vertex_id = VertexId(1024);
251
252 index.add_vertex(coord, vertex_id);
253
254 let row_results = index.vertices_in_row_range(5, 5);
256 assert_eq!(row_results, vec![vertex_id]);
257
258 let col_results = index.vertices_in_col_range(10, 10);
260 assert_eq!(col_results, vec![vertex_id]);
261
262 let row_results = index.vertices_in_row_range(3, 7);
264 assert_eq!(row_results, vec![vertex_id]);
265 }
266
267 #[test]
268 fn test_remove_vertex() {
269 let mut index = SheetIndex::new();
270 let coord = PackedCoord::new(5, 10);
271 let vertex_id = VertexId(1024);
272
273 index.add_vertex(coord, vertex_id);
274 assert_eq!(index.len(), 1);
275
276 index.remove_vertex(coord, vertex_id);
277 assert_eq!(index.len(), 0);
278
279 let row_results = index.vertices_in_row_range(5, 5);
281 assert!(row_results.is_empty());
282 }
283
284 #[test]
285 fn test_update_vertex_position() {
286 let mut index = SheetIndex::new();
287 let old_coord = PackedCoord::new(5, 10);
288 let new_coord = PackedCoord::new(15, 20);
289 let vertex_id = VertexId(1024);
290
291 index.add_vertex(old_coord, vertex_id);
292 index.update_vertex(old_coord, new_coord, vertex_id);
293
294 let old_row_results = index.vertices_in_row_range(5, 5);
296 assert!(old_row_results.is_empty());
297
298 let new_row_results = index.vertices_in_row_range(15, 15);
300 assert_eq!(new_row_results, vec![vertex_id]);
301
302 let new_col_results = index.vertices_in_col_range(20, 20);
303 assert_eq!(new_col_results, vec![vertex_id]);
304 }
305
306 #[test]
307 fn test_range_queries() {
308 let mut index = SheetIndex::new();
309
310 for row in 0..10 {
312 for col in 0..5 {
313 let coord = PackedCoord::new(row, col);
314 let vertex_id = VertexId(1024 + row * 5 + col);
315 index.add_vertex(coord, vertex_id);
316 }
317 }
318
319 let row_results = index.vertices_in_row_range(3, 5);
321 assert_eq!(row_results.len(), 15);
322
323 let col_results = index.vertices_in_col_range(1, 2);
325 assert_eq!(col_results.len(), 20);
326
327 let rect_results = index.vertices_in_rect(3, 5, 1, 2);
329 assert_eq!(rect_results.len(), 6);
330 }
331
332 #[test]
333 fn test_sparse_sheet_efficiency() {
334 let mut index = SheetIndex::new();
335
336 index.add_vertex(PackedCoord::new(100, 5), VertexId(1024));
338 index.add_vertex(PackedCoord::new(50_000, 10), VertexId(1025));
339 index.add_vertex(PackedCoord::new(100_000, 15), VertexId(1026));
340 index.add_vertex(PackedCoord::new(500_000, 20), VertexId(1027));
341 index.add_vertex(PackedCoord::new(999_999, 25), VertexId(1028));
342
343 assert_eq!(index.len(), 5);
344
345 let high_rows = index.vertices_in_row_range(100_000, u32::MAX);
347 assert_eq!(high_rows.len(), 3);
348
349 let col_range = index.vertices_in_col_range(10, 20);
351 assert_eq!(col_range.len(), 3); }
353
354 #[test]
355 fn test_shift_operation_query() {
356 let mut index = SheetIndex::new();
357
358 for row in [10, 20, 30, 40, 50] {
360 index.add_vertex(PackedCoord::new(row, 0), VertexId(1024 + row));
361 }
362
363 let vertices_to_shift = index.vertices_in_row_range(25, u32::MAX);
365 assert_eq!(vertices_to_shift.len(), 3); for col in 1..=3 {
369 index.add_vertex(PackedCoord::new(5, col), VertexId(2000 + col));
370 }
371
372 let vertices_to_delete = index.vertices_in_col_range(1, 3);
373 assert_eq!(vertices_to_delete.len(), 3);
374 }
375
376 #[test]
377 fn test_viewport_query() {
378 let mut index = SheetIndex::new();
379
380 for row in (0..10000).step_by(100) {
382 for col in 0..10 {
383 index.add_vertex(PackedCoord::new(row, col), VertexId(row * 10 + col));
384 }
385 }
386
387 let viewport = index.vertices_in_rect(500, 1500, 2, 7);
389
390 assert_eq!(viewport.len(), 66);
392 }
393}