1use super::interval_tree::IntervalTree;
2use super::vertex::VertexId;
3use formualizer_common::Coord as AbsCoord;
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: &[(AbsCoord, VertexId)]) {
59 self.add_vertices_batch(items);
60 }
61
62 pub fn add_vertex(&mut self, coord: AbsCoord, 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: &[(AbsCoord, 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: AbsCoord, 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(&mut self, old_coord: AbsCoord, new_coord: AbsCoord, vertex_id: VertexId) {
154 self.remove_vertex(old_coord, vertex_id);
155 self.add_vertex(new_coord, vertex_id);
156 }
157
158 pub fn vertices_in_row_range(&self, start: u32, end: u32) -> Vec<VertexId> {
163 let mut result = HashSet::new();
164
165 for (_low, _high, values) in self.row_tree.query(start, end) {
166 result.extend(values.iter());
167 }
168
169 result.into_iter().collect()
170 }
171
172 pub fn vertices_in_col_range(&self, start: u32, end: u32) -> Vec<VertexId> {
177 let mut result = HashSet::new();
178
179 for (_low, _high, values) in self.col_tree.query(start, end) {
180 result.extend(values.iter());
181 }
182
183 result.into_iter().collect()
184 }
185
186 pub fn vertices_in_rect(
191 &self,
192 start_row: u32,
193 end_row: u32,
194 start_col: u32,
195 end_col: u32,
196 ) -> Vec<VertexId> {
197 let row_vertices: HashSet<_> = self
199 .vertices_in_row_range(start_row, end_row)
200 .into_iter()
201 .collect();
202
203 let col_vertices: HashSet<_> = self
205 .vertices_in_col_range(start_col, end_col)
206 .into_iter()
207 .collect();
208
209 row_vertices.intersection(&col_vertices).copied().collect()
211 }
212
213 pub fn len(&self) -> usize {
215 let mut unique: HashSet<VertexId> = HashSet::new();
216
217 for (_low, _high, values) in self.row_tree.query(0, u32::MAX) {
219 unique.extend(values.iter());
220 }
221
222 unique.len()
223 }
224
225 pub fn is_empty(&self) -> bool {
227 self.row_tree.is_empty()
228 }
229
230 pub fn clear(&mut self) {
232 self.row_tree = IntervalTree::new();
233 self.col_tree = IntervalTree::new();
234 }
235}
236
237#[cfg(test)]
238mod tests {
239 use super::*;
240
241 #[test]
242 fn test_add_and_query_single_vertex() {
243 let mut index = SheetIndex::new();
244 let coord = AbsCoord::new(5, 10);
245 let vertex_id = VertexId(1024);
246
247 index.add_vertex(coord, vertex_id);
248
249 let row_results = index.vertices_in_row_range(5, 5);
251 assert_eq!(row_results, vec![vertex_id]);
252
253 let col_results = index.vertices_in_col_range(10, 10);
255 assert_eq!(col_results, vec![vertex_id]);
256
257 let row_results = index.vertices_in_row_range(3, 7);
259 assert_eq!(row_results, vec![vertex_id]);
260 }
261
262 #[test]
263 fn test_remove_vertex() {
264 let mut index = SheetIndex::new();
265 let coord = AbsCoord::new(5, 10);
266 let vertex_id = VertexId(1024);
267
268 index.add_vertex(coord, vertex_id);
269 assert_eq!(index.len(), 1);
270
271 index.remove_vertex(coord, vertex_id);
272 assert_eq!(index.len(), 0);
273
274 let row_results = index.vertices_in_row_range(5, 5);
276 assert!(row_results.is_empty());
277 }
278
279 #[test]
280 fn test_update_vertex_position() {
281 let mut index = SheetIndex::new();
282 let old_coord = AbsCoord::new(5, 10);
283 let new_coord = AbsCoord::new(15, 20);
284 let vertex_id = VertexId(1024);
285
286 index.add_vertex(old_coord, vertex_id);
287 index.update_vertex(old_coord, new_coord, vertex_id);
288
289 let old_row_results = index.vertices_in_row_range(5, 5);
291 assert!(old_row_results.is_empty());
292
293 let new_row_results = index.vertices_in_row_range(15, 15);
295 assert_eq!(new_row_results, vec![vertex_id]);
296
297 let new_col_results = index.vertices_in_col_range(20, 20);
298 assert_eq!(new_col_results, vec![vertex_id]);
299 }
300
301 #[test]
302 fn test_range_queries() {
303 let mut index = SheetIndex::new();
304
305 for row in 0..10 {
307 for col in 0..5 {
308 let coord = AbsCoord::new(row, col);
309 let vertex_id = VertexId(1024 + row * 5 + col);
310 index.add_vertex(coord, vertex_id);
311 }
312 }
313
314 let row_results = index.vertices_in_row_range(3, 5);
316 assert_eq!(row_results.len(), 15);
317
318 let col_results = index.vertices_in_col_range(1, 2);
320 assert_eq!(col_results.len(), 20);
321
322 let rect_results = index.vertices_in_rect(3, 5, 1, 2);
324 assert_eq!(rect_results.len(), 6);
325 }
326
327 #[test]
328 fn test_sparse_sheet_efficiency() {
329 let mut index = SheetIndex::new();
330
331 index.add_vertex(AbsCoord::new(100, 5), VertexId(1024));
333 index.add_vertex(AbsCoord::new(50_000, 10), VertexId(1025));
334 index.add_vertex(AbsCoord::new(100_000, 15), VertexId(1026));
335 index.add_vertex(AbsCoord::new(500_000, 20), VertexId(1027));
336 index.add_vertex(AbsCoord::new(999_999, 25), VertexId(1028));
337
338 assert_eq!(index.len(), 5);
339
340 let high_rows = index.vertices_in_row_range(100_000, u32::MAX);
342 assert_eq!(high_rows.len(), 3);
343
344 let col_range = index.vertices_in_col_range(10, 20);
346 assert_eq!(col_range.len(), 3); }
348
349 #[test]
350 fn test_shift_operation_query() {
351 let mut index = SheetIndex::new();
352
353 for row in [10, 20, 30, 40, 50] {
355 index.add_vertex(AbsCoord::new(row, 0), VertexId(1024 + row));
356 }
357
358 let vertices_to_shift = index.vertices_in_row_range(25, u32::MAX);
360 assert_eq!(vertices_to_shift.len(), 3); for col in 1..=3 {
364 index.add_vertex(AbsCoord::new(5, col), VertexId(2000 + col));
365 }
366
367 let vertices_to_delete = index.vertices_in_col_range(1, 3);
368 assert_eq!(vertices_to_delete.len(), 3);
369 }
370
371 #[test]
372 fn test_viewport_query() {
373 let mut index = SheetIndex::new();
374
375 for row in (0..10000).step_by(100) {
377 for col in 0..10 {
378 index.add_vertex(AbsCoord::new(row, col), VertexId(row * 10 + col));
379 }
380 }
381
382 let viewport = index.vertices_in_rect(500, 1500, 2, 7);
384
385 assert_eq!(viewport.len(), 66);
387 }
388}