Skip to main content

vtk_pure_rs/data/
multi_block.rs

1use crate::data::{FieldData, PolyData, ImageData, UnstructuredGrid, RectilinearGrid, StructuredGrid};
2use crate::data::traits::DataObject;
3
4/// A block in a MultiBlockDataSet, which can hold different dataset types.
5#[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/// A composite dataset containing named blocks of heterogeneous data types.
54///
55/// Analogous to VTK's `vtkMultiBlockDataSet`. Each block can be any dataset
56/// type, including nested MultiBlockDataSets.
57#[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    /// Add a named block.
69    pub fn add_block(&mut self, name: impl Into<String>, block: Block) {
70        self.blocks.push((Some(name.into()), block));
71    }
72
73    /// Add an unnamed block.
74    pub fn add_unnamed_block(&mut self, block: Block) {
75        self.blocks.push((None, block));
76    }
77
78    /// Number of blocks.
79    pub fn num_blocks(&self) -> usize {
80        self.blocks.len()
81    }
82
83    /// Get block by index.
84    pub fn block(&self, idx: usize) -> Option<&Block> {
85        self.blocks.get(idx).map(|(_, b)| b)
86    }
87
88    /// Get block name by index.
89    pub fn block_name(&self, idx: usize) -> Option<&str> {
90        self.blocks.get(idx).and_then(|(n, _)| n.as_deref())
91    }
92
93    /// Get block by name (first match).
94    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    /// Iterate over all blocks with their optional names.
101    pub fn iter(&self) -> impl Iterator<Item = (Option<&str>, &Block)> {
102        self.blocks.iter().map(|(n, b)| (n.as_deref(), b))
103    }
104
105    /// Remove a block by index. Returns the removed block.
106    pub fn remove_block(&mut self, idx: usize) -> (Option<String>, Block) {
107        self.blocks.remove(idx)
108    }
109
110    /// Remove a block by name. Returns the removed block if found.
111    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    /// Get all PolyData blocks.
120    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    /// Get all ImageData blocks.
128    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    /// Get all UnstructuredGrid blocks.
136    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    /// Flatten all nested MultiBlockDataSets into a single level.
144    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    /// Builder: add a named block.
168    pub fn with_block(mut self, name: impl Into<String>, block: Block) -> Self {
169        self.add_block(name, block);
170        self
171    }
172
173    /// Total number of blocks at all levels (recursive).
174    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}