sroa: Handle looking through chains of KeyValue instances

base/compiler/ssair/passes.jl

@@ -866,36 +866,38 @@ function lift_leaves_keyvalue(compact::IncrementalCompact, @nospecialize(key),
     for i = 1:length(leaves)
         leaf = leaves[i]
         cache_key = leaf
-        if isa(leaf, AnySSAValue)
-            (def, leaf) = walk_to_def(compact, leaf)
-            if is_known_invoke_or_call(def, Core.OptimizedGenerics.KeyValue.set, compact)
-                @assert isexpr(def, :invoke)
-                if length(def.args) in (5, 6)
-                    collection = def.args[end-2]
-                    set_key = def.args[end-1]
-                    set_val_idx = length(def.args)
-                elseif length(def.args) == 4
-                    collection = def.args[end-1]
-                    # Key is deleted
-                    # TODO: Model this
-                    return nothing
-                elseif length(def.args) == 3
-                    collection = def.args[end]
-                    # The whole collection is deleted
-                    # TODO: Model this
-                    return nothing
-                else
-                    return nothing
-                end
-                if set_key === key || (egal_tfunc(𝕃ₒ, argextype(key, compact), argextype(set_key, compact)) == Const(true))
-                    lift_arg!(compact, leaf, cache_key, def, set_val_idx, lifted_leaves)
+        while true
+            if isa(leaf, AnySSAValue)
+                (def, leaf) = walk_to_def(compact, leaf)
+                if is_known_invoke_or_call(def, Core.OptimizedGenerics.KeyValue.set, compact)
+                    @assert isexpr(def, :invoke)
+                    if length(def.args) in (5, 6)
+                        collection = def.args[end-2]
+                        set_key = def.args[end-1]
+                        set_val_idx = length(def.args)
+                    elseif length(def.args) == 4
+                        collection = def.args[end-1]
+                        # Key is deleted
+                        # TODO: Model this
+                        return nothing
+                    elseif length(def.args) == 3
+                        collection = def.args[end]
+                        # The whole collection is deleted
+                        # TODO: Model this
+                        return nothing
+                    else
+                        return nothing
+                    end
+                    if set_key === key || (egal_tfunc(𝕃ₒ, argextype(key, compact), argextype(set_key, compact)) == Const(true))
+                        lift_arg!(compact, leaf, cache_key, def, set_val_idx, lifted_leaves)
+                        break
+                    end
+                    leaf = collection
                     continue
                 end
-                # TODO: Continue walking the chain
-                return nothing
             end
+            return nothing
         end
-        return nothing
     end
     return lifted_leaves
 end
@@ -919,11 +921,11 @@ function lift_keyvalue_get!(compact::IncrementalCompact, idx::Int, stmt::Expr, 
     (lifted_val, nest) = perform_lifting!(compact,
         visited_philikes, key, result_t, lifted_leaves, collection, nothing)
 
-    compact[idx] = lifted_val === nothing ? nothing : Expr(:call, Core.tuple, lifted_val.val)
+    compact[idx] = lifted_val === nothing ? nothing : Expr(:call, GlobalRef(Core, :tuple), lifted_val.val)
     finish_phi_nest!(compact, nest)
     if lifted_val !== nothing
-        if !⊑(𝕃ₒ, compact[SSAValue(idx)][:type], result_t)
-            compact[SSAValue(idx)][:flag] |= IR_FLAG_REFINED
+        if !⊑(𝕃ₒ, compact[SSAValue(idx)][:type], tuple_tfunc(𝕃ₒ, Any[result_t]))
+            add_flag!(compact[SSAValue(idx)], IR_FLAG_REFINED)
         end
     end
 

base/scopedvalues.jl

@@ -70,11 +70,14 @@ function Scope(parent::Union{Nothing, Scope}, key::ScopedValue{T}, value) where
     return Scope(ScopeStorage(parent.values, key=>val))
 end
 
-function Scope(scope, pairs::Pair{<:ScopedValue}...)
-    for pair in pairs
-        scope = Scope(scope, pair...)
-    end
-    return scope::Scope
+function Scope(scope, pair::Pair{<:ScopedValue})
+    return Scope(scope, pair...)
+end
+
+function Scope(scope, pair1::Pair{<:ScopedValue}, pair2::Pair{<:ScopedValue}, pairs::Pair{<:ScopedValue}...)
+    # Unroll this loop through recursion to make sure that
+    # our compiler optimization support works
+    return Scope(Scope(scope, pair1...), pair2, pairs...)
 end
 Scope(::Nothing) = nothing
 

test/compiler/irpasses.jl

@@ -1549,11 +1549,20 @@ function persistent_dict_elim()
     a = Base.PersistentDict(:a => 1)
     return a[:a]
 end
+
 # Ideally we would be able to fully eliminate this,
 # but currently this would require an extra round of constprop
 @test_broken fully_eliminated(persistent_dict_elim)
 @test code_typed(persistent_dict_elim)[1][1].code[end] == Core.ReturnNode(1)
 
+function persistent_dict_elim_multiple()
+    a = Base.PersistentDict(:a => 1)
+    b = Base.PersistentDict(a, :b => 2)
+    return b[:a]
+end
+@test_broken fully_eliminated(persistent_dict_elim_multiple)
+@test code_typed(persistent_dict_elim_multiple)[1][1].code[end] == Core.ReturnNode(1)
+
 # Test CFG simplify with try/catch blocks
 let code = Any[
         # Block 1