Skip to main content

dial9_trace_format/
encoder.rs

1//! High-level encoder for writing trace files.
2//!
3//! [`Encoder`] writes the file header, registers schemas, interns strings, and
4//! encodes events with delta-compressed timestamps. It is the primary entry
5//! point for producing trace data.
6
7use crate::TraceEvent;
8use crate::codec::{self, PoolEntry, StackPoolEntry, WireTypeId};
9use crate::schema::{SchemaEntry, SchemaRegistry};
10use crate::types::{
11    CountingWriter, EncodeState, EventEncoder, InternedStackFrames, InternedString,
12};
13use std::any::TypeId;
14use std::collections::{HashMap, HashSet};
15use std::hash::{BuildHasherDefault, Hasher};
16use std::io::{self, Write};
17use std::sync::Arc;
18
19/// A fast, non-cryptographic hasher using FxHash's multiply-shift strategy.
20///
21/// For HashMap keys that are already well-distributed (TypeId, Arc<str>), this
22/// avoids hash collisions.
23#[doc(hidden)]
24#[derive(Default)]
25pub struct FxHasher(u64);
26
27impl FxHasher {
28    #[inline]
29    fn hash_word(&mut self, word: u64) {
30        self.0 = (self.0.rotate_left(5) ^ word).wrapping_mul(0x517cc1b727220a95);
31    }
32}
33
34impl Hasher for FxHasher {
35    #[inline]
36    fn write(&mut self, mut bytes: &[u8]) {
37        while bytes.len() >= 8 {
38            self.hash_word(u64::from_ne_bytes(bytes[..8].try_into().unwrap()));
39            bytes = &bytes[8..];
40        }
41        if bytes.len() >= 4 {
42            self.hash_word(u32::from_ne_bytes(bytes[..4].try_into().unwrap()) as u64);
43            bytes = &bytes[4..];
44        }
45        for &b in bytes {
46            self.hash_word(b as u64);
47        }
48    }
49
50    #[inline]
51    fn write_u8(&mut self, i: u8) {
52        self.hash_word(i as u64);
53    }
54
55    #[inline]
56    fn write_u16(&mut self, i: u16) {
57        self.hash_word(i as u64);
58    }
59
60    #[inline]
61    fn write_u32(&mut self, i: u32) {
62        self.hash_word(i as u64);
63    }
64
65    #[inline]
66    fn write_u64(&mut self, i: u64) {
67        self.hash_word(i);
68    }
69
70    #[inline]
71    fn write_usize(&mut self, i: usize) {
72        self.hash_word(i as u64);
73    }
74
75    #[inline]
76    fn write_u128(&mut self, i: u128) {
77        self.hash_word(i as u64);
78        self.hash_word((i >> 64) as u64);
79    }
80
81    #[inline]
82    fn finish(&self) -> u64 {
83        self.0
84    }
85}
86
87#[doc(hidden)]
88pub type FxBuildHasher = BuildHasherDefault<FxHasher>;
89#[doc(hidden)]
90pub type FxHashMap<K, V> = HashMap<K, V, FxBuildHasher>;
91#[doc(hidden)]
92pub type FxHashSet<T> = HashSet<T, FxBuildHasher>;
93
94/// A schema handle returned by [`Encoder::register_schema`] or created via
95/// [`Schema::new`].
96///
97/// Carries the full schema definition (name + fields) so it can auto-register
98/// itself with any encoder on first use. This means a `Schema` created on one
99/// encoder can be passed to a different encoder and it will just work.
100///
101/// `Schema` is cheap to clone (internally `Arc`-backed). Create it once and
102/// reuse it across events; see [`Encoder::write_event`].
103#[derive(Clone, Debug)]
104pub struct Schema {
105    pub(crate) entry: Arc<SchemaEntry>,
106    /// Pre-computed `Arc<str>` of the schema name, used as a cheap HashMap key
107    /// (clone is a pointer bump instead of a String allocation).
108    name_key: Arc<str>,
109}
110
111impl Schema {
112    /// Create a schema handle without an encoder.
113    ///
114    /// The schema will be lazily registered the first time it is passed to
115    /// [`Encoder::write_event`].
116    pub fn new(name: &str, fields: Vec<crate::schema::FieldDef>) -> Self {
117        let name_key: Arc<str> = Arc::from(name);
118        Self {
119            entry: Arc::new(SchemaEntry {
120                name: name.to_string(),
121                has_timestamp: true,
122                fields,
123                annotations: Vec::new(),
124            }),
125            name_key,
126        }
127    }
128
129    /// Create a schema handle from a complete [`SchemaEntry`].
130    pub fn from_entry(entry: SchemaEntry) -> Self {
131        let name_key: Arc<str> = Arc::from(entry.name.as_str());
132        Self {
133            entry: Arc::new(entry),
134            name_key,
135        }
136    }
137
138    /// Schema name.
139    pub fn name(&self) -> &str {
140        &self.entry.name
141    }
142
143    /// Schema field definitions.
144    pub fn fields(&self) -> &[crate::schema::FieldDef] {
145        &self.entry.fields
146    }
147}
148
149/// Key for schema lookup — either by name (manual registration) or by Rust
150/// `TypeId` (derive macro path).
151#[derive(Clone, PartialEq, Eq, Hash)]
152enum SchemaKey {
153    Name(Arc<str>),
154    RustType(TypeId),
155}
156
157/// Trace file encoder.
158///
159/// Writes the binary file header, registers event schemas, interns strings
160/// into a pool, and encodes events with delta-compressed timestamps.
161///
162/// The default type parameter (`Vec<u8>`) buffers everything in memory;
163/// use [`Encoder::new_to`] to write to an arbitrary [`Write`] sink.
164/// Upper bound on [`Encoder::dynamic_schema_cache`]. Far above any sane
165/// number of distinct live schema handles, small enough that pinned
166/// `SchemaEntry` Arcs stay negligible.
167const DYNAMIC_SCHEMA_CACHE_LIMIT: usize = 1024;
168
169pub struct Encoder<W: Write = Vec<u8>> {
170    state: EncodeState<W>,
171    registry: SchemaRegistry,
172    string_pool: FxHashMap<String, u32>,
173    next_pool_id: u32,
174    stack_pool: FxHashMap<Box<[u64]>, u32>,
175    next_stack_pool_id: u32,
176    schema_ids: FxHashMap<SchemaKey, WireTypeId>,
177    /// Identity fast path for dynamic [`Schema`] handles: wire ids keyed by
178    /// the address of the schema's shared `SchemaEntry` allocation. The
179    /// `Arc` is kept alive in the value so the address cannot be reused
180    /// while cached. Repeated `write_event` calls with the same handle skip
181    /// the name hash and deep schema comparison in `ensure_registered`.
182    ///
183    /// Bounded to [`DYNAMIC_SCHEMA_CACHE_LIMIT`] entries and cleared when
184    /// full: a caller that mints a fresh `Schema` per event would otherwise
185    /// grow it (and pin the `Arc`s) without bound. Clearing only costs the
186    /// fast path; registration stays correct through the name-keyed slow
187    /// path.
188    dynamic_schema_cache: FxHashMap<usize, (Arc<SchemaEntry>, WireTypeId)>,
189    /// Per-type dense cache keyed by `TraceEvent::type_slot()`.
190    /// Stores `wire_id + 1` so that `0` means "unset".
191    slot_cache: Vec<u32>,
192    /// Bitset over `0..STATIC_WIRE_ID_LIMIT`: which fast-path wire IDs (type
193    /// slots) have had their schema frame emitted on this encoder. 256 bits =
194    /// 32 bytes inline.
195    registered_ids: [u64; (crate::STATIC_WIRE_ID_LIMIT as usize) / 64],
196}
197
198impl Default for Encoder<Vec<u8>> {
199    fn default() -> Self {
200        Self::new()
201    }
202}
203
204impl Encoder<Vec<u8>> {
205    pub fn new() -> Self {
206        let mut buf = Vec::new();
207        codec::encode_header(&mut buf).expect("Vec::write_all cannot fail");
208        Self {
209            state: EncodeState::new(buf),
210            registry: SchemaRegistry::new(),
211            string_pool: FxHashMap::default(),
212            next_pool_id: 0,
213            stack_pool: FxHashMap::default(),
214            next_stack_pool_id: 0,
215            schema_ids: FxHashMap::default(),
216            dynamic_schema_cache: FxHashMap::default(),
217            slot_cache: Vec::new(),
218            registered_ids: [0; (crate::STATIC_WIRE_ID_LIMIT as usize) / 64],
219        }
220    }
221
222    /// Consume the encoder and return the encoded bytes.
223    pub fn finish(self) -> Vec<u8> {
224        self.state.writer.into_inner()
225    }
226}
227
228impl<W: Write> Encoder<W> {
229    /// Create an encoder that writes to an arbitrary writer.
230    /// Writes the file header immediately.
231    pub fn new_to(mut writer: W) -> io::Result<Self> {
232        codec::encode_header(&mut writer)?;
233        Ok(Self {
234            state: EncodeState::new(writer),
235            registry: SchemaRegistry::new(),
236            string_pool: FxHashMap::default(),
237            next_pool_id: 0,
238            stack_pool: FxHashMap::default(),
239            next_stack_pool_id: 0,
240            schema_ids: FxHashMap::default(),
241            dynamic_schema_cache: FxHashMap::default(),
242            slot_cache: Vec::new(),
243            registered_ids: [0; (crate::STATIC_WIRE_ID_LIMIT as usize) / 64],
244        })
245    }
246
247    /// Create an encoder seeded from decoded state. Used by
248    /// [`Decoder::into_encoder`](crate::decoder::Decoder::into_encoder).
249    pub(crate) fn from_decoder(
250        mut registry: SchemaRegistry,
251        string_pool: crate::decoder::StringPool,
252        stack_pool: crate::decoder::StackPool,
253        timestamp_base_ns: u64,
254        writer: W,
255    ) -> Self {
256        let mut pool = FxHashMap::default();
257        let mut next_pool_id: u32 = 0;
258        for (id, value) in string_pool.0.into_iter() {
259            pool.insert(value, id.raw_id());
260            if id.raw_id() >= next_pool_id {
261                next_pool_id = id.raw_id() + 1;
262            }
263        }
264
265        let mut new_stack_pool: FxHashMap<Box<[u64]>, u32> = FxHashMap::default();
266        let mut next_stack_pool_id: u32 = 0;
267        for (id, frames) in stack_pool.0.into_iter() {
268            new_stack_pool.insert(frames.into_boxed_slice(), id.raw_id());
269            if id.raw_id() >= next_stack_pool_id {
270                next_stack_pool_id = id.raw_id() + 1;
271            }
272        }
273
274        let mut schema_ids = FxHashMap::default();
275        for (wire_id, entry) in registry.entries() {
276            schema_ids.insert(SchemaKey::Name(Arc::from(entry.name.as_str())), wire_id);
277        }
278        registry.sync_next_id();
279
280        let mut state = EncodeState::new(writer);
281        state.set_ts_base_unchecked(timestamp_base_ns);
282
283        Self {
284            state,
285            registry,
286            string_pool: pool,
287            next_pool_id,
288            stack_pool: new_stack_pool,
289            next_stack_pool_id,
290            schema_ids,
291            dynamic_schema_cache: FxHashMap::default(),
292            slot_cache: Vec::new(),
293            registered_ids: [0; (crate::STATIC_WIRE_ID_LIMIT as usize) / 64],
294        }
295    }
296
297    /// Consume the encoder and return the inner writer.
298    pub fn into_inner(self) -> W {
299        self.state.writer.into_inner()
300    }
301
302    /// Borrow the inner writer.
303    pub fn as_inner(&self) -> &W {
304        self.state.writer.inner()
305    }
306
307    /// Total bytes written through this encoder (including the file header).
308    pub fn bytes_written(&self) -> u64 {
309        self.state.writer.bytes_written()
310    }
311
312    /// Reset the encoder to a new writer, preserving internal allocations.
313    /// Returns the old writer. Writes a file header to the new writer.
314    pub fn reset_to(&mut self, mut new_writer: W) -> io::Result<W> {
315        codec::encode_header(&mut new_writer)?;
316        self.string_pool.clear();
317        self.next_pool_id = 0;
318        self.stack_pool.clear();
319        self.next_stack_pool_id = 0;
320        self.registry.clear();
321        self.schema_ids.clear();
322        self.dynamic_schema_cache.clear();
323        self.slot_cache.fill(0);
324        self.registered_ids.fill(0);
325        // creating a new EncodeState resets the timestamp delta
326        let old_state = std::mem::replace(&mut self.state, EncodeState::new(new_writer));
327        Ok(old_state.writer.into_inner())
328    }
329
330    /// Ensure a schema is registered with this encoder. Returns the wire type
331    /// ID for this encoder's output stream.
332    ///
333    /// Idempotent if the schema matches. Errors if a different schema was
334    /// already registered under the same name.
335    fn ensure_registered(&mut self, schema: &Schema) -> io::Result<WireTypeId> {
336        let identity = Arc::as_ptr(&schema.entry) as usize;
337        if let Some((_, wire_id)) = self.dynamic_schema_cache.get(&identity) {
338            return Ok(*wire_id);
339        }
340        let wire_id = self.ensure_registered_slow(schema)?;
341        if self.dynamic_schema_cache.len() >= DYNAMIC_SCHEMA_CACHE_LIMIT {
342            // Pathological usage (a fresh handle per event); drop the cache
343            // rather than the memory. Well-behaved callers re-enter their
344            // entry on the next event.
345            self.dynamic_schema_cache.clear();
346        }
347        self.dynamic_schema_cache
348            .insert(identity, (Arc::clone(&schema.entry), wire_id));
349        Ok(wire_id)
350    }
351
352    /// Name-keyed registration with collision validation; the slow path
353    /// behind the identity cache above.
354    fn ensure_registered_slow(&mut self, schema: &Schema) -> io::Result<WireTypeId> {
355        let key = SchemaKey::Name(Arc::clone(&schema.name_key));
356        if let Some(&wire_id) = self.schema_ids.get(&key) {
357            // TODO: unify registry and schema_ids to avoid this error case
358            let Some(existing) = self.registry.get(wire_id) else {
359                return Err(io::Error::other(format!(
360                    "corrupted internal state. {wire_id:?} in schema_ids but not in registry."
361                )));
362            };
363            if *existing == *schema.entry {
364                return Ok(wire_id);
365            }
366            return Err(io::Error::new(
367                io::ErrorKind::InvalidInput,
368                format!(
369                    "schema already registered with different definition: {}",
370                    schema.name()
371                ),
372            ));
373        }
374        let id = self.registry.next_type_id();
375        codec::encode_schema(id, &schema.entry, &mut self.state.writer)?;
376        if !schema.entry.annotations.is_empty() {
377            codec::encode_schema_annotations(
378                id,
379                &schema.entry.annotations,
380                &mut self.state.writer,
381            )?;
382        }
383        self.registry
384            .register(id, (*schema.entry).clone())
385            .expect("schema registration failed");
386        self.schema_ids.insert(key, id);
387        Ok(id)
388    }
389
390    /// Register a schema by name. Returns a [`Schema`] handle that can be
391    /// passed to [`write_event`](Self::write_event) (on this or any other
392    /// encoder).
393    ///
394    /// All schemas have timestamps. When writing events, the first element of
395    /// `values` must be `FieldValue::Varint(timestamp_ns)`. It is extracted and
396    /// encoded in the event header (not as a regular field).
397    ///
398    /// Eagerly writes the schema frame. Idempotent if the definition matches.
399    pub fn register_schema(
400        &mut self,
401        name: &str,
402        fields: Vec<crate::schema::FieldDef>,
403    ) -> io::Result<Schema> {
404        let schema = Schema::new(name, fields);
405        self.ensure_registered(&schema)?;
406        Ok(schema)
407    }
408
409    /// Register a pre-built [`Schema`] handle with this encoder.
410    ///
411    /// Eagerly writes the schema frame (and annotation frame if annotations
412    /// are present). Idempotent if the definition matches.
413    pub fn register_existing(&mut self, schema: &Schema) -> io::Result<WireTypeId> {
414        self.ensure_registered(schema)
415    }
416
417    /// Write an event for a schema.
418    ///
419    /// The first element of `values` must be `FieldValue::Varint(timestamp_ns)`
420    /// — it is extracted and encoded in the event header, not as a regular
421    /// field. The remaining values must match the schema's field count.
422    ///
423    /// If this encoder hasn't seen `schema` before, it is auto-registered
424    /// (the schema frame is written before the event).
425    ///
426    /// # Performance
427    ///
428    /// Create the `Schema` once and reuse it (or clones of it) across events.
429    /// Reused handles hit an identity cache; a fresh handle per event falls
430    /// back to registration by name (hash and compare) on every write.
431    pub fn write_event(
432        &mut self,
433        schema: &Schema,
434        values: &[crate::types::FieldValue],
435    ) -> io::Result<()> {
436        use crate::types::FieldValue;
437
438        let type_id = self.ensure_registered(schema)?;
439        let expected_fields = schema.entry.fields.len();
440
441        let ts_ns = match values.first() {
442            Some(FieldValue::Varint(ns)) => *ns,
443            _ => {
444                return Err(io::Error::new(
445                    io::ErrorKind::InvalidInput,
446                    "first value must be FieldValue::Varint(timestamp_ns)",
447                ));
448            }
449        };
450        let field_values = &values[1..];
451
452        if field_values.len() != expected_fields {
453            return Err(io::Error::new(
454                io::ErrorKind::InvalidInput,
455                format!(
456                    "value count ({}) does not match schema field count ({}) for schema '{}'",
457                    field_values.len(),
458                    expected_fields,
459                    schema.name(),
460                ),
461            ));
462        }
463
464        let ts_delta = self.state.encode_timestamp_delta(ts_ns)?;
465        self.state.writer.write_all(&[codec::TAG_EVENT])?;
466        self.state.writer.write_all(&type_id.0.to_le_bytes())?;
467        codec::encode_u24_le(ts_delta, &mut self.state.writer)?;
468        let mut enc = EventEncoder::new(&mut self.state);
469        for (i, v) in field_values.iter().enumerate() {
470            enc.write_field_value(v, schema.entry.fields[i].field_type)?;
471        }
472        Ok(())
473    }
474
475    /// Write a derived TraceEvent. Auto-registers the schema on first call for this type.
476    /// Handles timestamp encoding: emits TimestampReset if needed, packs u24 delta in header.
477    pub fn write<T: TraceEvent + 'static>(&mut self, event: &T) -> io::Result<()> {
478        let slot = T::type_slot();
479        let tid = if slot != 0 && slot < crate::STATIC_WIRE_ID_LIMIT {
480            let word = (slot >> 6) as usize;
481            let bit = 1u64 << (slot & 63);
482            if self.registered_ids[word] & bit == 0 {
483                self.register_fast_id::<T>(slot)?;
484            }
485            WireTypeId(slot)
486        } else {
487            let s = slot as usize;
488            let cached = self.slot_cache.get(s).copied().unwrap_or(0);
489            if cached != 0 {
490                WireTypeId((cached - 1) as u16)
491            } else {
492                self.resolve_dynamic_wire_id::<T>(s)?
493            }
494        };
495        let ts_ns = event.timestamp();
496        let ts_delta = self.state.encode_timestamp_delta(ts_ns)?;
497        self.state.writer.write_all(&[codec::TAG_EVENT])?;
498        self.state.writer.write_all(&tid.0.to_le_bytes())?;
499        codec::encode_u24_le(ts_delta, &mut self.state.writer)?;
500        let mut enc = EventEncoder::new(&mut self.state);
501        event.encode_fields(&mut enc)
502    }
503
504    /// Slow path for `write::<T>`: resolve the wire ID via the schema-ids
505    /// hashmap (registering the schema if needed) and populate the slot cache
506    /// so the next call for the same type takes the fast path.
507    #[cold]
508    fn resolve_dynamic_wire_id<T: TraceEvent + 'static>(
509        &mut self,
510        slot: usize,
511    ) -> io::Result<WireTypeId> {
512        let key = SchemaKey::RustType(TypeId::of::<T>());
513        let tid = if let Some(&existing) = self.schema_ids.get(&key) {
514            existing
515        } else {
516            let schema = Schema::from_entry(T::schema_entry());
517            let id = self.ensure_registered(&schema)?;
518            self.schema_ids.insert(key, id);
519            id
520        };
521        if slot != 0 {
522            if self.slot_cache.len() <= slot {
523                self.slot_cache.resize(slot + 1, 0);
524            }
525            self.slot_cache[slot] = (tid.0 as u32) + 1;
526        }
527        Ok(tid)
528    }
529
530    /// First write of a slot in `1..STATIC_WIRE_ID_LIMIT`: emit the schema frame
531    /// at the slot `id` and mark the bitset, so later writes skip registration.
532    #[cold]
533    fn register_fast_id<T: TraceEvent + 'static>(&mut self, id: u16) -> io::Result<()> {
534        let entry = T::schema_entry();
535        let wire = WireTypeId(id);
536        codec::encode_schema(wire, &entry, &mut self.state.writer)?;
537        if !entry.annotations.is_empty() {
538            codec::encode_schema_annotations(wire, &entry.annotations, &mut self.state.writer)?;
539        }
540        self.registry.register(wire, entry).map_err(|e| {
541            io::Error::new(
542                io::ErrorKind::InvalidInput,
543                format!("wire id {id} collision: {e}"),
544            )
545        })?;
546        // mark id registered: set bit (id % 64) in word (id / 64)
547        self.registered_ids[(id >> 6) as usize] |= 1u64 << (id & 63);
548        Ok(())
549    }
550
551    /// Intern a string, emitting a pool frame if new. Returns an [`InternedString`] handle.
552    pub fn intern_string(&mut self, s: &str) -> io::Result<InternedString> {
553        if let Some(&id) = self.string_pool.get(s) {
554            return Ok(InternedString(id));
555        }
556        let id = self.next_pool_id;
557        self.next_pool_id += 1;
558        self.string_pool.insert(s.to_string(), id);
559        codec::encode_string_pool(
560            &[PoolEntry {
561                pool_id: id,
562                data: s.as_bytes().to_vec(),
563            }],
564            &mut self.state.writer,
565        )?;
566        Ok(InternedString(id))
567    }
568
569    pub fn write_string_pool(&mut self, entries: &[PoolEntry]) -> io::Result<()> {
570        codec::encode_string_pool(entries, &mut self.state.writer)
571    }
572
573    /// Intern a stack-frame vector, emitting a stack-pool frame if new.
574    /// Returns an [`InternedStackFrames`] handle.
575    pub fn intern_stack_frames(&mut self, frames: &[u64]) -> io::Result<InternedStackFrames> {
576        if let Some(&id) = self.stack_pool.get(frames) {
577            return Ok(InternedStackFrames(id));
578        }
579        let id = self.next_stack_pool_id;
580        self.next_stack_pool_id += 1;
581        self.stack_pool.insert(frames.into(), id);
582        codec::encode_stack_pool(
583            &[StackPoolEntry {
584                pool_id: id,
585                // TODO: allow `StackPoolEntry` to have borrowed frames avoiding the unecessary clone here
586                // https://github.com/dial9-rs/dial9-tokio-telemetry/issues/358
587                frames: frames.to_vec(),
588            }],
589            &mut self.state.writer,
590        )?;
591        Ok(InternedStackFrames(id))
592    }
593
594    pub fn write_stack_pool(&mut self, entries: &[StackPoolEntry]) -> io::Result<()> {
595        codec::encode_stack_pool(entries, &mut self.state.writer)
596    }
597
598    /// Flush the underlying writer.
599    pub fn flush(&mut self) -> io::Result<()> {
600        self.state.writer.flush()
601    }
602
603    /// Convert this encoder into a [`RawEncoder`] that only supports writing
604    /// pre-encoded bytes. The byte count is preserved so rotation decisions
605    /// remain correct.
606    ///
607    /// Use this after writing any structured data (headers, segment metadata)
608    /// to switch to a raw-only mode for appending pre-encoded batches.
609    pub fn into_raw_encoder(self) -> RawEncoder<W> {
610        RawEncoder {
611            writer: self.state.writer,
612        }
613    }
614}
615
616/// A write-only encoder that accepts pre-encoded bytes.
617///
618/// Created by [`Encoder::into_raw_encoder`] after the file header and any
619/// structured metadata have been written. Carries no schema registry, string
620/// pool, or timestamp state — it simply forwards bytes to the underlying
621/// writer while tracking the total byte count.
622pub struct RawEncoder<W> {
623    writer: CountingWriter<W>,
624}
625
626impl<W: Write> RawEncoder<W> {
627    /// Write pre-encoded bytes to the underlying writer.
628    pub fn write_raw(&mut self, bytes: &[u8]) -> io::Result<()> {
629        self.writer.write_all(bytes)
630    }
631
632    /// Total bytes written (including bytes written by the [`Encoder`] before
633    /// conversion).
634    pub fn bytes_written(&self) -> u64 {
635        self.writer.bytes_written()
636    }
637
638    /// Flush the underlying writer.
639    pub fn flush(&mut self) -> io::Result<()> {
640        self.writer.flush()
641    }
642
643    /// Consume the raw encoder and return the inner writer.
644    pub fn into_inner(self) -> W {
645        self.writer.into_inner()
646    }
647}
648
649impl Encoder<Vec<u8>> {
650    pub fn write_infallible<T: TraceEvent + 'static>(&mut self, event: &T) {
651        self.write(event).expect("writing to Vec<u8> is infallible")
652    }
653
654    pub fn intern_string_infallible(&mut self, s: &str) -> InternedString {
655        self.intern_string(s)
656            .expect("interning into Vec<u8> is infallible")
657    }
658
659    pub fn intern_stack_frames_infallible(&mut self, frames: &[u64]) -> InternedStackFrames {
660        self.intern_stack_frames(frames)
661            .expect("interning into Vec<u8> is infallible")
662    }
663
664    /// Resets the encoder to point to a new backing Vec returning the old one
665    pub fn reset_to_infallible(&mut self, data: Vec<u8>) -> Vec<u8> {
666        self.reset_to(data)
667            .expect("writing to Vec<u8> is infallible")
668    }
669}
670
671#[cfg(test)]
672mod tests {
673    use super::*;
674    use crate::schema::FieldDef;
675    use crate::types::{FieldType, FieldValue};
676
677    #[test]
678    fn encoder_writes_header() {
679        let enc = Encoder::new();
680        let data = enc.finish();
681        assert_eq!(&data[..5], &[0x54, 0x52, 0x43, 0x00, 1]);
682    }
683
684    #[test]
685    fn encoder_register_and_write_event() {
686        let mut enc = Encoder::new();
687        let schema = enc
688            .register_schema(
689                "Ev",
690                vec![FieldDef {
691                    name: "v".into(),
692                    field_type: FieldType::Varint,
693                }],
694            )
695            .unwrap();
696        enc.write_event(
697            &schema,
698            &[FieldValue::Varint(1_000), FieldValue::Varint(42)],
699        )
700        .unwrap();
701        let data = enc.finish();
702        assert!(data.len() > 5);
703    }
704
705    #[test]
706    fn idempotent_re_registration() {
707        let mut enc = Encoder::new();
708        let fields = vec![FieldDef {
709            name: "v".into(),
710            field_type: FieldType::Varint,
711        }];
712        let _s1 = enc.register_schema("Ev", fields.clone()).unwrap();
713        let _s2 = enc.register_schema("Ev", fields).unwrap();
714        // Both succeed — same schema, same name
715    }
716
717    #[test]
718    fn re_registration_different_schema_errors() {
719        let mut enc = Encoder::new();
720        enc.register_schema(
721            "Ev",
722            vec![FieldDef {
723                name: "v".into(),
724                field_type: FieldType::Varint,
725            }],
726        )
727        .unwrap();
728        let result = enc.register_schema(
729            "Ev",
730            vec![FieldDef {
731                name: "different".into(),
732                field_type: FieldType::Bool,
733            }],
734        );
735        assert!(result.is_err());
736    }
737
738    #[test]
739    fn schema_auto_registers_on_write() {
740        use crate::decoder::{DecodedFrame, Decoder};
741
742        // Create a schema without an encoder
743        let schema = Schema::new(
744            "Lazy",
745            vec![FieldDef {
746                name: "v".into(),
747                field_type: FieldType::Varint,
748            }],
749        );
750
751        // Write to an encoder that hasn't seen this schema — auto-registers
752        let mut enc = Encoder::new();
753        enc.write_event(
754            &schema,
755            &[FieldValue::Varint(1_000), FieldValue::Varint(42)],
756        )
757        .unwrap();
758
759        let bytes = enc.finish();
760        let mut dec = Decoder::new(&bytes).unwrap();
761        let frames = dec.decode_all();
762        assert!(matches!(&frames[0], DecodedFrame::Schema(s) if s.name == "Lazy"));
763        if let DecodedFrame::Event { values, .. } = &frames[1] {
764            assert_eq!(*values, vec![FieldValue::Varint(42)]);
765        } else {
766            panic!("expected event");
767        }
768    }
769
770    #[test]
771    fn schema_portable_across_encoders() {
772        use crate::decoder::{DecodedFrame, Decoder};
773
774        let mut enc1 = Encoder::new();
775        let schema = enc1
776            .register_schema(
777                "Shared",
778                vec![FieldDef {
779                    name: "v".into(),
780                    field_type: FieldType::Varint,
781                }],
782            )
783            .unwrap();
784        enc1.write_event(&schema, &[FieldValue::Varint(1_000), FieldValue::Varint(1)])
785            .unwrap();
786
787        // Pass the same Schema to a different encoder
788        let mut enc2 = Encoder::new();
789        enc2.write_event(&schema, &[FieldValue::Varint(2_000), FieldValue::Varint(2)])
790            .unwrap();
791
792        // Both encoders produce valid output
793        for (enc, expected_val) in [(enc1, 1u64), (enc2, 2u64)] {
794            let bytes = enc.finish();
795            let mut dec = Decoder::new(&bytes).unwrap();
796            let frames = dec.decode_all();
797            let event = frames
798                .iter()
799                .find(|f| matches!(f, DecodedFrame::Event { .. }))
800                .unwrap();
801            if let DecodedFrame::Event { values, .. } = event {
802                assert_eq!(values[0], FieldValue::Varint(expected_val));
803            }
804        }
805    }
806
807    #[test]
808    fn encoder_intern_string_deduplicates() {
809        let mut enc = Encoder::new();
810        let id1 = enc.intern_string("hello").unwrap();
811        let id2 = enc.intern_string("hello").unwrap();
812        let id3 = enc.intern_string("world").unwrap();
813        assert_eq!(id1, id2);
814        assert_ne!(id1, id3);
815    }
816
817    #[test]
818    fn encoder_intern_stack_frames_deduplicates() {
819        let mut enc = Encoder::new();
820        let stack_a: &[u64] = &[0x1000, 0x2000, 0x3000];
821        let stack_b: &[u64] = &[0x4000, 0x5000];
822        let id1 = enc.intern_stack_frames(stack_a).unwrap();
823        let id2 = enc.intern_stack_frames(stack_a).unwrap();
824        let id3 = enc.intern_stack_frames(stack_b).unwrap();
825        assert_eq!(id1, id2);
826        assert_ne!(id1, id3);
827    }
828
829    #[test]
830    fn stack_pool_round_trip_via_decoder() {
831        use crate::decoder::Decoder;
832        use crate::types::InternedStackFrames;
833
834        let mut enc = Encoder::new();
835        let stack_a: &[u64] = &[0xdead, 0xbeef, 0xcafe];
836        let stack_b: &[u64] = &[0x1, 0x2];
837        let id_a = enc.intern_stack_frames(stack_a).unwrap();
838        let id_b = enc.intern_stack_frames(stack_b).unwrap();
839        let bytes = enc.finish();
840
841        let mut dec = Decoder::new(&bytes).unwrap();
842        let _ = dec.decode_all();
843        assert_eq!(
844            dec.stack_pool().get(InternedStackFrames(id_a.raw_id())),
845            Some(stack_a)
846        );
847        assert_eq!(
848            dec.stack_pool().get(InternedStackFrames(id_b.raw_id())),
849            Some(stack_b)
850        );
851    }
852
853    #[test]
854    fn for_each_event_populates_stack_pool() {
855        use crate::decoder::Decoder;
856        use crate::schema::FieldDef;
857        use crate::types::{FieldType, FieldValue, InternedStackFrames};
858
859        let mut enc = Encoder::new();
860        let schema = enc
861            .register_schema(
862                "CpuSampleEvent",
863                vec![FieldDef {
864                    name: "callchain".into(),
865                    field_type: FieldType::PooledStackFrames,
866                }],
867            )
868            .unwrap();
869        let stack: &[u64] = &[0x1234, 0x5678, 0x9abc];
870        let id = enc.intern_stack_frames(stack).unwrap();
871        enc.write_event(
872            &schema,
873            &[
874                FieldValue::Varint(1_000_000),
875                FieldValue::PooledStackFrames(id),
876            ],
877        )
878        .unwrap();
879        let bytes = enc.finish();
880
881        let mut dec = Decoder::new(&bytes).unwrap();
882        let mut event_count = 0;
883        dec.for_each_event(|_ev| {
884            event_count += 1;
885        })
886        .unwrap();
887        assert_eq!(event_count, 1);
888        assert_eq!(
889            dec.stack_pool().get(InternedStackFrames(id.raw_id())),
890            Some(stack),
891        );
892    }
893
894    #[test]
895    fn encoder_intern_empty_stack_frames() {
896        use crate::decoder::Decoder;
897        use crate::types::InternedStackFrames;
898
899        let mut enc = Encoder::new();
900        let id1 = enc.intern_stack_frames(&[]).unwrap();
901        let id2 = enc.intern_stack_frames(&[]).unwrap();
902        assert_eq!(id1, id2);
903        let bytes = enc.finish();
904
905        let mut dec = Decoder::new(&bytes).unwrap();
906        let _ = dec.decode_all();
907        assert_eq!(
908            dec.stack_pool().get(InternedStackFrames(id1.raw_id())),
909            Some(&[][..])
910        );
911    }
912
913    #[test]
914    fn write_stack_pool_multi_entry_round_trip() {
915        use crate::decoder::Decoder;
916        use crate::types::InternedStackFrames;
917
918        let mut enc = Encoder::new();
919        let entries = vec![
920            StackPoolEntry {
921                pool_id: 0,
922                frames: vec![0xaaaa, 0xbbbb, 0xcccc],
923            },
924            StackPoolEntry {
925                pool_id: 1,
926                frames: vec![0x1111],
927            },
928            StackPoolEntry {
929                pool_id: 2,
930                frames: vec![],
931            },
932        ];
933        enc.write_stack_pool(&entries).unwrap();
934        let bytes = enc.finish();
935
936        let mut dec = Decoder::new(&bytes).unwrap();
937        let _ = dec.decode_all();
938        assert_eq!(
939            dec.stack_pool().get(InternedStackFrames(0)),
940            Some(&[0xaaaa, 0xbbbb, 0xcccc][..])
941        );
942        assert_eq!(
943            dec.stack_pool().get(InternedStackFrames(1)),
944            Some(&[0x1111][..])
945        );
946        assert_eq!(dec.stack_pool().get(InternedStackFrames(2)), Some(&[][..]));
947    }
948
949    #[test]
950    fn decoder_into_encoder_deduplicates_interned_stack_frames() {
951        use crate::decoder::Decoder;
952
953        let mut enc = Encoder::new();
954        let id1 = enc.intern_stack_frames(&[0x10, 0x20]).unwrap();
955        let base = enc.finish();
956
957        let mut decoder = Decoder::new(&base).unwrap();
958        while decoder.next_frame_ref().ok().flatten().is_some() {}
959        let mut output = Vec::new();
960        let mut ext = decoder.into_encoder(&mut output);
961        let id2 = ext.intern_stack_frames(&[0x10, 0x20]).unwrap();
962        let id3 = ext.intern_stack_frames(&[0x30]).unwrap();
963        assert_eq!(id1.raw_id(), id2.raw_id());
964        assert_ne!(id2.raw_id(), id3.raw_id());
965    }
966
967    #[test]
968    fn timestamp_round_trip() {
969        use crate::decoder::{DecodedFrame, Decoder};
970
971        let mut enc = Encoder::new();
972        let schema = enc
973            .register_schema(
974                "TS",
975                vec![FieldDef {
976                    name: "v".into(),
977                    field_type: FieldType::Varint,
978                }],
979            )
980            .unwrap();
981
982        let ts1 = 100_000u64;
983        let ts2 = 50_000u64;
984        let ts3 = 200_000_000u64;
985        let ts4 = 100_000_000u64;
986        enc.write_event(&schema, &[FieldValue::Varint(ts1), FieldValue::Varint(1)])
987            .unwrap();
988        enc.write_event(&schema, &[FieldValue::Varint(ts2), FieldValue::Varint(2)])
989            .unwrap();
990        enc.write_event(&schema, &[FieldValue::Varint(ts3), FieldValue::Varint(3)])
991            .unwrap();
992        enc.write_event(&schema, &[FieldValue::Varint(ts4), FieldValue::Varint(4)])
993            .unwrap();
994
995        let bytes = enc.finish();
996        let mut dec = Decoder::new(&bytes).unwrap();
997        let events: Vec<_> = dec
998            .decode_all()
999            .into_iter()
1000            .filter_map(|f| match f {
1001                DecodedFrame::Event {
1002                    timestamp_ns,
1003                    values,
1004                    ..
1005                } => Some((timestamp_ns, values)),
1006                _ => None,
1007            })
1008            .collect();
1009
1010        assert_eq!(events.len(), 4);
1011        assert_eq!(events[0].0, Some(ts1));
1012        assert_eq!(events[0].1, vec![FieldValue::Varint(1)]);
1013        assert_eq!(events[1].0, Some(ts2));
1014        assert_eq!(events[1].1, vec![FieldValue::Varint(2)]);
1015        assert_eq!(events[2].0, Some(ts3));
1016        assert_eq!(events[2].1, vec![FieldValue::Varint(3)]);
1017        assert_eq!(events[3].0, Some(ts4));
1018        assert_eq!(events[3].1, vec![FieldValue::Varint(4)]);
1019    }
1020
1021    #[test]
1022    fn encoder_new_to_writer() {
1023        let mut buf = Vec::new();
1024        let enc = Encoder::new_to(&mut buf).unwrap();
1025        drop(enc);
1026        assert!(buf.len() >= 5);
1027        assert_eq!(&buf[..5], &[0x54, 0x52, 0x43, 0x00, 1]);
1028    }
1029
1030    #[test]
1031    fn decoder_into_encoder_appends_without_header() {
1032        use crate::decoder::{DecodedFrame, Decoder};
1033
1034        // Create a trace with a header, a schema, and an event
1035        let mut enc = Encoder::new();
1036        let schema = enc
1037            .register_schema(
1038                "Ev",
1039                vec![FieldDef {
1040                    name: "v".into(),
1041                    field_type: FieldType::Varint,
1042                }],
1043            )
1044            .unwrap();
1045        enc.write_event(&schema, &[FieldValue::Varint(1_000), FieldValue::Varint(1)])
1046            .unwrap();
1047        let base = enc.finish();
1048
1049        // Decode all frames, then convert into an encoder that appends to output
1050        let mut decoder = Decoder::new(&base).unwrap();
1051        while decoder.next_frame_ref().ok().flatten().is_some() {}
1052        let mut output = Vec::new();
1053        let mut ext = decoder.into_encoder(&mut output);
1054        // Schema "Ev" is already known — no duplicate schema frame emitted
1055        ext.write_event(&schema, &[FieldValue::Varint(2_000), FieldValue::Varint(2)])
1056            .unwrap();
1057        drop(ext);
1058
1059        // Concatenate and decode
1060        let mut combined = base.clone();
1061        combined.extend_from_slice(&output);
1062        let mut dec = Decoder::new(&combined).unwrap();
1063        let events: Vec<_> = dec
1064            .decode_all()
1065            .into_iter()
1066            .filter_map(|f| match f {
1067                DecodedFrame::Event {
1068                    timestamp_ns,
1069                    values,
1070                    ..
1071                } => Some((timestamp_ns, values)),
1072                _ => None,
1073            })
1074            .collect();
1075        assert_eq!(events.len(), 2);
1076        assert_eq!(events[0].0, Some(1_000));
1077        assert_eq!(events[1].0, Some(2_000));
1078    }
1079
1080    #[test]
1081    fn decoder_into_encoder_deduplicates_interned_strings() {
1082        use crate::decoder::{DecodedFrame, Decoder};
1083
1084        // Create a trace with an interned string
1085        let mut enc = Encoder::new();
1086        let id1 = enc.intern_string("hello").unwrap();
1087        let base = enc.finish();
1088
1089        // Decode all frames, then convert into an encoder
1090        let mut decoder = Decoder::new(&base).unwrap();
1091        while decoder.next_frame_ref().ok().flatten().is_some() {}
1092        let mut output = Vec::new();
1093        let mut ext = decoder.into_encoder(&mut output);
1094        // "hello" is already interned, should reuse the same ID
1095        let id2 = ext.intern_string("hello").unwrap();
1096        let id3 = ext.intern_string("world").unwrap();
1097        drop(ext);
1098
1099        assert_eq!(id1, id2, "existing string should reuse pool ID");
1100        assert_ne!(id2, id3);
1101
1102        // "hello" should not produce a new StringPool frame; "world" should
1103        let mut combined = base.clone();
1104        combined.extend_from_slice(&output);
1105        let mut dec = Decoder::new(&combined).unwrap();
1106        let frames = dec.decode_all();
1107        let pool_frames: Vec<_> = frames
1108            .iter()
1109            .filter(|f| matches!(f, DecodedFrame::StringPool(_)))
1110            .collect();
1111        // One from the base trace ("hello"), one from extend ("world")
1112        assert_eq!(pool_frames.len(), 2);
1113    }
1114
1115    /// Minimal hand-rolled `TraceEvent` with a fixed fast-path slot, so the
1116    /// encoder tests don't depend on the derive crate.
1117    struct FastSlot {
1118        ts: u64,
1119    }
1120    impl TraceEvent for FastSlot {
1121        fn type_slot() -> u16 {
1122            5
1123        }
1124        fn event_name() -> &'static str {
1125            "FastSlot"
1126        }
1127        fn field_defs() -> Vec<FieldDef> {
1128            Vec::new()
1129        }
1130        fn timestamp(&self) -> u64 {
1131            self.ts
1132        }
1133        fn encode_fields<W: Write>(&self, _enc: &mut EventEncoder<'_, W>) -> io::Result<()> {
1134            Ok(())
1135        }
1136    }
1137
1138    #[test]
1139    fn fast_slot_registers_at_slot_id() {
1140        use crate::decoder::Decoder;
1141
1142        let mut enc = Encoder::new();
1143        enc.write(&FastSlot { ts: 1_000_000 }).unwrap();
1144        // A plain dynamic schema must land in the dynamic range, above slots.
1145        let dynamic = enc
1146            .register_schema(
1147                "Dyn",
1148                vec![FieldDef {
1149                    name: "v".into(),
1150                    field_type: FieldType::Varint,
1151                }],
1152            )
1153            .unwrap();
1154        enc.write_event(
1155            &dynamic,
1156            &[FieldValue::Varint(2_000), FieldValue::Varint(1)],
1157        )
1158        .unwrap();
1159        let bytes = enc.finish();
1160
1161        let mut dec = Decoder::new(&bytes).unwrap();
1162        let _ = dec.decode_all();
1163        // Fast-path event registered at its slot id.
1164        assert_eq!(
1165            dec.registry().get(WireTypeId(5)).unwrap().name(),
1166            "FastSlot"
1167        );
1168        // Dynamic schema sits at STATIC_WIRE_ID_LIMIT, not colliding with slots.
1169        assert_eq!(
1170            dec.registry()
1171                .get(WireTypeId(crate::STATIC_WIRE_ID_LIMIT))
1172                .unwrap()
1173                .name(),
1174            "Dyn"
1175        );
1176    }
1177
1178    #[test]
1179    fn register_and_write() {
1180        use crate::decoder::{DecodedFrame, Decoder};
1181
1182        let mut enc = Encoder::new();
1183        let schema = enc
1184            .register_schema(
1185                "MyEvent",
1186                vec![
1187                    FieldDef {
1188                        name: "count".into(),
1189                        field_type: FieldType::Varint,
1190                    },
1191                    FieldDef {
1192                        name: "name".into(),
1193                        field_type: FieldType::String,
1194                    },
1195                ],
1196            )
1197            .unwrap();
1198
1199        enc.write_event(
1200            &schema,
1201            &[
1202                FieldValue::Varint(1_000_000),
1203                FieldValue::Varint(42),
1204                FieldValue::String("hello".into()),
1205            ],
1206        )
1207        .unwrap();
1208
1209        let bytes = enc.finish();
1210        let mut dec = Decoder::new(&bytes).unwrap();
1211        let frames = dec.decode_all();
1212        let events: Vec<_> = frames
1213            .into_iter()
1214            .filter_map(|f| match f {
1215                DecodedFrame::Event {
1216                    timestamp_ns,
1217                    values,
1218                    ..
1219                } => Some((timestamp_ns, values)),
1220                _ => None,
1221            })
1222            .collect();
1223        assert_eq!(events.len(), 1);
1224        assert_eq!(events[0].0, Some(1_000_000));
1225        assert_eq!(events[0].1[0], FieldValue::Varint(42));
1226        assert_eq!(events[0].1[1], FieldValue::String("hello".into()));
1227    }
1228
1229    #[test]
1230    fn register_conflict_errors() {
1231        let mut enc = Encoder::new();
1232        enc.register_schema(
1233            "Ev",
1234            vec![FieldDef {
1235                name: "v".into(),
1236                field_type: FieldType::Varint,
1237            }],
1238        )
1239        .unwrap();
1240        let result = enc.register_schema(
1241            "Ev",
1242            vec![FieldDef {
1243                name: "other".into(),
1244                field_type: FieldType::Bool,
1245            }],
1246        );
1247        assert!(result.is_err());
1248    }
1249
1250    #[test]
1251    fn write_wrong_field_count_errors() {
1252        let mut enc = Encoder::new();
1253        let schema = enc
1254            .register_schema(
1255                "Ev",
1256                vec![FieldDef {
1257                    name: "v".into(),
1258                    field_type: FieldType::Varint,
1259                }],
1260            )
1261            .unwrap();
1262        // Pass 3 values (ts + 2 fields) for a 1-field schema
1263        let result = enc.write_event(
1264            &schema,
1265            &[
1266                FieldValue::Varint(0),
1267                FieldValue::Varint(1),
1268                FieldValue::Varint(2),
1269            ],
1270        );
1271        assert!(result.is_err());
1272    }
1273
1274    /// Verify that the encoder advances the timestamp base after each event,
1275    /// producing inter-event deltas rather than base-relative deltas.
1276    #[test]
1277    fn timestamp_base_advances_per_event() {
1278        use crate::decoder::{DecodedFrame, Decoder};
1279
1280        let mut enc = Encoder::new();
1281        let schema = enc
1282            .register_schema(
1283                "Ev",
1284                vec![FieldDef {
1285                    name: "v".into(),
1286                    field_type: FieldType::Varint,
1287                }],
1288            )
1289            .unwrap();
1290
1291        let ts1 = 12_000_000u64;
1292        let ts2 = 24_000_000u64;
1293        enc.write_event(&schema, &[FieldValue::Varint(ts1), FieldValue::Varint(1)])
1294            .unwrap();
1295        enc.write_event(&schema, &[FieldValue::Varint(ts2), FieldValue::Varint(2)])
1296            .unwrap();
1297
1298        let bytes = enc.finish();
1299
1300        let reset_count = bytes.iter().filter(|&&b| b == 0x05).count();
1301        assert_eq!(
1302            reset_count, 0,
1303            "base should advance per event, avoiding unnecessary resets"
1304        );
1305
1306        let mut dec = Decoder::new(&bytes).unwrap();
1307        let events: Vec<_> = dec
1308            .decode_all()
1309            .into_iter()
1310            .filter_map(|f| match f {
1311                DecodedFrame::Event { timestamp_ns, .. } => timestamp_ns,
1312                _ => None,
1313            })
1314            .collect();
1315        assert_eq!(events, vec![ts1, ts2]);
1316    }
1317
1318    #[test]
1319    fn reset_to_preserves_capacity() {
1320        let mut enc = Encoder::new();
1321        for i in 0..100 {
1322            enc.intern_string(&format!("string_{}", i)).unwrap();
1323        }
1324        let cap_before = enc.string_pool.capacity();
1325        let _bytes = enc.reset_to(Vec::new());
1326        let cap_after = enc.string_pool.capacity();
1327        assert_eq!(
1328            cap_before, cap_after,
1329            "string_pool capacity should be preserved after reset_to"
1330        );
1331    }
1332
1333    #[test]
1334    fn reset_to_returns_old_data_and_clears_state() {
1335        use crate::decoder::{DecodedFrame, Decoder};
1336
1337        let mut enc = Encoder::new();
1338        let schema = enc
1339            .register_schema(
1340                "Ev",
1341                vec![FieldDef {
1342                    name: "v".into(),
1343                    field_type: FieldType::Varint,
1344                }],
1345            )
1346            .unwrap();
1347        enc.write_event(
1348            &schema,
1349            &[FieldValue::Varint(1_000), FieldValue::Varint(42)],
1350        )
1351        .unwrap();
1352        let _s = enc.intern_string("hello").unwrap();
1353
1354        let old_bytes_written = enc.bytes_written();
1355        assert!(old_bytes_written > 0);
1356
1357        // --- reset ---
1358        let old = enc.reset_to_infallible(Vec::new());
1359
1360        // Invariant 1: old writer contains the data we wrote (decodable)
1361        let mut dec = Decoder::new(&old).unwrap();
1362        let frames = dec.decode_all();
1363        assert!(frames.iter().any(|f| matches!(f, DecodedFrame::Schema(_))));
1364        assert!(
1365            frames
1366                .iter()
1367                .any(|f| matches!(f, DecodedFrame::Event { .. }))
1368        );
1369        assert!(
1370            frames
1371                .iter()
1372                .any(|f| matches!(f, DecodedFrame::StringPool(_)))
1373        );
1374
1375        // Invariant 2: bytes_written resets to just the header size
1376        assert!(
1377            enc.bytes_written() < old_bytes_written,
1378            "bytes_written should reset (got {} vs old {})",
1379            enc.bytes_written(),
1380            old_bytes_written
1381        );
1382
1383        // Invariant 3: schemas are cleared — same schema must re-register
1384        // (write_event auto-registers, so we verify a new schema frame appears)
1385        enc.write_event(
1386            &schema,
1387            &[FieldValue::Varint(2_000), FieldValue::Varint(99)],
1388        )
1389        .unwrap();
1390
1391        // Invariant 4: string pool is cleared — re-interning emits a new pool frame
1392        let _s2 = enc.intern_string("hello").unwrap();
1393
1394        // Invariant 5: new output is a valid standalone trace
1395        let new_bytes = enc.reset_to_infallible(Vec::new());
1396        let mut dec2 = Decoder::new(&new_bytes).unwrap();
1397        let new_frames = dec2.decode_all();
1398        // Must have its own schema definition (not relying on old encoder state)
1399        assert!(
1400            new_frames
1401                .iter()
1402                .any(|f| matches!(f, DecodedFrame::Schema(s) if s.name == "Ev")),
1403            "new trace must contain schema definition"
1404        );
1405        // Must have its own string pool entry
1406        assert!(
1407            new_frames
1408                .iter()
1409                .any(|f| matches!(f, DecodedFrame::StringPool(_))),
1410            "new trace must contain string pool"
1411        );
1412        // Event must decode with correct timestamp (timestamp_base was reset)
1413        let event = new_frames
1414            .iter()
1415            .find_map(|f| match f {
1416                DecodedFrame::Event {
1417                    timestamp_ns,
1418                    values,
1419                    ..
1420                } => Some((timestamp_ns, values)),
1421                _ => None,
1422            })
1423            .expect("new trace must contain event");
1424        assert_eq!(*event.0, Some(2_000));
1425        assert_eq!(event.1[0], FieldValue::Varint(99));
1426    }
1427
1428    #[test]
1429    fn into_raw_encoder_preserves_byte_count() {
1430        let mut enc = Encoder::new();
1431        let schema = enc
1432            .register_schema(
1433                "Ev",
1434                vec![FieldDef {
1435                    name: "v".into(),
1436                    field_type: FieldType::Varint,
1437                }],
1438            )
1439            .unwrap();
1440        enc.write_event(
1441            &schema,
1442            &[FieldValue::Varint(1_000), FieldValue::Varint(42)],
1443        )
1444        .unwrap();
1445
1446        let bytes_before = enc.bytes_written();
1447        assert!(bytes_before > 0);
1448
1449        let raw = enc.into_raw_encoder();
1450        assert_eq!(
1451            raw.bytes_written(),
1452            bytes_before,
1453            "byte count must be preserved across conversion"
1454        );
1455    }
1456
1457    #[test]
1458    fn raw_encoder_write_raw_and_bytes_written() {
1459        let enc = Encoder::new();
1460        let initial = enc.bytes_written();
1461        let mut raw = enc.into_raw_encoder();
1462
1463        let payload = [0xAA; 100];
1464        raw.write_raw(&payload).unwrap();
1465
1466        assert_eq!(
1467            raw.bytes_written(),
1468            initial + payload.len() as u64,
1469            "bytes_written must include raw payload"
1470        );
1471    }
1472
1473    #[test]
1474    fn raw_encoder_into_inner_returns_all_data() {
1475        use crate::decoder::{DecodedFrame, Decoder};
1476
1477        // Write a structured event via Encoder, then append a raw batch
1478        // via RawEncoder, and verify the combined output decodes correctly.
1479        let mut enc = Encoder::new();
1480        let schema = enc
1481            .register_schema(
1482                "Ev",
1483                vec![FieldDef {
1484                    name: "v".into(),
1485                    field_type: FieldType::Varint,
1486                }],
1487            )
1488            .unwrap();
1489        enc.write_event(&schema, &[FieldValue::Varint(1_000), FieldValue::Varint(1)])
1490            .unwrap();
1491
1492        // Build a raw batch with the same schema
1493        let raw_batch = {
1494            let mut batch_enc = Encoder::new();
1495            batch_enc
1496                .write_event(&schema, &[FieldValue::Varint(2_000), FieldValue::Varint(2)])
1497                .unwrap();
1498            batch_enc.finish()
1499        };
1500
1501        let mut raw = enc.into_raw_encoder();
1502        raw.write_raw(&raw_batch).unwrap();
1503        let combined = raw.into_inner();
1504
1505        let mut dec = Decoder::new(&combined).unwrap();
1506        let events: Vec<_> = dec
1507            .decode_all()
1508            .into_iter()
1509            .filter_map(|f| match f {
1510                DecodedFrame::Event {
1511                    timestamp_ns,
1512                    values,
1513                    ..
1514                } => Some((timestamp_ns, values)),
1515                _ => None,
1516            })
1517            .collect();
1518
1519        assert_eq!(events.len(), 2);
1520        assert_eq!(events[0].0, Some(1_000));
1521        assert_eq!(events[0].1, vec![FieldValue::Varint(1)]);
1522        assert_eq!(events[1].0, Some(2_000));
1523        assert_eq!(events[1].1, vec![FieldValue::Varint(2)]);
1524    }
1525}
1526
1527#[cfg(test)]
1528mod dynamic_schema_cache_tests {
1529    use super::*;
1530    use crate::schema::{FieldDef, SchemaEntry};
1531    use crate::types::{FieldType, FieldValue};
1532
1533    fn schema(name: &str) -> Schema {
1534        Schema::from_entry(SchemaEntry::new(
1535            name,
1536            /* has_timestamp */ true,
1537            vec![FieldDef::new("v", FieldType::Varint)],
1538        ))
1539    }
1540
1541    /// The same handle re-registers through the identity cache, and a fresh
1542    /// handle for the same name resolves to the same wire id through the
1543    /// slow path (then caches its own identity).
1544    #[test]
1545    fn identity_cache_agrees_with_name_registration() {
1546        let mut enc = Encoder::new();
1547        let a = schema("Ev");
1548        let id1 = enc.ensure_registered(&a).unwrap();
1549        let id2 = enc.ensure_registered(&a).unwrap();
1550        assert_eq!(id1, id2, "same handle must reuse its wire id");
1551
1552        let b = schema("Ev"); // distinct allocation, same layout and name
1553        let id3 = enc.ensure_registered(&b).unwrap();
1554        assert_eq!(id1, id3, "same name must resolve to the same wire id");
1555        assert_eq!(enc.dynamic_schema_cache.len(), 2);
1556    }
1557
1558    /// A caller minting a fresh handle per event must not grow the cache
1559    /// (and pin schema Arcs) without bound.
1560    #[test]
1561    fn identity_cache_is_bounded() {
1562        let mut enc = Encoder::new();
1563        for i in 0..(DYNAMIC_SCHEMA_CACHE_LIMIT * 2 + 7) {
1564            // Cycle a few names so both fresh-per-event and fresh-name
1565            // shapes are covered; every handle is a distinct allocation.
1566            let s = schema(&format!("Ev{}", i % 3));
1567            enc.ensure_registered(&s).unwrap();
1568            assert!(
1569                enc.dynamic_schema_cache.len() <= DYNAMIC_SCHEMA_CACHE_LIMIT,
1570                "cache exceeded its bound at iteration {i}"
1571            );
1572        }
1573    }
1574
1575    /// Clearing the cache must not affect decodability: events written
1576    /// before and after the flush decode against one schema.
1577    #[test]
1578    fn events_across_cache_clears_decode() {
1579        let mut enc = Encoder::new();
1580        for i in 0..(DYNAMIC_SCHEMA_CACHE_LIMIT + 3) {
1581            let s = schema("Ev");
1582            // Timestamp plus the one schema field.
1583            enc.write_event(
1584                &s,
1585                &[FieldValue::Varint(i as u64), FieldValue::Varint(i as u64)],
1586            )
1587            .unwrap();
1588        }
1589        let data = enc.finish();
1590        let mut decoder = crate::decoder::Decoder::new(&data).unwrap();
1591        let mut count = 0u64;
1592        decoder
1593            .for_each_event(|ev| {
1594                assert_eq!(ev.name, "Ev");
1595                count += 1;
1596            })
1597            .unwrap();
1598        assert_eq!(count, (DYNAMIC_SCHEMA_CACHE_LIMIT + 3) as u64);
1599    }
1600}