From f818e86f48eeb7455b0dae3aeda1379fd3b63108 Mon Sep 17 00:00:00 2001
From: greg <greg.shuflin@protonmail.com>
Date: Tue, 15 Oct 2019 00:53:21 -0700
Subject: [PATCH] More visitor work

---
 schala-lang/language/src/ast/visitor.rs |  2 +
 schala-lang/language/src/ast/walker.rs  | 58 ++++++++++++++++++++++++-
 2 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/schala-lang/language/src/ast/visitor.rs b/schala-lang/language/src/ast/visitor.rs
index 9936cf4..9f0b783 100644
--- a/schala-lang/language/src/ast/visitor.rs
+++ b/schala-lang/language/src/ast/visitor.rs
@@ -22,6 +22,7 @@ pub trait ASTVisitor: Sized {
   fn call(&mut self, _f: &Expression, _arguments: &Vec<InvocationArgument>) {}
   fn index(&mut self, _indexee: &Expression, _indexers: &Vec<Expression>) {}
   fn if_expression(&mut self, _discrim: Option<&Expression>, _body: &IfExpressionBody) {}
+  fn condition_arm(&mut self, _arm: &ConditionArm) {}
   fn while_expression(&mut self, _condition: Option<&Expression>, _body: &Block) {}
   fn for_expression(&mut self, _enumerators: &Vec<Enumerator>, _body: &ForBody) {}
   fn lambda(&mut self, _params: &Vec<FormalParam>, _type_anno: Option<&TypeIdentifier>, _body: &Block) {}
@@ -35,4 +36,5 @@ pub trait ASTVisitor: Sized {
   fn bool_literal(&mut self, _b: bool) {}
   fn binexp(&mut self, _op: &BinOp, _lhs: &Expression, _rhs: &Expression) {}
   fn prefix_exp(&mut self, _op: &PrefixOp, _arg: &Expression) {}
+  fn pattern(&mut self, _pat: &Pattern) {}
 }
diff --git a/schala-lang/language/src/ast/walker.rs b/schala-lang/language/src/ast/walker.rs
index d1462c0..6298110 100644
--- a/schala-lang/language/src/ast/walker.rs
+++ b/schala-lang/language/src/ast/walker.rs
@@ -128,6 +128,7 @@ fn expression_kind<V: ASTVisitor>(v: &mut V, expression_kind: &ExpressionKind) {
     TupleLiteral(exprs) => {
       for expr in exprs {
         v.expression(expr);
+        expression(v, expr);
       }
     },
     Value(name) => v.qualified_name(name),
@@ -143,7 +144,11 @@ fn expression_kind<V: ASTVisitor>(v: &mut V, expression_kind: &ExpressionKind) {
       v.index(indexee, indexers);
       index(v, indexee, indexers);
     },
-    IfExpression { discriminator, body } => v.if_expression(deref_optional_box(discriminator), body),
+    IfExpression { discriminator, body } => {
+      v.if_expression(deref_optional_box(discriminator), body);
+      discriminator.as_ref().map(|d| expression(v, d));
+      if_expression_body(v, body);
+    },
     WhileExpression { condition, body } => v.while_expression(deref_optional_box(condition), body),
     ForExpression { enumerators, body } => v.for_expression(enumerators, body),
     Lambda { params , type_anno, body } => {
@@ -153,8 +158,59 @@ fn expression_kind<V: ASTVisitor>(v: &mut V, expression_kind: &ExpressionKind) {
     ListLiteral(exprs) => {
       for expr in exprs {
         v.expression(expr);
+        expression(v, expr);
       }
     },
   }
 }
 
+fn if_expression_body<V: ASTVisitor>(v: &mut V, body: &IfExpressionBody) {
+  use IfExpressionBody::*;
+  match body {
+    SimpleConditional { then_case, else_case } => {
+      walk_block(v, then_case);
+      else_case.as_ref().map(|block| walk_block(v, block));
+    },
+    SimplePatternMatch { pattern, then_case, else_case } => {
+      v.pattern(pattern);
+      walk_pattern(v, pattern);
+      walk_block(v, then_case);
+      else_case.as_ref().map(|block| walk_block(v, block));
+    },
+    CondList(arms) => {
+      for arm in arms {
+        v.condition_arm(arm);
+        condition_arm(v, arm);
+      }
+    }
+  }
+}
+
+fn condition_arm<V: ASTVisitor>(v: &mut V, arm: &ConditionArm) {
+  v.condition_arm(arm);
+}
+
+fn walk_pattern<V: ASTVisitor>(v: &mut V, pat: &Pattern) {
+  use Pattern::*;
+  match pat {
+    TuplePattern(patterns) => {
+      for pat in patterns {
+        v.pattern(pat);
+        walk_pattern(v, pat);
+      }
+    },
+    TupleStruct(_, patterns) => {
+      for pat in patterns {
+        v.pattern(pat);
+        walk_pattern(v, pat);
+      }
+    },
+    Record(_, name_and_patterns) => {
+      for (_, pat) in name_and_patterns {
+        v.pattern(pat);
+        walk_pattern(v, pat);
+      }
+    },
+    _ => ()
+  }
+}