2 files changed, 100 insertions, 0 deletions
diff --git a/src/matcher/implementation.lua b/src/matcher/implementation.lua
index 3e823d5b..4431d2ed 100644
--- a/src/matcher/implementation.lua
+++ b/src/matcher/implementation.lua
@@ -6,6 +6,7 @@ local scopes
 local result
 local namePos
 local colonPos
+local maybeCount
 
 local DUMMY_TABLE = {}
 
@@ -24,6 +25,14 @@ local function logicGet()
     return scope and scope.logicId
 end
 
+local function maybePush()
+    maybeCount = maybeCount + 1
+end
+
+local function maybePop()
+    maybeCount = maybeCount - 1
+end
+
 local function scopeInit()
     scopes = {{}}
 end
@@ -61,6 +70,7 @@ end
 local function globalSet(obj)
     obj.logicId = logicGet()
     obj.scopeId = #scopes
+    obj.maybeId = maybeCount
     local name = obj[1]
     for i = #scopes, 1, -1 do
         local scope = scopes[i]
@@ -97,6 +107,9 @@ local function checkImplementation(name, p)
                 goto CONTINUE
             end
             result[#result+1] = {start, finish}
+            if obj.maybeId and obj.maybeId > maybeCount then
+                goto CONTINUE
+            end
             do break end
             ::CONTINUE::
         end
@@ -232,6 +245,7 @@ end
 
 function defs.IfDef()
     logicPush()
+    maybePush()
     scopePush()
 end
 
@@ -259,15 +273,18 @@ end
 
 function defs.EndIf()
     logicFlush()
+    maybePop()
 end
 
 function defs.LoopDef(name)
+    maybePush()
     scopePush()
     scopeSet(name)
 end
 
 function defs.Loop()
     scopePop()
+    maybePop()
 end
 
 function defs.LoopStart(name, exp)
@@ -283,6 +300,7 @@ function defs.SimpleList(...)
 end
 
 function defs.InDef(names)
+    maybePush()
     scopePush()
     for _, name in ipairs(names) do
         scopeSet(name)
@@ -291,27 +309,33 @@ end
 
 function defs.In()
     scopePop()
+    maybePop()
 end
 
 function defs.WhileDef()
+    maybePush()
     scopePush()
 end
 
 function defs.While()
     scopePop()
+    maybePop()
 end
 
 function defs.RepeatDef()
+    maybePush()
     scopePush()
 end
 
 function defs.Until()
     scopePop()
+    maybePop()
 end
 
 return function (buf, pos_)
     pos = pos_
     result = nil
+    maybeCount = 0
     scopeInit()
 
     local suc, err = parser.grammar(buf, 'Lua', defs)
diff --git a/test/implementation/if.lua b/test/implementation/if.lua
index 0c7eefc8..a9ab8b22 100644
--- a/test/implementation/if.lua
+++ b/test/implementation/if.lua
@@ -6,3 +6,79 @@ else
     <?x?> = 1
 end
 ]]
+
+TEST [[
+<!x!> = 1
+if 1 then
+    <!x!> = 1
+else
+    <!x!> = 1
+end
+<?x?> = 1
+]]
+
+TEST [[
+<!x!> = 1
+if 1 then
+    <!x!> = 1
+elseif 1 then
+    <!x!> = 1
+else
+    <!x!> = 1
+end
+<?x?> = 1
+]]
+
+TEST [[
+<!x!> = 1
+if 1 then
+    <!x!> = 1
+elseif 1 then
+    <!x!> = 1
+    if 1 then
+        <!x!> = 1
+    end
+else
+    <!x!> = 1
+end
+<?x?> = 1
+]]
+
+TEST [[
+<!x!> = 1
+while true do
+    <!x!> = 1
+end
+<?x?> = 1
+]]
+
+TEST [[
+<!x!> = 1
+for _ in _ do
+    <!x!> = 1
+end
+<?x?> = 1
+]]
+
+TEST [[
+<!x!> = 1
+for _ = 1, 1 do
+    <!x!> = 1
+end
+<?x?> = 1
+]]
+
+TEST [[
+x = 1
+repeat
+    <!x!> = 1
+until <?x?> == 1
+]]
+
+TEST [[
+<!x!> = 1
+repeat
+    <!x!> = 1
+until 1
+<?x?> = 1
+]]