1use crate::{
5    de::Deserializer,
6    error::{Error, Result},
7    format::*,
8    ser::Serializer,
9    value::Value,
10};
11use erased_discriminant::Discriminant;
12use once_cell::sync::Lazy;
13use serde::{de::DeserializeSeed, Deserialize, Serialize};
14use std::any::TypeId;
15use std::collections::BTreeMap;
16
17pub type Registry = BTreeMap<String, ContainerFormat>;
19
20#[derive(Debug)]
23pub struct Tracer {
24    pub(crate) config: TracerConfig,
26
27    pub(crate) registry: Registry,
30
31    pub(crate) incomplete_enums: BTreeMap<String, EnumProgress>,
34
35    pub(crate) discriminants: BTreeMap<(TypeId, VariantId<'static>), Discriminant>,
37}
38
39#[derive(Copy, Clone, Debug)]
40pub(crate) enum EnumProgress {
41    NamedVariantsRemaining,
43    IndexedVariantsRemaining,
45}
46
47#[derive(Eq, PartialEq, Ord, PartialOrd, Debug)]
48pub(crate) enum VariantId<'a> {
49    Index(u32),
50    Name(&'a str),
51}
52
53#[derive(Debug, Default)]
56pub struct Samples {
57    pub(crate) values: BTreeMap<&'static str, Value>,
58}
59
60impl Samples {
61    pub fn new() -> Self {
63        Self::default()
64    }
65
66    pub fn value(&self, name: &'static str) -> Option<&Value> {
68        self.values.get(name)
69    }
70}
71
72#[derive(Debug)]
74pub struct TracerConfig {
75    pub(crate) is_human_readable: bool,
76    pub(crate) record_samples_for_newtype_structs: bool,
77    pub(crate) record_samples_for_tuple_structs: bool,
78    pub(crate) record_samples_for_structs: bool,
79    pub(crate) default_bool_value: bool,
80    pub(crate) default_u8_value: u8,
81    pub(crate) default_u16_value: u16,
82    pub(crate) default_u32_value: u32,
83    pub(crate) default_u64_value: u64,
84    pub(crate) default_u128_value: u128,
85    pub(crate) default_i8_value: i8,
86    pub(crate) default_i16_value: i16,
87    pub(crate) default_i32_value: i32,
88    pub(crate) default_i64_value: i64,
89    pub(crate) default_i128_value: i128,
90    pub(crate) default_f32_value: f32,
91    pub(crate) default_f64_value: f64,
92    pub(crate) default_char_value: char,
93    pub(crate) default_borrowed_str_value: &'static str,
94    pub(crate) default_string_value: String,
95    pub(crate) default_borrowed_bytes_value: &'static [u8],
96    pub(crate) default_byte_buf_value: Vec<u8>,
97}
98
99impl Default for TracerConfig {
100    fn default() -> Self {
102        Self {
103            is_human_readable: false,
104            record_samples_for_newtype_structs: true,
105            record_samples_for_tuple_structs: false,
106            record_samples_for_structs: false,
107            default_bool_value: false,
108            default_u8_value: 0,
109            default_u16_value: 0,
110            default_u32_value: 0,
111            default_u64_value: 0,
112            default_u128_value: 0,
113            default_i8_value: 0,
114            default_i16_value: 0,
115            default_i32_value: 0,
116            default_i64_value: 0,
117            default_i128_value: 0,
118            default_f32_value: 0.0,
119            default_f64_value: 0.0,
120            default_char_value: 'A',
121            default_borrowed_str_value: "",
122            default_string_value: String::new(),
123            default_borrowed_bytes_value: b"",
124            default_byte_buf_value: Vec::new(),
125        }
126    }
127}
128
129macro_rules! define_default_value_setter {
130    ($method:ident, $ty:ty) => {
131        pub fn $method(mut self, value: $ty) -> Self {
133            self.$method = value;
134            self
135        }
136    };
137}
138
139impl TracerConfig {
140    #[allow(clippy::wrong_self_convention)]
142    pub fn is_human_readable(mut self, value: bool) -> Self {
143        self.is_human_readable = value;
144        self
145    }
146
147    pub fn record_samples_for_newtype_structs(mut self, value: bool) -> Self {
149        self.record_samples_for_newtype_structs = value;
150        self
151    }
152
153    pub fn record_samples_for_tuple_structs(mut self, value: bool) -> Self {
155        self.record_samples_for_tuple_structs = value;
156        self
157    }
158
159    pub fn record_samples_for_structs(mut self, value: bool) -> Self {
161        self.record_samples_for_structs = value;
162        self
163    }
164
165    define_default_value_setter!(default_bool_value, bool);
166    define_default_value_setter!(default_u8_value, u8);
167    define_default_value_setter!(default_u16_value, u16);
168    define_default_value_setter!(default_u32_value, u32);
169    define_default_value_setter!(default_u64_value, u64);
170    define_default_value_setter!(default_u128_value, u128);
171    define_default_value_setter!(default_i8_value, i8);
172    define_default_value_setter!(default_i16_value, i16);
173    define_default_value_setter!(default_i32_value, i32);
174    define_default_value_setter!(default_i64_value, i64);
175    define_default_value_setter!(default_i128_value, i128);
176    define_default_value_setter!(default_f32_value, f32);
177    define_default_value_setter!(default_f64_value, f64);
178    define_default_value_setter!(default_char_value, char);
179    define_default_value_setter!(default_borrowed_str_value, &'static str);
180    define_default_value_setter!(default_string_value, String);
181    define_default_value_setter!(default_borrowed_bytes_value, &'static [u8]);
182    define_default_value_setter!(default_byte_buf_value, Vec<u8>);
183}
184
185impl Tracer {
186    pub fn new(config: TracerConfig) -> Self {
188        Self {
189            config,
190            registry: BTreeMap::new(),
191            incomplete_enums: BTreeMap::new(),
192            discriminants: BTreeMap::new(),
193        }
194    }
195
196    pub fn trace_value<T>(&mut self, samples: &mut Samples, value: &T) -> Result<(Format, Value)>
202    where
203        T: ?Sized + Serialize,
204    {
205        let serializer = Serializer::new(self, samples);
206        let (mut format, sample) = value.serialize(serializer)?;
207        format.reduce();
208        Ok((format, sample))
209    }
210
211    pub fn trace_type_once<'de, T>(&mut self, samples: &'de Samples) -> Result<(Format, T)>
220    where
221        T: Deserialize<'de>,
222    {
223        let mut format = Format::unknown();
224        let deserializer = Deserializer::new(self, samples, &mut format);
225        let value = T::deserialize(deserializer)?;
226        format.reduce();
227        Ok((format, value))
228    }
229
230    pub fn trace_type_once_with_seed<'de, S>(
232        &mut self,
233        samples: &'de Samples,
234        seed: S,
235    ) -> Result<(Format, S::Value)>
236    where
237        S: DeserializeSeed<'de>,
238    {
239        let mut format = Format::unknown();
240        let deserializer = Deserializer::new(self, samples, &mut format);
241        let value = seed.deserialize(deserializer)?;
242        format.reduce();
243        Ok((format, value))
244    }
245
246    pub fn trace_type<'de, T>(&mut self, samples: &'de Samples) -> Result<(Format, Vec<T>)>
250    where
251        T: Deserialize<'de>,
252    {
253        let mut values = Vec::new();
254        loop {
255            let (format, value) = self.trace_type_once::<T>(samples)?;
256            values.push(value);
257            if let Format::TypeName(name) = &format {
258                if let Some(&progress) = self.incomplete_enums.get(name) {
259                    self.incomplete_enums.remove(name);
261                    if let EnumProgress::NamedVariantsRemaining = progress {
262                        values.pop().unwrap();
263                    }
264                    continue;
265                }
266            }
267            return Ok((format, values));
268        }
269    }
270
271    pub fn trace_simple_type<'de, T>(&mut self) -> Result<(Format, Vec<T>)>
276    where
277        T: Deserialize<'de>,
278    {
279        static SAMPLES: Lazy<Samples> = Lazy::new(Samples::new);
280        self.trace_type(&SAMPLES)
281    }
282
283    pub fn trace_type_with_seed<'de, S>(
285        &mut self,
286        samples: &'de Samples,
287        seed: S,
288    ) -> Result<(Format, Vec<S::Value>)>
289    where
290        S: DeserializeSeed<'de> + Clone,
291    {
292        let mut values = Vec::new();
293        loop {
294            let (format, value) = self.trace_type_once_with_seed(samples, seed.clone())?;
295            values.push(value);
296            if let Format::TypeName(name) = &format {
297                if let Some(&progress) = self.incomplete_enums.get(name) {
298                    self.incomplete_enums.remove(name);
300                    if let EnumProgress::NamedVariantsRemaining = progress {
301                        values.pop().unwrap();
302                    }
303                    continue;
304                }
305            }
306            return Ok((format, values));
307        }
308    }
309
310    pub fn registry(self) -> Result<Registry> {
317        let mut registry = self.registry;
318        for (name, format) in registry.iter_mut() {
319            format
320                .normalize()
321                .map_err(|_| Error::UnknownFormatInContainer(name.clone()))?;
322        }
323        if self.incomplete_enums.is_empty() {
324            Ok(registry)
325        } else {
326            Err(Error::MissingVariants(
327                self.incomplete_enums.into_keys().collect(),
328            ))
329        }
330    }
331
332    pub fn registry_unchecked(self) -> Registry {
335        let mut registry = self.registry;
336        for format in registry.values_mut() {
337            format.normalize().unwrap_or(());
338        }
339        registry
340    }
341
342    pub(crate) fn record_container(
343        &mut self,
344        samples: &mut Samples,
345        name: &'static str,
346        format: ContainerFormat,
347        value: Value,
348        record_value: bool,
349    ) -> Result<(Format, Value)> {
350        self.registry.entry(name.to_string()).unify(format)?;
351        if record_value {
352            samples.values.insert(name, value.clone());
353        }
354        Ok((Format::TypeName(name.into()), value))
355    }
356
357    pub(crate) fn record_variant(
358        &mut self,
359        samples: &mut Samples,
360        name: &'static str,
361        variant_index: u32,
362        variant_name: &'static str,
363        variant: VariantFormat,
364        variant_value: Value,
365    ) -> Result<(Format, Value)> {
366        let mut variants = BTreeMap::new();
367        variants.insert(
368            variant_index,
369            Named {
370                name: variant_name.into(),
371                value: variant,
372            },
373        );
374        let format = ContainerFormat::Enum(variants);
375        let value = Value::Variant(variant_index, Box::new(variant_value));
376        self.record_container(samples, name, format, value, false)
377    }
378
379    pub(crate) fn get_sample<'de, 'a>(
380        &'a self,
381        samples: &'de Samples,
382        name: &'static str,
383    ) -> Option<(&'a ContainerFormat, &'de Value)> {
384        match samples.value(name) {
385            Some(value) => {
386                let format = self
387                    .registry
388                    .get(name)
389                    .expect("recorded containers should have a format already");
390                Some((format, value))
391            }
392            None => None,
393        }
394    }
395}