3 files changed, 97 insertions, 27 deletions

diff --git a/src/matcher/definition.lua b/src/matcher/definition.lua
index 41ad9f8a..91b5090f 100644
--- a/src/matcher/definition.lua
+++ b/src/matcher/definition.lua
@@ -13,9 +13,9 @@ local function scopeInit()
     scopes = {{}}
 end
 
-local function scopeSet(name, p)
+local function scopeSet(name)
     local scope = scopes[#scopes]
-    scope[name] = p
+    scope[name[1]] = name[2]
 end
 
 local function scopeGet(name)
@@ -50,23 +50,22 @@ function defs.Name(p, str)
 end
 
 function defs.LocalVar(name)
-    scopeSet(name[1], name[2])
+    scopeSet(name)
 end
 
 function defs.LocalSet(name)
-    scopeSet(name[1], name[2])
+    scopeSet(name)
 end
 
 function defs.Function(func)
     local names = func.name
     if names and #names == 1 then
-        local name = names[1]
-        scopeSet(name[1], name[2])
+        scopeSet(names[1])
     end
     return func
 end
 
-function defs.DoStart()
+function defs.DoDef()
     scopePush()
 end
 
@@ -74,7 +73,7 @@ function defs.Do()
     scopePop()
 end
 
-function defs.IfStart()
+function defs.IfDef()
     scopePush()
 end
 
@@ -82,7 +81,7 @@ function defs.If()
     scopePop()
 end
 
-function defs.ElseIfStart()
+function defs.ElseIfDef()
     scopePush()
 end
 
@@ -90,7 +89,7 @@ function defs.ElseIf()
     scopePop()
 end
 
-function defs.ElseStart()
+function defs.ElseDef()
     scopePush()
 end
 
@@ -98,6 +97,34 @@ function defs.Else()
     scopePop()
 end
 
+function defs.LoopDef(name)
+    scopePush()
+    scopeSet(name)
+end
+
+function defs.Loop()
+    scopePop()
+end
+
+function defs.LoopStart(name, exp)
+    return name
+end
+
+function defs.InList(...)
+    return {...}
+end
+
+function defs.InDef(names)
+    scopePush()
+    for _, name in ipairs(names) do
+        scopeSet(name)
+    end
+end
+
+function defs.In()
+    scopePop()
+end
+
 return function (buf, pos_)
     pos = pos_
     result = nil
diff --git a/src/parser/grammar.lua b/src/parser/grammar.lua
index a2e4faa6..e95d2d7f 100644
--- a/src/parser/grammar.lua
+++ b/src/parser/grammar.lua
@@ -252,8 +252,8 @@ Action      <-  SEMICOLON / Do / Break / Return / Label / GoTo / If / For / Whil
 
 ExpList     <-  Exp (COMMA Exp)*
 
-Do          <-  DO -> DoStart
-                    (!END Action)* -> Do
+Do          <-  DO                  -> DoDef
+                    (!END Action)*  -> Do
                 END
 
 Break       <-  BREAK
@@ -268,24 +268,26 @@ If          <-  IfPart
                 ElseIfPart*
                 ElsePart?
                 END
-IfPart      <-  (IF Exp THEN) -> IfStart
-                    (!ELSEIF !ELSE !END Action)* -> If
-ElseIfPart  <-  (ELSEIF Exp THEN) -> ElseIfStart
-                    (!ELSE !END Action)* -> ElseIf
-ElsePart    <-  ELSE -> ElseStart
-                    (!END Action)* -> Else
+IfPart      <-  (IF Exp THEN)                       -> IfDef
+                    (!ELSEIF !ELSE !END Action)*    -> If
+ElseIfPart  <-  (ELSEIF Exp THEN)                   -> ElseIfDef
+                    (!ELSE !END Action)*            -> ElseIf
+ElsePart    <-  ELSE                                -> ElseDef
+                    (!END Action)*                  -> Else
+
 For         <-  Loop / In
-Loop        <-  FOR LoopStart LoopFinish LoopStep? DO
-                    (!END Action)*
+
+Loop        <-  (FOR LoopStart LoopFinish LoopStep? DO) -> LoopDef
+                    (!END Action)*                      -> Loop
                 END
-LoopStart   <-  Name ASSIGN Exp
+LoopStart   <-  (Name ASSIGN Exp)                       -> LoopStart
 LoopFinish  <-  COMMA Exp
 LoopStep    <-  COMMA Exp
 
-In          <-  FOR NameList IN ExpList DO
-                    (!END Action)*
+In          <-  (FOR InList IN ExpList DO)  -> InDef
+                    (!END Action)*          -> In
                 END
-NameList    <-  Name (COMMA Name)*
+InList      <-  (Name (COMMA Name)*)        -> InList
 
 While       <-  WHILE Exp DO
                     (!END Action)*
diff --git a/test/definition/init.lua b/test/definition/init.lua
index c2f2b974..61eccfbf 100644
--- a/test/definition/init.lua
+++ b/test/definition/init.lua
@@ -60,7 +60,7 @@ if <?x?> then
 end
 ]]
 
-test[[
+test [[
 local <!x!>
 if x then
     local x
@@ -69,7 +69,7 @@ elseif <?x?> then
 end
 ]]
 
-test[[
+test [[
 local <!x!>
 if x then
     local x
@@ -81,7 +81,7 @@ end
 <?x?> = 1
 ]]
 
-test[[
+test [[
 local <!x!>
 if x then
     <?x?> = 1
@@ -91,3 +91,44 @@ else
     local x
 end
 ]]
+
+test [[
+local <!x!>
+for x = 1, 10 do
+end
+<?x?> = 1
+]]
+
+test [[
+local x
+for <!x!> = 1, 10 do
+    <?x?> = 1
+end
+]]
+
+test [[
+local <!x!>
+for x in x do
+end
+<?x?> = 1
+]]
+
+test [[
+local <!x!>
+for x in <?x?> do
+end
+]]
+
+test [[
+local x
+for <!x!> in x do
+    <?x?> = 1
+end
+]]
+
+test [[
+local x
+for z, y, <!x!> in x do
+    <?x?> = 1
+end
+]]