1use std::collections::HashMap;
2use std::mem;
3
4use geo_types::Coord;
5use probabilistic_collections::SipHasherBuilder;
6use probabilistic_collections::hyperloglog::HyperLogLog;
7use usize_cast::{FromUsize as _, IntoUsize as _};
8
9use super::model::VertexBufferType;
10use crate::MltResult;
11use crate::codecs::hilbert::hilbert_sort_key;
12use crate::codecs::zigzag::encode_componentwise_delta_vec2s;
13use crate::decoder::GeometryType::{LineString, Point, Polygon};
14use crate::decoder::{
15 ColumnType, DictionaryType, GeometryType, GeometryValues, LengthType, LogicalEncoding, Morton,
16 OffsetType, PhysicalEncoding, StreamMeta, StreamType,
17};
18use crate::encoder::model::{CurveParams, StreamCtx};
19use crate::encoder::{Codecs, Encoder, PhysicalCodecs, write_stream_payload};
20
21#[hotpath::measure]
29fn build_morton_dict(vertices: &[i32], meta: Morton) -> MltResult<(Vec<u32>, Vec<u32>)> {
30 let codes: Vec<u32> = vertices
31 .chunks_exact(2)
32 .map(|c| meta.encode_morton(c[0], c[1]))
33 .collect::<Result<_, _>>()?;
34
35 let mut dict = codes.clone();
36 dict.sort_unstable();
37 dict.dedup();
38
39 #[expect(
40 clippy::cast_possible_truncation,
41 reason = "dict.len() <= u32::MAX (deduped u32 codes)"
42 )]
43 let code_to_idx: HashMap<u32, u32> = dict
44 .iter()
45 .enumerate()
46 .map(|(i, &c)| (c, i as u32))
47 .collect();
48 let offsets: Vec<u32> = codes.iter().map(|code| code_to_idx[code]).collect();
49
50 Ok((dict, offsets))
51}
52
53#[hotpath::measure]
64fn build_hilbert_dict(
65 vertices: &[i32],
66 params: CurveParams,
67 offsets: &mut Vec<u32>,
68 indexed: &mut Vec<u64>,
69 dict_xy: &mut Vec<i32>,
70 remap: &mut HashMap<u32, u32>,
71) {
72 offsets.clear();
73 indexed.clear();
74 dict_xy.clear();
75 remap.clear();
76
77 let coord_count = vertices.len() / 2;
78 if coord_count == 0 {
79 return;
80 }
81 offsets.reserve(coord_count);
82 indexed.reserve(coord_count);
83 dict_xy.reserve(coord_count * 2);
84 remap.reserve(coord_count);
85
86 for (i, c) in vertices.chunks_exact(2).enumerate() {
87 let k = hilbert_sort_key(Coord { x: c[0], y: c[1] }, params);
88 offsets.push(k);
89 let packed = (u64::from(k) << 32) | u64::from_usize(i);
92 indexed.push(packed);
93 }
94 indexed.sort_unstable();
95
96 let mut last_key: Option<u32> = None;
97 for &packed in &*indexed {
98 let key = (packed >> 32) as u32;
99 let src_idx = ((packed & 0xFFFF_FFFF) as u32).into_usize();
100 if last_key != Some(key) {
101 #[expect(
102 clippy::cast_possible_truncation,
103 reason = "dict.len() <= coord_count <= u32::MAX"
104 )]
105 let slot = (dict_xy.len() / 2) as u32;
106 dict_xy.push(vertices[src_idx * 2]);
107 dict_xy.push(vertices[src_idx * 2 + 1]);
108 remap.insert(key, slot);
109 last_key = Some(key);
110 }
111 }
112
113 for k in offsets.iter_mut() {
114 *k = remap[k];
115 }
116}
117
118#[inline]
123fn extend_offsets(lengths: &mut Vec<u32>, offsets: &[u32]) -> usize {
124 lengths.extend(offsets.windows(2).map(|w| w[1] - w[0]));
125 offsets.len() - 1
126}
127
128fn encode_root_length_stream(
138 geom_types: &[GeometryType],
139 geom_offsets: &[u32],
140 buffer_id: GeometryType,
141) -> Vec<u32> {
142 if geom_offsets.len() == geom_types.len() + 1 {
143 geom_types
145 .iter()
146 .zip(geom_offsets.windows(2))
147 .filter(|&(&t, _)| t > buffer_id)
148 .map(|(_, w)| w[1] - w[0])
149 .collect()
150 } else {
151 geom_types
153 .iter()
154 .filter(|&&t| t > buffer_id)
155 .zip(geom_offsets.windows(2))
156 .map(|(_, w)| w[1] - w[0])
157 .collect()
158 }
159}
160
161fn encode_level1_length_stream(
163 geom_types: &[GeometryType],
164 geom_offsets: &[u32],
165 part_offsets: &[u32],
166 is_line_string_present: bool,
167) -> Vec<u32> {
168 let mut lengths = Vec::new();
169 let mut part_idx = 0;
170
171 for (i, &geom_type) in geom_types.iter().enumerate() {
172 if geom_type.is_polygon() || (is_line_string_present && geom_type.is_linestring()) {
173 let n = (geom_offsets[i + 1] - geom_offsets[i]).into_usize();
174 part_idx += extend_offsets(&mut lengths, &part_offsets[part_idx..=part_idx + n]);
175 }
176 }
179
180 lengths
181}
182
183fn encode_ring_lengths_for_mixed(
191 geom_types: &[GeometryType],
192 part_offsets: &[u32],
193 ring_offsets: &[u32],
194 has_line_string: bool,
195) -> Vec<u32> {
196 let mut lengths = Vec::new();
197 for (i, &geom_type) in geom_types.iter().enumerate() {
198 if geom_type.is_polygon() || (has_line_string && geom_type.is_linestring()) {
199 let s = part_offsets[i].into_usize();
200 let e = part_offsets[i + 1].into_usize();
201 extend_offsets(&mut lengths, &ring_offsets[s..=e]);
202 }
203 }
204 lengths
205}
206
207fn encode_level2_length_stream(
213 geom_types: &[GeometryType],
214 geom_offsets: &[u32],
215 part_offsets: &[u32],
216 ring_offsets: &[u32],
217) -> Vec<u32> {
218 let mut lengths = Vec::new();
219 let mut part_idx = 0;
220 let mut ring_idx = 0;
221
222 for (i, &geom_type) in geom_types.iter().enumerate() {
223 let count = (geom_offsets[i + 1] - geom_offsets[i]).into_usize();
224
225 if geom_type.is_polygon() {
229 for _ in 0..count {
231 let n = (part_offsets[part_idx + 1] - part_offsets[part_idx]).into_usize();
232 ring_idx += extend_offsets(&mut lengths, &ring_offsets[ring_idx..=ring_idx + n]);
233 part_idx += 1;
234 }
235 } else if geom_type.is_linestring() {
236 ring_idx += extend_offsets(&mut lengths, &ring_offsets[ring_idx..=ring_idx + count]);
238 }
239 }
241
242 lengths
243}
244
245fn encode_level1_without_ring_buffer_length_stream(
252 geom_types: &[GeometryType],
253 geom_offsets: &[u32],
254 part_offsets: &[u32],
255) -> Vec<u32> {
256 let mut lengths = Vec::new();
257 let mut part_idx = 0;
258
259 for (i, &geom_type) in geom_types.iter().enumerate() {
260 if geom_type.is_linestring() {
261 let n = (geom_offsets[i + 1] - geom_offsets[i]).into_usize();
262 part_idx += extend_offsets(&mut lengths, &part_offsets[part_idx..=part_idx + n]);
263 }
264 }
266
267 lengths
268}
269
270fn normalize_geometry_offsets(vector_types: &[GeometryType], geom_offsets: &[u32]) -> Vec<u32> {
272 let mut normalized = Vec::with_capacity(vector_types.len() + 1);
273 let mut offset = 0_u32;
274 let mut sparse_idx = 0_usize; for &geom_type in vector_types {
277 normalized.push(offset);
278
279 if geom_type.is_multi() {
280 if sparse_idx + 1 < geom_offsets.len() {
282 let start = geom_offsets[sparse_idx];
283 let end = geom_offsets[sparse_idx + 1];
284 offset += end - start;
285 sparse_idx += 1;
286 }
287 } else {
288 offset += 1;
290 }
291 }
292
293 normalized.push(offset);
294 normalized
295}
296
297fn normalize_part_offsets_for_rings(
309 vector_types: &[GeometryType],
310 part_offsets: &[u32],
311 ring_offsets: &[u32],
312) -> Vec<u32> {
313 let mut normalized = Vec::with_capacity(vector_types.len() + 1);
314 let mut ring_idx = 0_u32;
315 let mut part_idx = 0_usize;
316
317 for &geom_type in vector_types {
318 normalized.push(ring_idx);
319
320 if geom_type == Point {
321 } else if geom_type.is_linestring() {
323 ring_idx += 1;
325 } else if geom_type.is_polygon() && part_idx + 1 < part_offsets.len() {
326 let ring_count = part_offsets[part_idx + 1] - part_offsets[part_idx];
328 ring_idx += ring_count;
329 part_idx += 1;
330 }
331 }
333
334 debug_assert_eq!(
336 ring_idx.into_usize(),
337 ring_offsets.len().saturating_sub(1),
338 "ring index mismatch after normalization"
339 );
340 normalized.push(ring_idx);
341 normalized
342}
343
344#[hotpath::measure]
353fn dict_may_be_beneficial(vertices: &[i32], enc: &Encoder) -> bool {
354 const MAXIMUM_UNIQUENESS_THRESHOLD_FOR_DICT: f64 = 0.66;
355
356 let coord_count = vertices.len() / 2;
357 if coord_count == 0 || enc.morton_cache.is_none() {
358 return false;
359 }
360
361 let mut hll = HyperLogLog::<Coord<i32>>::with_hasher(0.03, SipHasherBuilder::from_seed(0, 0));
362 for c in vertices.chunks_exact(2) {
363 hll.insert(&Coord::<i32> { x: c[0], y: c[1] });
364 }
365 #[expect(clippy::cast_precision_loss)]
366 let estimated_unique = hll.len().clamp(0.0, coord_count as f64);
367 #[expect(clippy::cast_precision_loss)]
368 let uniqueness_ratio = estimated_unique / coord_count as f64;
369 uniqueness_ratio < MAXIMUM_UNIQUENESS_THRESHOLD_FOR_DICT
370}
371
372fn get_morton(enc: &Encoder) -> Morton {
376 enc.morton_cache.expect(
377 "morton_cache populated by StagedLayer::encode_into; gated by dict_may_be_beneficial",
378 )
379}
380
381fn get_hilbert_params(enc: &Encoder) -> CurveParams {
383 enc.hilbert_cache
384 .expect("hilbert_cache populated by StagedLayer::encode_into")
385}
386
387fn encode_vec2_vertex_stream(
390 vertices: &[i32],
391 enc: &mut Encoder,
392 codecs: &mut Codecs,
393) -> MltResult<u8> {
394 let delta = encode_componentwise_delta_vec2s(vertices, &mut codecs.logical.u32_tmp);
395 let ctx = StreamCtx::geom(StreamType::Data(DictionaryType::Vertex), "vertex");
396 let logical = LogicalEncoding::ComponentwiseDelta;
397 write_geo_precomputed_stream(delta, ctx, logical, enc, &mut codecs.physical)
398}
399
400fn encode_morton_vertex_streams(
403 vertices: &[i32],
404 enc: &mut Encoder,
405 codecs: &mut Codecs,
406) -> MltResult<u8> {
407 let morton = get_morton(enc);
408 let (dict, offsets) = build_morton_dict(vertices, morton)?;
409 let mut n: u8 = 0;
410
411 let ctx = StreamCtx::geom(StreamType::Offset(OffsetType::Vertex), "vertex_offsets");
412 n += write_geo_u32_stream(&offsets, ctx, enc, codecs)?;
413
414 let delta = encode_morton_deltas(&dict, &mut codecs.logical.u32_tmp);
415 let ctx = StreamCtx::geom(StreamType::Data(DictionaryType::Morton), "vertex");
416 let logical = LogicalEncoding::MortonDelta(morton);
417 n += write_geo_precomputed_stream(delta, ctx, logical, enc, &mut codecs.physical)?;
418 Ok(n)
419}
420
421fn encode_hilbert_vertex_streams(
425 vertices: &[i32],
426 enc: &mut Encoder,
427 codecs: &mut Codecs,
428) -> MltResult<u8> {
429 let params = get_hilbert_params(enc);
430 let mut n: u8 = 0;
431
432 let mut offsets = mem::take(&mut codecs.logical.hilbert_offsets);
435 let mut indexed = mem::take(&mut codecs.logical.hilbert_indexed);
436 let mut dict_xy = mem::take(&mut codecs.logical.hilbert_dict_xy);
437 let mut remap = mem::take(&mut codecs.logical.hilbert_remap);
438
439 build_hilbert_dict(
440 vertices,
441 params,
442 &mut offsets,
443 &mut indexed,
444 &mut dict_xy,
445 &mut remap,
446 );
447 codecs.logical.hilbert_indexed = indexed;
450 codecs.logical.hilbert_remap = remap;
451
452 let ctx = StreamCtx::geom(StreamType::Offset(OffsetType::Vertex), "vertex_offsets");
453 n += write_geo_u32_stream(&offsets, ctx, enc, codecs)?;
454
455 encode_componentwise_delta_vec2s(&dict_xy, &mut offsets);
458 let ctx = StreamCtx::geom(StreamType::Data(DictionaryType::Vertex), "vertex");
459 let logical = LogicalEncoding::ComponentwiseDelta;
460 n += write_geo_precomputed_stream(&offsets, ctx, logical, enc, &mut codecs.physical)?;
461
462 codecs.logical.hilbert_offsets = offsets;
463 codecs.logical.hilbert_dict_xy = dict_xy;
464 Ok(n)
465}
466
467fn write_geo_u32_stream(
473 data: &[u32],
474 ctx: StreamCtx,
475 enc: &mut Encoder,
476 codecs: &mut Codecs,
477) -> MltResult<u8> {
478 Ok(if data.is_empty() && !enc.force_stream(&ctx) {
479 0
480 } else {
481 codecs.write_int_stream(data, &ctx, enc)?;
482 1
483 })
484}
485
486fn write_geo_precomputed_stream(
491 data: &[u32],
492 ctx: StreamCtx,
493 logical: LogicalEncoding,
494 enc: &mut Encoder,
495 physical: &mut PhysicalCodecs,
496) -> MltResult<u8> {
497 use PhysicalEncoding as PE;
498
499 Ok(if data.is_empty() && !enc.force_stream(&ctx) {
500 0
501 } else {
502 if let Some(int_enc) = enc.override_int_enc(&ctx) {
503 physical.write_encoded_as::<[u32]>(&ctx, enc, logical, data, int_enc.physical)?;
504 } else if data.is_empty() {
505 let meta = StreamMeta::new2(ctx.stream_type, logical, PE::None, 0)?;
506 write_stream_payload(enc.data_mut(), meta, false, &[])?;
507 } else {
508 let allow_fastpfor = enc.config().allow_fastpfor();
509 let mut alt = enc.try_alternatives();
510 if allow_fastpfor {
511 alt.with(|enc| {
512 let vals = physical.fastpfor(data)?;
513 let meta =
514 StreamMeta::new2(ctx.stream_type, logical, PE::FastPFor256, data.len())?;
515 write_stream_payload(enc.data_mut(), meta, false, vals)
516 })?;
517 }
518 alt.with(|enc| {
519 let vals = physical.varint(data);
520 let meta = StreamMeta::new2(ctx.stream_type, logical, PE::VarInt, data.len())?;
521 write_stream_payload(enc.data_mut(), meta, false, vals)
522 })?;
523 }
524 1
525 })
526}
527
528impl GeometryValues {
529 #[hotpath::measure]
531 pub fn write_to(self, enc: &mut Encoder, codecs: &mut Codecs) -> MltResult<()> {
532 let Self {
533 vector_types,
534 geometry_offsets,
535 part_offsets,
536 ring_offsets,
537 index_buffer,
538 triangles,
539 vertices,
540 } = self;
541
542 let geom_offsets = geometry_offsets.unwrap_or_default();
547 let part_offsets = part_offsets.unwrap_or_default();
548 let ring_offsets = ring_offsets.unwrap_or_default();
549 let index_buffer = index_buffer.unwrap_or_default();
550 let triangles = triangles.unwrap_or_default();
551 let vertices = vertices.unwrap_or_default();
552
553 if enc.hilbert_cache.is_none() {
557 enc.hilbert_cache = Some(CurveParams::from_vertices(&vertices));
558 }
559 if enc.morton_cache.is_none() {
560 let p = enc.hilbert_cache.expect("populated above");
561 enc.morton_cache = Morton::new(p.bits, p.shift).ok();
562 }
563
564 let meta: Vec<u32> = vector_types.iter().map(|t| *t as u32).collect();
565
566 let part_offsets = if geom_offsets.is_empty()
567 && !ring_offsets.is_empty()
568 && !part_offsets.is_empty()
569 && part_offsets.len() != vector_types.len() + 1
570 {
571 normalize_part_offsets_for_rings(&vector_types, &part_offsets, &ring_offsets)
573 } else {
574 part_offsets
575 };
576
577 enc.write_column_type(ColumnType::Geometry)?;
580 let stream_count_pos = enc.data().len();
581 enc.data_mut().push(0); let mut n: u8 = 0;
583
584 let ctx = StreamCtx::geom(StreamType::Length(LengthType::VarBinary), "meta");
586 codecs.write_int_stream(&meta, &ctx, enc)?;
587 n += 1;
588
589 if !geom_offsets.is_empty() {
591 let geom_offsets = if geom_offsets.len() == vector_types.len() + 1 {
592 geom_offsets
593 } else {
594 normalize_geometry_offsets(&vector_types, &geom_offsets)
595 };
596 let data = encode_root_length_stream(&vector_types, &geom_offsets, Polygon);
597 let ctx = StreamCtx::geom(StreamType::Length(LengthType::Geometries), "geometries");
598 n += write_geo_u32_stream(&data, ctx, enc, codecs)?;
599
600 if !part_offsets.is_empty() {
606 if ring_offsets.is_empty() {
607 let data = encode_level1_without_ring_buffer_length_stream(
609 &vector_types,
610 &geom_offsets,
611 &part_offsets,
612 );
613 let ctx = StreamCtx::geom(StreamType::Length(LengthType::Parts), "no_rings");
614 n += write_geo_u32_stream(&data, ctx, enc, codecs)?;
615 } else {
616 let data = encode_level1_length_stream(
619 &vector_types,
620 &geom_offsets,
621 &part_offsets,
622 false,
623 );
624 let ctx = StreamCtx::geom(StreamType::Length(LengthType::Parts), "rings");
625 n += write_geo_u32_stream(&data, ctx, enc, codecs)?;
626
627 let data = encode_level2_length_stream(
628 &vector_types,
629 &geom_offsets,
630 &part_offsets,
631 &ring_offsets,
632 );
633 let ctx = StreamCtx::geom(StreamType::Length(LengthType::Rings), "rings2");
634 n += write_geo_u32_stream(&data, ctx, enc, codecs)?;
635 }
636 }
637 } else if !part_offsets.is_empty() {
638 if ring_offsets.is_empty() {
639 let data = encode_root_length_stream(&vector_types, &part_offsets, Point);
640 let ctx = StreamCtx::geom(StreamType::Length(LengthType::Parts), "no_rings");
641 n += write_geo_u32_stream(&data, ctx, enc, codecs)?;
642 } else {
643 let ctx = StreamCtx::geom(StreamType::Length(LengthType::Geometries), "geometries");
647 n += write_geo_u32_stream(&[], ctx, enc, codecs)?;
648
649 let data = encode_root_length_stream(&vector_types, &part_offsets, LineString);
650 let ctx = StreamCtx::geom(StreamType::Length(LengthType::Parts), "parts");
651 n += write_geo_u32_stream(&data, ctx, enc, codecs)?;
652
653 let has_line_string = vector_types
657 .iter()
658 .copied()
659 .any(GeometryType::is_linestring);
660 let data = encode_ring_lengths_for_mixed(
661 &vector_types,
662 &part_offsets,
663 &ring_offsets,
664 has_line_string,
665 );
666 let ctx = StreamCtx::geom(StreamType::Length(LengthType::Rings), "parts_ring");
667 n += write_geo_u32_stream(&data, ctx, enc, codecs)?;
668 }
669 }
670
671 let ctx = StreamCtx::geom(StreamType::Length(LengthType::Triangles), "triangles");
672 n += write_geo_u32_stream(&triangles, ctx, enc, codecs)?;
673 let ctx = StreamCtx::geom(StreamType::Offset(OffsetType::Index), "triangles_indexes");
674 n += write_geo_u32_stream(&index_buffer, ctx, enc, codecs)?;
675
676 if let Some(forced) = enc.override_vertex_buffer_type() {
677 n += match forced {
678 VertexBufferType::Vec2 => encode_vec2_vertex_stream(&vertices, enc, codecs)?,
679 VertexBufferType::Morton => encode_morton_vertex_streams(&vertices, enc, codecs)?,
680 VertexBufferType::Hilbert => encode_hilbert_vertex_streams(&vertices, enc, codecs)?,
681 };
682 } else if dict_may_be_beneficial(&vertices, enc) {
683 let mut winner_size: usize = usize::MAX;
685 let mut winner_stream_cnt: u8 = 0;
686 let mut alt = enc.try_alternatives();
687 alt.with(|e| {
688 let ds = e.data().len();
689 let ms = e.meta().len();
690 winner_stream_cnt = encode_vec2_vertex_stream(&vertices, e, codecs)?;
691 winner_size = (e.data().len() - ds) + (e.meta().len() - ms);
692 Ok(())
693 })?;
694 alt.with(|e| {
695 let ds = e.data().len();
696 let ms = e.meta().len();
697 let cnt = encode_hilbert_vertex_streams(&vertices, e, codecs)?;
698 let size = (e.data().len() - ds) + (e.meta().len() - ms);
699 if size < winner_size {
700 winner_stream_cnt = cnt;
701 winner_size = size;
702 }
703 Ok(())
704 })?;
705 alt.with(|e| {
706 let ds = e.data().len();
707 let ms = e.meta().len();
708 let cnt = encode_morton_vertex_streams(&vertices, e, codecs)?;
709 let size = (e.data().len() - ds) + (e.meta().len() - ms);
710 if size < winner_size {
711 winner_stream_cnt = cnt;
712 }
713 Ok(())
714 })?;
715 drop(alt);
716 n += winner_stream_cnt;
717 } else {
718 n += encode_vec2_vertex_stream(&vertices, enc, codecs)?;
719 }
720
721 debug_assert!(n <= 127, "geometry stream count must fit in one byte");
723 enc.data_mut()[stream_count_pos] = n;
724 Ok(())
725 }
726}
727
728fn encode_morton_deltas<'a>(codes: &[u32], buffer: &'a mut Vec<u32>) -> &'a mut Vec<u32> {
729 buffer.clear();
730 if let Some(&first) = codes.first() {
731 buffer.reserve(codes.len());
732 buffer.extend(std::iter::once(first).chain(codes.windows(2).map(|w| w[1] - w[0])));
733 }
734 buffer
735}
736
737#[cfg(test)]
738mod tests {
739 use super::*;
740
741 #[test]
742 fn test_build_morton_dict() {
743 let meta = Morton { bits: 4, shift: 0 };
744 let vertices = [1, 2, 3, 4, 1, 2, 0, 0];
746 let (dict, offsets) = build_morton_dict(&vertices, meta).unwrap();
747
748 assert!(
749 dict.windows(2).all(|w| w[0] < w[1]),
750 "dict not sorted/unique"
751 );
752 assert_eq!(offsets.len(), 4, "offsets length == number of vertex pairs");
753 assert_eq!(offsets[0], offsets[2], "duplicate (1,2) should share index");
754 assert!(offsets.iter().all(|&o| o.into_usize() < dict.len()));
755 }
756
757 #[test]
758 fn test_encode_root_length_stream() {
759 let types = vec![Polygon];
761 let offsets = vec![0, 1]; let lengths = encode_root_length_stream(&types, &offsets, Polygon);
764 assert!(lengths.is_empty());
766
767 let types = vec![GeometryType::MultiPolygon];
769 let offsets = vec![0, 2]; let lengths = encode_root_length_stream(&types, &offsets, Polygon);
772 assert_eq!(lengths, vec![2]);
773 }
774}