Display groups more nicely

1 files changed, 21 insertions, 19 deletions

diff --git a/rprt-engine/src/explainer.rs b/rprt-engine/src/explainer.rs
index b62bee4..9770c72 100644
--- a/rprt-engine/src/explainer.rs
+++ b/rprt-engine/src/explainer.rs
@@ -32,13 +32,19 @@ fn main() {
     print!("{}", tree);
 }
 
+enum NodeKind {
+    Hook(HookKind),
+    Train(u32),
+    Group,
+    Plain,
+}
+
 struct TreeNode {
     pos: usize,
     depth: usize,
     children: Vec<usize>,
     is_leaf: bool,
-    hook_kind: Option<HookKind>,
-    train_arity: Option<u8>,
+    node_kind: NodeKind,
 }
 
 fn generate_tree(formatted: &str, expr: &Composite) -> String {
@@ -111,12 +117,13 @@ fn generate_tree(formatted: &str, expr: &Composite) -> String {
             let node_at_pos = interior.iter().find(|n| n.pos == i);
 
             if let Some(node) = node_at_pos {
-                if let Some(hook_kind) = node.hook_kind {
-                    line[i] = hook_kind.to_string().chars().next().unwrap();
-                } else if let Some(arity) = node.train_arity {
-                    line[i] = char::from_digit(arity as u32, 10).unwrap();
-                } else if up || down || in_span {
-                    line[i] = box_char(up, down, left, right);
+                match node.node_kind {
+                    NodeKind::Hook(hook_kind) => line[i] = hook_kind.to_string().chars().next().unwrap(),
+                    NodeKind::Train(arity) => line[i] = char::from_digit(arity as u32, 10).unwrap(),
+                    NodeKind::Group => line[i] = '+',
+                    NodeKind::Plain => if up || down || in_span {
+                        line[i] = box_char(up, down, left, right);
+                    }
                 }
             } else if up || down || in_span {
                 line[i] = box_char(up, down, left, right);
@@ -161,8 +168,7 @@ fn build_tree_from_expr(
                 depth,
                 children: vec![],
                 is_leaf: true,
-                hook_kind: None,
-                train_arity: None,
+                node_kind: NodeKind::Plain,
             });
             offset
         }
@@ -177,8 +183,7 @@ fn build_tree_from_expr(
                 depth,
                 children: vec![f_pos, g_pos],
                 is_leaf: false,
-                hook_kind: None,
-                train_arity: Some(2),
+                node_kind: NodeKind::Train(2),
             });
             f_pos
         }
@@ -195,8 +200,7 @@ fn build_tree_from_expr(
                 depth,
                 children: vec![f_pos, g_pos, h_pos],
                 is_leaf: false,
-                hook_kind: None,
-                train_arity: Some(3),
+                node_kind: NodeKind::Train(3),
             });
             g_pos
         }
@@ -211,14 +215,13 @@ fn build_tree_from_expr(
                 depth,
                 children: vec![left_pos, right_pos],
                 is_leaf: false,
-                hook_kind: Some(*kind),
-                train_arity: None,
+                node_kind: NodeKind::Hook(*kind),
             });
             left_pos
         }
         Composite::Group { operations } => {
             let mut children = Vec::new();
-            let mut current_offset = offset + 1;
+            let mut current_offset = offset + 2;
 
             for (i, op) in operations.iter().enumerate() {
                 if i > 0 {
@@ -235,8 +238,7 @@ fn build_tree_from_expr(
                 depth,
                 children,
                 is_leaf: false,
-                hook_kind: None,
-                train_arity: None,
+                node_kind: NodeKind::Group,
             });
             node_pos
         }