IR: compress consecutive pop_loc nodes (#24109)

Author: vtjnash

base/inference.jl

@@ -5625,30 +5625,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

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

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
 

src/codegen.cpp

@@ -5503,15 +5503,20 @@ 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;
                 cur_prop.loc = DI.loc;
                 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;