sroa: Fix incorrect scope counting (#53630)

Sroa was incorrectly assuming that every :leave leaves exactly one
scope. In reality, it leaves as many scopes as the corresponding :leave
references. Fix that to fix #53521.

Author: Keno

base/compiler/ssair/passes.jl

@@ -1225,10 +1225,21 @@ function sroa_pass!(ir::IRCode, inlining::Union{Nothing,InliningState}=nothing)
             bb = compact.active_result_bb - 1
             bbs = scope_mapping[bb]
             if isexpr(stmt, :leave) && bbs != SSAValue(0)
-                update_scope_mapping!(scope_mapping, bb+1, scope_mapping[block_for_inst(compact, bbs)])
-            else
-                update_scope_mapping!(scope_mapping, bb+1, bbs)
+                # Here we want to count the number of scopes that we're leaving,
+                # which is the same as the number of EnterNodes being referenced
+                # by `stmt.args`. Which have :scope set. In practice, the frontend
+                # does emit these in order, so we could simply go to the last one,
+                # but we want to avoid making that semantic assumption.
+                for i = 1:length(stmt.args)
+                    scope = stmt.args[i]
+                    scope === nothing && continue
+                    enter = compact[scope][:inst]
+                    @assert isa(enter, EnterNode)
+                    isdefined(enter, :scope) || continue
+                    bbs = scope_mapping[block_for_inst(compact, bbs)]
+                end
             end
+            update_scope_mapping!(scope_mapping, bb+1, bbs)
         end
         # check whether this statement is `getfield` / `setfield!` (or other "interesting" statement)
         is_setfield = is_isdefined = is_finalizer = is_keyvalue_get = false

test/compiler/irpasses.jl

@@ -1801,3 +1801,22 @@ let n = 1000
     end
 end
 @test f_1000_blocks() == 0
+
+# https://github.com/JuliaLang/julia/issues/53521
+# Incorrect scope counting in :leave
+using Base.ScopedValues
+function f53521()
+    VALUE = ScopedValue(1)
+    @with VALUE => 2 begin
+        for i = 1
+            @with VALUE => 3 begin
+                try
+                    foo()
+                catch
+                    nothing
+                end
+            end
+        end
+    end
+end
+@test code_typed(f53521)[1][2] === Nothing