1use 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#[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 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 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#[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 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 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 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 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#[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 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 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 pub fn is_empty(&self) -> bool {
165 self.len() == 0
166 }
167}
168
169pub trait MixedScalar: Sized + 'static {
171 const SCALAR_TYPE: ScalarType;
173
174 fn wrap(section: Section<Self, VecStorage<Self>>) -> TaggedSection;
176 fn unwrap(section: &TaggedSection) -> Option<&Section<Self, VecStorage<Self>>>;
178 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#[derive(Clone, Debug, Default)]
328pub struct MixedSectionStore {
329 sections: BTreeMap<String, TaggedSection>,
330}
331
332impl MixedSectionStore {
333 pub fn new() -> Self {
335 Self::default()
336 }
337
338 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 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 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 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 pub fn get_tagged(&self, name: &str) -> Option<&TaggedSection> {
371 self.sections.get(name)
372 }
373
374 pub fn iter(&self) -> impl Iterator<Item = (&String, &TaggedSection)> {
376 self.sections.iter()
377 }
378
379 pub fn iter_mut(&mut self) -> impl Iterator<Item = (&String, &mut TaggedSection)> {
381 self.sections.iter_mut()
382 }
383
384 pub fn is_empty(&self) -> bool {
386 self.sections.is_empty()
387 }
388
389 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 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}