From 226f14bfef937f3d7884a447811f590ea2b97a1e Mon Sep 17 00:00:00 2001
From: Jameson Nash <vtjnash@gmail.com>
Date: Thu, 12 Oct 2017 14:42:18 -0400
Subject: [PATCH] IR: compress consecutive pop_loc nodes

---
 base/inference.jl      | 64 +++++++++++++++++++++++++++++++-----------
 base/show.jl           |  3 ++
 doc/src/devdocs/ast.md |  2 ++
 src/codegen.cpp        |  7 ++++-
 4 files changed, 59 insertions(+), 17 deletions(-)

diff --git a/base/inference.jl b/base/inference.jl
index c8d60eac7c86e..00bab2c0f215c 100644
--- a/base/inference.jl
+++ b/base/inference.jl
@@ -5622,30 +5622,62 @@ function meta_elim_pass!(code::Array{Any,1}, do_coverage::Bool)
             push!(prev_dbg_stack, 0)
             push!(push_loc_pos_stack, i)
         elseif arg1 === :pop_loc
-            prev_dbg = if length(prev_dbg_stack) > 1
-                pop!(prev_dbg_stack)
-            else
-                prev_dbg_stack[end]
-            end
-            if prev_dbg > 0
-                code[prev_dbg] = nothing
+            npops = (nargs > 1 ? args[2]::Int : 1)
+            for pop in 1:npops
+                prev_dbg = if length(prev_dbg_stack) > 1
+                    pop!(prev_dbg_stack)
+                else
+                    prev_dbg_stack[end]
+                end
+                if prev_dbg > 0
+                    code[prev_dbg] = nothing
+                end
+                push_loc = if length(push_loc_pos_stack) > 1
+                    pop!(push_loc_pos_stack)
+                else
+                    push_loc_pos_stack[end]
+                end
+                if push_loc > 0
+                    code[push_loc] = nothing
+                    npops -= 1
+                else
+                    prev_dbg_stack[end] = 0
+                    push_loc_pos_stack[end] = 0
+                end
             end
-            push_loc = if length(push_loc_pos_stack) > 1
-                pop!(push_loc_pos_stack)
+            if npops > 1
+                code[i] = Expr(:meta, :pop_loc, npops)
+            elseif npops == 1
+                code[i] = Expr(:meta, :pop_loc)
             else
-                push_loc_pos_stack[end]
-            end
-            if push_loc > 0
-                code[push_loc] = nothing
                 code[i] = nothing
-            else
-                prev_dbg_stack[end] = 0
-                push_loc_pos_stack[end] = 0
             end
         else
             continue
         end
     end
+
+    # combine consecutive :pop_loc instructions
+    lastpop = nothing
+    npops = 0
+    for i in 1:length(code)
+        ex = code[i]
+        if isa(ex, Expr) && ex.head === :meta && length(ex.args) > 0 && ex.args[1] == :pop_loc
+            npops += (length(ex.args) > 1 ? ex.args[2]::Int : 1)
+            if lastpop === nothing
+                lastpop = ex
+            else
+                code[i] = nothing
+            end
+        elseif ex !== nothing && lastpop !== nothing
+            if npops > 1
+                resize!(lastpop.args, 2)
+                lastpop.args[2] = npops
+            end
+            lastpop = nothing
+            npops = 0
+        end
+    end
 end
 
 # does the same job as alloc_elim_pass for allocations inline in getfields
diff --git a/base/show.jl b/base/show.jl
index c61dd830340da..130b5f992f8cc 100644
--- a/base/show.jl
+++ b/base/show.jl
@@ -1202,6 +1202,9 @@ function show_unquoted(io::IO, ex::Expr, indent::Int, prec::Int)
     elseif head === :meta && length(args) == 1 && args[1] === :pop_loc
         print(io, "# meta: pop location")
         show_type = false
+    elseif head === :meta && length(args) == 2 && args[1] === :pop_loc
+        print(io, "# meta: pop locations ($(args[2]))")
+        show_type = false
     # print anything else as "Expr(head, args...)"
     else
         if head !== :invoke
diff --git a/doc/src/devdocs/ast.md b/doc/src/devdocs/ast.md
index fb6d10dd3bd80..a0f9e4257ad9d 100644
--- a/doc/src/devdocs/ast.md
+++ b/doc/src/devdocs/ast.md
@@ -189,6 +189,8 @@ These symbols appear in the `head` field of `Expr`s in lowered form.
 
       * `:pop_loc`: returns to the source location before the matching `:push_loc`.
 
+          * `args[2]::Int` (optional) specifies the number of `push_loc` to pop
+
 
 ### Method
 
diff --git a/src/codegen.cpp b/src/codegen.cpp
index 7f13e04fcec9d..0ec40acbfba3f 100644
--- a/src/codegen.cpp
+++ b/src/codegen.cpp
@@ -5503,6 +5503,11 @@ static std::unique_ptr<Module> emit_function(
                 cur_prop.loc_changed = true;
             }
             else if (meta_arg == (jl_value_t*)jl_symbol("pop_loc")) {
+                unsigned npops = 1;
+                if (jl_expr_nargs(expr) > 1)
+                    npops = jl_unbox_long(jl_exprarg(expr, 1));
+                for (unsigned i = 1; i < npops; i++)
+                    DI_stack.pop_back();
                 cur_prop.is_poploc = true;
                 auto &DI = DI_stack.back();
                 SP = DI.sp;
@@ -5510,8 +5515,8 @@ static std::unique_ptr<Module> emit_function(
                 cur_prop.file = DI.file;
                 cur_prop.line = DI.line;
                 cur_prop.in_user_code = DI.in_user_code;
-                DI_stack.pop_back();
                 cur_prop.loc_changed = true;
+                DI_stack.pop_back();
             }
         }
         stmtprops[i] = cur_prop;