luci/spatial/
shape.rs

1//! Geo shape storage, GeoJSON parsing, and serialization.
2//!
3//! Provides [`GeoShapeStore`] for per-segment geo_shape field storage with a
4//! packed R-tree for spatial candidate selection and serialized shapes for
5//! exact predicate evaluation.
6//!
7//! See [[feature-geo-shape]] and [[geospatial]].
8
9use ::geo::algorithm::bounding_rect::BoundingRect;
10use ::geo::geometry::{
11    Coord, Geometry, GeometryCollection, LineString, MultiLineString, MultiPoint, MultiPolygon,
12    Point, Polygon, Rect,
13};
14use geo_index::rtree::sort::HilbertSort;
15use geo_index::rtree::{RTreeBuilder, RTreeIndex, RTreeRef};
16
17// ---------------------------------------------------------------------------
18// GeoJSON parsing
19// ---------------------------------------------------------------------------
20
21/// Parse a GeoJSON value into a `geo::Geometry<f64>`.
22///
23/// Supports all standard GeoJSON types plus Elasticsearch's `envelope`
24/// extension. Coordinates use GeoJSON order: `[longitude, latitude]`.
25pub fn parse_geojson(value: &serde_json::Value) -> Option<Geometry<f64>> {
26    let obj = value.as_object()?;
27    let shape_type = obj.get("type").and_then(|v| v.as_str())?;
28
29    match shape_type {
30        "Point" | "point" => {
31            let coords = obj.get("coordinates")?;
32            let c = parse_coord(coords)?;
33            Some(Geometry::Point(Point(c)))
34        }
35        "MultiPoint" | "multipoint" => {
36            let coords = obj.get("coordinates")?.as_array()?;
37            let points: Option<Vec<Point<f64>>> =
38                coords.iter().map(|c| parse_coord(c).map(Point)).collect();
39            Some(Geometry::MultiPoint(MultiPoint(points?)))
40        }
41        "LineString" | "linestring" => {
42            let coords = obj.get("coordinates")?;
43            let line = parse_linestring(coords)?;
44            Some(Geometry::LineString(line))
45        }
46        "MultiLineString" | "multilinestring" => {
47            let coords = obj.get("coordinates")?.as_array()?;
48            let lines: Option<Vec<LineString<f64>>> = coords.iter().map(parse_linestring).collect();
49            Some(Geometry::MultiLineString(MultiLineString(lines?)))
50        }
51        "Polygon" | "polygon" => {
52            let coords = obj.get("coordinates")?;
53            let poly = parse_polygon(coords)?;
54            Some(Geometry::Polygon(poly))
55        }
56        "MultiPolygon" | "multipolygon" => {
57            let coords = obj.get("coordinates")?.as_array()?;
58            let polys: Option<Vec<Polygon<f64>>> = coords.iter().map(parse_polygon).collect();
59            Some(Geometry::MultiPolygon(MultiPolygon(polys?)))
60        }
61        "GeometryCollection" | "geometrycollection" => {
62            let geoms = obj.get("geometries")?.as_array()?;
63            let parsed: Option<Vec<Geometry<f64>>> = geoms.iter().map(parse_geojson).collect();
64            Some(Geometry::GeometryCollection(GeometryCollection(parsed?)))
65        }
66        "envelope" | "Envelope" => {
67            // ES extension: [[minLon, maxLat], [maxLon, minLat]]
68            let coords = obj.get("coordinates")?.as_array()?;
69            if coords.len() != 2 {
70                return None;
71            }
72            let tl = parse_coord(&coords[0])?; // [minLon, maxLat]
73            let br = parse_coord(&coords[1])?; // [maxLon, minLat]
74            Some(Geometry::Rect(Rect::new(
75                Coord { x: tl.x, y: br.y }, // min corner
76                Coord { x: br.x, y: tl.y }, // max corner
77            )))
78        }
79        _ => None,
80    }
81}
82
83/// Parse a single GeoJSON coordinate `[lon, lat]`.
84fn parse_coord(value: &serde_json::Value) -> Option<Coord<f64>> {
85    let arr = value.as_array()?;
86    if arr.len() < 2 {
87        return None;
88    }
89    let x = arr[0].as_f64()?; // longitude
90    let y = arr[1].as_f64()?; // latitude
91    Some(Coord { x, y })
92}
93
94/// Parse a GeoJSON coordinate array into a `LineString`.
95fn parse_linestring(value: &serde_json::Value) -> Option<LineString<f64>> {
96    let arr = value.as_array()?;
97    let coords: Option<Vec<Coord<f64>>> = arr.iter().map(parse_coord).collect();
98    Some(LineString(coords?))
99}
100
101/// Parse a GeoJSON polygon (array of rings).
102fn parse_polygon(value: &serde_json::Value) -> Option<Polygon<f64>> {
103    let rings = value.as_array()?;
104    if rings.is_empty() {
105        return None;
106    }
107    let exterior = parse_linestring(&rings[0])?;
108    let interiors: Option<Vec<LineString<f64>>> = rings[1..].iter().map(parse_linestring).collect();
109    Some(Polygon::new(exterior, interiors?))
110}
111
112// ---------------------------------------------------------------------------
113// Bounding box computation
114// ---------------------------------------------------------------------------
115
116/// Compute (min_x, min_y, max_x, max_y) for a geometry.
117/// Returns `(min_lon, min_lat, max_lon, max_lat)`.
118pub fn compute_bbox(geom: &Geometry<f64>) -> Option<(f64, f64, f64, f64)> {
119    let rect = geom.bounding_rect()?;
120    Some((rect.min().x, rect.min().y, rect.max().x, rect.max().y))
121}
122
123// ---------------------------------------------------------------------------
124// Shape serialization
125// ---------------------------------------------------------------------------
126
127const TAG_POINT: u8 = 0;
128const TAG_MULTI_POINT: u8 = 1;
129const TAG_LINE_STRING: u8 = 2;
130const TAG_MULTI_LINE_STRING: u8 = 3;
131const TAG_POLYGON: u8 = 4;
132const TAG_MULTI_POLYGON: u8 = 5;
133const TAG_GEOMETRY_COLLECTION: u8 = 6;
134const TAG_RECT: u8 = 7;
135
136/// Serialize a geometry into a compact binary format.
137pub fn serialize_shape(geom: &Geometry<f64>) -> Vec<u8> {
138    let mut buf = Vec::new();
139    write_shape(geom, &mut buf);
140    buf
141}
142
143fn write_coord(c: &Coord<f64>, buf: &mut Vec<u8>) {
144    buf.extend_from_slice(&c.x.to_le_bytes());
145    buf.extend_from_slice(&c.y.to_le_bytes());
146}
147
148fn write_coords(coords: &[Coord<f64>], buf: &mut Vec<u8>) {
149    buf.extend_from_slice(&(coords.len() as u32).to_le_bytes());
150    for c in coords {
151        write_coord(c, buf);
152    }
153}
154
155fn write_shape(geom: &Geometry<f64>, buf: &mut Vec<u8>) {
156    match geom {
157        Geometry::Point(p) => {
158            buf.push(TAG_POINT);
159            write_coord(&p.0, buf);
160        }
161        Geometry::MultiPoint(mp) => {
162            buf.push(TAG_MULTI_POINT);
163            buf.extend_from_slice(&(mp.0.len() as u32).to_le_bytes());
164            for p in &mp.0 {
165                write_coord(&p.0, buf);
166            }
167        }
168        Geometry::LineString(ls) => {
169            buf.push(TAG_LINE_STRING);
170            write_coords(&ls.0, buf);
171        }
172        Geometry::MultiLineString(mls) => {
173            buf.push(TAG_MULTI_LINE_STRING);
174            buf.extend_from_slice(&(mls.0.len() as u32).to_le_bytes());
175            for ls in &mls.0 {
176                write_coords(&ls.0, buf);
177            }
178        }
179        Geometry::Polygon(poly) => {
180            buf.push(TAG_POLYGON);
181            let num_rings = 1 + poly.interiors().len();
182            buf.extend_from_slice(&(num_rings as u32).to_le_bytes());
183            write_coords(&poly.exterior().0, buf);
184            for ring in poly.interiors() {
185                write_coords(&ring.0, buf);
186            }
187        }
188        Geometry::MultiPolygon(mp) => {
189            buf.push(TAG_MULTI_POLYGON);
190            buf.extend_from_slice(&(mp.0.len() as u32).to_le_bytes());
191            for poly in &mp.0 {
192                let num_rings = 1 + poly.interiors().len();
193                buf.extend_from_slice(&(num_rings as u32).to_le_bytes());
194                write_coords(&poly.exterior().0, buf);
195                for ring in poly.interiors() {
196                    write_coords(&ring.0, buf);
197                }
198            }
199        }
200        Geometry::GeometryCollection(gc) => {
201            buf.push(TAG_GEOMETRY_COLLECTION);
202            buf.extend_from_slice(&(gc.0.len() as u32).to_le_bytes());
203            for g in &gc.0 {
204                let child = serialize_shape(g);
205                buf.extend_from_slice(&(child.len() as u32).to_le_bytes());
206                buf.extend_from_slice(&child);
207            }
208        }
209        Geometry::Rect(r) => {
210            buf.push(TAG_RECT);
211            write_coord(&r.min(), buf);
212            write_coord(&r.max(), buf);
213        }
214        // geo crate has Line and Triangle — treat as LineString
215        Geometry::Line(l) => {
216            buf.push(TAG_LINE_STRING);
217            let coords = vec![l.start, l.end];
218            write_coords(&coords, buf);
219        }
220        Geometry::Triangle(t) => {
221            buf.push(TAG_POLYGON);
222            buf.extend_from_slice(&1u32.to_le_bytes()); // 1 ring
223            let coords = vec![t.0, t.1, t.2, t.0]; // closed ring
224            write_coords(&coords, buf);
225        }
226    }
227}
228
229/// Deserialize a geometry from binary data. Returns (geometry, bytes_consumed).
230pub fn deserialize_shape(data: &[u8]) -> Option<(Geometry<f64>, usize)> {
231    if data.is_empty() {
232        return None;
233    }
234    let tag = data[0];
235    let mut pos = 1;
236
237    match tag {
238        TAG_POINT => {
239            let c = read_coord(data, &mut pos)?;
240            Some((Geometry::Point(Point(c)), pos))
241        }
242        TAG_MULTI_POINT => {
243            let count = read_u32(data, &mut pos)? as usize;
244            let mut points = Vec::with_capacity(count);
245            for _ in 0..count {
246                points.push(Point(read_coord(data, &mut pos)?));
247            }
248            Some((Geometry::MultiPoint(MultiPoint(points)), pos))
249        }
250        TAG_LINE_STRING => {
251            let coords = read_coords(data, &mut pos)?;
252            Some((Geometry::LineString(LineString(coords)), pos))
253        }
254        TAG_MULTI_LINE_STRING => {
255            let count = read_u32(data, &mut pos)? as usize;
256            let mut lines = Vec::with_capacity(count);
257            for _ in 0..count {
258                lines.push(LineString(read_coords(data, &mut pos)?));
259            }
260            Some((Geometry::MultiLineString(MultiLineString(lines)), pos))
261        }
262        TAG_POLYGON => {
263            let poly = read_polygon(data, &mut pos)?;
264            Some((Geometry::Polygon(poly), pos))
265        }
266        TAG_MULTI_POLYGON => {
267            let count = read_u32(data, &mut pos)? as usize;
268            let mut polys = Vec::with_capacity(count);
269            for _ in 0..count {
270                polys.push(read_polygon(data, &mut pos)?);
271            }
272            Some((Geometry::MultiPolygon(MultiPolygon(polys)), pos))
273        }
274        TAG_GEOMETRY_COLLECTION => {
275            let count = read_u32(data, &mut pos)? as usize;
276            let mut geoms = Vec::with_capacity(count);
277            for _ in 0..count {
278                let len = read_u32(data, &mut pos)? as usize;
279                let (g, _) = deserialize_shape(&data[pos..pos + len])?;
280                geoms.push(g);
281                pos += len;
282            }
283            Some((Geometry::GeometryCollection(GeometryCollection(geoms)), pos))
284        }
285        TAG_RECT => {
286            let min = read_coord(data, &mut pos)?;
287            let max = read_coord(data, &mut pos)?;
288            Some((Geometry::Rect(Rect::new(min, max)), pos))
289        }
290        _ => None,
291    }
292}
293
294fn read_f64(data: &[u8], pos: &mut usize) -> Option<f64> {
295    if *pos + 8 > data.len() {
296        return None;
297    }
298    let val = f64::from_le_bytes(data[*pos..*pos + 8].try_into().ok()?);
299    *pos += 8;
300    Some(val)
301}
302
303fn read_u32(data: &[u8], pos: &mut usize) -> Option<u32> {
304    if *pos + 4 > data.len() {
305        return None;
306    }
307    let val = u32::from_le_bytes(data[*pos..*pos + 4].try_into().ok()?);
308    *pos += 4;
309    Some(val)
310}
311
312fn read_coord(data: &[u8], pos: &mut usize) -> Option<Coord<f64>> {
313    let x = read_f64(data, pos)?;
314    let y = read_f64(data, pos)?;
315    Some(Coord { x, y })
316}
317
318fn read_coords(data: &[u8], pos: &mut usize) -> Option<Vec<Coord<f64>>> {
319    let count = read_u32(data, pos)? as usize;
320    let mut coords = Vec::with_capacity(count);
321    for _ in 0..count {
322        coords.push(read_coord(data, pos)?);
323    }
324    Some(coords)
325}
326
327fn read_polygon(data: &[u8], pos: &mut usize) -> Option<Polygon<f64>> {
328    let num_rings = read_u32(data, pos)? as usize;
329    if num_rings == 0 {
330        return None;
331    }
332    let exterior = LineString(read_coords(data, pos)?);
333    let mut interiors = Vec::with_capacity(num_rings.saturating_sub(1));
334    for _ in 1..num_rings {
335        interiors.push(LineString(read_coords(data, pos)?));
336    }
337    Some(Polygon::new(exterior, interiors))
338}
339
340// ---------------------------------------------------------------------------
341// GeoShapeStore
342// ---------------------------------------------------------------------------
343
344/// Per-segment geo shape storage with packed R-tree for candidate selection.
345///
346/// Stores bounding boxes for R-tree indexing and serialized shape data for
347/// exact predicate evaluation. Indexed by doc_id.
348///
349/// See [[feature-geo-shape]].
350#[derive(Clone)]
351pub struct GeoShapeStore {
352    count: usize,
353    /// Per-doc bounding boxes (min_x, min_y, max_x, max_y). NaN = null.
354    bbox_min_xs: Vec<f64>,
355    bbox_min_ys: Vec<f64>,
356    bbox_max_xs: Vec<f64>,
357    bbox_max_ys: Vec<f64>,
358    /// Per-doc offsets into shape_data. u32::MAX = null.
359    shape_offsets: Vec<u32>,
360    /// Concatenated serialized shapes.
361    shape_data: Vec<u8>,
362    /// Pre-built packed R-tree bytes (built during to_bytes, loaded from from_bytes).
363    rtree_data: Vec<u8>,
364}
365
366impl GeoShapeStore {
367    pub fn new() -> Self {
368        Self {
369            count: 0,
370            bbox_min_xs: Vec::new(),
371            bbox_min_ys: Vec::new(),
372            bbox_max_xs: Vec::new(),
373            bbox_max_ys: Vec::new(),
374            shape_offsets: Vec::new(),
375            shape_data: Vec::new(),
376            rtree_data: Vec::new(),
377        }
378    }
379
380    /// Add a geometry for the current doc.
381    pub fn add(&mut self, geom: &Geometry<f64>) {
382        let offset = self.shape_data.len() as u32;
383        let serialized = serialize_shape(geom);
384        self.shape_data.extend_from_slice(&serialized);
385        self.shape_offsets.push(offset);
386
387        if let Some((min_x, min_y, max_x, max_y)) = compute_bbox(geom) {
388            self.bbox_min_xs.push(min_x);
389            self.bbox_min_ys.push(min_y);
390            self.bbox_max_xs.push(max_x);
391            self.bbox_max_ys.push(max_y);
392        } else {
393            self.bbox_min_xs.push(f64::NAN);
394            self.bbox_min_ys.push(f64::NAN);
395            self.bbox_max_xs.push(f64::NAN);
396            self.bbox_max_ys.push(f64::NAN);
397        }
398        self.count += 1;
399    }
400
401    /// Add a null entry for a doc without a geo_shape value.
402    pub fn add_null(&mut self) {
403        self.bbox_min_xs.push(f64::NAN);
404        self.bbox_min_ys.push(f64::NAN);
405        self.bbox_max_xs.push(f64::NAN);
406        self.bbox_max_ys.push(f64::NAN);
407        self.shape_offsets.push(u32::MAX);
408        self.count += 1;
409    }
410
411    pub fn len(&self) -> usize {
412        self.count
413    }
414
415    pub fn is_empty(&self) -> bool {
416        self.count == 0
417    }
418
419    /// Access the raw shape offsets (for R-tree doc ID mapping).
420    pub fn shape_offsets_ref(&self) -> &[u32] {
421        &self.shape_offsets
422    }
423
424    /// Get the pre-built R-tree bytes (empty if not yet built).
425    pub fn rtree_data(&self) -> &[u8] {
426        &self.rtree_data
427    }
428
429    /// Check if a doc's shape is a Rect (bbox == exact geometry).
430    pub fn is_rect_shape(&self, doc_id: u32) -> bool {
431        let i = doc_id as usize;
432        if i >= self.count {
433            return false;
434        }
435        let offset = self.shape_offsets[i];
436        if offset == u32::MAX {
437            return false;
438        }
439        let o = offset as usize;
440        o < self.shape_data.len() && self.shape_data[o] == TAG_RECT
441    }
442
443    /// Get the bounding box for a doc, or None if null.
444    pub fn get_bbox(&self, doc_id: u32) -> Option<(f64, f64, f64, f64)> {
445        let i = doc_id as usize;
446        if i >= self.count {
447            return None;
448        }
449        let min_x = self.bbox_min_xs[i];
450        if min_x.is_nan() {
451            return None;
452        }
453        Some((
454            min_x,
455            self.bbox_min_ys[i],
456            self.bbox_max_xs[i],
457            self.bbox_max_ys[i],
458        ))
459    }
460
461    /// Deserialize and return the full geometry for a doc.
462    pub fn get_shape(&self, doc_id: u32) -> Option<Geometry<f64>> {
463        let i = doc_id as usize;
464        if i >= self.count {
465            return None;
466        }
467        let offset = self.shape_offsets[i];
468        if offset == u32::MAX {
469            return None;
470        }
471        let (geom, _) = deserialize_shape(&self.shape_data[offset as usize..])?;
472        Some(geom)
473    }
474
475    /// Build a packed R-tree from the stored bounding boxes.
476    pub fn build_rtree(&self) -> Vec<u8> {
477        let valid_count = self
478            .shape_offsets
479            .iter()
480            .filter(|&&o| o != u32::MAX)
481            .count();
482        if valid_count == 0 {
483            return Vec::new();
484        }
485        let mut builder = RTreeBuilder::<f64>::new(valid_count as u32);
486        for i in 0..self.count {
487            if self.shape_offsets[i] != u32::MAX {
488                builder.add(
489                    self.bbox_min_xs[i],
490                    self.bbox_min_ys[i],
491                    self.bbox_max_xs[i],
492                    self.bbox_max_ys[i],
493                );
494            }
495        }
496        let tree = builder.finish::<HilbertSort>();
497        tree.into_inner()
498    }
499
500    /// Search the R-tree for candidates whose bboxes intersect the query bbox.
501    /// Returns doc IDs (not R-tree insertion indices).
502    pub fn search_rtree(
503        rtree_data: &[u8],
504        query_bbox: (f64, f64, f64, f64),
505        shape_offsets: &[u32],
506    ) -> Vec<u32> {
507        if rtree_data.is_empty() {
508            return Vec::new();
509        }
510        let tree = match RTreeRef::<f64>::try_new(&rtree_data) {
511            Ok(t) => t,
512            Err(_) => return Vec::new(),
513        };
514        let (min_x, min_y, max_x, max_y) = query_bbox;
515        let rtree_indices = tree.search(min_x, min_y, max_x, max_y);
516
517        // Map R-tree insertion indices back to doc IDs.
518        // Insertion order = iteration order of non-null docs.
519        let doc_id_map: Vec<u32> = shape_offsets
520            .iter()
521            .enumerate()
522            .filter(|(_, o)| **o != u32::MAX)
523            .map(|(i, _)| i as u32)
524            .collect();
525
526        rtree_indices
527            .iter()
528            .filter_map(|&idx| doc_id_map.get(idx as usize).copied())
529            .collect()
530    }
531
532    // -----------------------------------------------------------------------
533    // Binary serialization
534    // -----------------------------------------------------------------------
535
536    /// Serialize to bytes for segment storage.
537    pub fn to_bytes(&self) -> Vec<u8> {
538        let rtree_data = self.build_rtree();
539        let mut buf = Vec::new();
540
541        // Count
542        buf.extend_from_slice(&(self.count as u32).to_le_bytes());
543
544        // Bbox arrays
545        for v in &self.bbox_min_xs {
546            buf.extend_from_slice(&v.to_le_bytes());
547        }
548        for v in &self.bbox_min_ys {
549            buf.extend_from_slice(&v.to_le_bytes());
550        }
551        for v in &self.bbox_max_xs {
552            buf.extend_from_slice(&v.to_le_bytes());
553        }
554        for v in &self.bbox_max_ys {
555            buf.extend_from_slice(&v.to_le_bytes());
556        }
557
558        // Shape offsets
559        for v in &self.shape_offsets {
560            buf.extend_from_slice(&v.to_le_bytes());
561        }
562
563        // Shape data
564        buf.extend_from_slice(&(self.shape_data.len() as u32).to_le_bytes());
565        buf.extend_from_slice(&self.shape_data);
566
567        // R-tree
568        buf.extend_from_slice(&(rtree_data.len() as u32).to_le_bytes());
569        buf.extend_from_slice(&rtree_data);
570
571        buf
572    }
573
574    /// Deserialize from segment bytes.
575    pub fn from_bytes(data: &[u8]) -> Self {
576        if data.len() < 4 {
577            return Self::new();
578        }
579        let count = u32::from_le_bytes(data[0..4].try_into().unwrap()) as usize;
580        let mut pos = 4;
581
582        let read_f64_vec = |data: &[u8], pos: &mut usize, n: usize| -> Vec<f64> {
583            let mut v = Vec::with_capacity(n);
584            for _ in 0..n {
585                let val = f64::from_le_bytes(data[*pos..*pos + 8].try_into().unwrap());
586                v.push(val);
587                *pos += 8;
588            }
589            v
590        };
591
592        let bbox_min_xs = read_f64_vec(data, &mut pos, count);
593        let bbox_min_ys = read_f64_vec(data, &mut pos, count);
594        let bbox_max_xs = read_f64_vec(data, &mut pos, count);
595        let bbox_max_ys = read_f64_vec(data, &mut pos, count);
596
597        let mut shape_offsets = Vec::with_capacity(count);
598        for _ in 0..count {
599            let v = u32::from_le_bytes(data[pos..pos + 4].try_into().unwrap());
600            shape_offsets.push(v);
601            pos += 4;
602        }
603
604        let shape_data_len = u32::from_le_bytes(data[pos..pos + 4].try_into().unwrap()) as usize;
605        pos += 4;
606        let shape_data = data[pos..pos + shape_data_len].to_vec();
607        pos += shape_data_len;
608
609        let rtree_data = if pos + 4 <= data.len() {
610            let rtree_len = u32::from_le_bytes(data[pos..pos + 4].try_into().unwrap()) as usize;
611            pos += 4;
612            if rtree_len > 0 && pos + rtree_len <= data.len() {
613                data[pos..pos + rtree_len].to_vec()
614            } else {
615                Vec::new()
616            }
617        } else {
618            Vec::new()
619        };
620
621        Self {
622            count,
623            bbox_min_xs,
624            bbox_min_ys,
625            bbox_max_xs,
626            bbox_max_ys,
627            shape_offsets,
628            shape_data,
629            rtree_data,
630        }
631    }
632
633    /// Get the R-tree bytes from already-serialized store data.
634    /// Used by the reader to pass R-tree data directly to search_rtree().
635    pub fn rtree_offset_in_bytes(data: &[u8]) -> Option<(usize, usize)> {
636        if data.len() < 4 {
637            return None;
638        }
639        let count = u32::from_le_bytes(data[0..4].try_into().unwrap()) as usize;
640        // Skip: count(4) + 4 bbox arrays(count*8 each) + offsets(count*4) + shape_data_len(4) + shape_data
641        let mut pos = 4 + count * 8 * 4 + count * 4;
642        if pos + 4 > data.len() {
643            return None;
644        }
645        let shape_data_len = u32::from_le_bytes(data[pos..pos + 4].try_into().unwrap()) as usize;
646        pos += 4 + shape_data_len;
647        if pos + 4 > data.len() {
648            return None;
649        }
650        let rtree_len = u32::from_le_bytes(data[pos..pos + 4].try_into().unwrap()) as usize;
651        pos += 4;
652        if rtree_len == 0 || pos + rtree_len > data.len() {
653            return None;
654        }
655        Some((pos, rtree_len))
656    }
657}
658
659// ---------------------------------------------------------------------------
660// Tests
661// ---------------------------------------------------------------------------
662
663#[cfg(test)]
664mod tests {
665    use super::*;
666    use serde_json::json;
667
668    #[test]
669    fn parse_point() {
670        let v = json!({"type": "Point", "coordinates": [-73.98, 40.75]});
671        let g = parse_geojson(&v).unwrap();
672        if let Geometry::Point(p) = g {
673            assert!((p.x() - (-73.98)).abs() < 1e-10);
674            assert!((p.y() - 40.75).abs() < 1e-10);
675        } else {
676            panic!("expected Point");
677        }
678    }
679
680    #[test]
681    fn parse_polygon() {
682        let v = json!({
683            "type": "Polygon",
684            "coordinates": [[[0.0,0.0],[10.0,0.0],[10.0,10.0],[0.0,10.0],[0.0,0.0]]]
685        });
686        let g = parse_geojson(&v).unwrap();
687        assert!(matches!(g, Geometry::Polygon(_)));
688    }
689
690    #[test]
691    fn parse_envelope() {
692        let v = json!({"type": "envelope", "coordinates": [[-77.0, 39.0], [-75.0, 38.0]]});
693        let g = parse_geojson(&v).unwrap();
694        if let Geometry::Rect(r) = g {
695            assert!((r.min().x - (-77.0)).abs() < 1e-10);
696            assert!((r.min().y - 38.0).abs() < 1e-10);
697            assert!((r.max().x - (-75.0)).abs() < 1e-10);
698            assert!((r.max().y - 39.0).abs() < 1e-10);
699        } else {
700            panic!("expected Rect");
701        }
702    }
703
704    #[test]
705    fn parse_invalid_returns_none() {
706        assert!(parse_geojson(&json!({"type": "Unknown"})).is_none());
707        assert!(parse_geojson(&json!(42)).is_none());
708    }
709
710    #[test]
711    fn serialize_roundtrip_point() {
712        let geom = Geometry::Point(Point::new(-73.98, 40.75));
713        let bytes = serialize_shape(&geom);
714        let (decoded, consumed) = deserialize_shape(&bytes).unwrap();
715        assert_eq!(consumed, bytes.len());
716        assert_eq!(geom, decoded);
717    }
718
719    #[test]
720    fn serialize_roundtrip_polygon() {
721        let poly = Polygon::new(
722            LineString(vec![
723                Coord { x: 0.0, y: 0.0 },
724                Coord { x: 10.0, y: 0.0 },
725                Coord { x: 10.0, y: 10.0 },
726                Coord { x: 0.0, y: 10.0 },
727                Coord { x: 0.0, y: 0.0 },
728            ]),
729            vec![],
730        );
731        let geom = Geometry::Polygon(poly);
732        let bytes = serialize_shape(&geom);
733        let (decoded, _) = deserialize_shape(&bytes).unwrap();
734        assert_eq!(geom, decoded);
735    }
736
737    #[test]
738    fn serialize_roundtrip_rect() {
739        let geom = Geometry::Rect(Rect::new(
740            Coord { x: -77.0, y: 38.0 },
741            Coord { x: -75.0, y: 39.0 },
742        ));
743        let bytes = serialize_shape(&geom);
744        let (decoded, _) = deserialize_shape(&bytes).unwrap();
745        assert_eq!(geom, decoded);
746    }
747
748    #[test]
749    fn store_add_and_get() {
750        let mut store = GeoShapeStore::new();
751        let poly = Geometry::Polygon(Polygon::new(
752            LineString(vec![
753                Coord { x: 0.0, y: 0.0 },
754                Coord { x: 10.0, y: 0.0 },
755                Coord { x: 10.0, y: 10.0 },
756                Coord { x: 0.0, y: 10.0 },
757                Coord { x: 0.0, y: 0.0 },
758            ]),
759            vec![],
760        ));
761        store.add(&poly);
762        store.add_null();
763
764        assert_eq!(store.len(), 2);
765        assert!(store.get_bbox(0).is_some());
766        assert!(store.get_bbox(1).is_none());
767        assert_eq!(store.get_shape(0).unwrap(), poly);
768        assert!(store.get_shape(1).is_none());
769    }
770
771    #[test]
772    fn store_serialization_roundtrip() {
773        let mut store = GeoShapeStore::new();
774        store.add(&Geometry::Point(Point::new(1.0, 2.0)));
775        store.add(&Geometry::Polygon(Polygon::new(
776            LineString(vec![
777                Coord { x: 0.0, y: 0.0 },
778                Coord { x: 5.0, y: 0.0 },
779                Coord { x: 5.0, y: 5.0 },
780                Coord { x: 0.0, y: 0.0 },
781            ]),
782            vec![],
783        )));
784
785        let bytes = store.to_bytes();
786        let restored = GeoShapeStore::from_bytes(&bytes);
787
788        assert_eq!(restored.len(), 2);
789        assert_eq!(restored.get_shape(0).unwrap(), store.get_shape(0).unwrap());
790        assert_eq!(restored.get_shape(1).unwrap(), store.get_shape(1).unwrap());
791    }
792
793    #[test]
794    fn rtree_search() {
795        let mut store = GeoShapeStore::new();
796        // Doc 0: polygon at (0,0)-(10,10)
797        store.add(&Geometry::Polygon(Polygon::new(
798            LineString(vec![
799                Coord { x: 0.0, y: 0.0 },
800                Coord { x: 10.0, y: 0.0 },
801                Coord { x: 10.0, y: 10.0 },
802                Coord { x: 0.0, y: 10.0 },
803                Coord { x: 0.0, y: 0.0 },
804            ]),
805            vec![],
806        )));
807        // Doc 1: polygon at (20,20)-(30,30)
808        store.add(&Geometry::Polygon(Polygon::new(
809            LineString(vec![
810                Coord { x: 20.0, y: 20.0 },
811                Coord { x: 30.0, y: 20.0 },
812                Coord { x: 30.0, y: 30.0 },
813                Coord { x: 20.0, y: 30.0 },
814                Coord { x: 20.0, y: 20.0 },
815            ]),
816            vec![],
817        )));
818
819        let rtree_data = store.build_rtree();
820        assert!(!rtree_data.is_empty());
821
822        // Search overlapping doc 0 only
823        let hits =
824            GeoShapeStore::search_rtree(&rtree_data, (5.0, 5.0, 15.0, 15.0), &store.shape_offsets);
825        assert_eq!(hits, vec![0]);
826
827        // Search overlapping doc 1 only
828        let hits = GeoShapeStore::search_rtree(
829            &rtree_data,
830            (25.0, 25.0, 35.0, 35.0),
831            &store.shape_offsets,
832        );
833        assert_eq!(hits, vec![1]);
834
835        // Search overlapping both
836        let hits =
837            GeoShapeStore::search_rtree(&rtree_data, (0.0, 0.0, 30.0, 30.0), &store.shape_offsets);
838        assert!(hits.contains(&0) && hits.contains(&1));
839    }
840}
luci/spatial/shape.rs

luci/spatial/
shape.rs
