Improve inference for common reflection operations

Keno · Keno · commit 3b5c7e1e877b · 2017-01-31T14:51:51.000-05:00
Make inference be able to infer the type constructor. This is a little
tricky, since we don't really have a good way to represent this. It's
not quite `Const(TypeVar(:T,lb,ub))`, because a) the bounds may only
be accurate up to type equality and b) TypeVar is not `isbits`, so
it's not actually egal to the value we'll have at runtime. Additionally
Type{T} already has meaning as a partially constructed type (e.g. an
unwrapped UnionAll), so using T::Type{T} runs the risk of confusing
this with an unwrapped type. Instead, introduce a new Const-like type,
`PartialTypeVar`, which carries the type var and also keeps track of
whether the bounds were egal or only typequal (we don't take advantage
of that yet, but we could in the future).

Additionally, improve the inference of `typename` and allow constant folding
of field accesses on TypeName and SimpleVectors (to be able to constant
fold T.parameters[1]).
diff --git a/base/inference.jl b/base/inference.jl
@@ -82,6 +82,15 @@ type Conditional
     end
 end
 
+immutable PartialTypeVar
+    tv::TypeVar
+    # N.B.: Currently unused, but would allow turning something back
+    # into Const, if the bounds are pulled out of this TypeVar
+    lb_certain::Bool
+    ub_certain::Bool
+    PartialTypeVar(tv::TypeVar, lb_certain::Bool, ub_certain::Bool) = new(tv, lb_certain, ub_certain)
+end
+
 function rewrap(t::ANY, u::ANY)
     isa(t, Const) && return t
     isa(t, Conditional) && return t
