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::property::group_string_properties;
9use crate::encoder::sort::{curve_params, sort_order};
10use crate::encoder::source::LayerSource;
11use crate::encoder::tile::stage_layer;
12use crate::encoder::{Codecs, Encoder, EncoderConfig, SortStrategy, encode01};
13#[cfg(feature = "unstable-v2")]
14use crate::encoder::{WireVersion, encode02};
15use crate::tile::TileLayer;
16use crate::{MltError, MltResult, PropValueRef};
17
18impl StagedLayer {
19    /// Encode and serialize the layer directly into `enc`, without creating any
20    /// intermediate representation.
21    ///
22    /// This is the hot path inside `TileLayer::encode`: each sort-strategy
23    /// trial calls this method on its own fresh `Encoder`, and only the
24    /// `Encoder` with the smallest `total_len()` is kept.
25    #[hotpath::measure]
26    pub fn encode_into(self, enc: Encoder, codecs: &mut Codecs) -> MltResult<Encoder> {
27        if self.name.is_empty() {
28            return Err(MltError::MissingLayerName);
29        }
30        #[cfg(feature = "unstable-v2")]
31        match enc.config().wire_version() {
32            WireVersion::V01 => encode01::encode_into01(self, enc, codecs),
33            WireVersion::V02 => encode02::encode_into02(self, enc, codecs),
34        }
35        #[cfg(not(feature = "unstable-v2"))]
36        encode01::encode_into01(self, enc, codecs)
37    }
38}
39
40/// Scratch capacity above which a [`Codecs`] is dropped instead of kept for the next layer.
41const POOLED_CODECS_MAX_BYTES: usize = 16 << 20;
42
43thread_local! {
44    static CODECS_POOL: RefCell<Option<Codecs>> = const { RefCell::new(None) };
45}
46
47/// A [`Codecs`] borrowed from this thread's pool, which keeps scratch buffers warm between layers.
48struct PooledCodecs(Option<Codecs>);
49
50impl PooledCodecs {
51    fn take() -> Self {
52        let pooled = CODECS_POOL.try_with(|pool| pool.borrow_mut().take());
53        Self(Some(pooled.ok().flatten().unwrap_or_default()))
54    }
55}
56
57impl Deref for PooledCodecs {
58    type Target = Codecs;
59
60    fn deref(&self) -> &Codecs {
61        self.0.as_ref().expect("codecs are only taken on drop")
62    }
63}
64
65impl DerefMut for PooledCodecs {
66    fn deref_mut(&mut self) -> &mut Codecs {
67        self.0.as_mut().expect("codecs are only taken on drop")
68    }
69}
70
71impl Drop for PooledCodecs {
72    fn drop(&mut self) {
73        if let Some(codecs) = self.0.take()
74            && codecs.scratch_bytes() <= POOLED_CODECS_MAX_BYTES
75        {
76            let _ = CODECS_POOL.try_with(|pool| *pool.borrow_mut() = Some(codecs));
77        }
78    }
79}
80
81/// Seed the encoder's curve-derived caches so the Hilbert/Morton dictionary
82/// builders skip their min/max scan. `Morton::new` returns `Err` when bits > 16;
83/// `dict_may_be_beneficial` reads `morton_cache.is_none()` and falls back to
84/// a Vec2-only path in that case.
85fn seed_curve_caches(enc: &mut Encoder, curve_params: CurveParams) {
86    enc.hilbert_cache = Some(curve_params);
87    enc.morton_cache = Morton::new(curve_params.bits, curve_params.shift).ok();
88}
89
90impl TileLayer {
91    /// Encode a [`TileLayer`] to bytes, automatically optimizing all encoding choices.
92    ///
93    /// This is the primary encoding entry point. It:
94    /// 1. Determines which sort strategies to try based on `cfg`
95    /// 2. Tries each sort strategy, encoding and measuring the output size
96    /// 3. Returns the smallest encoding as a complete layer record (including tag and length prefix)
97    ///
98    /// All encoding choices - sort order, per-stream integer encodings, string compression,
99    /// vertex buffer layout - are selected automatically to minimize output size.
100    pub fn encode(&self, cfg: EncoderConfig) -> MltResult<Vec<u8>> {
101        encode_layer(self, cfg)
102    }
103}
104
105/// Encode the layer `source` provides, as [`TileLayer::encode`] describes.
106#[hotpath::measure]
107pub(crate) fn encode_layer(source: &impl LayerSource, cfg: EncoderConfig) -> MltResult<Vec<u8>> {
108    if source.name().is_empty() {
109        return Err(MltError::MissingLayerName);
110    }
111    let len = source.feature_count();
112    if len == 0 {
113        return Ok(Vec::new());
114    }
115
116    let mut sort_by = vec![SortStrategy::Unsorted];
117    if cfg.attempt_spatial_morton_sort() {
118        sort_by.push(SortStrategy::SpatialMorton);
119    }
120    if cfg.attempt_spatial_hilbert_sort() {
121        sort_by.push(SortStrategy::SpatialHilbert);
122    }
123    // A stable sort of ascending (or absent) IDs is the source order.
124    if cfg.attempt_id_sort() && !(0..len).map(|f| source.id(f)).is_sorted() {
125        sort_by.push(SortStrategy::Id);
126    }
127
128    let stats = analyze(source, cfg.allow_shared_dict())?;
129    // Bounds are order-invariant, so this scan is shared across every
130    // sort trial and the encoder's Hilbert/Morton dictionary builders.
131    let curve_params = curve_params(source);
132    let stage = |sort| {
133        let order = sort_order(source, sort, curve_params);
134        stage_layer(source, &order, &stats, cfg.into())
135    };
136
137    // `Encoder::preserve_results` clears caches only on the moved-out
138    // archive, so a single seeding here serves every trial that reuses
139    // `enc`.
140    let mut enc = Encoder::new(cfg);
141    seed_curve_caches(&mut enc, curve_params);
142    let mut codecs = PooledCodecs::take();
143    if let [sort] = sort_by[..] {
144        return stage(sort)
145            .encode_into(enc, &mut codecs)?
146            .into_layer_bytes();
147    }
148    let mut best: Option<Encoder> = None;
149    for &sort in &sort_by {
150        enc = stage(sort).encode_into(enc, &mut codecs)?;
151        if best
152            .as_ref()
153            .is_none_or(|best| enc.total_len() < best.total_len())
154        {
155            best = Some(enc.preserve_results());
156        } else {
157            // Drop the losing trial's bytes, or the next trial would append to them and overcount its total_len.
158            enc.clear_results();
159        }
160    }
161    best.expect("at least two trials ran").into_layer_bytes()
162}
163
164/// Row-order-independent presence classification for IDs and properties.
165#[derive(Debug, Clone, Copy, PartialEq, Eq)]
166pub enum Presence {
167    /// No feature has a value for this logical column.
168    AllNull,
169    /// Every feature has a value for this logical column.
170    AllPresent,
171    /// Some, but not all, features have a value for this logical column.
172    Mixed,
173    /// Mixed presence with the same per-feature mask as an earlier property column.
174    ///
175    /// Only tells the stager the column is nullable, same as [`Self::Mixed`]. Which
176    /// columns actually end up sharing one presence bitfield is decided per wire
177    /// format at write time, from the staged masks - see `SharedPresence` in
178    /// `encoder::encode02`.
179    SameAsProp(usize),
180}
181impl Presence {
182    /// Create presence value
183    #[must_use]
184    pub fn from_bits(bits: &BitVec<u8>, existing: &[(BitVec<u8>, usize)]) -> Self {
185        if bits.not_any() {
186            Self::AllNull
187        } else if bits.all() {
188            Self::AllPresent
189        } else if let Some((_, idx)) = existing.iter().find(|(v, _)| v == bits) {
190            Self::SameAsProp(*idx)
191        } else {
192            Self::Mixed
193        }
194    }
195}
196
197/// How a property participates in a shared dictionary group.
198#[derive(Debug, Clone, PartialEq, Eq)]
199pub enum SharedDictRole {
200    /// The property is encoded as a standalone column.
201    None,
202    /// The property is the first column in this group and emits the shared dictionary with this prefix.
203    Owner(String),
204    /// The property is emitted by the group owner at this property index.
205    Member(usize),
206}
207
208/// Row-order-independent facts for a single property column.
209#[derive(Debug, Clone, PartialEq, Eq)]
210pub struct PropertyStats {
211    pub presence: Presence,
212    pub stats: PropertyTypedStats,
213}
214
215/// Row-order-independent layer facts computed once before sort trials.
216#[derive(Debug, Clone, PartialEq, Eq)]
217pub struct LayerStats {
218    pub id: Option<PropertyStats>,
219    pub properties: Vec<PropertyStats>,
220}
221
222/// Row-order-independent value statistics for a property column.
223#[derive(Debug, Clone, Default, PartialEq, Eq)]
224pub enum PropertyTypedStats {
225    /// No present values.
226    #[default]
227    None,
228    Bool,
229    Signed {
230        min: i64,
231        max: i64,
232    },
233    Unsigned {
234        min: u64,
235        max: u64,
236    },
237    F32,
238    F64,
239    String {
240        shared_dict: SharedDictRole,
241    },
242}
243
244impl PropertyTypedStats {
245    #[must_use]
246    pub fn values_fit_u32(&self) -> bool {
247        match self {
248            Self::None | Self::Bool | Self::F32 | Self::F64 | Self::String { .. } => false,
249            Self::Signed { min, max } => *min >= 0 && u32::try_from(*max).is_ok(),
250            Self::Unsigned { max, .. } => u32::try_from(*max).is_ok(),
251        }
252    }
253
254    /// Returns `true` if every value fits in an `i32`.
255    /// Unlike [`Self::values_fit_u32`] this admits negative values.
256    #[must_use]
257    pub fn values_fit_i32(&self) -> bool {
258        match self {
259            Self::None | Self::Bool | Self::F32 | Self::F64 | Self::String { .. } => false,
260            Self::Signed { min, max } => i32::try_from(*min).is_ok() && i32::try_from(*max).is_ok(),
261            Self::Unsigned { max, .. } => i32::try_from(*max).is_ok(),
262        }
263    }
264
265    #[must_use]
266    pub fn shared_dict(&self) -> SharedDictRole {
267        match self {
268            Self::String { shared_dict, .. } => shared_dict.clone(),
269            Self::None
270            | Self::Bool
271            | Self::Signed { .. }
272            | Self::Unsigned { .. }
273            | Self::F32
274            | Self::F64 => SharedDictRole::None,
275        }
276    }
277
278    pub(crate) fn set_shared_dict(&mut self, role: SharedDictRole) {
279        match self {
280            Self::String { shared_dict, .. } => *shared_dict = role,
281            Self::None
282            | Self::Bool
283            | Self::Signed { .. }
284            | Self::Unsigned { .. }
285            | Self::F32
286            | Self::F64 => debug_assert_eq!(role, SharedDictRole::None),
287        }
288    }
289
290    pub(crate) fn push(
291        &mut self,
292        prop: PropValueRef<'_>,
293        column_idx: usize,
294        property_name: &str,
295    ) -> MltResult<()> {
296        match prop {
297            PropValueRef::Bool(_) => self.merge_same_kind(Self::Bool, column_idx, property_name),
298            PropValueRef::I8(v) => self.merge_signed(i64::from(v), column_idx, property_name),
299            PropValueRef::U8(v) => self.merge_unsigned(u64::from(v), column_idx, property_name),
300            PropValueRef::I32(v) => self.merge_signed(i64::from(v), column_idx, property_name),
301            PropValueRef::U32(v) => self.merge_unsigned(u64::from(v), column_idx, property_name),
302            PropValueRef::I64(v) => self.merge_signed(v, column_idx, property_name),
303            PropValueRef::U64(v) => self.merge_unsigned(v, column_idx, property_name),
304            PropValueRef::F32(_) => self.merge_same_kind(Self::F32, column_idx, property_name),
305            PropValueRef::F64(_) => self.merge_same_kind(Self::F64, column_idx, property_name),
306            PropValueRef::Str(_) => self.merge_string(column_idx, property_name),
307        }
308    }
309
310    fn merge_signed(
311        &mut self,
312        value: i64,
313        column_idx: usize,
314        property_name: &str,
315    ) -> MltResult<()> {
316        match self {
317            Self::None => {
318                *self = Self::Signed {
319                    min: value,
320                    max: value,
321                };
322            }
323            Self::Signed { min, max } => {
324                *min = (*min).min(value);
325                *max = (*max).max(value);
326            }
327            Self::Bool | Self::Unsigned { .. } | Self::F32 | Self::F64 | Self::String { .. } => {
328                return mixed_prop_err(column_idx, property_name);
329            }
330        }
331        Ok(())
332    }
333
334    fn merge_unsigned(
335        &mut self,
336        value: u64,
337        column_idx: usize,
338        property_name: &str,
339    ) -> MltResult<()> {
340        match self {
341            Self::None => {
342                *self = Self::Unsigned {
343                    min: value,
344                    max: value,
345                };
346            }
347            Self::Unsigned { min, max } => {
348                *min = (*min).min(value);
349                *max = (*max).max(value);
350            }
351            Self::Bool | Self::Signed { .. } | Self::F32 | Self::F64 | Self::String { .. } => {
352                return mixed_prop_err(column_idx, property_name);
353            }
354        }
355        Ok(())
356    }
357
358    fn merge_string(&mut self, column_idx: usize, property_name: &str) -> MltResult<()> {
359        match self {
360            Self::None => {
361                *self = Self::String {
362                    shared_dict: SharedDictRole::None,
363                };
364            }
365            Self::String { .. } => {}
366            Self::Bool | Self::Signed { .. } | Self::Unsigned { .. } | Self::F32 | Self::F64 => {
367                return mixed_prop_err(column_idx, property_name);
368            }
369        }
370        Ok(())
371    }
372
373    fn merge_same_kind(
374        &mut self,
375        kind: Self,
376        column_idx: usize,
377        property_name: &str,
378    ) -> MltResult<()> {
379        match self {
380            Self::None => *self = kind,
381            Self::Bool if matches!(kind, Self::Bool) => {}
382            Self::F32 if matches!(kind, Self::F32) => {}
383            Self::F64 if matches!(kind, Self::F64) => {}
384            Self::Bool
385            | Self::Signed { .. }
386            | Self::Unsigned { .. }
387            | Self::F32
388            | Self::F64
389            | Self::String { .. } => return mixed_prop_err(column_idx, property_name),
390        }
391        Ok(())
392    }
393}
394
395#[cfg(any(test, feature = "__private"))]
396impl TileLayer {
397    /// Analyze a [`TileLayer`] and return reusable ID/property facts for the optimizer.
398    pub(crate) fn analyze(&self, allow_shared_dict: bool) -> MltResult<LayerStats> {
399        analyze(self, allow_shared_dict)
400    }
401}
402
403/// Analyze the layer `source` provides and return reusable ID/property facts for the optimizer.
404#[hotpath::measure]
405pub(crate) fn analyze(source: &impl LayerSource, allow_shared_dict: bool) -> MltResult<LayerStats> {
406    let mut property_bits = Vec::with_capacity(source.property_count());
407    let mut properties = analyze_properties(source, &mut property_bits)?;
408    let id = analyze_ids(source, &property_bits);
409    if allow_shared_dict {
410        group_string_properties(source, &mut properties);
411    }
412    Ok(LayerStats { id, properties })
413}
414
415fn analyze_ids(
416    source: &impl LayerSource,
417    property_bits: &[(BitVec<u8>, usize)],
418) -> Option<PropertyStats> {
419    let mut min = u64::MAX;
420    let mut max = 0u64;
421    let mut bits = BitVec::<u8>::with_capacity(source.feature_count());
422    for feature in 0..source.feature_count() {
423        if let Some(id) = source.id(feature) {
424            min = min.min(id);
425            max = max.max(id);
426            bits.push(true);
427        } else {
428            bits.push(false);
429        }
430    }
431    let presence = Presence::from_bits(&bits, property_bits);
432    if presence == Presence::AllNull {
433        None
434    } else {
435        Some(PropertyStats {
436            presence,
437            stats: PropertyTypedStats::Unsigned { min, max },
438        })
439    }
440}
441
442fn analyze_properties(
443    source: &impl LayerSource,
444    property_bits: &mut Vec<(BitVec<u8>, usize)>,
445) -> MltResult<Vec<PropertyStats>> {
446    let kinds = source.property_kinds();
447    (0..source.property_count())
448        .map(|col_idx| -> MltResult<PropertyStats> {
449            let mut stats = PropertyTypedStats::default();
450            let mut bits = BitVec::<u8>::with_capacity(source.feature_count());
451            for feature in 0..source.feature_count() {
452                if let Some(prop) = source.property(feature, col_idx) {
453                    if Some(&prop.kind()) != kinds.get(col_idx) {
454                        return mixed_prop_err(col_idx, source.property_name(col_idx));
455                    }
456                    stats.push(prop, col_idx, source.property_name(col_idx))?;
457                    bits.push(true);
458                } else {
459                    bits.push(false);
460                }
461            }
462
463            let presence = Presence::from_bits(&bits, property_bits);
464            if presence == Presence::Mixed {
465                property_bits.push((bits, col_idx));
466            }
467            Ok(PropertyStats { presence, stats })
468        })
469        .collect()
470}
471
472#[inline]
473fn mixed_prop_err<T>(column_idx: usize, property_name: &str) -> MltResult<T> {
474    Err(MltError::MixedPropertyTypes(
475        column_idx,
476        property_name.to_owned(),
477    ))
478}
mlt_core/encoder/optimizer.rs

mlt_core/encoder/
optimizer.rs
