Skip to main content

mlt_core/encoder/
writer.rs

1use std::collections::HashMap;
2use std::{io, mem};
3
4use fsst::Compressor;
5use integer_encoding::VarIntWriter as _;
6
7#[cfg(feature = "unstable-v2")]
8use crate::decoder::stream::header02::Family;
9use crate::decoder::{ColumnType, Morton};
10#[cfg(feature = "unstable-v2")]
11use crate::encoder::model::FloatEncoding;
12use crate::encoder::model::{CurveParams, ExplicitEncoder, StrEncoding, StreamCtx};
13use crate::encoder::{EncoderConfig, IntEncoder, VertexBufferType};
14use crate::utils::BinarySerializer as _;
15use crate::{MltError, MltResult};
16
17/// Stateful encoder that accumulates encoded layer bytes.
18///
19/// Logical temporary buffers live in `Codecs` and are passed alongside
20/// the encoder while a stream is being transformed and serialized. Physical
21/// encoders live here with their own scratch buffers, then copy complete
22/// payloads into [`data`](Encoder::data).
23///
24/// # Buffer layout
25///
26/// The MLT layer wire format is:
27///
28/// ```text
29/// [varint(body_len + 1)] [tag = 1]
30/// [name: string] [extent: varint] [column_count: varint]   <- hdr
31/// [col_type₁] [col_type₂] … [col_typeN]                    <- meta
32/// [col₁ stream data] [col₂ stream data] … [colN stream data] <- data
33/// ```
34///
35/// The three sections are accumulated into separate buffers so they can be
36/// combined at the end *without* any in-place insertion or extra copies:
37///
38/// * `hdr` - layer header (name, extent, `column_count`).
39/// * [`meta`] - column-type bytes (one byte + optional name per column).
40/// * [`data`] - encoded stream data; also the target of [`impl Write`].
41///
42/// # Sort-strategy trialing
43///
44/// Create one `Encoder` per sort-strategy trial, encode the layer into it,
45/// and keep the one whose `total_len()` is smallest:
46///
47/// ```rust,ignore
48/// let mut codecs = Codecs::default();
49/// let mut best: Option<Encoder> = None;
50/// for strategy in strategies {
51///     let mut enc = Encoder::new(cfg);
52///     layer.write_to(&mut enc, &mut codecs)?;
53///     if best.as_ref().is_none_or(|b| enc.total_len() < b.total_len()) {
54///         best = Some(enc);
55///     }
56/// }
57/// return best.unwrap().into_layer_bytes();
58/// ```
59///
60/// # Stream-level encoding alternatives
61///
62/// Use [`Encoder::try_alternatives`] to open a competition,
63/// then submit each candidate via `AltSession::with`.  The guard's `Drop`
64/// impl finalises the competition automatically:
65///
66/// ```rust,ignore
67/// let mut alt = enc.try_alternatives();
68/// alt.with(|enc| write_stream_as_varint(data, enc))?;
69/// alt.with(|enc| write_stream_as_fastpfor(data, enc))?;
70/// // alt drops -> keeps whichever was shorter
71/// ```
72///
73/// [`meta`]: Encoder::meta
74/// [`data`]: Encoder::data
75/// [`impl Write`]: Encoder#impl-Write
76#[derive(Default)]
77pub struct Encoder {
78    /// Encoding configuration: controls which optimization strategies are tried
79    /// (sort orders, compression algorithms, etc.).
80    ///
81    /// Set once at construction time via [`Encoder::new`]; propagated
82    /// automatically to all sub-encoders so individual encode methods do not
83    /// need a separate `cfg` argument.
84    cfg: EncoderConfig,
85
86    /// When [`Some`], property / ID / geometry encoders use `ExplicitEncoder`
87    /// callbacks instead of trying candidate encodings. When [`None`], the
88    /// automatic optimization path runs.
89    pub(crate) explicit: Option<ExplicitEncoder>,
90
91    /// Layer header bytes: `name`, `extent`, `column_count`.
92    ///
93    /// This section comes first in the wire format and is never subject to alternatives.
94    hdr: Vec<u8>,
95
96    /// Column-type metadata bytes.
97    ///
98    /// Each column contributes one type byte (plus a name string for property
99    /// columns).  Written by the `write_columns_meta_to` methods, which write
100    /// directly to `enc.meta`.  This section comes second in the wire format
101    /// and is never subject to alternatives (column types are fixed).
102    meta: Vec<u8>,
103
104    /// Encoded stream data.
105    ///
106    /// All stream counts, per-stream encoding-metadata bytes, and encoded
107    /// data bytes land here via [`impl Write`].  This section comes last in
108    /// the wire format and is where stream-level alternatives compete.
109    ///
110    /// [`impl Write`]: Encoder#impl-Write
111    data: Vec<u8>,
112
113    /// Morton parameters for this layer's vertex set; `None` if the extent
114    /// exceeds 16 bits per axis (Morton encoding is unusable in that case).
115    /// Pre-populated by [`StagedLayer::encode_into`](crate::encoder::StagedLayer::encode_into).
116    pub(crate) morton_cache: Option<Morton>,
117
118    /// Hilbert curve parameters for this layer's vertex set. Pre-populated by
119    /// [`StagedLayer::encode_into`](crate::encoder::StagedLayer::encode_into).
120    pub(crate) hilbert_cache: Option<CurveParams>,
121
122    /// Cached FSST compressor per string column, keyed by column name.
123    /// `None` means training found FSST not viable for that column.
124    /// Trained on deduplicated values on the first sort trial, reused on subsequent trials.
125    pub(crate) fsst_cache: HashMap<String, Option<Compressor>>,
126
127    /// The stream count a v2 decoder would infer from context at the current
128    /// write position: the layer's `feature_count`, or the presence popcount
129    /// while an optional column's data stream is being written.
130    ///
131    /// Read by the v2 stream-header codec to decide whether an explicit count
132    /// varint must be emitted; ignored entirely for v1 layers.
133    #[cfg(feature = "unstable-v2")]
134    pub(crate) count_context: u32,
135
136    /// The family the stream being written is numbered in, set by the v2 writers alongside [`Self::count_context`].
137    /// Ignored for v1 layers.
138    #[cfg(feature = "unstable-v2")]
139    pub(crate) family_context: Family,
140
141    // -----------------------------------------------------------------------
142    // Alternatives state - a stack that supports nested competitions.
143    //
144    // Invariant between candidates at any level:
145    //   data.len() == level.data_start + level.best_data_size.unwrap_or(0)
146    //   meta.len() == level.meta_start + level.best_meta_size.unwrap_or(0)
147    //
148    // Empty stack <-> no competition in progress.
149    // -----------------------------------------------------------------------
150    /// Stack of active encoding competitions, innermost last.
151    ///
152    /// Empty while no [`Encoder::try_alternatives`] session
153    /// is in progress.
154    alt_stack: Vec<AltLevel>,
155}
156
157impl Encoder {
158    /// Create a new encoder with the given [`EncoderConfig`].
159    ///
160    /// Use [`Encoder::default()`] when the default configuration is sufficient.
161    #[inline]
162    #[must_use]
163    pub fn new(cfg: EncoderConfig) -> Self {
164        Self {
165            cfg,
166            ..Self::default()
167        }
168    }
169
170    /// Like [`Self::new`] but with the explicit encoder set for deterministic encoding
171    /// (tests, synthetics). Use with `StagedLayer::encode_explicit`.
172    #[inline]
173    #[must_use]
174    pub fn with_explicit(cfg: EncoderConfig, explicit: ExplicitEncoder) -> Self {
175        Self {
176            cfg,
177            explicit: Some(explicit),
178            ..Self::default()
179        }
180    }
181
182    /// Ensure this encoder is in the good state, and moves results to a new instance.
183    /// This allows current instance to be reused for other experiment, avoiding repeat of some operations.
184    #[must_use]
185    pub(crate) fn preserve_results(&mut self) -> Self {
186        assert_eq!(self.alt_stack.len(), 0, "Alternatives stack is not empty");
187        Self {
188            // Keep the config: the archived result still needs it to frame the
189            // layer with the right tag in `into_layer_bytes`.
190            cfg: self.cfg,
191            explicit: None,
192            hdr: mem::take(&mut self.hdr),
193            meta: mem::take(&mut self.meta),
194            data: mem::take(&mut self.data),
195            morton_cache: None,
196            hilbert_cache: None,
197            fsst_cache: HashMap::new(),
198            #[cfg(feature = "unstable-v2")]
199            count_context: 0,
200            #[cfg(feature = "unstable-v2")]
201            family_context: Family::Int,
202            alt_stack: vec![],
203        }
204    }
205
206    #[inline]
207    pub(crate) fn write_column_type(&mut self, column_type: ColumnType) -> MltResult<()> {
208        column_type.write_to(&mut self.meta).map_err(MltError::from)
209    }
210
211    #[inline]
212    pub(crate) fn write_column_name(&mut self, name: &str) -> MltResult<()> {
213        self.meta.write_string(name).map_err(MltError::from)
214    }
215
216    #[inline]
217    #[must_use]
218    pub fn config(&self) -> EncoderConfig {
219        self.cfg
220    }
221
222    #[inline]
223    #[must_use]
224    pub fn data(&self) -> &[u8] {
225        &self.data
226    }
227
228    #[inline]
229    pub(crate) fn data_mut(&mut self) -> &mut Vec<u8> {
230        &mut self.data
231    }
232
233    #[inline]
234    #[must_use]
235    pub fn meta(&self) -> &[u8] {
236        &self.meta
237    }
238
239    #[inline]
240    pub(crate) fn meta_mut(&mut self) -> &mut Vec<u8> {
241        &mut self.meta
242    }
243
244    #[inline]
245    #[must_use]
246    pub fn section_lens(&self) -> (usize, usize, usize) {
247        (self.hdr.len(), self.meta.len(), self.data.len())
248    }
249
250    #[inline]
251    pub(crate) fn write_column_header(
252        &mut self,
253        column_type: ColumnType,
254        name: &str,
255    ) -> MltResult<()> {
256        self.write_column_type(column_type)?;
257        self.write_column_name(name)
258    }
259
260    /// Write the v1 layer header (`name`, `extent`, `column_count`) to `hdr`.
261    ///
262    /// Must be called exactly once per layer, after all column meta and data.
263    #[hotpath::measure]
264    pub fn write_header01(
265        &mut self,
266        name: &str,
267        extent: u32,
268        column_count: usize,
269    ) -> MltResult<()> {
270        if name.is_empty() {
271            return Err(MltError::MissingLayerName);
272        }
273        debug_assert!(
274            self.alt_stack.is_empty(),
275            "write_header called with an open alternatives session"
276        );
277        let name_len = u32::try_from(name.len())?;
278        let column_count = u32::try_from(column_count)?;
279        self.hdr.write_varint(name_len).map_err(MltError::from)?;
280        self.hdr.extend_from_slice(name.as_bytes());
281        self.hdr.write_varint(extent).map_err(MltError::from)?;
282        self.hdr
283            .write_varint(column_count)
284            .map_err(MltError::from)?;
285        Ok(())
286    }
287
288    /// Write the v2 layer header (`name`, `extent`, `feature_count`) to `hdr`.
289    ///
290    /// Unlike v1, `column_count` is not part of the header - it precedes the
291    /// counted columns in the data section, after the geometry section.
292    #[cfg(feature = "unstable-v2")]
293    #[hotpath::measure]
294    pub(crate) fn write_header02(
295        &mut self,
296        name: &str,
297        extent: u32,
298        feature_count: u32,
299    ) -> MltResult<()> {
300        if name.is_empty() {
301            return Err(MltError::MissingLayerName);
302        }
303        debug_assert!(
304            self.alt_stack.is_empty(),
305            "write_header02 called with an open alternatives session"
306        );
307        self.hdr.write_string(name).map_err(MltError::from)?;
308        self.hdr.write_varint(extent).map_err(MltError::from)?;
309        self.hdr
310            .write_varint(feature_count)
311            .map_err(MltError::from)?;
312        Ok(())
313    }
314
315    /// When [`Self::explicit`] is [`Some`], returns the callback-chosen [`IntEncoder`].
316    /// [`None`] means run automatic candidate selection for that stream.
317    #[inline]
318    pub(crate) fn override_int_enc(&self, ctx: &StreamCtx<'_>) -> Option<IntEncoder> {
319        self.explicit.as_ref().map(|e| (e.get_int_encoder)(ctx))
320    }
321
322    /// When [`Self::explicit`] is [`Some`], returns the callback-chosen [`StrEncoding`].
323    /// [`None`] means run automatic string / shared-dict corpus selection.
324    #[inline]
325    pub(crate) fn override_str_enc(&self, name: &str) -> Option<StrEncoding> {
326        self.explicit.as_ref().map(|e| (e.get_str_encoding)(name))
327    }
328
329    /// When [`Self::explicit`] is [`Some`], returns the callback-chosen [`FloatEncoding`].
330    /// [`None`] means cost the float encodings the config allows against each other.
331    #[inline]
332    #[cfg(feature = "unstable-v2")]
333    pub(crate) fn override_float_enc(&self, name: &str) -> Option<FloatEncoding> {
334        self.explicit.as_ref().map(|e| (e.get_float_encoding)(name))
335    }
336
337    /// Pinned vertex layout when an explicit encoder is active.
338    #[inline]
339    #[allow(clippy::unused_self)]
340    pub(crate) fn override_vertex_buffer_type(&self) -> Option<VertexBufferType> {
341        self.explicit.as_ref().map(|e| e.vertex_buffer_type)
342    }
343
344    /// Whether to force writing a geometry stream even when its data is empty.
345    ///
346    /// Delegates to [`ExplicitEncoder::force_stream`]; returns `false` when no explicit
347    /// encoder is active (the default "skip empty streams" behavior).
348    #[inline]
349    pub(crate) fn force_stream(&self, ctx: &StreamCtx<'_>) -> bool {
350        self.explicit
351            .as_ref()
352            .is_some_and(|e| (e.force_stream)(ctx))
353    }
354
355    /// Total encoded bytes across all three sections (`hdr + meta + data`).
356    #[inline]
357    #[must_use]
358    pub fn total_len(&self) -> usize {
359        self.hdr.len() + self.meta.len() + self.data.len()
360    }
361
362    /// Empty the output buffers so this encoder can be reused for the next sort trial.
363    /// Keeps allocated capacity and the seeded curve/FSST caches.
364    /// Used when a trial loses; [`Self::preserve_results`] handles the winning case instead.
365    pub(crate) fn clear_results(&mut self) {
366        debug_assert!(self.alt_stack.is_empty(), "Alternatives stack is not empty");
367        self.hdr.clear();
368        self.meta.clear();
369        self.data.clear();
370    }
371
372    /// Concatenate `hdr + meta + data` into a single buffer **without** a
373    /// tag/size prefix.
374    ///
375    /// Use this when the caller expects raw layer body bytes (without the size/tag framing)
376    /// rather than a complete framed wire record - see [`Self::into_layer_bytes`] for the framed form.
377    #[must_use]
378    pub fn into_raw_bytes(mut self) -> Vec<u8> {
379        if self.hdr.is_empty() && self.meta.is_empty() {
380            return self.data;
381        }
382        let mut out = Vec::with_capacity(self.hdr.len() + self.meta.len() + self.data.len());
383        out.append(&mut self.hdr);
384        out.append(&mut self.meta);
385        out.append(&mut self.data);
386        out
387    }
388
389    /// Assemble the complete layer record.
390    pub fn into_layer_bytes(self) -> MltResult<Vec<u8>> {
391        #[cfg(feature = "unstable-v2")]
392        let tag = self.cfg.wire_version().tag();
393        // v1 is the only format this build writes, and `0x01` is its layer tag.
394        #[cfg(not(feature = "unstable-v2"))]
395        let tag = 1;
396        self.into_layer_bytes_with_tag(tag)
397    }
398
399    /// Assemble a complete layer record for the given `tag`:
400    /// `[varint(body_len + 1)][tag][hdr][meta][data]`.
401    fn into_layer_bytes_with_tag(mut self, tag: u8) -> MltResult<Vec<u8>> {
402        debug_assert!(
403            self.alt_stack.is_empty(),
404            "into_layer_bytes_with_tag called with an open alternatives session"
405        );
406        let body_len = self.hdr.len() + self.meta.len() + self.data.len();
407        let size = u32::try_from(body_len + 1)?; // +1 for the tag byte
408        let mut out = Vec::with_capacity(5 + 1 + body_len);
409        out.write_varint(size).map_err(MltError::from)?;
410        out.push(tag);
411        out.append(&mut self.hdr);
412        out.append(&mut self.meta);
413        out.append(&mut self.data);
414        Ok(out)
415    }
416
417    /// Begin a new encoding competition.
418    ///
419    /// Returns an `AltSession` guard.  Submit each candidate via
420    /// `AltSession::with`; the guard's `Drop` impl finalises
421    /// the competition and retains the shortest candidate automatically.
422    ///
423    /// Nesting is supported: calling `try_alternatives` inside a
424    /// `with` closure opens an inner competition on the same stack,
425    /// resolved before the outer candidate is committed.
426    ///
427    /// # Example
428    ///
429    /// ```rust,ignore
430    /// let mut alt = enc.try_alternatives();
431    /// for cand in candidates {
432    ///     alt.with(|enc| write_candidate(cand, enc))?;
433    /// }
434    /// // alt drops -> finalises the competition
435    /// ```
436    pub fn try_alternatives(&mut self) -> AltSession<'_> {
437        self.alt_stack.push(AltLevel {
438            data_start: self.data.len(),
439            meta_start: self.meta.len(),
440            best_data: None,
441            best_meta: None,
442        });
443        AltSession { enc: self }
444    }
445
446    /// Commit the current candidate at the innermost competition level.
447    ///
448    /// Compares bytes written since the last commit against the running best
449    /// by **total** (`data + meta`) size; keeps the shorter one.
450    ///
451    /// Called internally by `AltSession::with` on `Ok`.
452    fn alt_commit(&mut self) {
453        debug_assert!(
454            !self.alt_stack.is_empty(),
455            "alt_commit called outside an active AltSession"
456        );
457        let (data, meta, stack) = (&mut self.data, &mut self.meta, &mut self.alt_stack);
458        let level = stack.last_mut().unwrap();
459        Self::close_candidate(data, meta, level);
460    }
461
462    /// Finalize the innermost competition and pop it from the stack.
463    ///
464    /// Any bytes written since the last `alt_commit` are evaluated as a
465    /// final candidate; if no pending bytes exist and a best is already
466    /// recorded this is a cheap stack-pop.
467    fn alt_pop(&mut self) {
468        debug_assert!(
469            !self.alt_stack.is_empty(),
470            "alt_pop called outside an active AltSession"
471        );
472        {
473            let (data, meta, stack) = (&mut self.data, &mut self.meta, &mut self.alt_stack);
474            let level = stack.last_mut().unwrap();
475            let data_pending = data.len() - (level.data_start + level.best_data.unwrap_or(0));
476            let meta_pending = meta.len() - (level.meta_start + level.best_meta.unwrap_or(0));
477            if data_pending > 0 || meta_pending > 0 || level.best_data.is_none() {
478                Self::close_candidate(data, meta, level);
479            }
480        }
481        self.alt_stack.pop();
482    }
483
484    /// Shared compare-and-keep logic used by both `alt_commit` and `alt_pop`.
485    ///
486    /// Compares the bytes written since the last committed candidate against
487    /// the current best by **total** (`data + meta`) size.
488    /// Keeps the shorter one; ties preserve the existing best.
489    fn close_candidate(data: &mut Vec<u8>, meta: &mut Vec<u8>, level: &mut AltLevel) {
490        let best_data_end = level.data_start + level.best_data.unwrap_or(0);
491        let best_meta_end = level.meta_start + level.best_meta.unwrap_or(0);
492        let cand_data = data.len() - best_data_end;
493        let cand_meta = meta.len() - best_meta_end;
494        let cand_total = cand_data + cand_meta;
495        let best_total = level.best_data.unwrap_or(0) + level.best_meta.unwrap_or(0);
496        if level.best_data.is_none_or(|_| cand_total < best_total) {
497            // New best: shift data candidate bytes to data_start.
498            if level.best_data.is_some() {
499                data.copy_within(best_data_end..best_data_end + cand_data, level.data_start);
500                meta.copy_within(best_meta_end..best_meta_end + cand_meta, level.meta_start);
501            }
502            data.truncate(level.data_start + cand_data);
503            meta.truncate(level.meta_start + cand_meta);
504            level.best_data = Some(cand_data);
505            level.best_meta = Some(cand_meta);
506        } else {
507            // Not an improvement: discard.
508            data.truncate(best_data_end);
509            meta.truncate(best_meta_end);
510        }
511    }
512}
513
514/// State for one level of an encoding competition.
515///
516/// Tracks the starting position in both the [`data`](Encoder::data) and
517/// [`meta`](Encoder::meta) buffers, and the byte count of the best candidate
518/// committed so far.
519///
520/// Candidates are compared by **total** bytes (`data + meta`); the shorter one
521/// wins, with ties resolved in favor of the earlier candidate.
522#[derive(Debug, Default, Clone)]
523struct AltLevel {
524    data_start: usize,
525    meta_start: usize,
526    /// Byte count appended to `data` by the current best candidate.
527    best_data: Option<usize>,
528    /// Byte count appended to `meta` by the current best candidate.
529    best_meta: Option<usize>,
530}
531
532/// RAII guard for a stream-encoding competition opened by [`Encoder::try_alternatives`].
533///
534/// Submit each candidate via [`with`](AltSession::with); on `Ok` the candidate is
535/// committed (compared against the running best and kept if shorter); on `Err`
536/// the partial write is rolled back and the error propagates.  The guard's
537/// `Drop` impl finalises the competition automatically, so the [`Encoder`] is
538/// always left in a consistent state even when an error exits the loop early.
539///
540/// Nesting is allowed: calling [`Encoder::try_alternatives`] inside a
541/// `with` closure opens an inner competition that is fully
542/// resolved before the outer candidate is committed.
543#[must_use = "AltSession must be used; drop it to finalise the competition"]
544pub struct AltSession<'a> {
545    enc: &'a mut Encoder,
546}
547
548impl AltSession<'_> {
549    /// Encode one candidate.
550    ///
551    /// - **`Ok`** - commits the candidate; replaces the running best if shorter.
552    /// - **`Err`** - truncates the partial write back to the pre-call checkpoint
553    ///   and returns the error.  The guard's `Drop` still finalises the
554    ///   competition cleanly using whichever candidates succeeded so far.
555    #[hotpath::measure]
556    pub fn with<F>(&mut self, f: F) -> MltResult<()>
557    where
558        F: FnOnce(&mut Encoder) -> MltResult<()>,
559    {
560        let data_cp = self.enc.data.len();
561        let meta_cp = self.enc.meta.len();
562        match f(self.enc) {
563            Ok(()) => {
564                self.enc.alt_commit();
565                Ok(())
566            }
567            Err(e) => {
568                self.enc.data.truncate(data_cp);
569                self.enc.meta.truncate(meta_cp);
570                Err(e)
571            }
572        }
573    }
574}
575
576impl Drop for AltSession<'_> {
577    fn drop(&mut self) {
578        self.enc.alt_pop();
579    }
580}
581
582/// Writes bytes to [`Encoder::data`].
583///
584/// This blanket implementation makes `Encoder` compatible with all
585/// `BinarySerializer`, `VarIntWriter`, and other `Write`-based utilities so that
586/// stream-data methods do not need a separate code path.
587impl io::Write for Encoder {
588    #[inline]
589    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
590        self.data.write(buf)
591    }
592
593    #[inline]
594    fn flush(&mut self) -> io::Result<()> {
595        Ok(())
596    }
597
598    #[inline]
599    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
600        self.data.write_all(buf)
601    }
602}
603
604#[cfg(test)]
605mod tests {
606    use super::*;
607
608    /// Helper: directly extend `enc.data` with raw bytes (simulates a stream write).
609    fn push(enc: &mut Encoder, bytes: &[u8]) {
610        enc.data.extend_from_slice(bytes);
611    }
612
613    // ── basic single-level behavior ──────────────────────────────────────
614
615    /// The shortest candidate wins.
616    #[test]
617    fn alternatives_keeps_shortest() {
618        let mut enc = Encoder::default();
619        push(&mut enc, b"prefix");
620
621        let mut alt = enc.try_alternatives();
622        alt.with(|enc| {
623            push(enc, b"longer");
624            Ok(())
625        })
626        .unwrap(); // 6 bytes
627        alt.with(|enc| {
628            push(enc, b"ab");
629            Ok(())
630        })
631        .unwrap(); // 2 bytes - shortest
632        alt.with(|enc| {
633            push(enc, b"xyz");
634            Ok(())
635        })
636        .unwrap(); // 3 bytes
637        drop(alt);
638
639        assert_eq!(enc.data, b"prefixab");
640    }
641
642    /// On a tie the first candidate is kept (strict `<`, not `<=`).
643    #[test]
644    fn alternatives_tie_keeps_first() {
645        let mut enc = Encoder::default();
646
647        let mut alt = enc.try_alternatives();
648        alt.with(|enc| {
649            push(enc, b"aaa");
650            Ok(())
651        })
652        .unwrap(); // 3 bytes
653        alt.with(|enc| {
654            push(enc, b"bbb");
655            Ok(())
656        })
657        .unwrap(); // 3 bytes - equal
658        drop(alt);
659
660        assert_eq!(enc.data, b"aaa");
661    }
662
663    /// A single candidate is unconditionally the winner.
664    #[test]
665    fn alternatives_single_candidate() {
666        let mut enc = Encoder::default();
667
668        let mut alt = enc.try_alternatives();
669        alt.with(|enc| {
670            push(enc, b"only");
671            Ok(())
672        })
673        .unwrap();
674        drop(alt);
675
676        assert_eq!(enc.data, b"only");
677    }
678
679    /// Bytes written before `try_alternatives` are left intact throughout.
680    #[test]
681    fn prefix_bytes_are_preserved() {
682        let mut enc = Encoder::default();
683        push(&mut enc, b"HDR");
684
685        let mut alt = enc.try_alternatives();
686        alt.with(|enc| {
687            push(enc, b"long_encoding");
688            Ok(())
689        })
690        .unwrap(); // 13 bytes
691        alt.with(|enc| {
692            push(enc, b"short");
693            Ok(())
694        })
695        .unwrap(); // 5 bytes - winner
696        drop(alt);
697
698        assert_eq!(&enc.data[..3], b"HDR");
699        assert_eq!(&enc.data[3..], b"short");
700    }
701
702    /// Dropping the guard after all candidates are committed is a cheap stack-pop.
703    #[test]
704    fn drop_after_all_committed_is_noop() {
705        let mut enc = Encoder::default();
706
707        let mut alt = enc.try_alternatives();
708        alt.with(|enc| {
709            push(enc, b"best");
710            Ok(())
711        })
712        .unwrap();
713        drop(alt); // all candidates committed; drop just pops the stack
714
715        assert!(enc.alt_stack.is_empty(), "stack empty after drop");
716        assert_eq!(enc.data, b"best");
717    }
718
719    // ── nesting ───────────────────────────────────────────────────────────
720
721    /// An inner competition is resolved before the outer candidate is committed.
722    #[test]
723    fn nested_alternatives() {
724        let mut enc = Encoder::default();
725
726        let mut outer = enc.try_alternatives();
727
728        // Outer candidate A: header bytes + inner competition.
729        outer
730            .with(|enc| {
731                push(enc, b"A:");
732                let mut inner = enc.try_alternatives(); // inner level pushed
733                inner.with(|enc| {
734                    push(enc, b"long_inner");
735                    Ok(())
736                })?; // 10 bytes
737                inner.with(|enc| {
738                    push(enc, b"in");
739                    Ok(())
740                })?; // 2 bytes - inner winner
741                drop(inner); // inner done; enc = b"A:in"
742                push(enc, b"!");
743                Ok(())
744            })
745            .unwrap(); // outer candidate A = b"A:in!" (5 bytes)
746
747        // Outer candidate B: shorter overall.
748        outer
749            .with(|enc| {
750                push(enc, b"B");
751                Ok(())
752            })
753            .unwrap(); // 1 byte - winner
754        drop(outer);
755
756        assert_eq!(enc.data, b"B");
757    }
758
759    /// Stack depth tracks nesting level; inner guard drops before outer closure returns.
760    #[test]
761    fn nesting_depth_reflected_in_stack() {
762        let mut enc = Encoder::default();
763
764        assert_eq!(enc.alt_stack.len(), 0);
765        let mut outer = enc.try_alternatives();
766
767        outer
768            .with(|enc| {
769                assert_eq!(enc.alt_stack.len(), 1); // outer level on stack
770                let mut inner = enc.try_alternatives();
771                inner.with(|enc| {
772                    assert_eq!(enc.alt_stack.len(), 2); // both levels on stack
773                    push(enc, b"x");
774                    Ok(())
775                })?;
776                drop(inner); // inner popped
777                assert_eq!(enc.alt_stack.len(), 1);
778                push(enc, b"y");
779                Ok(())
780            })
781            .unwrap();
782
783        drop(outer); // outer popped
784        assert_eq!(enc.alt_stack.len(), 0);
785    }
786
787    // ── meta buffer tracking ──────────────────────────────────────────────
788
789    /// Writes to both `data` and `meta` are rolled back for the losing
790    /// candidate and kept for the winner, measured by total bytes.
791    #[test]
792    fn alternatives_tracks_meta_and_data() {
793        let mut enc = Encoder::default();
794        enc.data.extend_from_slice(b"D");
795        enc.meta.extend_from_slice(b"M");
796
797        let mut alt = enc.try_alternatives();
798        // Candidate A: 4 data + 2 meta = 6 total
799        alt.with(|enc| {
800            push(enc, b"DDDD");
801            enc.meta.extend_from_slice(b"mm");
802            Ok(())
803        })
804        .unwrap();
805        // Candidate B: 1 data + 1 meta = 2 total - winner
806        alt.with(|enc| {
807            push(enc, b"d");
808            enc.meta.extend_from_slice(b"n");
809            Ok(())
810        })
811        .unwrap();
812        drop(alt);
813
814        assert_eq!(enc.data, b"Dd");
815        assert_eq!(enc.meta, b"Mn");
816    }
817
818    // ── error rollback ────────────────────────────────────────────────────
819
820    /// A failing candidate is rolled back; prior best is preserved.
821    #[test]
822    fn error_candidate_is_rolled_back() {
823        let mut enc = Encoder::default();
824
825        let mut alt = enc.try_alternatives();
826        alt.with(|enc| {
827            push(enc, b"ok");
828            Ok(())
829        })
830        .unwrap();
831        let _ = alt.with(|enc| {
832            push(enc, b"partial");
833            Err(MltError::IntegerOverflow) // simulated failure
834        });
835        drop(alt);
836
837        assert_eq!(enc.data, b"ok"); // "partial" was rolled back; "ok" kept
838    }
839}