@@ -678,6 +687,7 @@ function limit_type_depth(t::ANY, d::Int, cov::Bool=true, var::Union{Void,TypeVa
     return (cov && !stillcov) ? UnionAll(var, R) : R
 end
 
+const DataType_name_fieldindex = fieldindex(DataType, :name)
 const DataType_parameters_fieldindex = fieldindex(DataType, :parameters)
 const DataType_types_fieldindex = fieldindex(DataType, :types)
 const DataType_super_fieldindex = fieldindex(DataType, :super)
@@ -713,7 +723,8 @@ function getfield_tfunc(s00::ANY, name)
             if isa(sv, Module) && isa(nv, Symbol)
                 return abstract_eval_global(sv, nv)
             end
-            if (isa(sv, DataType) || isimmutable(sv)) && isdefined(sv, nv)
+            if (isa(sv, DataType) || isa(sv, SimpleVector) || isa(sv, TypeName)
+                || isimmutable(sv)) && isdefined(sv, nv)
                 return abstract_eval_constant(getfield(sv, nv))
             end
         end
@@ -774,7 +785,8 @@ function getfield_tfunc(s00::ANY, name)
         sp = nothing
     end
     if (sp !== nothing &&
-        (fld == DataType_parameters_fieldindex ||
+        (fld == DataType_name_fieldindex ||
+         fld == DataType_parameters_fieldindex ||
          fld == DataType_types_fieldindex ||
          fld == DataType_super_fieldindex))
         return Const(getfield(sp, fld))
@@ -905,15 +917,20 @@ function apply_type_tfunc(headtypetype::ANY, args::ANY...)
         return Any
     end
     uncertain = false
+    canconst = true
     tparams = Any[]
     outervars = Any[]
     for i = 1:largs
         ai = args[i]
         if isType(ai)
             aip1 = ai.parameters[1]
+            canconst &= !isleaftype(aip1)
             push!(tparams, aip1)
-        elseif isa(ai, Const) && (isa(ai.val, Type) || valid_tparam(ai.val))
+        elseif isa(ai, Const) && (isa(ai.val, Type) || isa(ai.val, TypeVar) || valid_tparam(ai.val))
             push!(tparams, ai.val)
+        elseif isa(ai, PartialTypeVar)
+            canconst = false
+            push!(tparams, ai.tv)
         else
             # TODO: return `Bottom` for trying to apply a non-UnionAll
             uncertain = true
@@ -956,11 +973,11 @@ function apply_type_tfunc(headtypetype::ANY, args::ANY...)
         # doesn't match, which could happen if a type estimate is too coarse
         return Type{_} where _<:headtype
     end
-    !uncertain && return Const(appl)
+    !uncertain && canconst && return Const(appl)
     if isvarargtype(headtype)
         return Type
     end
-    if type_too_complex(appl,0)
+    if uncertain && type_too_complex(appl,0)
         return Type{_} where _<:headtype
     end
     if istuple
@@ -1476,6 +1493,25 @@ function Pair_name()
     return _Pair_name
 end
 
+_typename(a) = Union{}
+_typename(a::Vararg) = Any
+_typename(a::TypeVar) = Any
+_typename(a::DataType) = Const(a.name)
+function _typename(a::Union)
+    ta = _typename(a.a)
+    tb = _typename(a.b)
+    ta === tb ? tb : (ta === Any || tb === Any) ? Any : Union{}
+end
+_typename(union::UnionAll) = _typename(union.body)
+function typename_static(t)
+    # N.B.: typename maps type equivalence classes to a single value
+    if isa(t, Const) || isType(t)
+        return _typename(isa(t, Const) ? t.val : t.parameters[1])
+    else
+        return Any
+    end
+end
+
 function abstract_call(f::ANY, fargs::Union{Tuple{},Vector{Any}}, argtypes::Vector{Any}, vtypes::VarTable, sv::InferenceState)
     if f === _apply
         length(fargs) > 1 || return Any
@@ -1557,19 +1593,65 @@ function abstract_call(f::ANY, fargs::Union{Tuple{},Vector{Any}}, argtypes::Vect
             end
         end
         return Any
-    elseif f === UnionAll
-        if length(fargs) == 3 && isa(argtypes[2], Const)
-            tv = argtypes[2].val
-            if isa(tv, TypeVar)
+    elseif f === TypeVar
+        lb = Union{}
+        ub = Any
+        ub_certain = lb_certain = true
+        if length(fargs) >= 2 && isa(argtypes[2], Const)
+            nv = argtypes[2].val
+            ubidx = 3
+            if length(fargs) >= 4
+                ubidx = 4
                 if isa(argtypes[3], Const)
-                    body = argtypes[3].val
+                    lb = argtypes[3].val
                 elseif isType(argtypes[3])
-                    body = argtypes[3].parameters[1]
+                    lb = argtypes[3].parameters[1]
+                    lb_certain = false
                 else
-                    return Any
+                    return TypeVar
                 end
-                return abstract_eval_constant(UnionAll(tv, body))
             end
+            if length(fargs) >= ubidx
+                if isa(argtypes[ubidx], Const)
+                    ub = argtypes[ubidx].val
+                elseif isType(argtypes[ubidx])
+                    ub = argtypes[ubidx].parameters[1]
+                    ub_certain = false
+                else
+                    return TypeVar
+                end
+            end
+            tv = TypeVar(nv, lb, ub)
+            return PartialTypeVar(tv, lb_certain, ub_certain)
+        end
+        return TypeVar
+    elseif f === UnionAll
+        if length(fargs) == 3
+            canconst = true
+            if isa(argtypes[3], Const)
+                body = argtypes[3].val
+            elseif isType(argtypes[3])
+                body = argtypes[3].parameters[1]
+                canconst = false
+            else
+                return Any
+            end
+            if isa(argtypes[2], Const)
+                tv = argtypes[2].val
+            elseif isa(argtypes[2], PartialTypeVar)
+                ptv = argtypes[2]
+                tv = ptv.tv
+                canconst = false
+            else
+                return Any
+            end
+            !isa(tv, TypeVar) && return Any
+            if !isa(body, Type) && !isa(body, TypeVar)
+                return Any
+            end
+            theunion = UnionAll(tv, body)
+            ret = canconst ? abstract_eval_constant(theunion) : Type{theunion}
+            return ret
         end
         return Any
     elseif f === return_type
@@ -1595,7 +1677,16 @@ function abstract_call(f::ANY, fargs::Union{Tuple{},Vector{Any}}, argtypes::Vect
 
     if length(argtypes)>2 && argtypes[3] ⊑ Int
         at2 = widenconst(argtypes[2])
-        if (at2 <: Tuple ||
+        if at2 <: SimpleVector && istopfunction(tm, f, :getindex)
+            if isa(argtypes[2], Const) && isa(argtypes[3], Const)
+                svecval = argtypes[2].val
+                idx = argtypes[3].val
+                if isa(idx, Int) && 1 <= idx <= length(svecval) &
+                      isassigned(svecval, idx)
+                    return Const(getindex(svecval, idx))
+                end
+            end
+        elseif (at2 <: Tuple ||
             (isa(at2, DataType) && (at2::DataType).name === Pair_name()))
             # allow tuple indexing functions to take advantage of constant
             # index arguments.
@@ -1617,6 +1708,12 @@ function abstract_call(f::ANY, fargs::Union{Tuple{},Vector{Any}}, argtypes::Vect
 
     if istopfunction(tm, f, :promote_type) || istopfunction(tm, f, :typejoin)
         return Type
+    elseif length(argtypes) == 2 && istopfunction(tm, f, :typename)
+        t = argtypes[2]
+        if isa(t, Const) || isType(t)
+            return typename_static(t)
+        end
+        return Any
     end
 
     if sv.params.inlining
@@ -1905,6 +2002,7 @@ function widenconst(c::Const)
         return typeof(c.val)
     end
 end
+widenconst(c::PartialTypeVar) = TypeVar
 widenconst(t::ANY) = t
 
 issubstate(a::VarState, b::VarState) = (a.typ ⊑ b.typ && a.undef <= b.undef)
@@ -3552,7 +3650,7 @@ function inlineable(f::ANY, ft::ANY, e::Expr, atypes::Vector{Any}, sv::Inference
             if method.name == :getindex || method.name == :next || method.name == :indexed_next
                 if length(atypes) > 2 && atypes[3] ⊑ Int
                     at2 = widenconst(atypes[2])
-                    if (at2 <: Tuple ||
+                    if (at2 <: Tuple || at2 <: SimpleVector ||
                         (isa(at2, DataType) && (at2::DataType).name === Pair_name()))
                         force_infer = true
                     end
diff --git a/test/inference.jl b/test/inference.jl
@@ -588,3 +588,14 @@ f11015(a::AT11015) = g11015(Base.fieldtype(typeof(a), :f), true)
 g11015(::Type{Bool}, ::Bool) = 2.0
 @test Int <: Base.return_types(f11015, (AT11015,))[1]
 @test f11015(AT11015(true)) === 1
+
+# Inference for some type-level computation
+fUnionAll{T}(::Type{T}) = Type{S} where S <: T
+@inferred fUnionAll(Real) == Type{T} where T <: Real
+@inferred fUnionAll(Rational{T} where T <: AbstractFloat) == Type{T} where T<:(Rational{S} where S <: AbstractFloat)
+
+fComplicatedUnionAll{T}(::Type{T}) = Type{Tuple{S,rand() >= 0.5 ? Int : Float64}} where S <: T
+let pub = Base.parameter_upper_bound, x = fComplicatedUnionAll(Real)
+    @test pub(pub(x, 1), 1) == Real
+    @test pub(pub(x, 1), 2) == Int || pub(pub(x, 1), 2) == Float64
+end
diff --git a/test/reflection.jl b/test/reflection.jl
@@ -629,9 +629,12 @@ end
 @test Base.parameter_upper_bound(ReflectionExample, 2) === Any
 @test Base.parameter_upper_bound(ReflectionExample{T, N} where T where N <: Real, 2) === Real
 
-@test Base.typename(ReflectionExample{Float64, Int64}).wrapper === ReflectionExample
-@test Base.typename(ReflectionExample{Float64, N} where N).wrapper === ReflectionExample
-@test Base.typename(ReflectionExample{T, Int64} where T).wrapper === ReflectionExample
-@test Base.typename(ReflectionExample).wrapper === ReflectionExample
-@test Base.typename(Union{ReflectionExample{Union{},1},ReflectionExample{Float64,1}}).wrapper === ReflectionExample
-@test_throws ErrorException Base.typename(Union{Int, Float64})
+let
+    wrapperT(T) = Base.typename(T).wrapper
+    @test @inferred wrapperT(ReflectionExample{Float64, Int64}) == ReflectionExample
+    @test @inferred wrapperT(ReflectionExample{Float64, N} where N) == ReflectionExample
+    @test @inferred wrapperT(ReflectionExample{T, Int64} where T) == ReflectionExample
+    @test @inferred wrapperT(ReflectionExample) == ReflectionExample
+    @test @inferred wrapperT(Union{ReflectionExample{Union{},1},ReflectionExample{Float64,1}}) == ReflectionExample
+    @test_throws ErrorException Base.typename(Union{Int, Float64})
+end
