1use std::cell::RefCell;
2use std::ops::{Deref, DerefMut};
3
4use bitvec::vec::BitVec;
5
6use crate::decoder::Morton;
7use crate::encoder::model::{CurveParams, StagedLayer};
8use crate::encoder::{
9 Codecs, Encoder, EncoderConfig, SortStrategy, encode01, spatial_sort_likely_to_help,
10};
11#[cfg(feature = "unstable-v2")]
12use crate::encoder::{WireVersion, encode02};
13use crate::tile::{PropKind, TileLayer};
14use crate::{MltError, MltResult, PropValue};
15
16impl StagedLayer {
17 #[hotpath::measure]
24 pub fn encode_into(self, enc: Encoder, codecs: &mut Codecs) -> MltResult<Encoder> {
25 if self.name.is_empty() {
26 return Err(MltError::MissingLayerName);
27 }
28 #[cfg(feature = "unstable-v2")]
29 match enc.config().wire_version() {
30 WireVersion::V01 => encode01::encode_into01(self, enc, codecs),
31 WireVersion::V02 => encode02::encode_into02(self, enc, codecs),
32 }
33 #[cfg(not(feature = "unstable-v2"))]
34 encode01::encode_into01(self, enc, codecs)
35 }
36}
37
38const POOLED_CODECS_MAX_BYTES: usize = 16 << 20;
40
41thread_local! {
42 static CODECS_POOL: RefCell<Option<Codecs>> = const { RefCell::new(None) };
43}
44
45struct PooledCodecs(Option<Codecs>);
47
48impl PooledCodecs {
49 fn take() -> Self {
50 let pooled = CODECS_POOL.try_with(|pool| pool.borrow_mut().take());
51 Self(Some(pooled.ok().flatten().unwrap_or_default()))
52 }
53}
54
55impl Deref for PooledCodecs {
56 type Target = Codecs;
57
58 fn deref(&self) -> &Codecs {
59 self.0.as_ref().expect("codecs are only taken on drop")
60 }
61}
62
63impl DerefMut for PooledCodecs {
64 fn deref_mut(&mut self) -> &mut Codecs {
65 self.0.as_mut().expect("codecs are only taken on drop")
66 }
67}
68
69impl Drop for PooledCodecs {
70 fn drop(&mut self) {
71 if let Some(codecs) = self.0.take()
72 && codecs.scratch_bytes() <= POOLED_CODECS_MAX_BYTES
73 {
74 let _ = CODECS_POOL.try_with(|pool| *pool.borrow_mut() = Some(codecs));
75 }
76 }
77}
78
79fn seed_curve_caches(enc: &mut Encoder, curve_params: CurveParams) {
84 enc.hilbert_cache = Some(curve_params);
85 enc.morton_cache = Morton::new(curve_params.bits, curve_params.shift).ok();
86}
87
88const SORT_TRIAL_THRESHOLD: usize = 512;
91
92impl TileLayer {
93 #[hotpath::measure]
103 pub fn encode(mut self, cfg: EncoderConfig) -> MltResult<Vec<u8>> {
104 if self.name().is_empty() {
105 return Err(MltError::MissingLayerName);
106 }
107 if self.features().is_empty() {
108 return Ok(Vec::new());
109 }
110 if !cfg.allow_8bit_ints() {
111 self.widen_8bit_ints();
112 }
113
114 let mut sort_by = vec![SortStrategy::Unsorted];
115 let try_spatial_sort =
116 cfg.attempt_spatial_morton_sort() || cfg.attempt_spatial_hilbert_sort();
117 if try_spatial_sort
118 && (self.feature_count() < SORT_TRIAL_THRESHOLD || spatial_sort_likely_to_help(&self))
119 {
120 if cfg.attempt_spatial_morton_sort() {
121 sort_by.push(SortStrategy::SpatialMorton);
122 }
123 if cfg.attempt_spatial_hilbert_sort() {
124 sort_by.push(SortStrategy::SpatialHilbert);
125 }
126 }
127 if cfg.attempt_id_sort() {
128 sort_by.push(SortStrategy::Id);
129 }
130
131 let stats = self.analyze(cfg.allow_shared_dict())?;
132 let curve_params = self.curve_params();
135
136 let mut enc = Encoder::new(cfg);
140 seed_curve_caches(&mut enc, curve_params);
141
142 let (last, init) = sort_by.split_last().expect("at least one strategy");
143 let mut codecs = PooledCodecs::take();
144 if init.is_empty() {
145 StagedLayer::from_tile(self, *last, &stats, cfg.tessellate(), curve_params)
146 .encode_into(enc, &mut codecs)?
147 } else {
148 enc = {
149 let first = init[0];
150 StagedLayer::from_tile(self.clone(), first, &stats, cfg.tessellate(), curve_params)
151 .encode_into(enc, &mut codecs)?
152 };
153 let mut best = enc.preserve_results();
154 for &sort in &init[1..] {
156 let layer = StagedLayer::from_tile(
157 self.clone(),
158 sort,
159 &stats,
160 cfg.tessellate(),
161 curve_params,
162 );
163 enc = layer.encode_into(enc, &mut codecs)?;
164 if enc.total_len() < best.total_len() {
165 best = enc.preserve_results();
166 } else {
167 enc.clear_results();
169 }
170 }
171 let layer = StagedLayer::from_tile(self, *last, &stats, cfg.tessellate(), curve_params);
173 enc = layer.encode_into(enc, &mut codecs)?;
174 if enc.total_len() < best.total_len() {
175 best = enc.preserve_results();
176 }
177 best
178 }
179 .into_layer_bytes()
180 }
181}
182
183#[derive(Debug, Clone, Copy, PartialEq, Eq)]
185pub enum Presence {
186 AllNull,
188 AllPresent,
190 Mixed,
192 SameAsProp(usize),
199}
200impl Presence {
201 #[must_use]
203 pub fn from_bits(bits: &BitVec<u8>, existing: &[(BitVec<u8>, usize)]) -> Self {
204 if bits.not_any() {
205 Self::AllNull
206 } else if bits.all() {
207 Self::AllPresent
208 } else if let Some((_, idx)) = existing.iter().find(|(v, _)| v == bits) {
209 Self::SameAsProp(*idx)
210 } else {
211 Self::Mixed
212 }
213 }
214}
215
216#[derive(Debug, Clone, PartialEq, Eq)]
218pub enum SharedDictRole {
219 None,
221 Owner(String),
223 Member(usize),
225}
226
227#[derive(Debug, Clone, PartialEq, Eq)]
229pub struct PropertyStats {
230 pub presence: Presence,
231 pub stats: PropertyTypedStats,
232}
233
234#[derive(Debug, Clone, PartialEq, Eq)]
236pub struct LayerStats {
237 pub id: Option<PropertyStats>,
238 pub properties: Vec<PropertyStats>,
239}
240
241#[derive(Debug, Clone, Default, PartialEq, Eq)]
243pub enum PropertyTypedStats {
244 #[default]
246 None,
247 Bool,
248 Signed {
249 min: i64,
250 max: i64,
251 },
252 Unsigned {
253 min: u64,
254 max: u64,
255 },
256 F32,
257 F64,
258 String {
259 shared_dict: SharedDictRole,
260 },
261}
262
263impl PropertyTypedStats {
264 #[must_use]
265 pub fn values_fit_u32(&self) -> bool {
266 match self {
267 Self::None | Self::Bool | Self::F32 | Self::F64 | Self::String { .. } => false,
268 Self::Signed { min, max } => *min >= 0 && u32::try_from(*max).is_ok(),
269 Self::Unsigned { max, .. } => u32::try_from(*max).is_ok(),
270 }
271 }
272
273 #[must_use]
276 pub fn values_fit_i32(&self) -> bool {
277 match self {
278 Self::None | Self::Bool | Self::F32 | Self::F64 | Self::String { .. } => false,
279 Self::Signed { min, max } => i32::try_from(*min).is_ok() && i32::try_from(*max).is_ok(),
280 Self::Unsigned { max, .. } => i32::try_from(*max).is_ok(),
281 }
282 }
283
284 #[must_use]
285 pub fn shared_dict(&self) -> SharedDictRole {
286 match self {
287 Self::String { shared_dict, .. } => shared_dict.clone(),
288 Self::None
289 | Self::Bool
290 | Self::Signed { .. }
291 | Self::Unsigned { .. }
292 | Self::F32
293 | Self::F64 => SharedDictRole::None,
294 }
295 }
296
297 pub(crate) fn set_shared_dict(&mut self, role: SharedDictRole) {
298 match self {
299 Self::String { shared_dict, .. } => *shared_dict = role,
300 Self::None
301 | Self::Bool
302 | Self::Signed { .. }
303 | Self::Unsigned { .. }
304 | Self::F32
305 | Self::F64 => debug_assert_eq!(role, SharedDictRole::None),
306 }
307 }
308
309 pub(crate) fn push(
310 &mut self,
311 prop: &PropValue,
312 column_idx: usize,
313 property_name: &str,
314 ) -> MltResult<bool> {
315 match prop {
316 PropValue::Bool(Some(_)) => {
317 self.merge_same_kind(Self::Bool, column_idx, property_name)?;
318 }
319 PropValue::I8(Some(v)) => {
320 self.merge_signed(i64::from(*v), column_idx, property_name)?;
321 }
322 PropValue::U8(Some(v)) => {
323 self.merge_unsigned(u64::from(*v), column_idx, property_name)?;
324 }
325 PropValue::I32(Some(v)) => {
326 self.merge_signed(i64::from(*v), column_idx, property_name)?;
327 }
328 PropValue::U32(Some(v)) => {
329 self.merge_unsigned(u64::from(*v), column_idx, property_name)?;
330 }
331 PropValue::I64(Some(v)) => self.merge_signed(*v, column_idx, property_name)?,
332 PropValue::U64(Some(v)) => self.merge_unsigned(*v, column_idx, property_name)?,
333 PropValue::F32(Some(_)) => {
334 self.merge_same_kind(Self::F32, column_idx, property_name)?;
335 }
336 PropValue::F64(Some(_)) => {
337 self.merge_same_kind(Self::F64, column_idx, property_name)?;
338 }
339 PropValue::Str(Some(_)) => self.merge_string(column_idx, property_name)?,
340 PropValue::Bool(None)
341 | PropValue::I8(None)
342 | PropValue::U8(None)
343 | PropValue::I32(None)
344 | PropValue::U32(None)
345 | PropValue::I64(None)
346 | PropValue::U64(None)
347 | PropValue::F32(None)
348 | PropValue::F64(None)
349 | PropValue::Str(None) => return Ok(false),
350 }
351 Ok(true)
352 }
353
354 fn merge_signed(
355 &mut self,
356 value: i64,
357 column_idx: usize,
358 property_name: &str,
359 ) -> MltResult<()> {
360 match self {
361 Self::None => {
362 *self = Self::Signed {
363 min: value,
364 max: value,
365 };
366 }
367 Self::Signed { min, max } => {
368 *min = (*min).min(value);
369 *max = (*max).max(value);
370 }
371 Self::Bool | Self::Unsigned { .. } | Self::F32 | Self::F64 | Self::String { .. } => {
372 return mixed_prop_err(column_idx, property_name);
373 }
374 }
375 Ok(())
376 }
377
378 fn merge_unsigned(
379 &mut self,
380 value: u64,
381 column_idx: usize,
382 property_name: &str,
383 ) -> MltResult<()> {
384 match self {
385 Self::None => {
386 *self = Self::Unsigned {
387 min: value,
388 max: value,
389 };
390 }
391 Self::Unsigned { min, max } => {
392 *min = (*min).min(value);
393 *max = (*max).max(value);
394 }
395 Self::Bool | Self::Signed { .. } | Self::F32 | Self::F64 | Self::String { .. } => {
396 return mixed_prop_err(column_idx, property_name);
397 }
398 }
399 Ok(())
400 }
401
402 fn merge_string(&mut self, column_idx: usize, property_name: &str) -> MltResult<()> {
403 match self {
404 Self::None => {
405 *self = Self::String {
406 shared_dict: SharedDictRole::None,
407 };
408 }
409 Self::String { .. } => {}
410 Self::Bool | Self::Signed { .. } | Self::Unsigned { .. } | Self::F32 | Self::F64 => {
411 return mixed_prop_err(column_idx, property_name);
412 }
413 }
414 Ok(())
415 }
416
417 fn merge_same_kind(
418 &mut self,
419 kind: Self,
420 column_idx: usize,
421 property_name: &str,
422 ) -> MltResult<()> {
423 match self {
424 Self::None => *self = kind,
425 Self::Bool if matches!(kind, Self::Bool) => {}
426 Self::F32 if matches!(kind, Self::F32) => {}
427 Self::F64 if matches!(kind, Self::F64) => {}
428 Self::Bool
429 | Self::Signed { .. }
430 | Self::Unsigned { .. }
431 | Self::F32
432 | Self::F64
433 | Self::String { .. } => return mixed_prop_err(column_idx, property_name),
434 }
435 Ok(())
436 }
437}
438
439impl TileLayer {
440 #[hotpath::measure]
442 pub(crate) fn analyze(&self, allow_shared_dict: bool) -> MltResult<LayerStats> {
443 let mut property_bits = Vec::with_capacity(self.property_names().len());
444 let mut properties = self.analyze_properties(&mut property_bits)?;
445 let id = self.analyze_ids(&property_bits);
446 if allow_shared_dict {
447 self.group_string_properties(&mut properties);
448 }
449 Ok(LayerStats { id, properties })
450 }
451
452 fn analyze_ids(&self, property_bits: &[(BitVec<u8>, usize)]) -> Option<PropertyStats> {
453 let mut min = u64::MAX;
454 let mut max = 0u64;
455 let mut bits = BitVec::<u8>::with_capacity(self.feature_count());
456 for feature in self.features() {
457 if let Some(id) = feature.id() {
458 min = min.min(id);
459 max = max.max(id);
460 bits.push(true);
461 } else {
462 bits.push(false);
463 }
464 }
465 let presence = Presence::from_bits(&bits, property_bits);
466 if presence == Presence::AllNull {
467 None
468 } else {
469 Some(PropertyStats {
470 presence,
471 stats: PropertyTypedStats::Unsigned { min, max },
472 })
473 }
474 }
475
476 fn analyze_properties(
477 &self,
478 property_bits: &mut Vec<(BitVec<u8>, usize)>,
479 ) -> MltResult<Vec<PropertyStats>> {
480 self.property_names()
481 .iter()
482 .enumerate()
483 .map(|(col_idx, name)| -> MltResult<PropertyStats> {
484 let mut kind = None;
485 let mut stats = PropertyTypedStats::default();
486 let mut bits = BitVec::<u8>::with_capacity(self.feature_count());
487 for feature in self.features() {
488 let prop = feature.properties().get(col_idx);
489 if let Some(prop_kind) = prop.map(PropKind::from) {
490 match kind {
491 Some(kind) if kind != prop_kind => {
492 return mixed_prop_err(col_idx, name.as_str());
493 }
494 None => kind = Some(prop_kind),
495 _ => {}
496 }
497 }
498 if let Some(prop) = prop
499 && stats.push(prop, col_idx, name)?
500 {
501 bits.push(true);
502 } else {
503 bits.push(false);
504 }
505 }
506
507 let presence = Presence::from_bits(&bits, property_bits);
508 if presence == Presence::Mixed {
509 property_bits.push((bits, col_idx));
510 }
511 Ok(PropertyStats { presence, stats })
512 })
513 .collect()
514 }
515}
516
517#[inline]
518fn mixed_prop_err<T>(column_idx: usize, property_name: &str) -> MltResult<T> {
519 Err(MltError::MixedPropertyTypes(
520 column_idx,
521 property_name.to_owned(),
522 ))
523}