From a484df12e3d495f8652607c36a94cba91a5830d0 Mon Sep 17 00:00:00 2001
From: Victor <melnikukr@mail.ru>
Date: Fri, 22 Jul 2016 22:44:46 +0300
Subject: [PATCH] =?UTF-8?q?=D0=98=D1=81=D0=BF=D1=80=D0=B0=D0=B2=D0=BB?=
 =?UTF-8?q?=D0=B5=D0=BD=20=D0=BF=D0=BE=D0=B4=D1=81=D1=87=D1=91=D1=82=20?=
 =?UTF-8?q?=D0=BC=D0=BE=D0=B4=D0=B8=D1=84=D0=B8=D0=BA=D0=B0=D1=86=D0=B8?=
 =?UTF-8?q?=D0=B9=20=D0=BF=D0=B5=D1=80=D0=B5=D0=BC=D0=B5=D0=BD=D0=BD=D1=8B?=
 =?UTF-8?q?=D1=85=20=D0=BF=D1=80=D0=B8=20=D1=83=D0=BD=D0=B0=D1=80=D0=BD?=
 =?UTF-8?q?=D1=8B=D1=85=20=D0=BE=D0=BF=D0=B5=D1=80=D0=B0=D1=86=D0=B8=D1=8F?=
 =?UTF-8?q?=D1=85?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../parser/optimization/VariablesGrabber.java | 44 ++++++++++++++++++-
 1 file changed, 42 insertions(+), 2 deletions(-)

diff --git a/src/com/annimon/ownlang/parser/optimization/VariablesGrabber.java b/src/com/annimon/ownlang/parser/optimization/VariablesGrabber.java
index d46f824..e980f68 100644
--- a/src/com/annimon/ownlang/parser/optimization/VariablesGrabber.java
+++ b/src/com/annimon/ownlang/parser/optimization/VariablesGrabber.java
@@ -1,11 +1,16 @@
 package com.annimon.ownlang.parser.optimization;
 
+import com.annimon.ownlang.parser.ast.Accessible;
 import com.annimon.ownlang.parser.ast.Argument;
 import com.annimon.ownlang.parser.ast.AssignmentExpression;
+import com.annimon.ownlang.parser.ast.ContainerAccessExpression;
 import com.annimon.ownlang.parser.ast.DestructuringAssignmentStatement;
+import com.annimon.ownlang.parser.ast.ForeachArrayStatement;
+import com.annimon.ownlang.parser.ast.ForeachMapStatement;
 import com.annimon.ownlang.parser.ast.FunctionDefineStatement;
 import com.annimon.ownlang.parser.ast.MatchExpression;
 import com.annimon.ownlang.parser.ast.Node;
+import com.annimon.ownlang.parser.ast.UnaryExpression;
 import com.annimon.ownlang.parser.ast.ValueExpression;
 import com.annimon.ownlang.parser.ast.VariableExpression;
 import static com.annimon.ownlang.parser.visitors.VisitorUtils.isValue;
@@ -46,6 +51,19 @@ public class VariablesGrabber extends OptimizationVisitor<Map<String, VariableIn
         return super.visit(s, t);
     }
 
+    @Override
+    public Node visit(ForeachArrayStatement s, Map<String, VariableInfo> t) {
+        t.put(s.variable, variableInfo(t, s.variable));
+        return super.visit(s, t);
+    }
+
+    @Override
+    public Node visit(ForeachMapStatement s, Map<String, VariableInfo> t) {
+        t.put(s.key, variableInfo(t, s.key));
+        t.put(s.value, variableInfo(t, s.value));
+        return super.visit(s, t);
+    }
+
     @Override
     public Node visit(FunctionDefineStatement s, Map<String, VariableInfo> t) {
         for (Argument argument : s.arguments) {
@@ -57,10 +75,32 @@ public class VariablesGrabber extends OptimizationVisitor<Map<String, VariableIn
 
     @Override
     public Node visit(MatchExpression s, Map<String, VariableInfo> t) {
-        // no visit match expression
-        return s;
+        for (MatchExpression.Pattern pattern : s.patterns) {
+            if (pattern instanceof MatchExpression.VariablePattern) {
+                final String variableName = ((MatchExpression.VariablePattern) pattern).variable;
+                t.put(variableName, variableInfo(t, variableName));
+            }
+        }
+        return super.visit(s, t);
     }
 
+    @Override
+    public Node visit(UnaryExpression s, Map<String, VariableInfo> t) {
+        if (s.expr1 instanceof Accessible) {
+            if (s.expr1 instanceof VariableExpression) {
+                final String variableName = ((VariableExpression) s.expr1).name;
+                t.put(variableName, variableInfo(t, variableName));
+            }
+            if (s.expr1 instanceof ContainerAccessExpression) {
+                final String variableName = ((ContainerAccessExpression) s.expr1).variable;
+                t.put(variableName, variableInfo(t, variableName));
+            }
+        }
+        return super.visit(s, t);
+    }
+
+
+
     private VariableInfo variableInfo(Map<String, VariableInfo> t, final String variableName) {
         final VariableInfo var;
         if (t.containsKey(variableName)) {