1 files changed, 117 insertions, 0 deletions
diff --git a/forest/examples/query_forest.rs b/forest/examples/query_forest.rs
new file mode 100644
index 0000000..fe70ca8
--- /dev/null
+++ b/forest/examples/query_forest.rs
@@ -0,0 +1,117 @@
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let mut file = std::fs::OpenOptions::new().read(true).open("forest.rep")?;
+
+    use std::io::Read;
+
+    let mut source: Vec<u8> = Vec::new();
+
+    file.read_to_end(&mut source)?;
+
+    use forest::bytes::BytesForest;
+
+    #[allow(unused_imports)]
+    use forest::iter::query::{preprocess_forest, QForest};
+
+    use graph::Graph;
+
+    let bytes_forest = BytesForest::try_from(source)?;
+
+    dbg!(
+        bytes_forest.nodes_len(),
+        bytes_forest.children_of(0)?.collect::<Vec<_>>()
+    );
+
+    #[cfg(feature = "test-print-viz")]
+    bytes_forest.print_viz("test forest.gv")?;
+
+    let preprocessed = preprocess_forest(bytes_forest, 0)?;
+
+    dbg!(&preprocessed);
+
+    use std::io::{BufRead, Write};
+
+    println!("Input a query:");
+    std::io::stdout().flush()?;
+
+    #[derive(Debug, Default, Copy, Clone)]
+    enum QueryType {
+        SubtreeRange(usize),
+        #[default]
+        Switches,
+        DFSNext(usize),
+        OrderToDFS(usize),
+        DFSToOrder(usize),
+    }
+
+    for line in std::io::stdin().lock().lines() {
+        let line = line?;
+
+        if line == "q" {
+            break;
+        }
+
+        let qtype = if line.starts_with('s') {
+            1u8
+        } else if line.starts_with('b') {
+            2u8
+        } else if line.starts_with('o') {
+            3u8
+        } else if line.starts_with('d') {
+            4u8
+        } else {
+            0u8
+        };
+
+        let type_str: &str;
+
+        let q = if line.starts_with('s') || line.starts_with('o') || line.starts_with('d') {
+            if let Ok(n) = line[1..].parse::<usize>() {
+                n
+            } else {
+                0usize
+            }
+        } else if let Ok(n) = line.parse::<usize>() {
+            n
+        } else {
+            continue;
+        };
+
+        let query: QueryType = match qtype {
+            0 => QueryType::DFSNext(q),
+            1 => QueryType::SubtreeRange(q),
+            2 => QueryType::default(),
+            3 => QueryType::OrderToDFS(q),
+            _ => QueryType::DFSToOrder(q),
+        };
+
+        let ans = match query {
+            QueryType::SubtreeRange(n) => {
+                type_str = "subtree range";
+                format!("{:?}", preprocessed.subtree_ranges.get(n))
+            }
+            QueryType::Switches => {
+                type_str = "branch switches";
+                format!("{:?}", preprocessed.branch_switches)
+            }
+            QueryType::DFSNext(n) => {
+                type_str = "dfs";
+                format!("{:?}", preprocessed.dfs_next.get(n).copied())
+            }
+            QueryType::OrderToDFS(n) => {
+                type_str = "order to dfs";
+                format!("{:?}", preprocessed.translate_to_dfs(n))
+            }
+            QueryType::DFSToOrder(n) => {
+                type_str = "dfs to order";
+                format!("{:?}", preprocessed.translate_from_dfs(n))
+            }
+        };
+
+        println!("answer for {type_str} = {ans}");
+
+        println!("Input a query:");
+        std::io::stdout().flush()?;
+    }
+
+    Ok(())
+}