3 files changed, 17 insertions, 13 deletions

diff --git a/src/ast.rs b/src/ast.rs
index 472172b..a23c688 100644
--- a/src/ast.rs
+++ b/src/ast.rs
@@ -1,3 +1,4 @@
+use crate::tokenizer::If;
 use std::iter::Peekable;
 
 use crate::parser::{Expr, ExprT, Metadata, Scope};
@@ -30,10 +31,8 @@ pub fn parse<'a>(vec: &Vec<Expr<'a>>) -> Vec<DepthExpr<'a>> {
 
                 let else_branch = if let Some(n) = it.peek() {
                     match n.expr {
-                        ExprT::Decision((s, _)) => {
-                            if s.trim_start()
-                                .starts_with("else")
-                            {
+                        ExprT::Decision((_, scope)) => {
+                            if let Scope::If(If::Else) = scope {
                                 it.next();
                                 Some(Box::new(parse(&entire_scopes(&mut it))))
                             } else {
@@ -75,9 +74,6 @@ fn entire_scopes<'a>(it: &mut Peekable<std::slice::Iter<'_, Expr<'a>>>) -> Vec<E
     let mut scope_count: u32 = 1; // assume we already entered the scope
 
     for i in it {
-        if scope_count == 0 {
-            break;
-        }
         if let Some(meta) = i.meta {
             match meta {
                 Metadata::StartScope => {
@@ -93,6 +89,9 @@ fn entire_scopes<'a>(it: &mut Peekable<std::slice::Iter<'_, Expr<'a>>>) -> Vec<E
             }
         }
         exprs.push(*i);
+        if scope_count == 0 {
+            break;
+        }
     }
 
     exprs
diff --git a/src/main.rs b/src/main.rs
index ecbb08c..baa44b5 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) {abcdsdf} else if {asd else if {}}}";
+    let test = "if (first) {firstthen; if (firstthenif) {firstthenifthen} else {firstthenelse}} else {firstelse}";
 
     println!("{:#?}", parse(&parser::parse(test)));
 }
diff --git a/src/parser.rs b/src/parser.rs
index 87d2470..fbeacc5 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -76,11 +76,16 @@ pub fn process(strs: SplitWithMetadata<'_>) -> Vec<Expr<'_>> {
         // TODO: fix one line if else as not delimited by {}
         for s in str.split_whitespace() {
             if let Some(token) = tokenizer::parse(s) {
-                exp = match token {
-                    If(t) => Some(ExprT::Decision((str, Scope::If(t.to_owned())))),
-                    Loop(t) => Some(ExprT::Decision((str, Scope::Loop(t.to_owned())))),
-                    Throw => Some(ExprT::IO(str)),
-                    IO => Some(ExprT::IO(str)),
+                match token {
+                    If(t) => {
+                        if let Some(ExprT::Decision((_, Scope::If(If::Else)))) = exp {
+                        } else {
+                            exp = Some(ExprT::Decision((str, Scope::If(t.to_owned()))));
+                        }
+                    }
+                    Loop(t) => exp = Some(ExprT::Decision((str, Scope::Loop(t.to_owned())))),
+                    Throw => exp = Some(ExprT::IO(str)),
+                    IO => exp = Some(ExprT::IO(str)),
                 }
             }
         }