early take fix

Author: aviatesk

base/compiler/ssair/EscapeAnalysis/EscapeAnalysis.jl

@@ -809,11 +809,6 @@ end
     yroot = find_root!(estate.aliasset, yidx)
     if xroot ≠ yroot
         union!(estate.aliasset, xroot, yroot)
-        xinfo = estate.escapes[xidx]
-        yinfo = estate.escapes[yidx]
-        xyinfo = xinfo ⊔ₑ yinfo
-        estate.escapes[xidx] = xyinfo
-        estate.escapes[yidx] = xyinfo
         return true
     end
     return false
@@ -851,6 +846,10 @@ function add_alias_change!(astate::AnalysisState, @nospecialize(x), @nospecializ
     yidx = iridx(y, estate)
     if xidx !== nothing && yidx !== nothing && !isaliased(xidx, yidx, astate.estate)
         pushfirst!(astate.changes, AliasChange(xidx, yidx))
+        # add new escape change here so that it's shared among the expanded `aliasset`
+        xinfo = estate.escapes[xidx]
+        yinfo = estate.escapes[yidx]
+        add_escape_change!(astate, x, xinfo ⊔ₑ yinfo)
     end
     return nothing
 end
@@ -937,6 +936,8 @@ function compute_frameinfo(ir::IRCode)
             end
             arrayinfo[idx] = dims
         elseif arrayinfo !== nothing
+            # TODO this super limited alias analysis is able to handle only very simple cases
+            # this should be replaced with a proper forward dimension analysis
             if isa(stmt, PhiNode)
                 values = stmt.values
                 local dims = nothing
@@ -946,12 +947,13 @@ function compute_frameinfo(ir::IRCode)
                         if isa(val, SSAValue) && haskey(arrayinfo, val.id)
                             if dims === nothing
                                 dims = arrayinfo[val.id]
-                            elseif dims ≠ arrayinfo[val.id]
-                                dims = nothing
-                                break
+                                continue
+                            elseif dims == arrayinfo[val.id]
+                                continue
                             end
                         end
                     end
+                    @goto next_stmt
                 end
                 if dims !== nothing
                     arrayinfo[idx] = dims
@@ -1107,28 +1109,30 @@ function escape_new!(astate::AnalysisState, pc::Int, args::Vector{Any})
         # fields are known precisely: propagate escape information imposed on recorded possibilities to the exact field values
         infos = AliasInfo.infos
         nf = length(infos)
-        objinfo = ignore_aliasinfo(objinfo)
+        objinfo′ = ignore_aliasinfo(objinfo)
         for i in 2:nargs
             i-1 > nf && break # may happen when e.g. ϕ-node merges values with different types
             arg = args[i]
             add_alias_escapes!(astate, arg, infos[i-1])
             push!(infos[i-1], -pc) # record def
             # propagate the escape information of this object ignoring field information
-            add_escape_change!(astate, arg, objinfo)
+            add_escape_change!(astate, arg, objinfo′)
             add_liveness_change!(astate, arg, pc)
         end
+        add_escape_change!(astate, obj, EscapeInfo(objinfo, AliasInfo)) # update with new AliasInfo
     elseif isa(AliasInfo, Unindexable) && !AliasInfo.array
         # fields are known partially: propagate escape information imposed on recorded possibilities to all fields values
         info = AliasInfo.info
-        objinfo = ignore_aliasinfo(objinfo)
+        objinfo′ = ignore_aliasinfo(objinfo)
         for i in 2:nargs
             arg = args[i]
             add_alias_escapes!(astate, arg, info)
             push!(info, -pc) # record def
             # propagate the escape information of this object ignoring field information
-            add_escape_change!(astate, arg, objinfo)
+            add_escape_change!(astate, arg, objinfo′)
             add_liveness_change!(astate, arg, pc)
         end
+        add_escape_change!(astate, obj, EscapeInfo(objinfo, AliasInfo)) # update with new AliasInfo
     else
         # this object has been used as array, but it is allocated as struct here (i.e. should throw)
         # update obj's field information and just handle this case conservatively
@@ -1405,13 +1409,11 @@ function escape_builtin!(::typeof(getfield), astate::AnalysisState, pc::Int, arg
         isa(AliasInfo, Unindexable) && @goto record_unindexable_use
         @label record_indexable_use
         push!(AliasInfo.infos[fidx], pc) # record use
-        objinfo = EscapeInfo(objinfo, AliasInfo)
-        add_escape_change!(astate, obj, objinfo)
+        add_escape_change!(astate, obj, EscapeInfo(objinfo, AliasInfo)) # update with new AliasInfo
     elseif isa(AliasInfo, Unindexable) && !AliasInfo.array
         @label record_unindexable_use
         push!(AliasInfo.info, pc) # record use
-        objinfo = EscapeInfo(objinfo, AliasInfo)
-        add_escape_change!(astate, obj, objinfo)
+        add_escape_change!(astate, obj, EscapeInfo(objinfo, AliasInfo)) # update with new AliasInfo
     else
         # this object has been used as array, but it is used as struct here (i.e. should throw)
         # update obj's field information and just handle this case conservatively
@@ -1422,6 +1424,7 @@ function escape_builtin!(::typeof(getfield), astate::AnalysisState, pc::Int, arg
         # as the most conservative propagation
         ssainfo = estate[SSAValue(pc)]
         add_escape_change!(astate, obj, ssainfo)
+        add_alias_change!(astate, obj, SSAValue(pc))
     end
     return false
 end
@@ -1457,7 +1460,7 @@ function escape_builtin!(::typeof(setfield!), astate::AnalysisState, pc::Int, ar
         add_alias_escapes!(astate, val, AliasInfo.infos[fidx])
         push!(AliasInfo.infos[fidx], -pc) # record def
         objinfo = EscapeInfo(objinfo, AliasInfo)
-        add_escape_change!(astate, obj, objinfo)
+        add_escape_change!(astate, obj, objinfo) # update with new AliasInfo
         # propagate the escape information of this object ignoring field information
         add_escape_change!(astate, val, ignore_aliasinfo(objinfo))
     elseif isa(AliasInfo, Unindexable) && !AliasInfo.array
@@ -1466,7 +1469,7 @@ function escape_builtin!(::typeof(setfield!), astate::AnalysisState, pc::Int, ar
         add_alias_escapes!(astate, val, AliasInfo.info)
         push!(AliasInfo.info, -pc) # record def
         objinfo = EscapeInfo(objinfo, AliasInfo)
-        add_escape_change!(astate, obj, objinfo)
+        add_escape_change!(astate, obj, objinfo) # update with new AliasInfo
         # propagate the escape information of this object ignoring field information
         add_escape_change!(astate, val, ignore_aliasinfo(objinfo))
     else
@@ -1546,6 +1549,7 @@ function escape_builtin!(::typeof(arrayref), astate::AnalysisState, pc::Int, arg
         @label conservative_propagation
         ssainfo = estate[SSAValue(pc)]
         add_escape_change!(astate, ary, ssainfo)
+        add_alias_change!(astate, ary, SSAValue(pc))
     end
     return true
 end