1 files changed, 74 insertions, 0 deletions
diff --git a/matroska/src/unflatten.rs b/matroska/src/unflatten.rs
new file mode 100644
index 0000000..03ca0f5
--- /dev/null
+++ b/matroska/src/unflatten.rs
@@ -0,0 +1,74 @@
+use crate::{matroska::MatroskaTag, Master};
+use anyhow::Result;
+
+pub struct Unflat<'a> {
+    pub item: MatroskaTag,
+    pub children: Option<Unflatten<'a>>,
+}
+
+pub struct Unflatten<'a> {
+    inner: &'a mut dyn Iterator<Item = Result<MatroskaTag>>,
+    stop: bool,
+    end: Option<MatroskaTag>,
+}
+
+impl<'a> Unflatten<'a> {
+    pub fn new(inner: &'a mut dyn Iterator<Item = Result<MatroskaTag>>) -> Self {
+        Self {
+            inner,
+            stop: false,
+            end: None,
+        }
+    }
+    pub fn new_with_end(
+        inner: &'a mut dyn Iterator<Item = Result<MatroskaTag>>,
+        start: MatroskaTag,
+    ) -> Self {
+        Self {
+            inner,
+            stop: false,
+            end: Some(MatroskaTag::construct_master(start.id(), Master::End).unwrap()),
+        }
+    }
+
+    pub fn next(&mut self) -> Option<Result<Unflat>> {
+        if self.stop {
+            return None;
+        }
+        match self.inner.next() {
+            None => None,
+            Some(Err(e)) => Some(Err(e)),
+            Some(Ok(item)) => {
+                let master = MatroskaTag::is_master(item.id()).unwrap();
+                if Some(&item) == self.end.as_ref() {
+                    self.stop = true;
+                    None
+                } else {
+                    Some(Ok(Unflat {
+                        children: if master {
+                            let end =
+                                MatroskaTag::construct_master(item.id(), Master::End).unwrap();
+                            if end == item {
+                                return None;
+                            }
+                            Some(Unflatten {
+                                inner: self.inner,
+                                stop: false,
+                                end: Some(end),
+                            })
+                        } else {
+                            None
+                        },
+                        item,
+                    }))
+                }
+            }
+        }
+    }
+}
+
+impl Drop for Unflatten<'_> {
+    fn drop(&mut self) {
+        while let Some(_) = self.next() {}
+    }
+}