Ellide immutable allocation in some simple cases

Extend tuple_elim_pass and getfield_elim_pass to handle immutable object allocations.

Author: carnaval

base/inference.jl

@@ -1436,7 +1436,7 @@ function typeinf(linfo::LambdaInfo, atypes::ANY, sparams::SimpleVector, def, cop
         end
     end
     # TODO: typeinf currently gets stuck without this
-    if linfo.name === :abstract_interpret || linfo.name === :tuple_elim_pass || linfo.name === :abstract_call_gf
+    if linfo.name === :abstract_interpret || linfo.name === :alloc_elim_pass || linfo.name === :abstract_call_gf
         return (linfo.ast, Any)
     end
 
@@ -1861,7 +1861,7 @@ function typeinf_uncached(linfo::LambdaInfo, atypes::ANY, sparams::SimpleVector,
                 sv.vars = append_any(f_argnames(fulltree), map(vi->vi[1], fulltree.args[2][1]))
                 inbounds_meta_elim_pass(fulltree.args[3])
             end
-            tuple_elim_pass(fulltree, sv)
+            alloc_elim_pass(fulltree, sv)
             getfield_elim_pass(fulltree.args[3], sv)
         end
         linfo.inferred = true
@@ -2429,15 +2429,15 @@ function inlineable(f::ANY, ft::ANY, e::Expr, atype::ANY, sv::VarInfo, enclosing
         vaname = args[na]
         len_argexprs = length(argexprs)
         valen = len_argexprs-na+1
-        if valen>0 && !occurs_outside_getfield(body, vaname, sv, valen)
+        if valen>0 && !occurs_outside_getfield(body, vaname, sv, valen, ())
             # argument tuple is not used as a whole, so convert function body
             # to one accepting the exact number of arguments we have.
             newnames = unique_names(ast,valen)
             if needcopy
                 body = astcopy(body)
                 needcopy = false
             end
-            replace_getfield!(ast, body, vaname, newnames, sv, 1)
+            replace_getfield!(ast, body, vaname, newnames, (), sv, 1)
             args = vcat(args[1:na-1], newnames)
             na = length(args)
 
@@ -3190,28 +3190,27 @@ symequal(x::Symbol    , y::SymbolNode) = is(x,y.name)
 symequal(x::GenSym    , y::GenSym)     = is(x.id,y.id)
 symequal(x::ANY       , y::ANY)        = is(x,y)
 
-function occurs_outside_getfield(e::ANY, sym::ANY, sv::VarInfo, tuplen::Int)
+function occurs_outside_getfield(e::ANY, sym::ANY, sv::VarInfo, field_count, field_names)
     if is(e, sym) || (isa(e, SymbolNode) && is(e.name, sym))
         return true
     end
     if isa(e,Expr)
         e = e::Expr
         if is_known_call(e, getfield, sv) && symequal(e.args[2],sym)
-            targ = e.args[2]
-            if !(exprtype(targ,sv) <: Tuple)
-                return true
-            end
             idx = e.args[3]
-            if !isa(idx,Int) || !(1 <= idx <= tuplen)
-                return true
+            if isa(idx,QuoteNode) && (idx.value in field_names)
+                return false
             end
-            return false
+            if isa(idx,Int) && (1 <= idx <= field_count)
+                return false
+            end
+            return true
         end
         if is(e.head,:(=))
-            return occurs_outside_getfield(e.args[2], sym, sv, tuplen)
+            return occurs_outside_getfield(e.args[2], sym, sv, field_count, field_names)
         else
             for a in e.args
-                if occurs_outside_getfield(a, sym, sv, tuplen)
+                if occurs_outside_getfield(a, sym, sv, field_count, field_names)
                     return true
                 end
             end
@@ -3230,46 +3229,80 @@ function inbounds_meta_elim_pass(e::Expr)
     end
 end
 
-# replace getfield(tuple(exprs...), i) with exprs[i]
+# does the same job as alloc_elim_pass for allocations inline in getfields
+# TODO can probably be removed when we switch to a linear IR
 function getfield_elim_pass(e::Expr, sv)
     for i = 1:length(e.args)
         ei = e.args[i]
         if isa(ei,Expr)
             getfield_elim_pass(ei, sv)
             if is_known_call(ei, getfield, sv) && length(ei.args)==3 &&
-                isa(ei.args[3],Int)
+                (isa(ei.args[3],Int) || isa(ei.args[3],QuoteNode))
                 e1 = ei.args[2]
                 j = ei.args[3]
                 if isa(e1,Expr)
-                    if is_known_call(e1, tuple, sv) && (1 <= j < length(e1.args))
-                        ok = true
-                        for k = 2:length(e1.args)
-                            k == j+1 && continue
-                            if !effect_free(e1.args[k], sv, true)
-                                ok = false; break
-                            end
+                    alloc = is_immutable_allocation(e1, sv)
+                    if !is(alloc, false)
+                        flen, fnames = alloc
+                        if isa(j,QuoteNode)
+                            j = findfirst(fnames, j.value)
                         end
-                        if ok
-                            e.args[i] = e1.args[j+1]
+                        if 1 <= j <= flen
+                            ok = true
+                            for k = 2:length(e1.args)
+                                k == j+1 && continue
+                                if !effect_free(e1.args[k], sv, true)
+                                    ok = false; break
+                                end
+                            end
+                            if ok
+                                e.args[i] = e1.args[j+1]
+                            end
                         end
                     end
-                elseif isa(e1,Tuple) && (1 <= j <= length(e1))
+                elseif isa(e1,Tuple) && isa(j,Int) && (1 <= j <= length(e1))
                     e1j = e1[j]
                     if !(isa(e1j,Number) || isa(e1j,AbstractString) || isa(e1j,Tuple) ||
                          isa(e1j,Type))
                         e1j = QuoteNode(e1j)
                     end
                     e.args[i] = e1j
-                elseif isa(e1,QuoteNode) && isa(e1.value,Tuple) && (1 <= j <= length(e1.value))
+                elseif isa(e1,QuoteNode) && isa(e1.value,Tuple) && isa(j,Int) && (1 <= j <= length(e1.value))
                     e.args[i] = QuoteNode(e1.value[j])
                 end
             end
         end
     end
 end
 
-# eliminate allocation of unnecessary tuples
-function tuple_elim_pass(ast::Expr, sv::VarInfo)
+# check if e is a successful allocation of an immutable struct
+# if it is, returns (n,f) such that it is always valid to call
+# getfield(..., 1 <= x <= n) or getfield(..., x in f) on the result
+function is_immutable_allocation(e :: ANY, sv::VarInfo)
+    isa(e, Expr) || return false
+    if is_known_call(e, tuple, sv)
+        return (length(e.args)-1,())
+    elseif e.head === :new
+        typ = exprtype(e, sv)
+        if isleaftype(typ) && !typ.mutable
+            @assert(isa(typ,DataType))
+            nf = length(e.args)-1
+            names = fieldnames(typ)
+            @assert(nf <= nfields(typ))
+            if nf < nfields(typ)
+                # some fields were left undef
+                # we could potentially propagate Bottom
+                # for pointer fields
+                names = names[1:nf]
+            end
+            return (nf, names)
+        end
+    end
+    false
+end
+# eliminate allocation of unnecessary immutables
+# that are only used as arguments to safe getfield calls
+function alloc_elim_pass(ast::Expr, sv::VarInfo)
     bexpr = ast.args[3]::Expr
     body = (ast.args[3].args)::Array{Any,1}
     vs = find_sa_vars(ast)
@@ -3283,10 +3316,11 @@ function tuple_elim_pass(ast::Expr, sv::VarInfo)
         end
         var = e.args[1]
         rhs = e.args[2]
-        if isa(rhs,Expr) && is_known_call(rhs, tuple, sv)
+        alloc = is_immutable_allocation(rhs, sv)
+        if !is(alloc,false)
+            nv, field_names = alloc
             tup = rhs.args
-            nv = length(tup)-1
-            if occurs_outside_getfield(bexpr, var, sv, nv) || !is_local(sv, var)
+            if occurs_outside_getfield(bexpr, var, sv, nv, field_names) || !is_local(sv, var)
                 i += 1
                 continue
             end
@@ -3309,22 +3343,29 @@ function tuple_elim_pass(ast::Expr, sv::VarInfo)
                 end
             end
             i += n_ins
-            replace_getfield!(ast, bexpr, var, vals, sv, i)
+            replace_getfield!(ast, bexpr, var, vals, field_names, sv, i)
         else
             i += 1
         end
     end
 end
 
-function replace_getfield!(ast, e::ANY, tupname, vals, sv, i0)
+function replace_getfield!(ast, e::ANY, tupname, vals, field_names, sv, i0)
     if !isa(e,Expr)
         return
     end
     for i = i0:length(e.args)
         a = e.args[i]
         if isa(a,Expr) && is_known_call(a, getfield, sv) &&
             symequal(a.args[2],tupname)
-            val = vals[a.args[3]]
+            idx = if isa(a.args[3], Int)
+                a.args[3]
+            else
+                @assert isa(a.args[3], QuoteNode)
+                findfirst(field_names, a.args[3].value)
+            end
+            @assert(idx > 0) # clients should check that all getfields are valid
+            val = vals[idx]
             # original expression might have better type info than
             # the tuple element expression that's replacing it.
             if isa(val,SymbolNode)
@@ -3347,7 +3388,7 @@ function replace_getfield!(ast, e::ANY, tupname, vals, sv, i0)
             end
             e.args[i] = val
         else
-            replace_getfield!(ast, a, tupname, vals, sv, 1)
+            replace_getfield!(ast, a, tupname, vals, field_names, sv, 1)
         end
     end
 end

src/codegen.cpp

@@ -3488,6 +3488,7 @@ static jl_cgval_t emit_expr(jl_value_t *expr, jl_codectx_t *ctx)
         if (jl_is_type_type(ty) &&
             jl_is_datatype(jl_tparam0(ty)) &&
             jl_is_leaf_type(jl_tparam0(ty))) {
+            assert(nargs <= jl_datatype_nfields(jl_tparam0(ty))+1);
             return emit_new_struct(jl_tparam0(ty),nargs,args,ctx);
         }
         Value *typ = boxed(emit_expr(args[0], ctx), ctx);

src/julia.h

@@ -692,6 +692,7 @@ STATIC_INLINE char *jl_symbol_name_(jl_sym_t *s)
 #define DEFINE_FIELD_ACCESSORS(f)                                       \
     static inline uint32_t jl_field_##f(jl_datatype_t *st, int i)       \
     {                                                                   \
+        assert(i >= 0 && (size_t)i < jl_datatype_nfields(st));          \
         if (st->fielddesc_type == 0) {                                  \
             return ((jl_fielddesc8_t*)jl_datatype_fields(st))[i].f;     \
         }                                                               \
@@ -705,6 +706,7 @@ STATIC_INLINE char *jl_symbol_name_(jl_sym_t *s)
     static inline void jl_field_set##f(jl_datatype_t *st, int i,        \
                                        uint32_t val)                    \
     {                                                                   \
+        assert(i >= 0 && (size_t)i < jl_datatype_nfields(st));          \
         if (st->fielddesc_type == 0) {                                  \
             ((jl_fielddesc8_t*)jl_datatype_fields(st))[i].f = val;      \
         }                                                               \

test/core.jl

@@ -3845,3 +3845,20 @@ let ary = Vector{Any}(10)
         check_undef_and_fill(ary, 1:(2n + 4))
     end
 end
+
+# pr #15259
+immutable A15259
+    x
+    y
+end
+# check that allocation was ellided
+@eval f15259(x,y) = (a = $(Expr(:new, :A15259, :x, :y)); (a.x, a.y, getfield(a,1), getfield(a, 2)))
+@test isempty(filter(x -> isa(x,Expr) && x.head === :(=) &&
+                          isa(x.args[2], Expr) && x.args[2].head === :new,
+                     code_typed(f15259, (Any,Int))[1].args[3].args))
+@test f15259(1,2) == (1,2,1,2)
+# check that error cases are still correct
+@eval g15259(x,y) = (a = $(Expr(:new, :A15259, :x, :y)); a.z)
+@test_throws ErrorException g15259(1,1)
+@eval h15259(x,y) = (a = $(Expr(:new, :A15259, :x, :y)); getfield(a, 3))
+@test_throws BoundsError h15259(1,1)