facet_reflect/wip/
mod.rs

1use crate::{ReflectError, ValueId};
2use alloc::{format, string::ToString};
3use alloc::{vec, vec::Vec};
4use bitflags::bitflags;
5use core::{fmt, marker::PhantomData};
6use facet_core::{Def, Facet, FieldError, PtrConst, PtrMut, PtrUninit, Shape, Variant};
7use flat_map::FlatMap;
8use log::trace;
9use owo_colors::OwoColorize;
10
11use alloc::string::String;
12
13mod iset;
14pub use iset::*;
15
16mod enum_;
17mod flat_map;
18
19mod heap_value;
20pub use heap_value::*;
21
22fn def_kind(def: &Def) -> &'static str {
23    match def {
24        Def::Scalar(_) => "scalar",
25        Def::Struct(_) => "struct",
26        Def::Map(_) => "map",
27        Def::List(_) => "list",
28        Def::Enum(_) => "enum",
29        Def::Option(_) => "option",
30        Def::SmartPointer(_) => "smart_ptr",
31        _ => "other",
32    }
33}
34
35/// Represents a frame in the initialization stack
36pub struct Frame {
37    /// The value we're initializing
38    data: PtrUninit<'static>,
39
40    /// The shape of the value
41    shape: &'static Shape,
42
43    /// If set, when we're initialized, we must mark the
44    /// parent's indexth field as initialized.
45    field_index_in_parent: Option<usize>,
46
47    /// Tracking which of our fields are initialized
48    /// TODO: I'm not sure we should track "ourselves" as initialized — we always have the
49    /// parent to look out for, right now we're tracking children in two states, which isn't ideal
50    istate: IState,
51}
52
53impl Frame {
54    /// Given a ValueId and an IState, recompose a Frame suitable for tracking
55    fn recompose(id: ValueId, istate: IState) -> Self {
56        Frame {
57            data: PtrUninit::new(id.ptr as *mut u8),
58            shape: id.shape,
59            field_index_in_parent: None,
60            istate,
61        }
62    }
63
64    /// Deallocates the memory used by this frame if it was heap-allocated.
65    fn dealloc_if_needed(&mut self) {
66        if self.istate.flags.contains(FrameFlags::ALLOCATED) {
67            trace!(
68                "[{}] {:p} => deallocating {}",
69                self.istate.depth,
70                self.data.as_mut_byte_ptr().magenta(),
71                self.shape.green(),
72            );
73            if self.shape.layout.size() != 0 {
74                trace!("Ok I swear we're deallocating it");
75                unsafe {
76                    alloc::alloc::dealloc(self.data.as_mut_byte_ptr(), self.shape.layout);
77                }
78            }
79            self.istate.flags.remove(FrameFlags::ALLOCATED);
80        } else {
81            trace!(
82                "[{}] {:p} => NOT deallocating {} (not ALLOCATED)",
83                self.istate.depth,
84                self.data.as_mut_byte_ptr().magenta(),
85                self.shape.green(),
86            );
87        }
88    }
89}
90
91struct DebugToDisplay<T>(T);
92
93impl<T> fmt::Debug for DebugToDisplay<T>
94where
95    T: fmt::Display,
96{
97    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
98        fmt::Display::fmt(&self.0, f)
99    }
100}
101
102impl fmt::Debug for Frame {
103    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
104        f.debug_struct("Frame")
105            .field("shape", &DebugToDisplay(&self.shape))
106            .field("kind", &def_kind(&self.shape.def))
107            .field("index", &self.field_index_in_parent)
108            .field("mode", &self.istate.mode)
109            .finish()
110    }
111}
112
113impl Frame {
114    /// Returns the value ID for a frame
115    fn id(&self) -> ValueId {
116        ValueId::new(self.shape, self.data.as_byte_ptr())
117    }
118
119    /// Returns true if the frame is fully initialized
120    fn is_fully_initialized(&self) -> bool {
121        match self.shape.def {
122            Def::Struct(sd) => self.istate.fields.are_all_set(sd.fields.len()),
123            Def::Enum(_) => match self.istate.variant.as_ref() {
124                None => false,
125                Some(v) => self.istate.fields.are_all_set(v.data.fields.len()),
126            },
127            _ => self.istate.fields.are_all_set(1),
128        }
129    }
130
131    // Safety: only call if is fully initialized
132    unsafe fn drop_and_dealloc_if_needed(mut self) {
133        trace!(
134            "[Frame::drop] Dropping frame for shape {} at {:p}",
135            self.shape.blue(),
136            self.data.as_byte_ptr()
137        );
138        if let Some(drop_in_place) = self.shape.vtable.drop_in_place {
139            unsafe {
140                trace!(
141                    "[Frame::drop] Invoking drop_in_place for shape {} at {:p}",
142                    self.shape.green(),
143                    self.data.as_byte_ptr()
144                );
145                drop_in_place(self.data.assume_init());
146            }
147        } else {
148            trace!(
149                "[Frame::drop] No drop_in_place function for shape {}",
150                self.shape.blue(),
151            );
152        }
153        self.dealloc_if_needed();
154    }
155
156    /// Marks the frame as fully initialized
157    unsafe fn mark_fully_initialized(&mut self) {
158        match self.shape.def {
159            Def::Struct(sd) => {
160                self.istate.fields = ISet::all(sd.fields);
161            }
162            Def::Enum(_) => {
163                if let Some(variant) = &self.istate.variant {
164                    self.istate.fields = ISet::all(variant.data.fields);
165                }
166            }
167            _ => {
168                self.istate.fields.set(0);
169            }
170        }
171    }
172}
173
174/// Initialization state
175struct IState {
176    /// Variant chosen — for everything except enums, this stays None
177    variant: Option<Variant>,
178
179    /// Fields that were initialized. For scalars, we only track 0
180    fields: ISet,
181
182    /// The depth of the frame in the stack
183    depth: usize,
184
185    /// The special mode of this frame (if any)
186    mode: FrameMode,
187
188    /// If true, must be freed when dropped
189    flags: FrameFlags,
190
191    /// The current index for list elements
192    list_index: Option<usize>,
193
194    /// The current key for map elements
195    #[allow(dead_code)]
196    map_key: Option<String>,
197}
198
199bitflags! {
200    /// Flags that can be applied to frames
201    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
202    pub struct FrameFlags: u64 {
203        /// An empty set of flags
204        const EMPTY = 0;
205
206        /// We allocated this frame on the heap, we need to deallocated it when popping
207        const ALLOCATED = 1 << 0;
208
209        /// This value was moved out of — it's not part of the value we're building and
210        /// we shouldn't error out when we build and we notice it's not initialized.
211        /// In fact, it should not be tracked at all.
212        const MOVED = 1 << 1;
213    }
214
215    // Note: there is no 'initialized' flag because initialization can be partial — it's tracked via `ISet`
216}
217
218impl IState {
219    /// Creates a new `IState` with the given depth.
220    pub fn new(depth: usize, mode: FrameMode, flags: FrameFlags) -> Self {
221        Self {
222            variant: None,
223            fields: Default::default(),
224            depth,
225            mode,
226            flags,
227            list_index: None,
228            map_key: None,
229        }
230    }
231
232    /// Sets the list index and returns self for method chaining
233    #[allow(dead_code)]
234    pub fn with_list_index(mut self, index: usize) -> Self {
235        self.list_index = Some(index);
236        self
237    }
238
239    /// Sets the map key and returns self for method chaining
240    #[allow(dead_code)]
241    pub fn with_map_key(mut self, key: String) -> Self {
242        self.map_key = Some(key);
243        self
244    }
245}
246
247/// Represents the special mode a frame can be in
248#[derive(Debug, Clone, Copy, PartialEq, Eq)]
249pub enum FrameMode {
250    /// Root frame
251    Root,
252    /// Struct field
253    Field,
254    /// Frame represents a list element
255    ListElement,
256    /// Frame represents a map key
257    MapKey,
258    /// Frame represents a map value with the given key frame index
259    MapValue {
260        /// The index of the key frame associated with this map value
261        index: usize,
262    },
263    /// Frame represents the Some variant of an option (that we allocated)
264    OptionSome,
265    /// Frame represents the None variant of an option (no allocation needed)
266    /// Any `put` should fail
267    OptionNone,
268}
269
270/// A work-in-progress heap-allocated value
271pub struct Wip<'a> {
272    /// stack of frames to keep track of deeply nested initialization
273    frames: alloc::vec::Vec<Frame>,
274
275    /// keeps track of initialization of out-of-tree frames
276    istates: FlatMap<ValueId, IState>,
277
278    /// lifetime of the shortest reference we hold
279    phantom: PhantomData<&'a ()>,
280}
281
282impl<'a> Wip<'a> {
283    /// Returns the number of frames on the stack
284    pub fn frames_count(&self) -> usize {
285        self.frames.len()
286    }
287
288    /// Allocates a new value of the given shape
289    pub fn alloc_shape(shape: &'static Shape) -> Self {
290        let data = shape.allocate();
291        Self {
292            frames: alloc::vec![Frame {
293                data,
294                shape,
295                field_index_in_parent: None,
296                istate: IState::new(0, FrameMode::Root, FrameFlags::ALLOCATED),
297            }],
298            istates: Default::default(),
299            phantom: PhantomData,
300        }
301    }
302
303    /// Allocates a new value of type `S`
304    pub fn alloc<S: Facet>() -> Self {
305        Self::alloc_shape(S::SHAPE)
306    }
307
308    fn track(&mut self, frame: Frame) {
309        // fields might be partially initialized (in-place) and then
310        // we might come back to them, so because they're popped off
311        // the stack, we still need to track them _somewhere_
312        //
313        // the root also relies on being tracked in the drop impl
314        if !frame.istate.flags.contains(FrameFlags::MOVED) {
315            self.istates.insert(frame.id(), frame.istate);
316        }
317    }
318
319    unsafe fn mark_moved_out_of(&mut self, frame: &mut Frame) {
320        frame.dealloc_if_needed();
321        ISet::clear(&mut frame.istate.fields);
322        frame.istate.variant = None;
323        frame.istate.flags.insert(FrameFlags::MOVED);
324        // make sure this isn't tracked here anymore — we don't want to have
325        // any metadata associated with it that gets restored by mistake
326        self.istates.remove(&frame.id());
327    }
328
329    /// Returns the shape of the current frame
330    pub fn shape(&self) -> &'static Shape {
331        self.frames.last().unwrap().shape
332    }
333
334    /// Return true if the last frame is in option mode
335    pub fn in_option(&self) -> bool {
336        let Some(frame) = self.frames.last() else {
337            return false;
338        };
339        matches!(frame.istate.mode, FrameMode::OptionSome)
340    }
341
342    /// Returns the mode of the current frame
343    pub fn mode(&self) -> FrameMode {
344        self.frames.last().unwrap().istate.mode
345    }
346
347    /// Asserts everything is initialized and that invariants are upheld (if any)
348    pub fn build(mut self) -> Result<HeapValue<'a>, ReflectError> {
349        trace!("[{}] ⚒️ It's BUILD time", self.frames.len());
350
351        // 1. Track all frames currently on the stack into istates
352        while let Some(frame) = self.pop_inner() {
353            self.track(frame);
354        }
355
356        // 2. Find the root frame
357        let Some((root_id, _)) = self.istates.iter().find(|(_k, istate)| istate.depth == 0) else {
358            trace!("No root found, possibly already built or empty WIP");
359            return Err(ReflectError::OperationFailed {
360                shape: <()>::SHAPE,
361                operation: "tried to build a value but there was no root frame tracked",
362            });
363        };
364
365        let root_id = *root_id;
366        // We need to *keep* the root istate for the check, so we clone it or get it immutably.
367        // Let's retrieve it immutably first. The `istates` map will be dropped with `Wip` anyway.
368        let root_istate = self
369            .istates
370            .remove(&root_id)
371            .expect("Root ID found but not present in istates, this is a bug"); // Clone needed to avoid borrowing issues later potentially
372
373        let root_frame = Frame::recompose(root_id, root_istate);
374        let root_shape = root_frame.shape;
375        let root_data_ptr = root_frame.data; // Keep the root pointer for the final HeapValue
376
377        // 6. Transfer ownership of the root data to the HeapValue
378        // The root frame should have had the ALLOCATED flag if it was heap allocated.
379        // We need to ensure the Guard takes ownership correctly.
380        // Find the original root istate again to check the flag.
381        let guard = Guard {
382            ptr: root_data_ptr.as_mut_byte_ptr(),
383            layout: root_shape.layout,
384        };
385
386        // 3. Initialize `to_check`
387        let mut to_check = alloc::vec![root_frame];
388
389        // 4. Traverse the tree
390        while let Some(frame) = to_check.pop() {
391            trace!(
392                "Checking frame: shape={} at {:p}, flags={:?}, mode={:?}",
393                frame.shape.blue(),
394                frame.data.as_byte_ptr(),
395                frame.istate.flags.bright_magenta(),
396                frame.istate.mode,
397            );
398
399            // Skip moved frames
400            if frame.istate.flags.contains(FrameFlags::MOVED) {
401                trace!(
402                    "{}",
403                    "Frame was moved out of, skipping initialization check".yellow()
404                );
405                continue;
406            }
407
408            // Check initialization for the current frame
409            match frame.shape.def {
410                Def::Struct(sd) => {
411                    if !frame.is_fully_initialized() {
412                        // find the field that's not initialized
413                        for i in 0..sd.fields.len() {
414                            if !frame.istate.fields.has(i) {
415                                let field = &sd.fields[i];
416                                return Err(ReflectError::UninitializedField {
417                                    shape: frame.shape,
418                                    field_name: field.name,
419                                });
420                            }
421                        }
422                        // Should be unreachable
423                        unreachable!(
424                            "Enum variant not fully initialized but couldn't find which field"
425                        );
426                    }
427
428                    // If initialized, push children to check stack
429                    for (i, field) in sd.fields.iter().enumerate() {
430                        let field_shape = field.shape();
431                        let field_ptr = unsafe { frame.data.field_init_at(field.offset) };
432                        let field_id = ValueId::new(field_shape, field_ptr.as_byte_ptr());
433
434                        if let Some(field_istate) = self.istates.remove(&field_id) {
435                            trace!(
436                                "Queueing struct field check: #{} '{}' of {}: shape={}, ptr={:p}",
437                                i.to_string().bright_cyan(),
438                                field.name.bright_blue(),
439                                frame.shape.blue(),
440                                field_shape.green(),
441                                field_ptr.as_byte_ptr()
442                            );
443                            let field_frame = Frame::recompose(field_id, field_istate);
444                            to_check.push(field_frame);
445                        }
446                    }
447                }
448                Def::Enum(_ed) => {
449                    if let Some(variant) = &frame.istate.variant {
450                        if !frame.istate.fields.are_all_set(variant.data.fields.len()) {
451                            // Find the uninitialized field
452                            for (i, field) in variant.data.fields.iter().enumerate() {
453                                if !frame.istate.fields.has(i) {
454                                    return Err(ReflectError::UninitializedEnumField {
455                                        shape: frame.shape,
456                                        variant_name: variant.name,
457                                        field_name: field.name,
458                                    });
459                                }
460                            }
461                            // Should be unreachable
462                            unreachable!(
463                                "Enum variant not fully initialized but couldn't find which field"
464                            );
465                        }
466
467                        // If initialized, push children to check stack
468                        for (i, field) in variant.data.fields.iter().enumerate() {
469                            let field_shape = field.shape();
470                            // Enum fields are potentially at different offsets depending on the variant layout.
471                            // We assume the `frame.data` points to the start of the enum's data payload
472                            // (after the discriminant if applicable and handled by layout).
473                            let field_ptr = unsafe { frame.data.field_init_at(field.offset) };
474                            let field_id = ValueId::new(field_shape, field_ptr.as_byte_ptr());
475
476                            if let Some(field_istate) = self.istates.remove(&field_id) {
477                                trace!(
478                                    "Queueing enum field check: #{} '{}' of variant '{}' of {}: shape={}, ptr={:p}",
479                                    i.to_string().bright_cyan(),
480                                    field.name.bright_blue(),
481                                    variant.name.yellow(),
482                                    frame.shape.blue(),
483                                    field_shape.green(),
484                                    field_ptr.as_byte_ptr()
485                                );
486                                let field_frame = Frame::recompose(field_id, field_istate);
487                                to_check.push(field_frame);
488                            }
489                        }
490                    } else {
491                        // No variant selected is an error during build
492                        return Err(ReflectError::NoVariantSelected { shape: frame.shape });
493                    }
494                }
495                // For types that manage their own contents (List, Map, Option, Scalar, etc.),
496                // we just need to check if the *container* itself is marked as initialized.
497                // The recursive check handles struct/enum *elements* within these containers if they exist.
498                Def::List(_)
499                | Def::Map(_)
500                | Def::Option(_)
501                | Def::Scalar(_)
502                | Def::SmartPointer(_)
503                | Def::Array(_)
504                | Def::Slice(_) => {
505                    if !frame.istate.fields.are_all_set(1) {
506                        // Check specific modes for better errors
507                        match frame.istate.mode {
508                            FrameMode::OptionNone => {
509                                // This should technically be marked initialized, but if not, treat as uninit Option
510                                return Err(ReflectError::UninitializedValue {
511                                    shape: frame.shape,
512                                });
513                            }
514                            // Add more specific checks if needed, e.g., for lists/maps that started but weren't finished?
515                            _ => {
516                                return Err(ReflectError::UninitializedValue {
517                                    shape: frame.shape,
518                                });
519                            }
520                        }
521                    }
522                    // No children to push onto `to_check` from the perspective of the *container* frame itself.
523                    // If a List contains Structs, those struct frames would have been pushed/popped
524                    // and their states tracked individually in `istates`, and checked when encountered via
525                    // `to_check` if they were fields of another struct/enum.
526                    // The `Drop` logic handles cleaning these contained items based on the container's drop_in_place.
527                    // For `build`, we trust that if the container is marked initialized, its contents are valid
528                    // according to its type's rules.
529                }
530                // Handle other Def variants if necessary
531                _ => {
532                    // Default: Check if initialized using the standard method
533                    if !frame.istate.fields.are_all_set(1) {
534                        return Err(ReflectError::UninitializedValue { shape: frame.shape });
535                    }
536                }
537            }
538        }
539
540        // If we finished the loop, all reachable and non-moved frames are initialized.
541        trace!("All reachable frames checked and initialized.");
542
543        // 5. Check invariants on the root
544        let data = unsafe { root_data_ptr.assume_init() };
545        if let Some(invariant_fn) = root_shape.vtable.invariants {
546            trace!(
547                "Checking invariants for root shape {} at {:p}",
548                root_shape.green(),
549                data.as_byte_ptr()
550            );
551            if !unsafe { invariant_fn(PtrConst::new(data.as_byte_ptr())) } {
552                return Err(ReflectError::InvariantViolation {
553                    invariant: "Custom validation function returned false",
554                });
555            }
556        } else {
557            trace!(
558                "No invariants to check for root shape {}",
559                root_shape.blue()
560            );
561        }
562
563        FlatMap::clear(&mut self.istates); // Prevent Drop from running on the successfully built value.
564
565        Ok(HeapValue {
566            guard: Some(guard),
567            shape: root_shape,
568            phantom: PhantomData,
569        })
570    }
571
572    /// Selects a field of a struct or enum variant by index and pushes it onto the frame stack.
573    ///
574    /// # Arguments
575    ///
576    /// * `index` - The index of the field to select.
577    ///
578    /// # Returns
579    ///
580    /// * `Ok(Self)` if the field was successfully selected and pushed.
581    /// * `Err(ReflectError)` if the current frame is not a struct or an enum with a selected variant,
582    ///   or if the field doesn't exist.
583    pub fn field(mut self, index: usize) -> Result<Self, ReflectError> {
584        let frame = self.frames.last_mut().unwrap();
585        let shape = frame.shape;
586
587        let (field, field_offset) = match shape.def {
588            Def::Struct(def) => {
589                if index >= def.fields.len() {
590                    return Err(ReflectError::FieldError {
591                        shape,
592                        field_error: FieldError::NoSuchField,
593                    });
594                }
595                let field = &def.fields[index];
596                (field, field.offset)
597            }
598            Def::Enum(_) => {
599                let Some(variant) = frame.istate.variant.as_ref() else {
600                    return Err(ReflectError::OperationFailed {
601                        shape,
602                        operation: "tried to access a field but no variant was selected",
603                    });
604                };
605
606                if index >= variant.data.fields.len() {
607                    return Err(ReflectError::FieldError {
608                        shape,
609                        field_error: FieldError::NoSuchField,
610                    });
611                }
612
613                let field = &variant.data.fields[index];
614                (field, field.offset)
615            }
616            _ => {
617                return Err(ReflectError::WasNotA {
618                    expected: "struct or enum",
619                    actual: shape,
620                });
621            }
622        };
623
624        let field_data = unsafe { frame.data.field_uninit_at(field_offset) };
625
626        let mut frame = Frame {
627            data: field_data,
628            shape: field.shape(),
629            field_index_in_parent: Some(index),
630            // we didn't have to allocate that field, it's a struct field, so it's not allocated
631            istate: IState::new(self.frames.len(), FrameMode::Field, FrameFlags::EMPTY),
632        };
633        trace!(
634            "[{}] Selecting field {} ({}#{}) of {}",
635            self.frames.len(),
636            field.name.blue(),
637            field.shape().green(),
638            index.yellow(),
639            shape.blue(),
640        );
641        if let Some(iset) = self.istates.remove(&frame.id()) {
642            trace!(
643                "[{}] Restoring saved state for {} (istate.mode = {:?}, istate.fields = {:?}, istate.flags = {:?}, istate.depth = {:?})",
644                self.frames.len(),
645                frame.id().shape.blue(),
646                iset.mode,
647                iset.fields,
648                iset.flags,
649                iset.depth
650            );
651            frame.istate = iset;
652        } else {
653            trace!(
654                "[{}] no saved state for field {} ({}#{}) of {}",
655                self.frames.len(),
656                field.name.blue(),
657                field.shape().green(),
658                index.yellow(),
659                shape.blue(),
660            );
661        }
662        self.frames.push(frame);
663        Ok(self)
664    }
665
666    /// Finds the index of a field in a struct or enum variant by name.
667    ///
668    /// # Arguments
669    ///
670    /// * `name` - The name of the field to find.
671    ///
672    /// # Returns
673    ///
674    /// * `Some(usize)` if the field was found.
675    /// * `None` if the current frame is not a struct or an enum with a selected variant,
676    ///   or if the field doesn't exist.
677    pub fn field_index(&self, name: &str) -> Option<usize> {
678        let frame = self.frames.last()?;
679        match frame.shape.def {
680            Def::Struct(def) => def.fields.iter().position(|f| f.name == name),
681            Def::Enum(_) => {
682                // Get the selected variant
683                let variant = frame.istate.variant.as_ref()?;
684                variant.data.fields.iter().position(|f| f.name == name)
685            }
686            _ => None,
687        }
688    }
689
690    /// Selects a field of a struct or enum variant by name and pushes it onto the frame stack.
691    ///
692    /// # Arguments
693    ///
694    /// * `name` - The name of the field to select.
695    ///
696    /// # Returns
697    ///
698    /// * `Ok(Self)` if the field was successfully selected and pushed.
699    /// * `Err(ReflectError)` if the current frame is not a struct or an enum with a selected variant,
700    ///   or if the field doesn't exist.
701    pub fn field_named(self, name: &str) -> Result<Self, ReflectError> {
702        let frame = self.frames.last().unwrap();
703        let shape = frame.shape;
704
705        // For enums, ensure a variant is selected
706        if let Def::Enum(_) = shape.def {
707            if frame.istate.variant.is_none() {
708                return Err(ReflectError::OperationFailed {
709                    shape,
710                    operation: "tried to access a field by name but no variant was selected",
711                });
712            }
713        }
714
715        let index = self.field_index(name).ok_or(ReflectError::FieldError {
716            shape,
717            field_error: FieldError::NoSuchField,
718        })?;
719
720        self.field(index)
721    }
722
723    /// Puts a value of type `T` into the current frame.
724    ///
725    /// # Arguments
726    ///
727    /// * `t` - The value to put into the frame.
728    ///
729    /// # Returns
730    ///
731    /// * `Ok(Self)` if the value was successfully put into the frame.
732    /// * `Err(ReflectError)` if there was an error putting the value into the frame.
733    pub fn put<'val, T: Facet + 'val>(mut self, t: T) -> Result<Wip<'val>, ReflectError>
734    where
735        'a: 'val,
736    {
737        let Some(frame) = self.frames.last_mut() else {
738            return Err(ReflectError::OperationFailed {
739                shape: T::SHAPE,
740                operation: "tried to put a T but there was no frame to put T into",
741            });
742        };
743
744        // check that the type matches
745        if !frame.shape.is_type::<T>() {
746            // maybe we're putting into an Option<T> ?
747            if frame.shape.is_type::<Option<T>>() {
748                trace!("Putting into an option!");
749                let Def::Option(od) = frame.shape.def else {
750                    unreachable!()
751                };
752
753                // unfortunately this all involves copying data over
754                let src = PtrConst::new(&raw const t);
755                if frame.istate.fields.is_any_set() {
756                    let data = unsafe { frame.data.assume_init() };
757                    unsafe { (od.vtable.replace_with_fn)(data, Some(src)) };
758                } else {
759                    let data = frame.data;
760                    unsafe { (od.vtable.init_some_fn)(data, src) };
761                }
762                unsafe {
763                    frame.mark_fully_initialized();
764                }
765                // t will drop naturally at the end of this function
766
767                let shape = frame.shape;
768                let index = frame.field_index_in_parent;
769
770                // mark the field as initialized
771                self.mark_field_as_initialized(shape, index)?;
772
773                trace!("[{}] Just put a {} value", self.frames.len(), shape.green());
774
775                return Ok(self);
776            }
777
778            return Err(ReflectError::WrongShape {
779                expected: frame.shape,
780                actual: T::SHAPE,
781            });
782        }
783
784        // de-initialize partially initialized fields
785        if frame.istate.variant.is_some() || frame.istate.fields.is_any_set() {
786            trace!(
787                "De-initializing partially initialized fields for {}",
788                frame.shape
789            );
790
791            match frame.shape.def {
792                Def::Struct(sd) => {
793                    for (i, field) in sd.fields.iter().enumerate() {
794                        if frame.istate.fields.has(i) {
795                            if let Some(drop_fn) = field.shape().vtable.drop_in_place {
796                                unsafe {
797                                    let field_ptr = frame.data.as_mut_byte_ptr().add(field.offset);
798                                    drop_fn(PtrMut::new(field_ptr));
799                                }
800                            }
801                        }
802                    }
803                }
804                Def::Enum(_) => {
805                    if let Some(variant) = &frame.istate.variant {
806                        for (i, field) in variant.data.fields.iter().enumerate() {
807                            if frame.istate.fields.has(i) {
808                                if let Some(drop_fn) = field.shape().vtable.drop_in_place {
809                                    unsafe {
810                                        let field_ptr =
811                                            frame.data.as_mut_byte_ptr().add(field.offset);
812                                        drop_fn(PtrMut::new(field_ptr));
813                                    }
814                                }
815                            }
816                        }
817                    }
818                }
819                _ => {
820                    // For scalar types, nothing to do if not fully initialized
821                }
822            }
823
824            // Reset initialization state
825            frame.istate.variant = None;
826            ISet::clear(&mut frame.istate.fields);
827        }
828
829        unsafe {
830            frame.data.put(t);
831            frame.mark_fully_initialized();
832        }
833
834        let shape = frame.shape;
835        let index = frame.field_index_in_parent;
836
837        // mark the field as initialized
838        self.mark_field_as_initialized(shape, index)?;
839
840        trace!("[{}] Just put a {} value", self.frames.len(), shape.green());
841
842        Ok(self)
843    }
844
845    /// Tries to parse the current frame's value from a string
846    pub fn parse(mut self, s: &str) -> Result<Self, ReflectError> {
847        let Some(frame) = self.frames.last_mut() else {
848            return Err(ReflectError::OperationFailed {
849                shape: <()>::SHAPE,
850                operation: "tried to parse value but there was no frame",
851            });
852        };
853
854        let shape = frame.shape;
855        let index = frame.field_index_in_parent;
856
857        let Some(parse_fn) = frame.shape.vtable.parse else {
858            return Err(ReflectError::OperationFailed {
859                shape: frame.shape,
860                operation: "type does not implement Parse",
861            });
862        };
863        match unsafe { (parse_fn)(s, frame.data) } {
864            Ok(_res) => {
865                unsafe {
866                    frame.mark_fully_initialized();
867                }
868
869                // mark the field as initialized
870                self.mark_field_as_initialized(shape, index)?;
871
872                Ok(self)
873            }
874            Err(_) => Err(ReflectError::OperationFailed {
875                shape,
876                operation: "parsing",
877            }),
878        }
879    }
880
881    /// Puts the default value in the current frame.
882    pub fn put_default(mut self) -> Result<Self, ReflectError> {
883        let Some(frame) = self.frames.last_mut() else {
884            return Err(ReflectError::OperationFailed {
885                shape: <()>::SHAPE,
886                operation: "tried to put default value but there was no frame",
887            });
888        };
889
890        let vtable = frame.shape.vtable;
891
892        let Some(default_in_place) = vtable.default_in_place else {
893            return Err(ReflectError::OperationFailed {
894                shape: frame.shape,
895                operation: "type does not implement Default",
896            });
897        };
898        unsafe {
899            default_in_place(frame.data);
900            frame.mark_fully_initialized();
901        }
902
903        let shape = frame.shape;
904        let index = frame.field_index_in_parent;
905
906        // mark the field as initialized
907        self.mark_field_as_initialized(shape, index)?;
908
909        Ok(self)
910    }
911
912    /// Marks a field as initialized in the parent frame.
913    fn mark_field_as_initialized(
914        &mut self,
915        shape: &'static Shape,
916        index: Option<usize>,
917    ) -> Result<(), ReflectError> {
918        if let Some(index) = index {
919            let parent_index = self.frames.len().saturating_sub(2);
920            let num_frames = self.frames.len();
921            let Some(parent) = self.frames.get_mut(parent_index) else {
922                return Err(ReflectError::OperationFailed {
923                    shape,
924                    operation: "was supposed to mark a field as initialized, but there was no parent frame",
925                });
926            };
927            let parent_shape = parent.shape;
928            trace!(
929                "[{}] {}.{} initialized with {}",
930                num_frames,
931                parent_shape.blue(),
932                index.yellow(),
933                shape.green()
934            );
935
936            if matches!(parent.shape.def, Def::Enum(_)) && parent.istate.variant.is_none() {
937                return Err(ReflectError::OperationFailed {
938                    shape,
939                    operation: "was supposed to mark a field as initialized, but the parent frame was an enum and didn't have a variant chosen",
940                });
941            }
942
943            if parent.istate.fields.has(index) {
944                return Err(ReflectError::OperationFailed {
945                    shape,
946                    operation: "was supposed to mark a field as initialized, but the parent frame already had it marked as initialized",
947                });
948            }
949
950            parent.istate.fields.set(index);
951        }
952        Ok(())
953    }
954
955    /// Returns the shape of the element type for a list/array
956    pub fn element_shape(&self) -> Result<&'static Shape, ReflectError> {
957        let frame = self.frames.last().unwrap();
958        let shape = frame.shape;
959
960        match shape.def {
961            Def::List(list_def) => Ok(list_def.t()),
962            _ => Err(ReflectError::WasNotA {
963                expected: "list or array",
964                actual: shape,
965            }),
966        }
967    }
968
969    /// Returns the shape of the key type for a map
970    pub fn key_shape(&self) -> Result<&'static Shape, ReflectError> {
971        let frame = self.frames.last().unwrap();
972        let shape = frame.shape;
973
974        match shape.def {
975            Def::Map(map_def) => Ok(map_def.k),
976            _ => Err(ReflectError::WasNotA {
977                expected: "map",
978                actual: shape,
979            }),
980        }
981    }
982
983    /// Creates an empty list without pushing any elements
984    pub fn put_empty_list(mut self) -> Result<Self, ReflectError> {
985        let Some(frame) = self.frames.last_mut() else {
986            return Err(ReflectError::OperationFailed {
987                shape: <()>::SHAPE,
988                operation: "tried to create empty list but there was no frame",
989            });
990        };
991
992        if !matches!(frame.shape.def, Def::List(_)) {
993            return Err(ReflectError::WasNotA {
994                expected: "list or array",
995                actual: frame.shape,
996            });
997        }
998
999        let vtable = frame.shape.vtable;
1000
1001        // Initialize an empty list
1002        let Some(default_in_place) = vtable.default_in_place else {
1003            return Err(ReflectError::OperationFailed {
1004                shape: frame.shape,
1005                operation: "list type does not implement Default",
1006            });
1007        };
1008
1009        unsafe {
1010            default_in_place(frame.data);
1011            frame.mark_fully_initialized();
1012        }
1013
1014        let shape = frame.shape;
1015        let index = frame.field_index_in_parent;
1016
1017        // Mark the field as initialized
1018        self.mark_field_as_initialized(shape, index)?;
1019
1020        Ok(self)
1021    }
1022
1023    /// Creates an empty map without pushing any entries
1024    pub fn put_empty_map(mut self) -> Result<Self, ReflectError> {
1025        let Some(frame) = self.frames.last_mut() else {
1026            return Err(ReflectError::OperationFailed {
1027                shape: <()>::SHAPE,
1028                operation: "tried to create empty map but there was no frame",
1029            });
1030        };
1031
1032        if !matches!(frame.shape.def, Def::Map(_)) {
1033            return Err(ReflectError::WasNotA {
1034                expected: "map or hash map",
1035                actual: frame.shape,
1036            });
1037        }
1038
1039        let vtable = frame.shape.vtable;
1040
1041        // Initialize an empty map
1042        let Some(default_in_place) = vtable.default_in_place else {
1043            return Err(ReflectError::OperationFailed {
1044                shape: frame.shape,
1045                operation: "map type does not implement Default",
1046            });
1047        };
1048
1049        unsafe {
1050            default_in_place(frame.data);
1051            frame.mark_fully_initialized();
1052        }
1053
1054        let shape = frame.shape;
1055        let index = frame.field_index_in_parent;
1056
1057        // Mark the field as initialized
1058        self.mark_field_as_initialized(shape, index)?;
1059
1060        Ok(self)
1061    }
1062
1063    /// Begins pushback mode for a list/array, allowing elements to be added one by one
1064    pub fn begin_pushback(mut self) -> Result<Self, ReflectError> {
1065        let Some(frame) = self.frames.last_mut() else {
1066            return Err(ReflectError::OperationFailed {
1067                shape: <()>::SHAPE,
1068                operation: "tried to begin pushback but there was no frame",
1069            });
1070        };
1071
1072        if !matches!(frame.shape.def, Def::List(_)) {
1073            return Err(ReflectError::WasNotA {
1074                expected: "list or array",
1075                actual: frame.shape,
1076            });
1077        }
1078
1079        let vtable = frame.shape.vtable;
1080
1081        // Initialize an empty list if it's not already initialized
1082        if !frame.istate.fields.has(0) {
1083            let Some(default_in_place) = vtable.default_in_place else {
1084                return Err(ReflectError::OperationFailed {
1085                    shape: frame.shape,
1086                    operation: "list type does not implement Default",
1087                });
1088            };
1089
1090            unsafe {
1091                default_in_place(frame.data);
1092                frame.istate.fields.set(0);
1093            }
1094        }
1095
1096        Ok(self)
1097    }
1098
1099    /// Begins insertion mode for a map, allowing key-value pairs to be added one by one
1100    pub fn begin_map_insert(mut self) -> Result<Self, ReflectError> {
1101        let Some(frame) = self.frames.last_mut() else {
1102            return Err(ReflectError::OperationFailed {
1103                shape: <()>::SHAPE,
1104                operation: "tried to begin map insertion but there was no frame",
1105            });
1106        };
1107
1108        if !matches!(frame.shape.def, Def::Map(_)) {
1109            return Err(ReflectError::WasNotA {
1110                expected: "map or hash map",
1111                actual: frame.shape,
1112            });
1113        }
1114
1115        let vtable = frame.shape.vtable;
1116
1117        // Initialize an empty map if it's not already initialized
1118        if !frame.istate.fields.has(0) {
1119            let Some(default_in_place) = vtable.default_in_place else {
1120                return Err(ReflectError::OperationFailed {
1121                    shape: frame.shape,
1122                    operation: "map type does not implement Default",
1123                });
1124            };
1125
1126            unsafe {
1127                default_in_place(frame.data);
1128                frame.istate.fields.set(0);
1129            }
1130        }
1131
1132        Ok(self)
1133    }
1134
1135    /// Pushes a new element onto the list/array
1136    ///
1137    /// This creates a new frame for the element. When this frame is popped,
1138    /// the element will be added to the list.
1139    pub fn push(mut self) -> Result<Self, ReflectError> {
1140        // Make sure we're initializing a list
1141        let frame = self.frames.last().unwrap();
1142        let list_shape = frame.shape;
1143
1144        if !matches!(list_shape.def, Def::List(_)) {
1145            return Err(ReflectError::WasNotA {
1146                expected: "list or array",
1147                actual: list_shape,
1148            });
1149        }
1150
1151        // If the list isn't initialized yet, initialize it
1152        if !frame.istate.fields.has(0) {
1153            self = self.begin_pushback()?;
1154        }
1155
1156        // Get the element type
1157        let element_shape = self.element_shape()?;
1158
1159        // Allocate memory for the element
1160        let element_data = element_shape.allocate();
1161
1162        // Create a new frame for the element
1163        let element_frame = Frame {
1164            data: element_data,
1165            shape: element_shape,
1166            field_index_in_parent: None, // No need for an index since we're using mode
1167            istate: IState::new(
1168                self.frames.len(),
1169                FrameMode::ListElement,
1170                FrameFlags::ALLOCATED,
1171            ),
1172        };
1173
1174        trace!(
1175            "[{}] Pushing element of type {} to list {}",
1176            self.frames.len(),
1177            element_shape.green(),
1178            list_shape.blue(),
1179        );
1180
1181        self.frames.push(element_frame);
1182        Ok(self)
1183    }
1184
1185    /// Prepare to push the `Some(T)` variant of an `Option<T>`.
1186    pub fn push_some(mut self) -> Result<Self, ReflectError> {
1187        // Make sure we're initializing an option
1188        let frame = self.frames.last().unwrap();
1189        let option_shape = frame.shape;
1190
1191        // Get the option definition
1192        let Def::Option(option_def) = option_shape.def else {
1193            return Err(ReflectError::WasNotA {
1194                expected: "option",
1195                actual: option_shape,
1196            });
1197        };
1198
1199        // Get the inner type of the option
1200        let inner_shape = option_def.t();
1201
1202        // Allocate memory for the inner value
1203        let inner_data = inner_shape.allocate();
1204
1205        // Create a new frame for the inner value
1206        let inner_frame = Frame {
1207            data: inner_data,
1208            shape: inner_shape,
1209            // this is only set when we pop
1210            field_index_in_parent: None,
1211            istate: IState::new(
1212                self.frames.len(),
1213                FrameMode::OptionSome,
1214                // TODO: we could lazy-allocate it when something like `field` is called, tbh
1215                FrameFlags::ALLOCATED,
1216            ),
1217        };
1218
1219        trace!(
1220            "[{}] Pushing option frame for {}",
1221            self.frames.len(),
1222            option_shape.blue(),
1223        );
1224
1225        self.frames.push(inner_frame);
1226        Ok(self)
1227    }
1228
1229    /// Pops a not-yet-initialized option frame, setting it to None in the parent
1230    ///
1231    /// This is used to set an option to None instead of Some.
1232    /// Steps:
1233    ///  1. Asserts the option frame is NOT initialized
1234    ///  2. Frees the memory for the pushed value
1235    ///  3. Pops the frame
1236    ///  4. Sets the parent option to its default value (i.e., None)
1237    ///  5. Pops the parent option (which is the actual `Option<T>`, but no longer in option mode)
1238    pub fn pop_some_push_none(mut self) -> Result<Self, ReflectError> {
1239        // 1. Option frame must exist
1240        let Some(frame) = self.frames.last_mut() else {
1241            return Err(ReflectError::OperationFailed {
1242                shape: <()>::SHAPE,
1243                operation: "tried to pop_some_push_none but there was no frame",
1244            });
1245        };
1246
1247        // 1. Make sure the current frame is an option inner frame in "Option" mode
1248        if frame.istate.mode != FrameMode::OptionSome {
1249            return Err(ReflectError::OperationFailed {
1250                shape: frame.shape,
1251                operation: "pop_some_push_none called, but frame was not in Option mode",
1252            });
1253        }
1254
1255        // 1. Check not initialized
1256        if frame.is_fully_initialized() {
1257            return Err(ReflectError::OperationFailed {
1258                shape: frame.shape,
1259                operation: "option frame already initialized, cannot pop_some_push_none",
1260            });
1261        }
1262
1263        frame.dealloc_if_needed();
1264
1265        // 3. Pop the frame (this discards, doesn't propagate up)
1266        let _frame = self.frames.pop().expect("frame already checked");
1267
1268        // 4. Set parent option (which we just popped into) to default (None)
1269        let parent_frame = self
1270            .frames
1271            .last_mut()
1272            .ok_or(ReflectError::OperationFailed {
1273                shape: <()>::SHAPE,
1274                operation: "tried to pop_some_push_none but there was no parent frame",
1275            })?;
1276
1277        // Safety: option frames are correctly sized, and data is valid
1278        unsafe {
1279            if let Some(default_fn) = parent_frame.shape.vtable.default_in_place {
1280                default_fn(parent_frame.data);
1281            } else {
1282                return Err(ReflectError::OperationFailed {
1283                    shape: parent_frame.shape,
1284                    operation: "option type does not implement Default",
1285                });
1286            }
1287            parent_frame.mark_fully_initialized();
1288        }
1289
1290        let Def::Option(od) = parent_frame.shape.def else {
1291            return Err(ReflectError::OperationFailed {
1292                shape: parent_frame.shape,
1293                operation: "pop_some_push_none and the parent isn't of type Option???",
1294            });
1295        };
1296
1297        // Now push a `None` frame
1298        let data = parent_frame.data;
1299
1300        let mut frame = Frame {
1301            data,
1302            shape: od.t(),
1303            field_index_in_parent: Some(0),
1304            istate: IState::new(self.frames.len(), FrameMode::OptionNone, FrameFlags::EMPTY),
1305        };
1306        unsafe {
1307            frame.mark_fully_initialized();
1308        }
1309
1310        self.frames.push(frame);
1311
1312        Ok(self)
1313    }
1314
1315    /// Pushes a new key frame for a map entry
1316    ///
1317    /// This creates a new frame for the key. After setting the key value,
1318    /// call `push_map_value` to create a frame for the corresponding value.
1319    pub fn push_map_key(mut self) -> Result<Self, ReflectError> {
1320        // Make sure we're initializing a map
1321        let frame = self.frames.last().unwrap();
1322        let map_shape = frame.shape;
1323
1324        if !matches!(map_shape.def, Def::Map(_)) {
1325            return Err(ReflectError::WasNotA {
1326                expected: "map or hash map",
1327                actual: map_shape,
1328            });
1329        }
1330
1331        // If the map isn't initialized yet, initialize it
1332        if !frame.istate.fields.has(0) {
1333            self = self.begin_map_insert()?;
1334        }
1335
1336        // Get the key type
1337        let key_shape = self.key_shape()?;
1338
1339        // Allocate memory for the key
1340        let key_data = key_shape.allocate();
1341
1342        // Create a new frame for the key
1343        let key_frame = Frame {
1344            data: key_data,
1345            shape: key_shape,
1346            field_index_in_parent: None,
1347            istate: IState::new(self.frames.len(), FrameMode::MapKey, FrameFlags::ALLOCATED),
1348        };
1349
1350        trace!(
1351            "[{}] Pushing key of type {} for map {}",
1352            self.frames.len(),
1353            key_shape.green(),
1354            map_shape.blue(),
1355        );
1356
1357        self.frames.push(key_frame);
1358        Ok(self)
1359    }
1360
1361    /// Pushes a new value frame for a map entry
1362    ///
1363    /// This should be called after pushing and initializing a key frame.
1364    /// When the value frame is popped, the key-value pair will be added to the map.
1365    pub fn push_map_value(mut self) -> Result<Self, ReflectError> {
1366        trace!("Wants to push map value. Frames = ");
1367        for (i, f) in self.frames.iter().enumerate() {
1368            trace!("Frame {}: {:?}", i, f);
1369        }
1370
1371        // First, ensure we have a valid key frame
1372        if self.frames.len() < 2 {
1373            return Err(ReflectError::OperationFailed {
1374                shape: <()>::SHAPE,
1375                operation: "tried to push map value but there was no key frame",
1376            });
1377        }
1378
1379        // Check the frame before the last to ensure it's a map key
1380        let key_frame_index = self.frames.len() - 1;
1381        let key_frame = &self.frames[key_frame_index];
1382
1383        // Verify the current frame is a key frame
1384        match key_frame.istate.mode {
1385            FrameMode::MapKey => {} // Valid - continue
1386            _ => {
1387                return Err(ReflectError::OperationFailed {
1388                    shape: key_frame.shape,
1389                    operation: "current frame is not a map key",
1390                });
1391            }
1392        }
1393
1394        // Check that the key is fully initialized
1395        if !key_frame.is_fully_initialized() {
1396            return Err(ReflectError::OperationFailed {
1397                shape: key_frame.shape,
1398                operation: "map key is not fully initialized",
1399            });
1400        }
1401
1402        // Get the parent map frame to verify we're working with a map
1403        let map_frame_index = self.frames.len() - 2;
1404        let map_frame = &self.frames[map_frame_index];
1405        let map_shape = map_frame.shape;
1406
1407        let Def::Map(map_def) = map_shape.def else {
1408            return Err(ReflectError::WasNotA {
1409                expected: "map",
1410                actual: map_frame.shape,
1411            });
1412        };
1413
1414        let value_shape = map_def.v;
1415
1416        // Allocate memory for the value
1417        let value_data = value_shape.allocate();
1418
1419        // Create a new frame for the value
1420        let value_frame = Frame {
1421            data: value_data,
1422            shape: value_shape,
1423            field_index_in_parent: None,
1424            istate: IState::new(
1425                self.frames.len(),
1426                FrameMode::MapValue {
1427                    index: key_frame_index,
1428                },
1429                FrameFlags::ALLOCATED,
1430            ),
1431        };
1432
1433        trace!(
1434            "[{}] Pushing value of type {} for map {} with key type {}",
1435            self.frames.len(),
1436            value_shape.green(),
1437            map_shape.blue(),
1438            key_frame.shape.yellow(),
1439        );
1440
1441        self.frames.push(value_frame);
1442        Ok(self)
1443    }
1444
1445    /// Pops the current frame — goes back up one level
1446    pub fn pop(mut self) -> Result<Self, ReflectError> {
1447        let Some(frame) = self.pop_inner() else {
1448            return Err(ReflectError::InvariantViolation {
1449                invariant: "No frame to pop",
1450            });
1451        };
1452        self.track(frame);
1453        Ok(self)
1454    }
1455
1456    fn pop_inner(&mut self) -> Option<Frame> {
1457        let mut frame = self.frames.pop()?;
1458        let frame_shape = frame.shape;
1459
1460        let init = frame.is_fully_initialized();
1461        trace!(
1462            "[{}] {} popped, {} initialized",
1463            self.frames.len(),
1464            frame_shape.blue(),
1465            if init {
1466                "✅ fully".style(owo_colors::Style::new().green())
1467            } else {
1468                "🚧 partially".style(owo_colors::Style::new().red())
1469            }
1470        );
1471        if init {
1472            if let Some(parent) = self.frames.last_mut() {
1473                if let Some(index) = frame.field_index_in_parent {
1474                    parent.istate.fields.set(index);
1475                }
1476            }
1477        }
1478
1479        // Handle special frame modes
1480        match frame.istate.mode {
1481            // Handle list element frames
1482            FrameMode::ListElement => {
1483                if frame.is_fully_initialized() {
1484                    // This was a list element, so we need to push it to the parent list
1485                    // Capture frame length and parent shape before mutable borrow
1486                    let frame_len = self.frames.len();
1487
1488                    // Get parent frame
1489                    let parent_frame = self.frames.last_mut().unwrap();
1490                    let parent_shape = parent_frame.shape;
1491
1492                    // Make sure the parent is a list
1493                    match parent_shape.def {
1494                        Def::List(_) => {
1495                            // Get the list vtable from the ListDef
1496                            if let Def::List(list_def) = parent_shape.def {
1497                                let list_vtable = list_def.vtable;
1498                                trace!(
1499                                    "[{}] Pushing element to list {}",
1500                                    frame_len,
1501                                    parent_shape.blue()
1502                                );
1503                                unsafe {
1504                                    // Convert the frame data pointer to Opaque and call push function from vtable
1505                                    (list_vtable.push)(
1506                                        PtrMut::new(parent_frame.data.as_mut_byte_ptr()),
1507                                        PtrMut::new(frame.data.as_mut_byte_ptr()),
1508                                    );
1509                                    self.mark_moved_out_of(&mut frame);
1510                                }
1511                            } else {
1512                                panic!("parent frame is not a list type");
1513                            }
1514                        }
1515                        _ => {
1516                            panic!("Expected list or array, got {}", frame.shape);
1517                        }
1518                    }
1519                }
1520            }
1521
1522            // Handle map value frames
1523            FrameMode::MapValue {
1524                index: key_frame_index,
1525            } if frame.is_fully_initialized() => {
1526                // This was a map value, so we need to insert the key-value pair into the map
1527
1528                // Now let's remove the key frame from the frames array
1529                let mut key_frame = self.frames.remove(key_frame_index);
1530
1531                // Make sure the key is fully initialized
1532                if !key_frame.istate.fields.is_any_set() {
1533                    panic!("key is not initialized when popping value frame");
1534                }
1535
1536                // Get parent map frame
1537                let frame_len = self.frames.len();
1538                let parent_frame = self.frames.last_mut().unwrap();
1539                let parent_shape = parent_frame.shape;
1540
1541                // Make sure the parent is a map
1542                match parent_shape.def {
1543                    Def::Map(_) => {
1544                        // Get the map vtable from the MapDef
1545                        if let Def::Map(map_def) = parent_shape.def {
1546                            trace!(
1547                                "[{}] Inserting key-value pair into map {}",
1548                                frame_len,
1549                                parent_shape.blue()
1550                            );
1551                            unsafe {
1552                                // Call the map's insert function with the key and value
1553                                (map_def.vtable.insert_fn)(
1554                                    parent_frame.data.assume_init(),
1555                                    key_frame.data.assume_init(),
1556                                    PtrMut::new(frame.data.as_mut_byte_ptr()),
1557                                );
1558                                self.mark_moved_out_of(&mut key_frame);
1559                                self.mark_moved_out_of(&mut frame);
1560                            }
1561                        } else {
1562                            panic!("parent frame is not a map type");
1563                        }
1564                    }
1565                    _ => {
1566                        panic!("Expected map or hash map, got {}", frame.shape);
1567                    }
1568                }
1569            }
1570
1571            // Handle option frames
1572            FrameMode::OptionSome => {
1573                if frame.is_fully_initialized() {
1574                    trace!("Popping OptionSome (fully init'd)");
1575
1576                    // This was an option Some value, so we need to set it in the parent option
1577                    let frame_len = self.frames.len();
1578
1579                    // Get parent frame
1580                    let parent_frame = self.frames.last_mut().unwrap();
1581                    let parent_shape = parent_frame.shape;
1582
1583                    // Make sure the parent is an option
1584                    match parent_shape.def {
1585                        Def::Option(option_def) => {
1586                            trace!(
1587                                "[{}] Setting Some value in option {}",
1588                                frame_len,
1589                                parent_shape.blue()
1590                            );
1591                            unsafe {
1592                                // Call the option's init_some function
1593                                (option_def.vtable.init_some_fn)(
1594                                    parent_frame.data,
1595                                    PtrConst::new(frame.data.as_byte_ptr()),
1596                                );
1597                                trace!("Marking parent frame as fully initialized");
1598                                parent_frame.mark_fully_initialized();
1599
1600                                self.mark_moved_out_of(&mut frame);
1601                            }
1602                        }
1603                        _ => {
1604                            panic!(
1605                                "Expected parent frame to be an option type, got {}",
1606                                frame.shape
1607                            );
1608                        }
1609                    }
1610                } else {
1611                    trace!("Popping OptionSome (not fully init'd)");
1612                }
1613            }
1614
1615            // Map keys are just tracked, they don't need special handling when popped
1616            // FIXME: that's not true, we need to deallocate them at least??
1617            FrameMode::MapKey => {}
1618
1619            // Field frame
1620            FrameMode::Field => {}
1621
1622            // Uninitialized special frames
1623            _ => {}
1624        }
1625
1626        Some(frame)
1627    }
1628
1629    /// Evict a frame from istates, along with all its children
1630    /// (because we're about to use `drop_in_place` on it — not
1631    /// yet though, we need to know the variant for enums, etc.)
1632    pub fn evict_tree(&mut self, frame: Frame) -> Frame {
1633        match frame.shape.def {
1634            Def::Struct(sd) => {
1635                for f in sd.fields {
1636                    let id = ValueId {
1637                        shape: f.shape(),
1638                        ptr: unsafe { frame.data.field_uninit_at(f.offset) }.as_byte_ptr(),
1639                    };
1640                    if let Some(istate) = self.istates.remove(&id) {
1641                        let frame = Frame::recompose(id, istate);
1642                        self.evict_tree(frame);
1643                    } else {
1644                        trace!("No istate found for field {}", f.name);
1645                    }
1646                }
1647            }
1648            Def::Enum(_ed) => {
1649                // Check if a variant is selected in the istate
1650                if let Some(variant) = &frame.istate.variant {
1651                    trace!(
1652                        "Evicting enum {} variant '{}' fields",
1653                        frame.shape.blue(),
1654                        variant.name.yellow()
1655                    );
1656                    // Iterate over the fields of the selected variant
1657                    for field in variant.data.fields {
1658                        // Calculate the pointer to the field within the enum's data payload
1659                        let field_ptr = unsafe { frame.data.field_uninit_at(field.offset) };
1660                        let field_shape = field.shape();
1661                        let field_id = ValueId::new(field_shape, field_ptr.as_byte_ptr());
1662
1663                        // Try to remove the field's state from istates
1664                        if let Some(field_istate) = self.istates.remove(&field_id) {
1665                            trace!(
1666                                "Evicting field '{}' (shape {}) of enum variant '{}'",
1667                                field.name.bright_blue(),
1668                                field_shape.green(),
1669                                variant.name.yellow()
1670                            );
1671                            // Recompose the frame for the field
1672                            let field_frame = Frame::recompose(field_id, field_istate);
1673                            // Recursively evict the field's subtree
1674                            self.evict_tree(field_frame);
1675                        } else {
1676                            trace!(
1677                                "Field '{}' (shape {}) of enum variant '{}' not found in istates, skipping eviction",
1678                                field.name.red(),
1679                                field_shape.red(),
1680                                variant.name.yellow()
1681                            );
1682                        }
1683                    }
1684                } else {
1685                    // No variant selected, nothing to evict within the enum
1686                    trace!(
1687                        "Enum {} has no variant selected, no fields to evict.",
1688                        frame.shape.blue()
1689                    );
1690                }
1691            }
1692            _ => {}
1693        }
1694        frame
1695    }
1696
1697    #[allow(rustdoc::broken_intra_doc_links)]
1698    /// Returns the current path in the JSON document as a string.
1699    /// For example: "$.users[0].name"
1700    pub fn path(&self) -> String {
1701        let mut path = String::from("$");
1702
1703        for (i, frame) in self.frames.iter().enumerate() {
1704            // Skip the root frame
1705            if i == 0 {
1706                continue;
1707            }
1708
1709            match frame.istate.mode {
1710                FrameMode::ListElement => {
1711                    // For arrays, we use bracket notation with index
1712                    if let Some(index) = frame.istate.list_index {
1713                        path.push_str(&format!("[{}]", index));
1714                    } else {
1715                        path.push_str("[?]");
1716                    }
1717                }
1718                FrameMode::MapKey => {
1719                    path.push_str(".key");
1720                }
1721                FrameMode::MapValue { index: _ } => {
1722                    path.push_str(".value");
1723                }
1724                FrameMode::OptionSome => {
1725                    path.push_str(".some");
1726                }
1727                FrameMode::OptionNone => {
1728                    path.push_str(".none");
1729                }
1730                FrameMode::Root => {
1731                    // Root doesn't add to the path
1732                }
1733                FrameMode::Field => {
1734                    // For struct fields, we use dot notation with field name
1735                    if let Some(index) = frame.field_index_in_parent {
1736                        // Find the parent frame to get the field name
1737                        if let Some(parent) = self.frames.get(i - 1) {
1738                            if let Def::Struct(sd) = parent.shape.def {
1739                                if index < sd.fields.len() {
1740                                    let field_name = sd.fields[index].name;
1741                                    path.push('.');
1742                                    path.push_str(field_name);
1743                                }
1744                            } else if let Def::Enum(_) = parent.shape.def {
1745                                if let Some(variant) = &parent.istate.variant {
1746                                    if index < variant.data.fields.len() {
1747                                        let field_name = variant.data.fields[index].name;
1748                                        path.push('.');
1749                                        path.push_str(field_name);
1750                                    }
1751                                }
1752                            }
1753                        }
1754                    }
1755                }
1756            }
1757        }
1758
1759        path
1760    }
1761}
1762
1763impl Drop for Wip<'_> {
1764    fn drop(&mut self) {
1765        while let Some(frame) = self.frames.pop() {
1766            self.track(frame);
1767        }
1768        trace!("🧹 Whole WIP is dropping",);
1769
1770        let Some((root_id, _)) = self.istates.iter().find(|(_k, istate)| istate.depth == 0) else {
1771            trace!("No root found, we probably built already");
1772            return;
1773        };
1774
1775        let root_id = *root_id;
1776        let root_istate = self.istates.remove(&root_id).unwrap();
1777        let root = Frame::recompose(root_id, root_istate);
1778        let mut to_clean = vec![root];
1779
1780        let mut _root_guard: Option<Guard> = None;
1781
1782        while let Some(mut frame) = to_clean.pop() {
1783            trace!(
1784                "Cleaning frame: shape={} at {:p}, flags={:?}, mode={:?}, fully_initialized={}",
1785                frame.shape.blue(),
1786                frame.data.as_byte_ptr(),
1787                frame.istate.flags.bright_magenta(),
1788                frame.istate.mode.yellow(),
1789                if frame.is_fully_initialized() {
1790                    "✅"
1791                } else {
1792                    "❌"
1793                }
1794            );
1795
1796            if frame.istate.flags.contains(FrameFlags::MOVED) {
1797                trace!(
1798                    "{}",
1799                    "Frame was moved out of, nothing to dealloc/drop_in_place".yellow()
1800                );
1801                continue;
1802            }
1803
1804            match frame.shape.def {
1805                Def::Struct(sd) => {
1806                    if frame.is_fully_initialized() {
1807                        trace!(
1808                            "Dropping fully initialized struct: {} at {:p}",
1809                            frame.shape.green(),
1810                            frame.data.as_byte_ptr()
1811                        );
1812                        let frame = self.evict_tree(frame);
1813                        unsafe { frame.drop_and_dealloc_if_needed() };
1814                    } else {
1815                        let num_fields = sd.fields.len();
1816                        trace!(
1817                            "De-initializing struct {} at {:p} field-by-field ({} fields)",
1818                            frame.shape.yellow(),
1819                            frame.data.as_byte_ptr(),
1820                            num_fields.to_string().bright_cyan()
1821                        );
1822                        for i in 0..num_fields {
1823                            if frame.istate.fields.has(i) {
1824                                let field = sd.fields[i];
1825                                let field_shape = field.shape();
1826                                let field_ptr = unsafe { frame.data.field_init_at(field.offset) };
1827                                let field_id = ValueId::new(field_shape, field_ptr.as_byte_ptr());
1828                                trace!(
1829                                    "Recursively cleaning field #{} '{}' of {}: field_shape={}, field_ptr={:p}",
1830                                    i.to_string().bright_cyan(),
1831                                    field.name.bright_blue(),
1832                                    frame.shape.blue(),
1833                                    field_shape.green(),
1834                                    field_ptr.as_byte_ptr()
1835                                );
1836                                let istate = self.istates.remove(&field_id).unwrap();
1837                                let field_frame = Frame::recompose(field_id, istate);
1838                                to_clean.push(field_frame);
1839                            } else {
1840                                trace!(
1841                                    "Field #{} '{}' of {} was NOT initialized, skipping",
1842                                    i.to_string().bright_cyan(),
1843                                    sd.fields[i].name.bright_red(),
1844                                    frame.shape.red()
1845                                );
1846                            }
1847                        }
1848
1849                        // we'll also need to clean up if we're root
1850                        if frame.istate.mode == FrameMode::Root {
1851                            _root_guard = Some(Guard {
1852                                ptr: frame.data.as_mut_byte_ptr(),
1853                                layout: frame.shape.layout,
1854                            });
1855                        }
1856                    }
1857                }
1858                Def::Enum(_ed) => {
1859                    trace!(
1860                        "{}",
1861                        format!(
1862                            "TODO: handle enum deallocation for {} at {:p}",
1863                            frame.shape.yellow(),
1864                            frame.data.as_byte_ptr()
1865                        )
1866                        .magenta()
1867                    );
1868
1869                    // we'll also need to clean up if we're root
1870                    if frame.istate.mode == FrameMode::Root {
1871                        _root_guard = Some(Guard {
1872                            ptr: frame.data.as_mut_byte_ptr(),
1873                            layout: frame.shape.layout,
1874                        });
1875                    }
1876                }
1877                Def::Array(_)
1878                | Def::Slice(_)
1879                | Def::List(_)
1880                | Def::Map(_)
1881                | Def::SmartPointer(_)
1882                | Def::Scalar(_)
1883                | Def::Option(_) => {
1884                    trace!(
1885                        "Can drop all at once for shape {} (def variant: {:?}, frame mode {:?}) at {:p}",
1886                        frame.shape.cyan(),
1887                        frame.shape.def,
1888                        frame.istate.mode.yellow(),
1889                        frame.data.as_byte_ptr(),
1890                    );
1891
1892                    if frame.is_fully_initialized() {
1893                        unsafe { frame.drop_and_dealloc_if_needed() }
1894                    } else {
1895                        frame.dealloc_if_needed();
1896                    }
1897                }
1898                _ => {}
1899            }
1900        }
1901
1902        // We might have some frames left over to deallocate for temporary allocations for keymap insertion etc.
1903        let mut all_ids = self.istates.keys().copied().collect::<Vec<_>>();
1904        for frame_id in all_ids.drain(..) {
1905            let frame_istate = self.istates.remove(&frame_id).unwrap();
1906
1907            trace!(
1908                "Checking leftover istate: id.shape={} id.ptr={:p} mode={:?}",
1909                frame_id.shape.cyan(),
1910                frame_id.ptr,
1911                frame_istate.mode.yellow()
1912            );
1913            let mut frame = Frame::recompose(frame_id, frame_istate);
1914
1915            if frame.is_fully_initialized() {
1916                trace!("It's fully initialized, we can drop it");
1917                unsafe { frame.drop_and_dealloc_if_needed() };
1918            } else if frame.istate.flags.contains(FrameFlags::ALLOCATED) {
1919                trace!("Not initialized but allocated, let's free it");
1920                frame.dealloc_if_needed();
1921            }
1922        }
1923    }
1924}