1use crate::data::{FieldData, PolyData, ImageData, UnstructuredGrid, RectilinearGrid, StructuredGrid};
2use crate::data::traits::DataObject;
3
4#[derive(Debug, Clone)]
6pub enum Block {
7 PolyData(PolyData),
8 ImageData(ImageData),
9 UnstructuredGrid(UnstructuredGrid),
10 RectilinearGrid(RectilinearGrid),
11 StructuredGrid(StructuredGrid),
12 MultiBlock(MultiBlockDataSet),
13}
14
15impl From<PolyData> for Block {
16 fn from(pd: PolyData) -> Self { Block::PolyData(pd) }
17}
18impl From<ImageData> for Block {
19 fn from(id: ImageData) -> Self { Block::ImageData(id) }
20}
21impl From<UnstructuredGrid> for Block {
22 fn from(ug: UnstructuredGrid) -> Self { Block::UnstructuredGrid(ug) }
23}
24impl From<RectilinearGrid> for Block {
25 fn from(rg: RectilinearGrid) -> Self { Block::RectilinearGrid(rg) }
26}
27impl From<StructuredGrid> for Block {
28 fn from(sg: StructuredGrid) -> Self { Block::StructuredGrid(sg) }
29}
30impl From<MultiBlockDataSet> for Block {
31 fn from(mb: MultiBlockDataSet) -> Self { Block::MultiBlock(mb) }
32}
33
34impl std::fmt::Display for Block {
35 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36 match self {
37 Block::PolyData(pd) => write!(f, "{pd}"),
38 Block::ImageData(id) => write!(f, "{id}"),
39 Block::UnstructuredGrid(ug) => write!(f, "{ug}"),
40 Block::RectilinearGrid(_) => write!(f, "RectilinearGrid"),
41 Block::StructuredGrid(_) => write!(f, "StructuredGrid"),
42 Block::MultiBlock(mb) => write!(f, "MultiBlock({} blocks)", mb.num_blocks()),
43 }
44 }
45}
46
47impl std::fmt::Display for MultiBlockDataSet {
48 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
49 write!(f, "MultiBlockDataSet: {} blocks", self.num_blocks())
50 }
51}
52
53#[derive(Debug, Clone, Default)]
58pub struct MultiBlockDataSet {
59 blocks: Vec<(Option<String>, Block)>,
60 field_data: FieldData,
61}
62
63impl MultiBlockDataSet {
64 pub fn new() -> Self {
65 Self::default()
66 }
67
68 pub fn add_block(&mut self, name: impl Into<String>, block: Block) {
70 self.blocks.push((Some(name.into()), block));
71 }
72
73 pub fn add_unnamed_block(&mut self, block: Block) {
75 self.blocks.push((None, block));
76 }
77
78 pub fn num_blocks(&self) -> usize {
80 self.blocks.len()
81 }
82
83 pub fn block(&self, idx: usize) -> Option<&Block> {
85 self.blocks.get(idx).map(|(_, b)| b)
86 }
87
88 pub fn block_name(&self, idx: usize) -> Option<&str> {
90 self.blocks.get(idx).and_then(|(n, _)| n.as_deref())
91 }
92
93 pub fn block_by_name(&self, name: &str) -> Option<&Block> {
95 self.blocks.iter().find_map(|(n, b)| {
96 if n.as_deref() == Some(name) { Some(b) } else { None }
97 })
98 }
99
100 pub fn iter(&self) -> impl Iterator<Item = (Option<&str>, &Block)> {
102 self.blocks.iter().map(|(n, b)| (n.as_deref(), b))
103 }
104
105 pub fn remove_block(&mut self, idx: usize) -> (Option<String>, Block) {
107 self.blocks.remove(idx)
108 }
109
110 pub fn remove_by_name(&mut self, name: &str) -> Option<Block> {
112 if let Some(idx) = self.blocks.iter().position(|(n, _)| n.as_deref() == Some(name)) {
113 Some(self.blocks.remove(idx).1)
114 } else {
115 None
116 }
117 }
118
119 pub fn poly_data_blocks(&self) -> Vec<(Option<&str>, &PolyData)> {
121 self.blocks.iter().filter_map(|(n, b)| match b {
122 Block::PolyData(pd) => Some((n.as_deref(), pd)),
123 _ => None,
124 }).collect()
125 }
126
127 pub fn image_data_blocks(&self) -> Vec<(Option<&str>, &ImageData)> {
129 self.blocks.iter().filter_map(|(n, b)| match b {
130 Block::ImageData(id) => Some((n.as_deref(), id)),
131 _ => None,
132 }).collect()
133 }
134
135 pub fn unstructured_grid_blocks(&self) -> Vec<(Option<&str>, &UnstructuredGrid)> {
137 self.blocks.iter().filter_map(|(n, b)| match b {
138 Block::UnstructuredGrid(ug) => Some((n.as_deref(), ug)),
139 _ => None,
140 }).collect()
141 }
142
143 pub fn flatten(&self) -> Vec<(Option<String>, Block)> {
145 let mut result = Vec::new();
146 for (name, block) in &self.blocks {
147 match block {
148 Block::MultiBlock(inner) => {
149 for (sub_name, sub_block) in inner.flatten() {
150 let combined_name = match (name.as_deref(), sub_name.as_deref()) {
151 (Some(p), Some(c)) => Some(format!("{p}/{c}")),
152 (Some(p), None) => Some(p.to_string()),
153 (None, Some(c)) => Some(c.to_string()),
154 (None, None) => None,
155 };
156 result.push((combined_name, sub_block));
157 }
158 }
159 other => {
160 result.push((name.clone(), other.clone()));
161 }
162 }
163 }
164 result
165 }
166
167 pub fn with_block(mut self, name: impl Into<String>, block: Block) -> Self {
169 self.add_block(name, block);
170 self
171 }
172
173 pub fn total_blocks_recursive(&self) -> usize {
175 let mut count = 0;
176 for (_, block) in &self.blocks {
177 match block {
178 Block::MultiBlock(inner) => count += inner.total_blocks_recursive(),
179 _ => count += 1,
180 }
181 }
182 count
183 }
184}
185
186impl DataObject for MultiBlockDataSet {
187 fn field_data(&self) -> &FieldData {
188 &self.field_data
189 }
190
191 fn field_data_mut(&mut self) -> &mut FieldData {
192 &mut self.field_data
193 }
194}
195
196#[cfg(test)]
197mod tests {
198 use super::*;
199
200 #[test]
201 fn multi_block_basics() {
202 let mut mb = MultiBlockDataSet::new();
203 let pd = PolyData::from_triangles(
204 vec![[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.5, 1.0, 0.0]],
205 vec![[0, 1, 2]],
206 );
207 mb.add_block("mesh", Block::PolyData(pd));
208 mb.add_block("image", Block::ImageData(ImageData::with_dimensions(2, 2, 2)));
209
210 assert_eq!(mb.num_blocks(), 2);
211 assert_eq!(mb.block_name(0), Some("mesh"));
212 assert!(mb.block_by_name("image").is_some());
213 }
214
215 #[test]
216 fn typed_getters() {
217 let mb = MultiBlockDataSet::new()
218 .with_block("tri", Block::PolyData(PolyData::new()))
219 .with_block("img", Block::ImageData(ImageData::with_dimensions(2, 2, 2)));
220 assert_eq!(mb.poly_data_blocks().len(), 1);
221 assert_eq!(mb.image_data_blocks().len(), 1);
222 assert_eq!(mb.unstructured_grid_blocks().len(), 0);
223 }
224
225 #[test]
226 fn remove_by_name() {
227 let mut mb = MultiBlockDataSet::new();
228 mb.add_block("a", Block::PolyData(PolyData::new()));
229 mb.add_block("b", Block::PolyData(PolyData::new()));
230 assert_eq!(mb.num_blocks(), 2);
231 mb.remove_by_name("a");
232 assert_eq!(mb.num_blocks(), 1);
233 assert_eq!(mb.block_name(0), Some("b"));
234 }
235
236 #[test]
237 fn flatten() {
238 let inner = MultiBlockDataSet::new()
239 .with_block("c1", Block::PolyData(PolyData::new()))
240 .with_block("c2", Block::PolyData(PolyData::new()));
241 let outer = MultiBlockDataSet::new()
242 .with_block("top", Block::PolyData(PolyData::new()))
243 .with_block("sub", Block::MultiBlock(inner));
244 let flat = outer.flatten();
245 assert_eq!(flat.len(), 3);
246 assert_eq!(flat[0].0.as_deref(), Some("top"));
247 assert_eq!(flat[1].0.as_deref(), Some("sub/c1"));
248 }
249
250 #[test]
251 fn total_recursive() {
252 let inner = MultiBlockDataSet::new()
253 .with_block("a", Block::PolyData(PolyData::new()))
254 .with_block("b", Block::PolyData(PolyData::new()));
255 let outer = MultiBlockDataSet::new()
256 .with_block("top", Block::PolyData(PolyData::new()))
257 .with_block("sub", Block::MultiBlock(inner));
258 assert_eq!(outer.total_blocks_recursive(), 3);
259 }
260
261 #[test]
262 fn nested_multi_block() {
263 let mut inner = MultiBlockDataSet::new();
264 inner.add_block("tri", Block::PolyData(PolyData::new()));
265
266 let mut outer = MultiBlockDataSet::new();
267 outer.add_block("sub", Block::MultiBlock(inner));
268
269 assert_eq!(outer.num_blocks(), 1);
270 if let Some(Block::MultiBlock(sub)) = outer.block(0) {
271 assert_eq!(sub.num_blocks(), 1);
272 } else {
273 panic!("expected nested MultiBlock");
274 }
275 }
276
277 #[test]
278 fn from_conversions() {
279 let pd = PolyData::new();
280 let block: Block = pd.into();
281 assert!(matches!(block, Block::PolyData(_)));
282
283 let img = ImageData::with_dimensions(2, 2, 2);
284 let block: Block = img.into();
285 assert!(matches!(block, Block::ImageData(_)));
286 }
287
288 #[test]
289 fn display() {
290 let mb = MultiBlockDataSet::new()
291 .with_block("tri", PolyData::new().into())
292 .with_block("img", ImageData::with_dimensions(2, 2, 2).into());
293 let s = format!("{mb}");
294 assert!(s.contains("2 blocks"));
295
296 let block: Block = PolyData::new().into();
297 let s = format!("{block}");
298 assert!(s.contains("PolyData"));
299 }
300}