Skip to main content

facet_reflect/partial/partial_api/
sets.rs

1use super::*;
2use crate::AllocatedShape;
3
4////////////////////////////////////////////////////////////////////////////////////////////////////
5// Sets
6////////////////////////////////////////////////////////////////////////////////////////////////////
7impl<const BORROW: bool> Partial<'_, BORROW> {
8    /// Initializes a set (HashSet, BTreeSet, etc.) if it hasn't been initialized before.
9    /// This is a prerequisite to `begin_set_item`/`set`/`end` or the shorthand `insert`.
10    ///
11    /// `init_set` does not clear the set if it was previously initialized.
12    /// `init_set` does not push a new frame to the stack, and thus does not
13    /// require `end` to be called afterwards.
14    pub fn init_set(mut self) -> Result<Self, ReflectError> {
15        crate::trace!("init_set()");
16        // Get shape upfront to avoid borrow conflicts
17        let shape = self.frames().last().unwrap().allocated.shape();
18        let frame = self.frames_mut().last_mut().unwrap();
19
20        match &frame.tracker {
21            Tracker::Scalar if !frame.is_init => {
22                // that's good, let's initialize it
23            }
24            Tracker::Scalar => {
25                // is_init is true - initialized (perhaps from a previous round?) but should be a set tracker, let's fix that:
26                frame.tracker = Tracker::Set {
27                    current_child: false,
28                };
29                return Ok(self);
30            }
31            Tracker::Set { .. } => {
32                if frame.is_init {
33                    // already initialized, nothing to do
34                    return Ok(self);
35                }
36            }
37            _ => {
38                let tracker_kind = frame.tracker.kind();
39                return Err(self.err(ReflectErrorKind::UnexpectedTracker {
40                    message: "init_set called but tracker isn't something set-like",
41                    current_tracker: tracker_kind,
42                }));
43            }
44        };
45
46        // Check that we have a Set
47        let set_def = match &shape.def {
48            Def::Set(set_def) => set_def,
49            _ => {
50                return Err(self.err(ReflectErrorKind::OperationFailed {
51                    shape,
52                    operation: "init_set can only be called on Set types",
53                }));
54            }
55        };
56
57        let init_fn = set_def.vtable.init_in_place_with_capacity;
58
59        // Initialize the set with default capacity (0)
60        unsafe {
61            init_fn(frame.data, 0);
62        }
63
64        // Update tracker to Set state and mark as initialized
65        frame.tracker = Tracker::Set {
66            current_child: false,
67        };
68        frame.is_init = true;
69
70        Ok(self)
71    }
72
73    /// Begins pushing an element to the set.
74    /// The element should be set using `set()` or similar methods, then `end()` to complete.
75    pub fn begin_set_item(mut self) -> Result<Self, ReflectError> {
76        crate::trace!("begin_set_item()");
77        // Get frame info immutably first
78        let (shape, parent_type_plan, set_def, is_init, has_current_child) = {
79            let frame = self.frames().last().unwrap();
80            let shape = frame.allocated.shape();
81            let set_def = match &shape.def {
82                Def::Set(set_def) => *set_def,
83                _ => {
84                    return Err(self.err(ReflectErrorKind::OperationFailed {
85                        shape,
86                        operation: "init_set_item can only be called on Set types",
87                    }));
88                }
89            };
90            let has_current_child = match &frame.tracker {
91                Tracker::Set { current_child } => *current_child,
92                _ => {
93                    return Err(self.err(ReflectErrorKind::OperationFailed {
94                        shape,
95                        operation: "must call init_set() before begin_set_item()",
96                    }));
97                }
98            };
99            (
100                shape,
101                frame.type_plan,
102                set_def,
103                frame.is_init,
104                has_current_child,
105            )
106        };
107
108        // Verify the tracker is in Set state and initialized
109        if !is_init {
110            return Err(self.err(ReflectErrorKind::OperationFailed {
111                shape,
112                operation: "must call init_set() before begin_set_item()",
113            }));
114        }
115        if has_current_child {
116            return Err(self.err(ReflectErrorKind::OperationFailed {
117                shape,
118                operation: "already pushing an element, call end() first",
119            }));
120        }
121
122        // Update tracker to indicate we're building a child
123        match &mut self.mode.stack_mut().last_mut().unwrap().tracker {
124            Tracker::Set { current_child } => *current_child = true,
125            _ => unreachable!(),
126        }
127
128        // Get the element shape
129        let element_shape = set_def.t();
130
131        // Allocate space for the new element
132        let element_layout = match element_shape.layout.sized_layout() {
133            Ok(layout) => layout,
134            Err(_) => {
135                return Err(self.err(ReflectErrorKind::Unsized {
136                    shape: element_shape,
137                    operation: "begin_set_item: calculating element layout",
138                }));
139            }
140        };
141        let element_ptr = facet_core::alloc_for_layout(element_layout);
142
143        // Push a new frame for the element
144        // Get child type plan NodeId for set items
145        let child_plan_id = self
146            .root_plan
147            .set_item_node_id(parent_type_plan)
148            .expect("TypePlan should have item node for Set");
149        self.mode.stack_mut().push(Frame::new(
150            element_ptr,
151            AllocatedShape::new(element_shape, element_layout.size()),
152            FrameOwnership::Owned,
153            child_plan_id,
154        ));
155
156        Ok(self)
157    }
158}