feat(ast): then_branch

2 files changed, 47 insertions, 18 deletions

diff --git a/src/ast.rs b/src/ast.rs
index a87b6a2..472172b 100644
--- a/src/ast.rs
+++ b/src/ast.rs
@@ -1,3 +1,5 @@
+use std::iter::Peekable;
+
 use crate::parser::{Expr, ExprT, Metadata, Scope};
 
 #[derive(Clone, Debug)]
@@ -5,13 +7,14 @@ pub enum DepthExpr<'a> {
     Decision {
         cond: Box<DepthExpr<'a>>,
         t: Scope,
+        then_branch: Box<Vec<DepthExpr<'a>>>,
         else_branch: Option<Box<Vec<DepthExpr<'a>>>>,
     },
     IO(&'a str),
     Process(&'a str),
 }
 
-pub fn parse<'a, 'b>(vec: &'b Vec<Expr<'a>>) -> Vec<DepthExpr<'a>> {
+pub fn parse<'a>(vec: &Vec<Expr<'a>>) -> Vec<DepthExpr<'a>> {
     let mut exprs: Vec<DepthExpr<'a>> = Vec::new();
 
     let mut it = vec
@@ -19,29 +22,55 @@ pub fn parse<'a, 'b>(vec: &'b Vec<Expr<'a>>) -> Vec<DepthExpr<'a>> {
         .peekable();
 
     while let Some(e) = it.next() {
-        if let ExprT::Decision((cond, scope)) = e.expr {
-            let cond = DepthExpr::Process(cond);
-            for e in it.clone() {
-                if e.meta == Some(Metadata::EndScope) {
-                    break;
+        match e.expr {
+            ExprT::Decision((cond, scope)) => {
+                let cond = DepthExpr::Process(cond);
+
+                let then_branch = parse(&entire_scopes(&mut it));
+
+                let else_branch = if let Some(n) = it.peek() {
+                    match n.expr {
+                        ExprT::Decision((s, _)) => {
+                            if s.trim_start()
+                                .starts_with("else")
+                            {
+                                it.next();
+                                Some(Box::new(parse(&entire_scopes(&mut it))))
+                            } else {
+                                None
+                            }
+                        }
+                        _ => None,
+                    }
+                } else {
+                    None
+                };
+
+                exprs.push(DepthExpr::Decision {
+                    cond: Box::new(cond),
+                    t: scope,
+                    then_branch: Box::new(then_branch),
+                    else_branch,
+                });
+            }
+            ExprT::IO(s) => {
+                if !s.is_empty() {
+                    exprs.push(DepthExpr::IO(s))
                 }
             }
-
-            exprs.push(DepthExpr::Decision {
-                cond: Box::new(cond),
-                t: scope,
-                else_branch: Some(Box::new(parse(&entire_scopes(
-                    it.clone()
-                        .cloned(),
-                )))),
-            });
+            ExprT::Process(s) => {
+                if !s.is_empty() {
+                    exprs.push(DepthExpr::Process(s))
+                }
+            }
+            ExprT::StartOrStop(_) => {}
         }
     }
 
     exprs
 }
 
-fn entire_scopes<'a>(it: impl Iterator<Item = Expr<'a>>) -> Vec<Expr<'a>> {
+fn entire_scopes<'a>(it: &mut Peekable<std::slice::Iter<'_, Expr<'a>>>) -> Vec<Expr<'a>> {
     let mut exprs: Vec<Expr<'a>> = Vec::new();
     let mut scope_count: u32 = 1; // assume we already entered the scope
 
@@ -63,7 +92,7 @@ fn entire_scopes<'a>(it: impl Iterator<Item = Expr<'a>>) -> Vec<Expr<'a>> {
                 }
             }
         }
-        exprs.push(i);
+        exprs.push(*i);
     }
 
     exprs
diff --git a/src/main.rs b/src/main.rs
index 5c638a7..ecbb08c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,7 +2,7 @@ use java2flowchart::ast::parse;
 use java2flowchart::parser;
 
 fn main() {
-    let test = "if (a) {if (a) {a} else if {asd else if {}}}";
+    let test = "if (a) {if (a) {abcdsdf} else if {asd else if {}}}";
 
     println!("{:#?}", parse(&parser::parse(test)));
 }