Skip to main content

mlt_core/encoder/
optimizer.rs

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    /// Encode and serialize the layer directly into `enc`, without creating any
18    /// intermediate representation.
19    ///
20    /// This is the hot path inside `TileLayer::encode`: each sort-strategy
21    /// trial calls this method on its own fresh `Encoder`, and only the
22    /// `Encoder` with the smallest `total_len()` is kept.
23    #[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
38/// Scratch capacity above which a [`Codecs`] is dropped instead of kept for the next layer.
39const POOLED_CODECS_MAX_BYTES: usize = 16 << 20;
40
41thread_local! {
42    static CODECS_POOL: RefCell<Option<Codecs>> = const { RefCell::new(None) };
43}
44
45/// A [`Codecs`] borrowed from this thread's pool, which keeps scratch buffers warm between layers.
46struct 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
79/// Seed the encoder's curve-derived caches so the Hilbert/Morton dictionary
80/// builders skip their min/max scan. `Morton::new` returns `Err` when bits > 16;
81/// `dict_may_be_beneficial` reads `morton_cache.is_none()` and falls back to
82/// a Vec2-only path in that case.
83fn 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
88/// Feature-count threshold above which the spatial trial is subject to the
89/// bounding-box pruning heuristic.
90const SORT_TRIAL_THRESHOLD: usize = 512;
91
92impl TileLayer {
93    /// Encode a [`TileLayer`] to bytes, automatically optimizing all encoding choices.
94    ///
95    /// This is the primary encoding entry point. It:
96    /// 1. Determines which sort strategies to try based on `cfg`
97    /// 2. Tries each sort strategy, encoding and measuring the output size
98    /// 3. Returns the smallest encoding as a complete layer record (including tag and length prefix)
99    ///
100    /// All encoding choices - sort order, per-stream integer encodings, string compression,
101    /// vertex buffer layout - are selected automatically to minimize output size.
102    #[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        // Bounds are order-invariant, so this scan is shared across every
133        // sort trial and the encoder's Hilbert/Morton dictionary builders.
134        let curve_params = self.curve_params();
135
136        // `Encoder::preserve_results` clears caches only on the moved-out
137        // archive, so a single seeding here serves every trial that reuses
138        // `enc`.
139        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            // Clone for all-but-last strategies
155            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                    // Drop the losing trial's bytes, or the next trial would append to them and overcount its total_len.
168                    enc.clear_results();
169                }
170            }
171            // Last strategy: consume self, no clone
172            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/// Row-order-independent presence classification for IDs and properties.
184#[derive(Debug, Clone, Copy, PartialEq, Eq)]
185pub enum Presence {
186    /// No feature has a value for this logical column.
187    AllNull,
188    /// Every feature has a value for this logical column.
189    AllPresent,
190    /// Some, but not all, features have a value for this logical column.
191    Mixed,
192    /// Mixed presence with the same per-feature mask as an earlier property column.
193    ///
194    /// Only tells the stager the column is nullable, same as [`Self::Mixed`]. Which
195    /// columns actually end up sharing one presence bitfield is decided per wire
196    /// format at write time, from the staged masks - see `SharedPresence` in
197    /// `encoder::encode02`.
198    SameAsProp(usize),
199}
200impl Presence {
201    /// Create presence value
202    #[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/// How a property participates in a shared dictionary group.
217#[derive(Debug, Clone, PartialEq, Eq)]
218pub enum SharedDictRole {
219    /// The property is encoded as a standalone column.
220    None,
221    /// The property is the first column in this group and emits the shared dictionary with this prefix.
222    Owner(String),
223    /// The property is emitted by the group owner at this property index.
224    Member(usize),
225}
226
227/// Row-order-independent facts for a single property column.
228#[derive(Debug, Clone, PartialEq, Eq)]
229pub struct PropertyStats {
230    pub presence: Presence,
231    pub stats: PropertyTypedStats,
232}
233
234/// Row-order-independent layer facts computed once before sort trials.
235#[derive(Debug, Clone, PartialEq, Eq)]
236pub struct LayerStats {
237    pub id: Option<PropertyStats>,
238    pub properties: Vec<PropertyStats>,
239}
240
241/// Row-order-independent value statistics for a property column.
242#[derive(Debug, Clone, Default, PartialEq, Eq)]
243pub enum PropertyTypedStats {
244    /// No present values.
245    #[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    /// Returns `true` if every value fits in an `i32`.
274    /// Unlike [`Self::values_fit_u32`] this admits negative values.
275    #[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    /// Analyze a [`TileLayer`] and return reusable ID/property facts for the optimizer.
441    #[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}