Skip to main content

mesh_sieve/data/
mixed_section.rs

1//! Mixed-type section storage with tagged scalar types.
2
3use crate::data::atlas::Atlas;
4use crate::data::section::Section;
5use crate::data::storage::VecStorage;
6use crate::mesh_error::MeshSieveError;
7use std::collections::BTreeMap;
8
9/// Scalar type tag for mixed sections.
10#[derive(Clone, Copy, Debug, Eq, PartialEq)]
11pub enum ScalarType {
12    F64,
13    F32,
14    I32,
15    I64,
16    U32,
17    U64,
18}
19
20impl ScalarType {
21    /// Returns a stable string label for the scalar type.
22    pub fn as_str(self) -> &'static str {
23        match self {
24            ScalarType::F64 => "f64",
25            ScalarType::F32 => "f32",
26            ScalarType::I32 => "i32",
27            ScalarType::I64 => "i64",
28            ScalarType::U32 => "u32",
29            ScalarType::U64 => "u64",
30        }
31    }
32
33    /// Parse a scalar type from a string label.
34    pub fn parse(tag: &str) -> Option<Self> {
35        match tag {
36            "f64" => Some(ScalarType::F64),
37            "f32" => Some(ScalarType::F32),
38            "i32" => Some(ScalarType::I32),
39            "i64" => Some(ScalarType::I64),
40            "u32" => Some(ScalarType::U32),
41            "u64" => Some(ScalarType::U64),
42            _ => None,
43        }
44    }
45}
46
47/// Tagged, type-erased section storage for mixed scalar types.
48#[derive(Clone, Debug)]
49pub enum TaggedSection {
50    F64(Section<f64, VecStorage<f64>>),
51    F32(Section<f32, VecStorage<f32>>),
52    I32(Section<i32, VecStorage<i32>>),
53    I64(Section<i64, VecStorage<i64>>),
54    U32(Section<u32, VecStorage<u32>>),
55    U64(Section<u64, VecStorage<u64>>),
56}
57
58impl TaggedSection {
59    /// Return the scalar type tag for this section.
60    pub fn scalar_type(&self) -> ScalarType {
61        match self {
62            TaggedSection::F64(_) => ScalarType::F64,
63            TaggedSection::F32(_) => ScalarType::F32,
64            TaggedSection::I32(_) => ScalarType::I32,
65            TaggedSection::I64(_) => ScalarType::I64,
66            TaggedSection::U32(_) => ScalarType::U32,
67            TaggedSection::U64(_) => ScalarType::U64,
68        }
69    }
70
71    /// Return the atlas backing this tagged section.
72    pub fn atlas(&self) -> &Atlas {
73        match self {
74            TaggedSection::F64(section) => section.atlas(),
75            TaggedSection::F32(section) => section.atlas(),
76            TaggedSection::I32(section) => section.atlas(),
77            TaggedSection::I64(section) => section.atlas(),
78            TaggedSection::U32(section) => section.atlas(),
79            TaggedSection::U64(section) => section.atlas(),
80        }
81    }
82
83    /// Gather values from this tagged section in atlas insertion order.
84    pub fn gather_in_order(&self) -> TaggedSectionBuffer {
85        match self {
86            TaggedSection::F64(section) => TaggedSectionBuffer::F64(section.gather_in_order()),
87            TaggedSection::F32(section) => TaggedSectionBuffer::F32(section.gather_in_order()),
88            TaggedSection::I32(section) => TaggedSectionBuffer::I32(section.gather_in_order()),
89            TaggedSection::I64(section) => TaggedSectionBuffer::I64(section.gather_in_order()),
90            TaggedSection::U32(section) => TaggedSectionBuffer::U32(section.gather_in_order()),
91            TaggedSection::U64(section) => TaggedSectionBuffer::U64(section.gather_in_order()),
92        }
93    }
94
95    /// Scatter values into this tagged section in atlas insertion order.
96    pub fn try_scatter_in_order(
97        &mut self,
98        buf: &TaggedSectionBuffer,
99    ) -> Result<(), MeshSieveError> {
100        match (self, buf) {
101            (TaggedSection::F64(section), TaggedSectionBuffer::F64(data)) => {
102                section.try_scatter_in_order(data)
103            }
104            (TaggedSection::F32(section), TaggedSectionBuffer::F32(data)) => {
105                section.try_scatter_in_order(data)
106            }
107            (TaggedSection::I32(section), TaggedSectionBuffer::I32(data)) => {
108                section.try_scatter_in_order(data)
109            }
110            (TaggedSection::I64(section), TaggedSectionBuffer::I64(data)) => {
111                section.try_scatter_in_order(data)
112            }
113            (TaggedSection::U32(section), TaggedSectionBuffer::U32(data)) => {
114                section.try_scatter_in_order(data)
115            }
116            (TaggedSection::U64(section), TaggedSectionBuffer::U64(data)) => {
117                section.try_scatter_in_order(data)
118            }
119            (section, buf) => Err(MeshSieveError::TaggedSectionTypeMismatch {
120                expected: section.scalar_type(),
121                found: buf.scalar_type(),
122            }),
123        }
124    }
125}
126
127/// Typed buffer for tagged section scatter/gather operations.
128#[derive(Clone, Debug)]
129pub enum TaggedSectionBuffer {
130    F64(Vec<f64>),
131    F32(Vec<f32>),
132    I32(Vec<i32>),
133    I64(Vec<i64>),
134    U32(Vec<u32>),
135    U64(Vec<u64>),
136}
137
138impl TaggedSectionBuffer {
139    /// Scalar type tag for this buffer.
140    pub fn scalar_type(&self) -> ScalarType {
141        match self {
142            TaggedSectionBuffer::F64(_) => ScalarType::F64,
143            TaggedSectionBuffer::F32(_) => ScalarType::F32,
144            TaggedSectionBuffer::I32(_) => ScalarType::I32,
145            TaggedSectionBuffer::I64(_) => ScalarType::I64,
146            TaggedSectionBuffer::U32(_) => ScalarType::U32,
147            TaggedSectionBuffer::U64(_) => ScalarType::U64,
148        }
149    }
150
151    /// Length of the underlying flat buffer.
152    pub fn len(&self) -> usize {
153        match self {
154            TaggedSectionBuffer::F64(data) => data.len(),
155            TaggedSectionBuffer::F32(data) => data.len(),
156            TaggedSectionBuffer::I32(data) => data.len(),
157            TaggedSectionBuffer::I64(data) => data.len(),
158            TaggedSectionBuffer::U32(data) => data.len(),
159            TaggedSectionBuffer::U64(data) => data.len(),
160        }
161    }
162
163    /// Return true if the buffer is empty.
164    pub fn is_empty(&self) -> bool {
165        self.len() == 0
166    }
167}
168
169/// Trait to map scalar types to tagged sections for typed accessors.
170pub trait MixedScalar: Sized + 'static {
171    /// Scalar type tag for this concrete type.
172    const SCALAR_TYPE: ScalarType;
173
174    /// Wrap a typed section into a tagged container.
175    fn wrap(section: Section<Self, VecStorage<Self>>) -> TaggedSection;
176    /// Borrow a typed section if the tag matches.
177    fn unwrap(section: &TaggedSection) -> Option<&Section<Self, VecStorage<Self>>>;
178    /// Mutably borrow a typed section if the tag matches.
179    fn unwrap_mut(section: &mut TaggedSection) -> Option<&mut Section<Self, VecStorage<Self>>>;
180}
181
182impl MixedScalar for f64 {
183    const SCALAR_TYPE: ScalarType = ScalarType::F64;
184
185    fn wrap(section: Section<Self, VecStorage<Self>>) -> TaggedSection {
186        TaggedSection::F64(section)
187    }
188
189    fn unwrap(section: &TaggedSection) -> Option<&Section<Self, VecStorage<Self>>> {
190        if let TaggedSection::F64(section) = section {
191            Some(section)
192        } else {
193            None
194        }
195    }
196
197    fn unwrap_mut(section: &mut TaggedSection) -> Option<&mut Section<Self, VecStorage<Self>>> {
198        if let TaggedSection::F64(section) = section {
199            Some(section)
200        } else {
201            None
202        }
203    }
204}
205
206impl MixedScalar for f32 {
207    const SCALAR_TYPE: ScalarType = ScalarType::F32;
208
209    fn wrap(section: Section<Self, VecStorage<Self>>) -> TaggedSection {
210        TaggedSection::F32(section)
211    }
212
213    fn unwrap(section: &TaggedSection) -> Option<&Section<Self, VecStorage<Self>>> {
214        if let TaggedSection::F32(section) = section {
215            Some(section)
216        } else {
217            None
218        }
219    }
220
221    fn unwrap_mut(section: &mut TaggedSection) -> Option<&mut Section<Self, VecStorage<Self>>> {
222        if let TaggedSection::F32(section) = section {
223            Some(section)
224        } else {
225            None
226        }
227    }
228}
229
230impl MixedScalar for i32 {
231    const SCALAR_TYPE: ScalarType = ScalarType::I32;
232
233    fn wrap(section: Section<Self, VecStorage<Self>>) -> TaggedSection {
234        TaggedSection::I32(section)
235    }
236
237    fn unwrap(section: &TaggedSection) -> Option<&Section<Self, VecStorage<Self>>> {
238        if let TaggedSection::I32(section) = section {
239            Some(section)
240        } else {
241            None
242        }
243    }
244
245    fn unwrap_mut(section: &mut TaggedSection) -> Option<&mut Section<Self, VecStorage<Self>>> {
246        if let TaggedSection::I32(section) = section {
247            Some(section)
248        } else {
249            None
250        }
251    }
252}
253
254impl MixedScalar for i64 {
255    const SCALAR_TYPE: ScalarType = ScalarType::I64;
256
257    fn wrap(section: Section<Self, VecStorage<Self>>) -> TaggedSection {
258        TaggedSection::I64(section)
259    }
260
261    fn unwrap(section: &TaggedSection) -> Option<&Section<Self, VecStorage<Self>>> {
262        if let TaggedSection::I64(section) = section {
263            Some(section)
264        } else {
265            None
266        }
267    }
268
269    fn unwrap_mut(section: &mut TaggedSection) -> Option<&mut Section<Self, VecStorage<Self>>> {
270        if let TaggedSection::I64(section) = section {
271            Some(section)
272        } else {
273            None
274        }
275    }
276}
277
278impl MixedScalar for u32 {
279    const SCALAR_TYPE: ScalarType = ScalarType::U32;
280
281    fn wrap(section: Section<Self, VecStorage<Self>>) -> TaggedSection {
282        TaggedSection::U32(section)
283    }
284
285    fn unwrap(section: &TaggedSection) -> Option<&Section<Self, VecStorage<Self>>> {
286        if let TaggedSection::U32(section) = section {
287            Some(section)
288        } else {
289            None
290        }
291    }
292
293    fn unwrap_mut(section: &mut TaggedSection) -> Option<&mut Section<Self, VecStorage<Self>>> {
294        if let TaggedSection::U32(section) = section {
295            Some(section)
296        } else {
297            None
298        }
299    }
300}
301
302impl MixedScalar for u64 {
303    const SCALAR_TYPE: ScalarType = ScalarType::U64;
304
305    fn wrap(section: Section<Self, VecStorage<Self>>) -> TaggedSection {
306        TaggedSection::U64(section)
307    }
308
309    fn unwrap(section: &TaggedSection) -> Option<&Section<Self, VecStorage<Self>>> {
310        if let TaggedSection::U64(section) = section {
311            Some(section)
312        } else {
313            None
314        }
315    }
316
317    fn unwrap_mut(section: &mut TaggedSection) -> Option<&mut Section<Self, VecStorage<Self>>> {
318        if let TaggedSection::U64(section) = section {
319            Some(section)
320        } else {
321            None
322        }
323    }
324}
325
326/// Store named sections with mixed scalar types.
327#[derive(Clone, Debug, Default)]
328pub struct MixedSectionStore {
329    sections: BTreeMap<String, TaggedSection>,
330}
331
332impl MixedSectionStore {
333    /// Create an empty mixed section store.
334    pub fn new() -> Self {
335        Self::default()
336    }
337
338    /// Insert a typed section into the store.
339    pub fn insert<T: MixedScalar>(
340        &mut self,
341        name: impl Into<String>,
342        section: Section<T, VecStorage<T>>,
343    ) -> Option<TaggedSection> {
344        self.sections.insert(name.into(), T::wrap(section))
345    }
346
347    /// Insert a tagged section into the store.
348    pub fn insert_tagged(
349        &mut self,
350        name: impl Into<String>,
351        section: TaggedSection,
352    ) -> Option<TaggedSection> {
353        self.sections.insert(name.into(), section)
354    }
355
356    /// Retrieve a typed section by name.
357    pub fn get<T: MixedScalar>(&self, name: &str) -> Option<&Section<T, VecStorage<T>>> {
358        self.sections.get(name).and_then(T::unwrap)
359    }
360
361    /// Retrieve a mutable typed section by name.
362    pub fn get_mut<T: MixedScalar>(
363        &mut self,
364        name: &str,
365    ) -> Option<&mut Section<T, VecStorage<T>>> {
366        self.sections.get_mut(name).and_then(T::unwrap_mut)
367    }
368
369    /// Retrieve a tagged section by name.
370    pub fn get_tagged(&self, name: &str) -> Option<&TaggedSection> {
371        self.sections.get(name)
372    }
373
374    /// Iterate over all named tagged sections.
375    pub fn iter(&self) -> impl Iterator<Item = (&String, &TaggedSection)> {
376        self.sections.iter()
377    }
378
379    /// Mutably iterate over all named tagged sections.
380    pub fn iter_mut(&mut self) -> impl Iterator<Item = (&String, &mut TaggedSection)> {
381        self.sections.iter_mut()
382    }
383
384    /// Return true if the store is empty.
385    pub fn is_empty(&self) -> bool {
386        self.sections.is_empty()
387    }
388
389    /// Gather all tagged sections into flat buffers in atlas insertion order.
390    pub fn gather_in_order(&self) -> BTreeMap<String, TaggedSectionBuffer> {
391        self.sections
392            .iter()
393            .map(|(name, section)| (name.clone(), section.gather_in_order()))
394            .collect()
395    }
396
397    /// Scatter all tagged sections from flat buffers in atlas insertion order.
398    pub fn try_scatter_in_order(
399        &mut self,
400        buffers: &BTreeMap<String, TaggedSectionBuffer>,
401    ) -> Result<(), MeshSieveError> {
402        for (name, section) in &mut self.sections {
403            let buf = buffers
404                .get(name)
405                .ok_or_else(|| MeshSieveError::MissingSectionName { name: name.clone() })?;
406            section.try_scatter_in_order(buf)?;
407        }
408        Ok(())
409    }
410}