From 77a997b85bc00e058193a3ee2f6ae6cf76505402 Mon Sep 17 00:00:00 2001
From: Tim Holy <tim.holy@gmail.com>
Date: Wed, 16 Sep 2020 02:47:45 -0500
Subject: [PATCH] Don't split scope-modifying blocks (#431)

Fixes #427
---
 src/construct.jl | 30 ++++++++++++++++++++++++------
 test/toplevel.jl | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/src/construct.jl b/src/construct.jl
index badd89a3af85a..6d4ecf9eb883d 100644
--- a/src/construct.jl
+++ b/src/construct.jl
@@ -437,11 +437,8 @@ mutable struct ExprSplitter
     lnn::Union{LineNumberNode,Nothing}
 end
 function ExprSplitter(mod::Module, ex::Expr; lnn=nothing)
-    index = Int[]
-    if ex.head === :block || ex.head === :toplevel
-        push!(index, 1)
-    end
-    iter = ExprSplitter([(mod,ex)], index, lnn)
+    iter = ExprSplitter(Tuple{Module,Expr}[], Int[], lnn)
+    push_modex!(iter, mod, ex)
     queuenext!(iter)
     return iter
 end
@@ -452,7 +449,15 @@ Base.eltype(::Type{ExprSplitter}) = Tuple{Module,Expr}
 function push_modex!(iter::ExprSplitter, mod::Module, ex::Expr)
     push!(iter.stack, (mod, ex))
     if ex.head === :toplevel || ex.head === :block
-        push!(iter.index, 1)
+        # Issue #427
+        modifies_scope = false
+        for a in ex.args
+            if isa(a, Expr) && a.head ∈ (:local, :global)
+                modifies_scope = true
+                break
+            end
+        end
+        push!(iter.index, modifies_scope ? 0 : 1)
     end
     return iter
 end
@@ -504,6 +509,10 @@ function queuenext!(iter::ExprSplitter)
     elseif head === :block || head === :toplevel
         # Container expression
         idx = iter.index[end]
+        if idx == 0
+            # return the whole block (issue #427)
+            return nothing
+        end
         while idx <= length(ex.args)
             a = ex.args[idx]
             if isa(a, LineNumberNode)
@@ -538,6 +547,15 @@ function Base.iterate(iter::ExprSplitter, state=nothing)
             push_modex!(iter, mod, body)
         end
     end
+    if ex.head === :block || ex.head === :toplevel
+        # This was a block that we couldn't safely descend into (issue #427)
+        if !isempty(iter.index) && iter.index[end] > length(iter.stack[end][2].args)
+            pop!(iter.stack)
+            pop!(iter.index)
+            queuenext!(iter)
+        end
+        return (mod, ex), nothing
+    end
     queuenext!(iter)
     # :global expressions can't be lowered. For debugging it might be nice
     # to still return the lnn, but then we have to work harder on detecting them.
diff --git a/test/toplevel.jl b/test/toplevel.jl
index eb8e32590a1e8..a9633640a4e83 100644
--- a/test/toplevel.jl
+++ b/test/toplevel.jl
@@ -480,3 +480,37 @@ end
     frame = Frame(ToplevelParameters, ex)
     @test JuliaInterpreter.finish!(frame, true) === nothing
 end
+
+@testset "Issue #427" begin
+    ex = :(begin
+        local foo = 10
+        sin(foo)
+    end)
+    for (mod, ex) in ExprSplitter(@__MODULE__, ex)
+        @test JuliaInterpreter.finish!(Frame(mod, ex), true) === nothing
+    end
+    @test length(collect(ExprSplitter(@__MODULE__, ex))) == 1
+    ex = :(begin
+        3 + 7
+        module Local
+            local foo = 10
+            sin(foo)
+        end
+    end)
+    modexs = collect(ExprSplitter(@__MODULE__, ex))
+    @test length(modexs) == 2
+    @test modexs[2][1] == getfield(@__MODULE__, :Local)
+    for (mod, ex) in modexs
+        @test JuliaInterpreter.finish!(Frame(mod, ex), true) === nothing
+    end
+    ex = :(begin
+        3 + 7
+        module Local
+            local foo = 10
+            sin(foo)
+        end
+        3 + 7
+    end)
+    modexs = collect(ExprSplitter(@__MODULE__, ex))
+    @test length(modexs) == 3
+end