pytorch_reader/nested/
de.rs

1use std::collections::HashMap;
2
3use super::data::NestedValue;
4use super::{adapter::BurnModuleAdapter, error::Error};
5
6use serde::de::{EnumAccess, VariantAccess};
7use serde::{
8    de::{self, DeserializeSeed, IntoDeserializer, MapAccess, SeqAccess, Visitor},
9    forward_to_deserialize_any,
10};
11
12const RECORD_ITEM_SUFFIX: &str = "RecordItem";
13
14#[inline]
15fn custom_err<T: core::fmt::Display>(msg: T) -> Error {
16    <Error as de::Error>::custom(msg)
17}
18
19/// Length of a `NestedValue` that represents a sequence, if it is one.
20fn sequence_len(value: &NestedValue) -> Option<usize> {
21    match value {
22        NestedValue::Vec(v) => Some(v.len()),
23        NestedValue::U8s(v) => Some(v.len()),
24        NestedValue::U16s(v) => Some(v.len()),
25        NestedValue::F32s(v) => Some(v.len()),
26        _ => None,
27    }
28}
29
30/// A deserializer for the nested value data structure.
31pub struct Deserializer<A: BurnModuleAdapter> {
32    // This string starts with the input data and characters are truncated off
33    // the beginning as data is parsed.
34    value: Option<NestedValue>,
35    default_for_missing_fields: bool,
36    phantom: std::marker::PhantomData<A>,
37}
38
39impl<A: BurnModuleAdapter> Deserializer<A> {
40    /// Creates a new deserializer with the given nested value.
41    ///
42    /// # Arguments
43    ///
44    /// * `value` - A nested value.
45    /// * `default_for_missing_fields` - A boolean indicating whether to add missing fields with default value.
46    pub fn new(value: NestedValue, default_for_missing_fields: bool) -> Self {
47        Self {
48            value: Some(value),
49            default_for_missing_fields,
50            phantom: std::marker::PhantomData,
51        }
52    }
53
54    fn extract_scalar<T>(
55        self,
56        expected: &'static str,
57        extractor: impl FnOnce(&NestedValue) -> Option<T>,
58    ) -> Result<T, Error> {
59        let value = self
60            .value
61            .ok_or_else(|| custom_err(format!("expected {expected}, found nothing")))?;
62
63        match extractor(&value) {
64            Some(val) => Ok(val),
65            None => Err(custom_err(format!("expected {expected} but got {value:?}"))),
66        }
67    }
68}
69
70impl<'de, A: BurnModuleAdapter> serde::Deserializer<'de> for Deserializer<A> {
71    type Error = Error;
72
73    fn deserialize_any<V>(self, visitor: V) -> Result<V::Value, Self::Error>
74    where
75        V: Visitor<'de>,
76    {
77        let default_for_missing_fields = self.default_for_missing_fields;
78        match self.value {
79            Some(NestedValue::Bool(v)) => visitor.visit_bool(v),
80            Some(NestedValue::String(v)) => visitor.visit_string(v),
81            Some(NestedValue::F32(v)) => visitor.visit_f32(v),
82            Some(NestedValue::F64(v)) => visitor.visit_f64(v),
83            Some(NestedValue::I16(v)) => visitor.visit_i16(v),
84            Some(NestedValue::I32(v)) => visitor.visit_i32(v),
85            Some(NestedValue::I64(v)) => visitor.visit_i64(v),
86            Some(NestedValue::U8(v)) => visitor.visit_u8(v),
87            Some(NestedValue::U16(v)) => visitor.visit_u16(v),
88            Some(NestedValue::U64(v)) => visitor.visit_u64(v),
89            Some(NestedValue::Map(map)) => {
90                visitor.visit_map(HashMapAccess::<A>::new(map, default_for_missing_fields))
91            }
92            Some(NestedValue::Vec(v)) => visitor.visit_seq(VecSeqAccess::<A, NestedValue>::new(
93                NestedValue::Vec(v),
94                default_for_missing_fields,
95            )?),
96            Some(NestedValue::U8s(v)) => visitor.visit_seq(VecSeqAccess::<A, u8>::new(
97                NestedValue::U8s(v),
98                default_for_missing_fields,
99            )?),
100            Some(NestedValue::U16s(v)) => visitor.visit_seq(VecSeqAccess::<A, u16>::new(
101                NestedValue::U16s(v),
102                default_for_missing_fields,
103            )?),
104            Some(NestedValue::F32s(v)) => visitor.visit_seq(VecSeqAccess::<A, f32>::new(
105                NestedValue::F32s(v),
106                default_for_missing_fields,
107            )?),
108            Some(NestedValue::Unsupported(type_name)) => Err(custom_err(format!(
109                "cannot deserialize unsupported source type: {type_name}"
110            ))),
111            // A placeholder has no concrete type, and `None` means nothing was
112            // available: neither can be turned into an untyped value.
113            Some(NestedValue::Default(_)) | None => Err(custom_err(
114                "cannot deserialize an untyped value: no concrete source value is available",
115            )),
116        }
117    }
118
119    fn deserialize_struct<V>(
120        self,
121        name: &'static str,
122        fields: &'static [&'static str],
123        visitor: V,
124    ) -> Result<V::Value, Self::Error>
125    where
126        V: Visitor<'de>,
127    {
128        let value = match self.value {
129            Some(value) => {
130                // Adapt modules
131                if let Some(name) = name.strip_suffix(RECORD_ITEM_SUFFIX) {
132                    A::adapt(name, value)
133                } else {
134                    value
135                }
136            }
137            None => {
138                return Err(custom_err(format!(
139                    "Expected some value but got {:?}",
140                    self.value
141                )));
142            }
143        };
144
145        match value {
146            NestedValue::Map(map) => {
147                // Add missing fields into the map with default value if needed.
148                let map = if self.default_for_missing_fields {
149                    let mut map = map;
150                    for field in fields.iter().map(|s| s.to_string()) {
151                        map.entry(field.clone())
152                            .or_insert(NestedValue::Default(Some(field)));
153                    }
154                    map
155                } else {
156                    map
157                };
158
159                visitor.visit_map(HashMapAccess::<A>::new(
160                    map,
161                    self.default_for_missing_fields,
162                ))
163            }
164
165            _ => Err(custom_err(format!("Expected struct but got {value:?}"))),
166        }
167    }
168
169    fn deserialize_string<V>(self, visitor: V) -> Result<V::Value, Self::Error>
170    where
171        V: Visitor<'de>,
172    {
173        match self.value {
174            Some(NestedValue::String(val)) => visitor.visit_string(val),
175            Some(other) => Err(custom_err(format!("expected string but got {other:?}"))),
176            None => Err(custom_err("expected string, found nothing")),
177        }
178    }
179
180    fn deserialize_ignored_any<V>(self, visitor: V) -> Result<V::Value, Self::Error>
181    where
182        V: Visitor<'de>,
183    {
184        visitor.visit_unit()
185    }
186
187    fn deserialize_map<V>(self, visitor: V) -> Result<V::Value, Self::Error>
188    where
189        V: Visitor<'de>,
190    {
191        match self.value {
192            Some(NestedValue::Map(map)) => visitor.visit_map(HashMapAccess::<A>::new(
193                map,
194                self.default_for_missing_fields,
195            )),
196
197            _ => Err(custom_err(format!(
198                "Expected map value but got {:?}",
199                self.value
200            ))),
201        }
202    }
203
204    fn deserialize_bool<V>(self, visitor: V) -> Result<V::Value, Self::Error>
205    where
206        V: Visitor<'de>,
207    {
208        let val = self.extract_scalar("bool", |v| v.clone().as_bool())?;
209        visitor.visit_bool(val)
210    }
211
212    fn deserialize_i8<V>(self, visitor: V) -> Result<V::Value, Self::Error>
213    where
214        V: Visitor<'de>,
215    {
216        let val = self.extract_scalar("i8", |v| {
217            v.clone().as_i16().and_then(|x| i8::try_from(x).ok())
218        })?;
219        visitor.visit_i8(val)
220    }
221
222    fn deserialize_i16<V>(self, visitor: V) -> Result<V::Value, Self::Error>
223    where
224        V: Visitor<'de>,
225    {
226        let val = self.extract_scalar("i16", |v| v.clone().as_i16())?;
227        visitor.visit_i16(val)
228    }
229
230    fn deserialize_i32<V>(self, visitor: V) -> Result<V::Value, Self::Error>
231    where
232        V: Visitor<'de>,
233    {
234        let val = self.extract_scalar("i32", |v| v.clone().as_i32())?;
235        visitor.visit_i32(val)
236    }
237
238    fn deserialize_i64<V>(self, visitor: V) -> Result<V::Value, Self::Error>
239    where
240        V: Visitor<'de>,
241    {
242        let val = self.extract_scalar("i64", |v| v.clone().as_i64())?;
243        visitor.visit_i64(val)
244    }
245
246    fn deserialize_u8<V>(self, visitor: V) -> Result<V::Value, Self::Error>
247    where
248        V: Visitor<'de>,
249    {
250        let val = self.extract_scalar("u8", |v| v.clone().as_u8())?;
251        visitor.visit_u8(val)
252    }
253
254    fn deserialize_u16<V>(self, visitor: V) -> Result<V::Value, Self::Error>
255    where
256        V: Visitor<'de>,
257    {
258        let val = self.extract_scalar("u16", |v| v.clone().as_u16())?;
259        visitor.visit_u16(val)
260    }
261
262    fn deserialize_u32<V>(self, visitor: V) -> Result<V::Value, Self::Error>
263    where
264        V: Visitor<'de>,
265    {
266        let val = self.extract_scalar("u32", |v| {
267            v.clone().as_u64().and_then(|x| u32::try_from(x).ok())
268        })?;
269        visitor.visit_u32(val)
270    }
271
272    fn deserialize_u64<V>(self, visitor: V) -> Result<V::Value, Self::Error>
273    where
274        V: Visitor<'de>,
275    {
276        let val = self.extract_scalar("u64", |v| v.clone().as_u64())?;
277        visitor.visit_u64(val)
278    }
279
280    fn deserialize_f32<V>(self, visitor: V) -> Result<V::Value, Self::Error>
281    where
282        V: Visitor<'de>,
283    {
284        let val = self.extract_scalar("f32", |v| v.clone().as_f32())?;
285        visitor.visit_f32(val)
286    }
287
288    fn deserialize_f64<V>(self, visitor: V) -> Result<V::Value, Self::Error>
289    where
290        V: Visitor<'de>,
291    {
292        let val = self.extract_scalar("f64", |v| v.clone().as_f64())?;
293        visitor.visit_f64(val)
294    }
295
296    fn deserialize_char<V>(self, visitor: V) -> Result<V::Value, Self::Error>
297    where
298        V: Visitor<'de>,
299    {
300        match self.value {
301            Some(NestedValue::String(s)) if s.chars().count() == 1 => {
302                visitor.visit_char(s.chars().next().expect("count checked above"))
303            }
304            Some(other) => Err(custom_err(format!(
305                "expected a single-character string but got {other:?}"
306            ))),
307            None => Err(custom_err("expected char, found nothing")),
308        }
309    }
310
311    fn deserialize_str<V>(self, visitor: V) -> Result<V::Value, Self::Error>
312    where
313        V: Visitor<'de>,
314    {
315        match self.value {
316            Some(NestedValue::String(val)) => visitor.visit_str(&val),
317            Some(other) => Err(custom_err(format!("expected str but got {other:?}"))),
318            None => Err(custom_err("expected str, found nothing")),
319        }
320    }
321
322    fn deserialize_bytes<V>(self, visitor: V) -> Result<V::Value, Self::Error>
323    where
324        V: Visitor<'de>,
325    {
326        match self.value {
327            Some(NestedValue::U8s(bytes)) => visitor.visit_bytes(&bytes),
328            Some(other) => Err(custom_err(format!("expected bytes but got {other:?}"))),
329            None => Err(custom_err("expected bytes, found nothing")),
330        }
331    }
332
333    fn deserialize_byte_buf<V>(self, visitor: V) -> Result<V::Value, Self::Error>
334    where
335        V: Visitor<'de>,
336    {
337        match self.value {
338            Some(NestedValue::U8s(bytes)) => visitor.visit_byte_buf(bytes),
339            Some(other) => Err(custom_err(format!(
340                "expected byte buffer but got {other:?}"
341            ))),
342            None => Err(custom_err("expected byte buffer, found nothing")),
343        }
344    }
345
346    fn deserialize_option<V>(self, visitor: V) -> Result<V::Value, Self::Error>
347    where
348        V: Visitor<'de>,
349    {
350        if let Some(value) = self.value {
351            visitor.visit_some(Deserializer::<A>::new(
352                value,
353                self.default_for_missing_fields,
354            ))
355        } else {
356            visitor.visit_none()
357        }
358    }
359
360    fn deserialize_unit<V>(self, visitor: V) -> Result<V::Value, Self::Error>
361    where
362        V: Visitor<'de>,
363    {
364        // A unit carries no data, so it is only valid when nothing (or an
365        // explicit placeholder) was provided. Anything else is a type error.
366        match self.value {
367            None | Some(NestedValue::Default(_)) => visitor.visit_unit(),
368            Some(other) => Err(custom_err(format!("expected unit but got {other:?}"))),
369        }
370    }
371
372    fn deserialize_unit_struct<V>(
373        self,
374        _name: &'static str,
375        visitor: V,
376    ) -> Result<V::Value, Self::Error>
377    where
378        V: Visitor<'de>,
379    {
380        self.deserialize_unit(visitor)
381    }
382
383    fn deserialize_newtype_struct<V>(
384        self,
385        _name: &'static str,
386        visitor: V,
387    ) -> Result<V::Value, Self::Error>
388    where
389        V: Visitor<'de>,
390    {
391        let value = self
392            .value
393            .ok_or_else(|| custom_err("expected value for newtype struct but got None"))?;
394        visitor.visit_newtype_struct(Deserializer::<A>::new(
395            value,
396            self.default_for_missing_fields,
397        ))
398    }
399
400    fn deserialize_seq<V>(self, visitor: V) -> Result<V::Value, Self::Error>
401    where
402        V: Visitor<'de>,
403    {
404        if let Some(value) = self.value {
405            match value {
406                NestedValue::Vec(_) => visitor.visit_seq(VecSeqAccess::<A, NestedValue>::new(
407                    value,
408                    self.default_for_missing_fields,
409                )?),
410                NestedValue::U8s(_) => visitor.visit_seq(VecSeqAccess::<A, u8>::new(
411                    value,
412                    self.default_for_missing_fields,
413                )?),
414                NestedValue::U16s(_) => visitor.visit_seq(VecSeqAccess::<A, u16>::new(
415                    value,
416                    self.default_for_missing_fields,
417                )?),
418                NestedValue::F32s(_) => visitor.visit_seq(VecSeqAccess::<A, f32>::new(
419                    value,
420                    self.default_for_missing_fields,
421                )?),
422                _ => Err(custom_err(format!("Expected Vec but got {value:?}"))),
423            }
424        } else {
425            Err(custom_err("Expected Vec but got None"))
426        }
427    }
428
429    fn deserialize_tuple<V>(self, len: usize, visitor: V) -> Result<V::Value, Self::Error>
430    where
431        V: Visitor<'de>,
432    {
433        // Tuples and fixed-size arrays are represented as sequences, and the
434        // declared length is enforced so trailing elements are never silently
435        // dropped (nor missing ones silently defaulted).
436        let actual = self.value.as_ref().and_then(sequence_len).ok_or_else(|| {
437            custom_err(format!(
438                "expected a sequence of length {len} but got {:?}",
439                self.value
440            ))
441        })?;
442        if actual != len {
443            return Err(<Error as de::Error>::invalid_length(
444                actual,
445                &format!("tuple of length {len}").as_str(),
446            ));
447        }
448        self.deserialize_seq(visitor)
449    }
450
451    fn deserialize_tuple_struct<V>(
452        self,
453        _name: &'static str,
454        len: usize,
455        visitor: V,
456    ) -> Result<V::Value, Self::Error>
457    where
458        V: Visitor<'de>,
459    {
460        let actual = self.value.as_ref().and_then(sequence_len).ok_or_else(|| {
461            custom_err(format!(
462                "expected a sequence of length {len} but got {:?}",
463                self.value
464            ))
465        })?;
466        if actual != len {
467            return Err(<Error as de::Error>::invalid_length(
468                actual,
469                &format!("tuple struct of length {len}").as_str(),
470            ));
471        }
472        self.deserialize_seq(visitor)
473    }
474
475    fn deserialize_enum<V>(
476        self,
477        _name: &'static str,
478        _variants: &'static [&'static str],
479        visitor: V,
480    ) -> Result<V::Value, Self::Error>
481    where
482        V: Visitor<'de>,
483    {
484        let value = self
485            .value
486            .ok_or_else(|| custom_err("expected value for enum but got None"))?;
487
488        // The nested encoding carries no separate variant tag: variants are
489        // stored as `NestedValue::Map { "DType": <variant name> }` (see
490        // `ProbeEnumAccess::unit_variant`). Require that tag and reject anything
491        // else rather than guessing a variant, since silently selecting one (for
492        // example the first declared variant) would decode the wrong data.
493        //
494        // The previous implementation additionally duplicated the caller's
495        // visitor with a raw `ptr::copy_nonoverlapping` so it could retry
496        // variants. `Visitor` is not `Copy`, so that bitwise copy aliased any
497        // heap-owned visitor state and was undefined behaviour (a double free,
498        // confirmed under Miri). The visitor is now consumed exactly once.
499        let variant = match &value {
500            NestedValue::Map(map) => match map.get("DType") {
501                Some(NestedValue::String(name)) => name.clone(),
502                other => {
503                    return Err(custom_err(format!(
504                        "cannot resolve the enum variant: expected a 'DType' string tag, got {other:?}"
505                    )));
506                }
507            },
508            other => {
509                return Err(custom_err(format!(
510                    "cannot resolve the enum variant: expected a map with a 'DType' tag, got {other:?}"
511                )));
512            }
513        };
514
515        visitor.visit_enum(ProbeEnumAccess::<A>::new(
516            value,
517            variant,
518            self.default_for_missing_fields,
519        ))
520    }
521
522    fn deserialize_identifier<V>(self, visitor: V) -> Result<V::Value, Self::Error>
523    where
524        V: Visitor<'de>,
525    {
526        match self.value {
527            Some(NestedValue::String(val)) => visitor.visit_str(&val),
528            Some(other) => Err(custom_err(format!(
529                "expected identifier string but got {other:?}"
530            ))),
531            None => Err(custom_err("expected identifier, found nothing")),
532        }
533    }
534}
535
536/// A sequence access for a vector in the nested value data structure.
537struct VecSeqAccess<A: BurnModuleAdapter, I> {
538    iter: Box<dyn Iterator<Item = I>>,
539    default_for_missing_fields: bool,
540    phantom: std::marker::PhantomData<A>,
541}
542
543impl<A: BurnModuleAdapter> VecSeqAccess<A, NestedValue> {
544    fn new(vec: NestedValue, default_for_missing_fields: bool) -> Result<Self, Error> {
545        match vec {
546            NestedValue::Vec(v) => Ok(VecSeqAccess {
547                iter: Box::new(v.into_iter()),
548                default_for_missing_fields,
549                phantom: std::marker::PhantomData,
550            }),
551            other => Err(custom_err(format!("expected Vec, found {other:?}"))),
552        }
553    }
554}
555
556impl<A: BurnModuleAdapter> VecSeqAccess<A, u8> {
557    fn new(vec: NestedValue, default_for_missing_fields: bool) -> Result<Self, Error> {
558        match vec {
559            NestedValue::U8s(v) => Ok(VecSeqAccess {
560                iter: Box::new(v.into_iter()),
561                default_for_missing_fields,
562                phantom: std::marker::PhantomData,
563            }),
564            other => Err(custom_err(format!("expected U8s, found {other:?}"))),
565        }
566    }
567}
568
569impl<A: BurnModuleAdapter> VecSeqAccess<A, u16> {
570    fn new(vec: NestedValue, default_for_missing_fields: bool) -> Result<Self, Error> {
571        match vec {
572            NestedValue::U16s(v) => Ok(VecSeqAccess {
573                iter: Box::new(v.into_iter()),
574                default_for_missing_fields,
575                phantom: std::marker::PhantomData,
576            }),
577            other => Err(custom_err(format!("expected U16s, found {other:?}"))),
578        }
579    }
580}
581
582impl<A: BurnModuleAdapter> VecSeqAccess<A, f32> {
583    fn new(vec: NestedValue, default_for_missing_fields: bool) -> Result<Self, Error> {
584        match vec {
585            NestedValue::F32s(v) => Ok(VecSeqAccess {
586                iter: Box::new(v.into_iter()),
587                default_for_missing_fields,
588                phantom: std::marker::PhantomData,
589            }),
590            other => Err(custom_err(format!("expected F32s, found {other:?}"))),
591        }
592    }
593}
594
595impl<'de, A> SeqAccess<'de> for VecSeqAccess<A, NestedValue>
596where
597    NestedValueWrapper<A>: IntoDeserializer<'de, Error>,
598    A: BurnModuleAdapter,
599{
600    type Error = Error;
601
602    fn next_element_seed<T>(&mut self, seed: T) -> Result<Option<T::Value>, Self::Error>
603    where
604        T: DeserializeSeed<'de>,
605    {
606        let item = match self.iter.next() {
607            Some(v) => v,
608            None => return Ok(None),
609        };
610
611        seed.deserialize(
612            NestedValueWrapper::<A>::new(item, self.default_for_missing_fields).into_deserializer(),
613        )
614        .map(Some)
615    }
616}
617
618impl<'de, A> SeqAccess<'de> for VecSeqAccess<A, u8>
619where
620    NestedValueWrapper<A>: IntoDeserializer<'de, Error>,
621    A: BurnModuleAdapter,
622{
623    type Error = Error;
624
625    fn next_element_seed<T>(&mut self, seed: T) -> Result<Option<T::Value>, Self::Error>
626    where
627        T: DeserializeSeed<'de>,
628    {
629        let item = match self.iter.next() {
630            Some(v) => v,
631            None => return Ok(None),
632        };
633
634        seed.deserialize(
635            NestedValueWrapper::<A>::new(NestedValue::U8(item), self.default_for_missing_fields)
636                .into_deserializer(),
637        )
638        .map(Some)
639    }
640}
641
642impl<'de, A> SeqAccess<'de> for VecSeqAccess<A, u16>
643where
644    NestedValueWrapper<A>: IntoDeserializer<'de, Error>,
645    A: BurnModuleAdapter,
646{
647    type Error = Error;
648
649    fn next_element_seed<T>(&mut self, seed: T) -> Result<Option<T::Value>, Self::Error>
650    where
651        T: DeserializeSeed<'de>,
652    {
653        let item = match self.iter.next() {
654            Some(v) => v,
655            None => return Ok(None),
656        };
657
658        seed.deserialize(
659            NestedValueWrapper::<A>::new(NestedValue::U16(item), self.default_for_missing_fields)
660                .into_deserializer(),
661        )
662        .map(Some)
663    }
664}
665
666impl<'de, A> SeqAccess<'de> for VecSeqAccess<A, f32>
667where
668    NestedValueWrapper<A>: IntoDeserializer<'de, Error>,
669    A: BurnModuleAdapter,
670{
671    type Error = Error;
672
673    fn next_element_seed<T>(&mut self, seed: T) -> Result<Option<T::Value>, Self::Error>
674    where
675        T: DeserializeSeed<'de>,
676    {
677        let item = match self.iter.next() {
678            Some(v) => v,
679            None => return Ok(None),
680        };
681
682        seed.deserialize(
683            NestedValueWrapper::<A>::new(NestedValue::F32(item), self.default_for_missing_fields)
684                .into_deserializer(),
685        )
686        .map(Some)
687    }
688}
689
690/// A map access for a map in the nested value data structure.
691struct HashMapAccess<A: BurnModuleAdapter> {
692    iter: std::collections::hash_map::IntoIter<String, NestedValue>,
693    next_value: Option<NestedValue>,
694    default_for_missing_fields: bool,
695    phantom: std::marker::PhantomData<A>,
696}
697
698impl<A: BurnModuleAdapter> HashMapAccess<A> {
699    fn new(map: HashMap<String, NestedValue>, default_for_missing_fields: bool) -> Self {
700        HashMapAccess {
701            iter: map.into_iter(),
702            next_value: None,
703            default_for_missing_fields,
704            phantom: std::marker::PhantomData,
705        }
706    }
707}
708
709impl<'de, A> MapAccess<'de> for HashMapAccess<A>
710where
711    String: IntoDeserializer<'de, Error>,
712    NestedValueWrapper<A>: IntoDeserializer<'de, Error>,
713    A: BurnModuleAdapter,
714{
715    type Error = Error;
716
717    fn next_key_seed<T>(&mut self, seed: T) -> Result<Option<T::Value>, Self::Error>
718    where
719        T: DeserializeSeed<'de>,
720    {
721        match self.iter.next() {
722            Some((k, v)) => {
723                self.next_value = Some(v);
724                seed.deserialize(k.into_deserializer()).map(Some)
725            }
726            None => Ok(None),
727        }
728    }
729
730    fn next_value_seed<T>(&mut self, seed: T) -> Result<T::Value, Self::Error>
731    where
732        T: DeserializeSeed<'de>,
733    {
734        match self.next_value.take() {
735            Some(NestedValue::Default(originator)) => {
736                seed.deserialize(DefaultDeserializer::new(originator))
737            }
738            Some(v) => seed.deserialize(
739                NestedValueWrapper::new(v, self.default_for_missing_fields).into_deserializer(),
740            ),
741            None => seed.deserialize(DefaultDeserializer::new(None)),
742        }
743    }
744}
745
746struct ProbeEnumAccess<A: BurnModuleAdapter> {
747    value: NestedValue,
748    current_variant: String,
749    default_for_missing_fields: bool,
750    phantom: std::marker::PhantomData<A>,
751}
752
753impl<A: BurnModuleAdapter> ProbeEnumAccess<A> {
754    fn new(value: NestedValue, current_variant: String, default_for_missing_fields: bool) -> Self {
755        ProbeEnumAccess {
756            value,
757            current_variant,
758            default_for_missing_fields,
759            phantom: std::marker::PhantomData,
760        }
761    }
762}
763
764impl<'de, A> EnumAccess<'de> for ProbeEnumAccess<A>
765where
766    A: BurnModuleAdapter,
767{
768    type Error = Error;
769    type Variant = Self;
770
771    fn variant_seed<V>(self, seed: V) -> Result<(V::Value, Self::Variant), Self::Error>
772    where
773        V: DeserializeSeed<'de>,
774    {
775        seed.deserialize(self.current_variant.clone().into_deserializer())
776            .map(|v| (v, self))
777    }
778}
779
780impl<'de, A> VariantAccess<'de> for ProbeEnumAccess<A>
781where
782    A: BurnModuleAdapter,
783{
784    type Error = Error;
785
786    fn newtype_variant_seed<T>(self, seed: T) -> Result<T::Value, Self::Error>
787    where
788        T: DeserializeSeed<'de>,
789    {
790        let value = seed.deserialize(
791            NestedValueWrapper::<A>::new(self.value, self.default_for_missing_fields)
792                .into_deserializer(),
793        )?;
794        Ok(value)
795    }
796
797    fn unit_variant(self) -> Result<(), Self::Error> {
798        match self.value {
799            NestedValue::Map(value) => match value.get("DType") {
800                Some(NestedValue::String(variant)) => {
801                    if *variant == self.current_variant {
802                        Ok(())
803                    } else {
804                        Err(Error::Other("Wrong variant".to_string()))
805                    }
806                }
807                Some(other) => Err(custom_err(format!(
808                    "expected DType variant as string, got {other:?}"
809                ))),
810                None => Err(custom_err("expected map containing 'DType' key for enum")),
811            },
812            other => Err(custom_err(format!(
813                "expected map for enum variant, got {other:?}"
814            ))),
815        }
816    }
817
818    fn tuple_variant<V>(self, _len: usize, visitor: V) -> Result<V::Value, Self::Error>
819    where
820        V: Visitor<'de>,
821    {
822        let default_for_missing_fields = self.default_for_missing_fields;
823        match self.value {
824            NestedValue::Vec(v) => visitor.visit_seq(VecSeqAccess::<A, NestedValue>::new(
825                NestedValue::Vec(v),
826                default_for_missing_fields,
827            )?),
828            NestedValue::U8s(v) => visitor.visit_seq(VecSeqAccess::<A, u8>::new(
829                NestedValue::U8s(v),
830                default_for_missing_fields,
831            )?),
832            NestedValue::U16s(v) => visitor.visit_seq(VecSeqAccess::<A, u16>::new(
833                NestedValue::U16s(v),
834                default_for_missing_fields,
835            )?),
836            NestedValue::F32s(v) => visitor.visit_seq(VecSeqAccess::<A, f32>::new(
837                NestedValue::F32s(v),
838                default_for_missing_fields,
839            )?),
840            other => Err(custom_err(format!(
841                "expected a sequence for a tuple enum variant but got {other:?}"
842            ))),
843        }
844    }
845
846    fn struct_variant<V>(
847        self,
848        _fields: &'static [&'static str],
849        visitor: V,
850    ) -> Result<V::Value, Self::Error>
851    where
852        V: Visitor<'de>,
853    {
854        match self.value {
855            NestedValue::Map(map) => visitor.visit_map(HashMapAccess::<A>::new(
856                map,
857                self.default_for_missing_fields,
858            )),
859            other => Err(custom_err(format!(
860                "expected a map for a struct enum variant but got {other:?}"
861            ))),
862        }
863    }
864}
865
866/// A wrapper for the nested value data structure with a burn module adapter.
867struct NestedValueWrapper<A: BurnModuleAdapter> {
868    value: NestedValue,
869    default_for_missing_fields: bool,
870    phantom: std::marker::PhantomData<A>,
871}
872
873impl<A: BurnModuleAdapter> NestedValueWrapper<A> {
874    fn new(value: NestedValue, default_for_missing_fields: bool) -> Self {
875        Self {
876            value,
877            default_for_missing_fields,
878            phantom: std::marker::PhantomData,
879        }
880    }
881}
882
883impl<A: BurnModuleAdapter> IntoDeserializer<'_, Error> for NestedValueWrapper<A> {
884    type Deserializer = Deserializer<A>;
885
886    fn into_deserializer(self) -> Self::Deserializer {
887        Deserializer::<A>::new(self.value, self.default_for_missing_fields)
888    }
889}
890
891/// A default deserializer that always returns the default value.
892struct DefaultDeserializer {
893    originator_field_name: Option<String>,
894}
895
896impl DefaultDeserializer {
897    fn new(originator_field_name: Option<String>) -> Self {
898        Self {
899            originator_field_name,
900        }
901    }
902}
903
904impl<'de> serde::Deserializer<'de> for DefaultDeserializer {
905    type Error = Error;
906
907    fn deserialize_any<V>(self, _visitor: V) -> Result<V::Value, Self::Error>
908    where
909        V: Visitor<'de>,
910    {
911        // A "default" deserializer has no concrete value to hand out; every
912        // typed entry point above supplies a type-specific zero instead.
913        Err(Error::Other(format!(
914            "Missing source values for the '{}' field of type '<unknown>'. Please verify the source data and ensure the field name is correct",
915            self.originator_field_name.unwrap_or("UNKNOWN".to_string()),
916        )))
917    }
918
919    fn deserialize_i32<V>(self, visitor: V) -> Result<V::Value, Self::Error>
920    where
921        V: Visitor<'de>,
922    {
923        visitor.visit_i32(Default::default())
924    }
925
926    fn deserialize_f32<V>(self, visitor: V) -> Result<V::Value, Self::Error>
927    where
928        V: Visitor<'de>,
929    {
930        visitor.visit_f32(Default::default())
931    }
932
933    fn deserialize_i16<V>(self, visitor: V) -> Result<V::Value, Self::Error>
934    where
935        V: Visitor<'de>,
936    {
937        visitor.visit_i16(Default::default())
938    }
939
940    fn deserialize_i64<V>(self, visitor: V) -> Result<V::Value, Self::Error>
941    where
942        V: Visitor<'de>,
943    {
944        visitor.visit_i64(Default::default())
945    }
946
947    fn deserialize_u16<V>(self, visitor: V) -> Result<V::Value, Self::Error>
948    where
949        V: Visitor<'de>,
950    {
951        visitor.visit_u16(Default::default())
952    }
953
954    fn deserialize_u64<V>(self, visitor: V) -> Result<V::Value, Self::Error>
955    where
956        V: Visitor<'de>,
957    {
958        visitor.visit_u64(Default::default())
959    }
960
961    fn deserialize_f64<V>(self, visitor: V) -> Result<V::Value, Self::Error>
962    where
963        V: Visitor<'de>,
964    {
965        visitor.visit_f64(Default::default())
966    }
967
968    fn deserialize_bool<V>(self, visitor: V) -> Result<V::Value, Self::Error>
969    where
970        V: Visitor<'de>,
971    {
972        visitor.visit_bool(Default::default())
973    }
974
975    fn deserialize_char<V>(self, visitor: V) -> Result<V::Value, Self::Error>
976    where
977        V: Visitor<'de>,
978    {
979        visitor.visit_char(Default::default())
980    }
981
982    fn deserialize_str<V>(self, visitor: V) -> Result<V::Value, Self::Error>
983    where
984        V: Visitor<'de>,
985    {
986        visitor.visit_str(Default::default())
987    }
988
989    fn deserialize_i8<V>(self, visitor: V) -> Result<V::Value, Self::Error>
990    where
991        V: Visitor<'de>,
992    {
993        visitor.visit_i8(Default::default())
994    }
995
996    fn deserialize_u8<V>(self, visitor: V) -> Result<V::Value, Self::Error>
997    where
998        V: Visitor<'de>,
999    {
1000        visitor.visit_u8(Default::default())
1001    }
1002
1003    fn deserialize_u32<V>(self, visitor: V) -> Result<V::Value, Self::Error>
1004    where
1005        V: Visitor<'de>,
1006    {
1007        visitor.visit_u32(Default::default())
1008    }
1009
1010    fn deserialize_option<V>(self, visitor: V) -> Result<V::Value, Self::Error>
1011    where
1012        V: Visitor<'de>,
1013    {
1014        visitor.visit_none()
1015    }
1016
1017    fn deserialize_seq<V>(self, visitor: V) -> Result<V::Value, Self::Error>
1018    where
1019        V: Visitor<'de>,
1020    {
1021        visitor.visit_seq(DefaultSeqAccess::new(None))
1022    }
1023
1024    fn deserialize_string<V>(self, visitor: V) -> Result<V::Value, Self::Error>
1025    where
1026        V: Visitor<'de>,
1027    {
1028        visitor.visit_string(Default::default())
1029    }
1030
1031    fn deserialize_struct<V>(
1032        self,
1033        name: &'static str,
1034        _fields: &'static [&'static str],
1035        _visitor: V,
1036    ) -> Result<V::Value, Self::Error>
1037    where
1038        V: Visitor<'de>,
1039    {
1040        Err(Error::Other(format!(
1041            "Missing source values for the '{}' field of type '{}'. Please verify the source data and ensure the field name is correct",
1042            self.originator_field_name.unwrap_or("UNKNOWN".to_string()),
1043            name,
1044        )))
1045    }
1046
1047    fn deserialize_tuple_struct<V>(
1048        self,
1049        _name: &'static str,
1050        len: usize,
1051        visitor: V,
1052    ) -> Result<V::Value, Self::Error>
1053    where
1054        V: Visitor<'de>,
1055    {
1056        visitor.visit_seq(DefaultSeqAccess::new(Some(len)))
1057    }
1058
1059    fn deserialize_tuple<V>(self, len: usize, visitor: V) -> Result<V::Value, Self::Error>
1060    where
1061        V: Visitor<'de>,
1062    {
1063        visitor.visit_seq(DefaultSeqAccess::new(Some(len)))
1064    }
1065
1066    fn deserialize_map<V>(self, visitor: V) -> Result<V::Value, Self::Error>
1067    where
1068        V: Visitor<'de>,
1069    {
1070        visitor.visit_map(DefaultMapAccess::new())
1071    }
1072
1073    fn deserialize_ignored_any<V>(self, visitor: V) -> Result<V::Value, Self::Error>
1074    where
1075        V: Visitor<'de>,
1076    {
1077        visitor.visit_unit()
1078    }
1079
1080    forward_to_deserialize_any! {
1081        u128 bytes byte_buf unit unit_struct newtype_struct
1082        enum identifier
1083    }
1084}
1085
1086/// A default sequence access that always returns None (empty sequence).
1087pub struct DefaultSeqAccess {
1088    size: Option<usize>,
1089}
1090
1091impl Default for DefaultSeqAccess {
1092    fn default() -> Self {
1093        Self::new(None)
1094    }
1095}
1096
1097impl DefaultSeqAccess {
1098    /// Creates a new default sequence access with the given size hint.
1099    pub fn new(size: Option<usize>) -> Self {
1100        DefaultSeqAccess { size }
1101    }
1102}
1103
1104impl<'de> SeqAccess<'de> for DefaultSeqAccess {
1105    type Error = Error;
1106
1107    fn next_element_seed<T>(&mut self, seed: T) -> Result<Option<T::Value>, Self::Error>
1108    where
1109        T: DeserializeSeed<'de>,
1110    {
1111        match self.size {
1112            Some(0) => Ok(None),
1113            Some(ref mut size) => {
1114                *size -= 1;
1115                seed.deserialize(DefaultDeserializer::new(None)).map(Some)
1116            }
1117            None => Ok(None),
1118        }
1119    }
1120
1121    fn size_hint(&self) -> Option<usize> {
1122        self.size
1123    }
1124}
1125
1126/// A default map access that always returns None (empty map).
1127pub struct DefaultMapAccess;
1128
1129impl Default for DefaultMapAccess {
1130    fn default() -> Self {
1131        Self::new()
1132    }
1133}
1134
1135impl DefaultMapAccess {
1136    /// Creates a new default map access.
1137    pub fn new() -> Self {
1138        DefaultMapAccess
1139    }
1140}
1141
1142impl<'de> MapAccess<'de> for DefaultMapAccess {
1143    type Error = Error;
1144
1145    fn next_key_seed<T>(&mut self, _seed: T) -> Result<Option<T::Value>, Self::Error>
1146    where
1147        T: DeserializeSeed<'de>,
1148    {
1149        Ok(None)
1150    }
1151
1152    fn next_value_seed<T>(&mut self, _seed: T) -> Result<T::Value, Self::Error>
1153    where
1154        T: DeserializeSeed<'de>,
1155    {
1156        // `next_key_seed` always returns `None`, so a well-behaved visitor never
1157        // asks for a value. Return an error instead of panicking if one does.
1158        Err(custom_err(
1159            "no value is available from the default (empty) map access",
1160        ))
1161    }
1162
1163    fn size_hint(&self) -> Option<usize> {
1164        None
1165    }
1166}
1167
1168#[cfg(test)]
1169mod tests {
1170    use super::*;
1171    use crate::nested::adapter::DefaultAdapter;
1172    use serde::Deserialize;
1173
1174    #[derive(Debug, Deserialize, PartialEq)]
1175    struct AllScalars {
1176        b: bool,
1177        i16_val: i16,
1178        i32_val: i32,
1179        i64_val: i64,
1180        u8_val: u8,
1181        u16_val: u16,
1182        u64_val: u64,
1183        f32_val: f32,
1184        f64_val: f64,
1185        s: String,
1186    }
1187
1188    #[derive(Debug, Deserialize, PartialEq)]
1189    struct Config {
1190        hidden_size: i32,
1191    }
1192
1193    #[derive(Debug, Deserialize, PartialEq)]
1194    struct Newtype(i32);
1195
1196    #[derive(Debug, Deserialize, PartialEq)]
1197    enum MockDType {
1198        Float32,
1199        Int32,
1200    }
1201
1202    #[test]
1203    fn should_deserialize_all_scalars_correctly() {
1204        let mut map = HashMap::new();
1205        map.insert("b".to_string(), NestedValue::Bool(true));
1206        map.insert("i16_val".to_string(), NestedValue::I16(16));
1207        map.insert("i32_val".to_string(), NestedValue::I32(32));
1208        map.insert("i64_val".to_string(), NestedValue::I64(64));
1209        map.insert("u8_val".to_string(), NestedValue::U8(8));
1210        map.insert("u16_val".to_string(), NestedValue::U16(16));
1211        map.insert("u64_val".to_string(), NestedValue::U64(64));
1212        map.insert("f32_val".to_string(), NestedValue::F32(1.25));
1213        map.insert("f64_val".to_string(), NestedValue::F64(2.5));
1214        map.insert("s".to_string(), NestedValue::String("burn".to_string()));
1215
1216        let de = Deserializer::<DefaultAdapter>::new(NestedValue::Map(map), false);
1217        let config = AllScalars::deserialize(de).unwrap();
1218
1219        assert_eq!(
1220            config,
1221            AllScalars {
1222                b: true,
1223                i16_val: 16,
1224                i32_val: 32,
1225                i64_val: 64,
1226                u8_val: 8,
1227                u16_val: 16,
1228                u64_val: 64,
1229                f32_val: 1.25,
1230                f64_val: 2.5,
1231                s: "burn".to_string(),
1232            }
1233        );
1234    }
1235
1236    #[test]
1237    fn should_deserialize_i64_as_i32_when_in_range() {
1238        let mut map = HashMap::new();
1239        map.insert("hidden_size".to_string(), NestedValue::I64(768));
1240
1241        let de = Deserializer::<DefaultAdapter>::new(NestedValue::Map(map), false);
1242        let config = Config::deserialize(de).unwrap();
1243
1244        assert_eq!(config, Config { hidden_size: 768 });
1245    }
1246
1247    #[test]
1248    fn should_return_err_on_type_mismatch() {
1249        let mut map = HashMap::new();
1250        map.insert(
1251            "hidden_size".to_string(),
1252            NestedValue::String("768".to_string()),
1253        );
1254
1255        let de = Deserializer::<DefaultAdapter>::new(NestedValue::Map(map), false);
1256        let result = Config::deserialize(de);
1257
1258        assert!(result.is_err());
1259    }
1260
1261    #[test]
1262    fn should_return_err_on_missing_field() {
1263        let map = HashMap::new();
1264        let de = Deserializer::<DefaultAdapter>::new(NestedValue::Map(map), false);
1265        let result = Config::deserialize(de);
1266
1267        assert!(result.is_err());
1268    }
1269
1270    #[test]
1271    fn should_deserialize_newtype_struct() {
1272        let de = Deserializer::<DefaultAdapter>::new(NestedValue::I32(42), false);
1273        let result = Newtype::deserialize(de).unwrap();
1274        assert_eq!(result, Newtype(42));
1275    }
1276
1277    #[test]
1278    fn should_return_err_on_invalid_seq() {
1279        #[derive(Debug, Deserialize)]
1280        #[allow(dead_code)]
1281        struct SeqHolder(Vec<i32>);
1282
1283        let de = Deserializer::<DefaultAdapter>::new(NestedValue::I32(42), false);
1284        let result = SeqHolder::deserialize(de);
1285        assert!(result.is_err());
1286    }
1287
1288    #[test]
1289    fn should_deserialize_dtype_enum() {
1290        let mut map = HashMap::new();
1291        map.insert(
1292            "DType".to_string(),
1293            NestedValue::String("Float32".to_string()),
1294        );
1295        let de = Deserializer::<DefaultAdapter>::new(NestedValue::Map(map), false);
1296        let result = MockDType::deserialize(de).unwrap();
1297        assert_eq!(result, MockDType::Float32);
1298    }
1299
1300    #[test]
1301    fn should_return_err_on_missing_dtype_key() {
1302        let mut map = HashMap::new();
1303        map.insert(
1304            "OtherKey".to_string(),
1305            NestedValue::String("Float32".to_string()),
1306        );
1307        let de = Deserializer::<DefaultAdapter>::new(NestedValue::Map(map), false);
1308        let result = MockDType::deserialize(de);
1309        assert!(result.is_err());
1310    }
1311
1312    #[test]
1313    fn should_return_err_on_invalid_enum_shape() {
1314        let de = Deserializer::<DefaultAdapter>::new(NestedValue::I32(42), false);
1315        let result = MockDType::deserialize(de);
1316        assert!(result.is_err());
1317    }
1318
1319    #[test]
1320    fn should_return_err_on_invalid_dtype_variant_type() {
1321        let mut map = HashMap::new();
1322        map.insert("DType".to_string(), NestedValue::I32(42));
1323        let de = Deserializer::<DefaultAdapter>::new(NestedValue::Map(map), false);
1324        let result = MockDType::deserialize(de);
1325        assert!(result.is_err());
1326    }
1327
1328    #[test]
1329    fn should_skip_a_field_the_target_does_not_name() {
1330        for extra in [
1331            NestedValue::Default(None),
1332            NestedValue::Unsupported("torch.device".to_string()),
1333        ] {
1334            let mut map = HashMap::new();
1335            map.insert("hidden_size".to_string(), NestedValue::I32(1));
1336            map.insert("extra".to_string(), extra);
1337
1338            let de = Deserializer::<DefaultAdapter>::new(NestedValue::Map(map), false);
1339            assert_eq!(Config::deserialize(de).unwrap(), Config { hidden_size: 1 });
1340        }
1341    }
1342
1343    #[test]
1344    fn should_return_err_on_unsupported_value() {
1345        let mut map = HashMap::new();
1346        map.insert(
1347            "hidden_size".to_string(),
1348            NestedValue::Unsupported("numpy.int64".to_string()),
1349        );
1350        let de = Deserializer::<DefaultAdapter>::new(NestedValue::Map(map), false);
1351        let err = Config::deserialize(de).unwrap_err();
1352        assert!(err.to_string().contains("numpy.int64"), "{err}");
1353    }
1354
1355    #[test]
1356    fn should_deserialize_u8s_as_byte_buf() {
1357        #[derive(Debug, Deserialize, PartialEq)]
1358        struct ByteHolder(Vec<u8>);
1359
1360        let de = Deserializer::<DefaultAdapter>::new(NestedValue::U8s(vec![1, 2, 3, 4]), false);
1361        let result = ByteHolder::deserialize(de).unwrap();
1362        assert_eq!(result, ByteHolder(vec![1, 2, 3, 4]));
1363    }
1364}
pytorch_reader/nested/de.rs

pytorch_reader/nested/
de.rs
