Miscellaneous type inferability fixes

src/Adaptivity/AdaptivityGlues.jl

@@ -235,6 +235,7 @@ function get_d_to_fface_to_cface(glue::AdaptivityGlue{<:RefinementGlue},
   # Local data for each coarse cell, at the RefinementRule level
   rrules = get_old_cell_refinement_rules(glue)
   ccell_to_d_to_faces = lazy_map(rr->map(d->Geometry.get_faces(get_grid_topology(rr.ref_grid),Dc,d),0:Dc),rrules)
+  #ccell_to_d_to_faces = lazy_map(rr->map(d->Geometry.get_faces(get_grid_topology(rr.ref_grid),Dc,d)::Table{Int32, Vector{Int32}, Vector{Int32}},0:Dc),rrules)
   ccell_to_d_to_fface_to_parent_face = lazy_map(get_d_to_face_to_parent_face,rrules)
   fcell_to_child_id = glue.n2o_cell_to_child_id
 

src/Algebra/LinearSolvers.jl

@@ -147,11 +147,11 @@ function test_linear_solver(
   ss = symbolic_setup(ls,A)
   ns = numerical_setup(ss,A)
   numerical_setup!(ns,A)
-  solve!(y,ns,b)
+  @inferred solve!(y,ns,b)
   @test x ≈ y
 
   y .= zero(eltype(y))
-  solve!(y,ls,A,b)
+  @inferred solve!(y,ls,A,b)
   @test x ≈ y
 
   op = AffineOperator(A,b)

src/Algebra/NonlinearSolvers.jl

@@ -59,15 +59,15 @@ function test_nonlinear_solver(
   pred::Function=isapprox)
 
   x1 = copy(x0)
-  cache = solve!(x1,nls,op)
+  cache = @inferred solve!(x1,nls,op)
   @test pred(x1,x)
 
   x1 = copy(x0)
-  cache = solve!(x1,nls,op,cache)
+  cache = @inferred solve!(x1,nls,op,cache)
   @test pred(x1,x)
 
   x1 = copy(x0)
-  cache = solve!(x1,nls,op,cache)
+  cache = @inferred solve!(x1,nls,op,cache)
   @test pred(x1,x)
 
 end

src/Arrays/Autodiff.jl

@@ -99,13 +99,15 @@ end
 ConfigMap(f) = ConfigMap(f,nothing)
 
 # TODO Prescribing long chunk size can lead to slow compilation times!
-function return_cache(k::ConfigMap{typeof(ForwardDiff.gradient)},x)
-  cfg = ForwardDiff.GradientConfig(k.tag,x,ForwardDiff.Chunk{length(x)}())
+function return_cache(k::ConfigMap{typeof(ForwardDiff.gradient)}, x)
+  chunk_x = ForwardDiff.Chunk{length(x)}() # TODO fix type instability
+  cfg = ForwardDiff.GradientConfig(k.tag, x, chunk_x)
   return cfg
 end
 
-function return_cache(k::ConfigMap{typeof(ForwardDiff.jacobian)},x)
-  cfg = ForwardDiff.JacobianConfig(k.tag,x,ForwardDiff.Chunk{length(x)}())
+function return_cache(k::ConfigMap{typeof(ForwardDiff.jacobian)}, x)
+  chunk_x = ForwardDiff.Chunk{length(x)}() # TODO fix type instability
+  cfg = ForwardDiff.JacobianConfig(k.tag, x, chunk_x)
   return cfg
 end
 

src/Arrays/KeyToValMaps.jl

@@ -13,12 +13,8 @@ function return_cache(m::KeyToValMap,key)
   d = Dict{K,V}()
 end
 
-function evaluate!(cache,m::KeyToValMap,key)
-  if haskey(cache,key)
-    return get(cache,key,nothing)
-  else
-    val = m.key_to_val(key)
-    cache[key] = val
-    return val
-  end
+function evaluate!(cache::Dict{K,V},m::KeyToValMap,key) where {K,V}
+  get!(cache, key) do
+    m.key_to_val(key)
+  end::V
 end

src/CellData/CellStates.jl

@@ -1,23 +1,27 @@
 
 """
+    CellState{T,P<:CellPoint,V<:AbstractArray} <: CellField
+
 This can be used as a CellField as long as one evaluates it
 on the stored CellPoint.
 """
-struct CellState{T,P<:CellPoint} <: CellField
+struct CellState{T,P<:CellPoint,V<:AbstractArray} <: CellField
   points::P
-  values::AbstractArray
+  values::V
 
   function CellState{T}(::UndefInitializer,points::CellPoint) where T
     values = _init_values(T,get_data(points))
     P = typeof(points)
-    new{T,P}(points,values)
+    V = typeof(values)
+    new{T,P,V}(points,values)
   end
 
   function CellState(v::Number,points::CellPoint)
     values = _init_values(v,get_data(points))
     T = typeof(v)
     P = typeof(points)
-    new{T,P}(points,values)
+    V = typeof(values)
+    new{T,P,V}(points,values)
   end
 end
 
@@ -50,7 +54,7 @@ function get_data(f::CellState)
 end
 
 get_triangulation(f::CellState) = get_triangulation(f.points)
-DomainStyle(::Type{CellState{T,P}}) where {T,P} = DomainStyle(P)
+DomainStyle(::Type{<:CellState{T,P}}) where {T,P} = DomainStyle(P)
 
 _get_cell_points(a::CellState) = a.points
 

src/FESpaces/CLagrangianFESpaces.jl

@@ -86,7 +86,7 @@ function CLagrangianFESpace(
 end
 
 # Allowed configurations of CLagrangianFESpace constructors
-# We need: 
+# We need:
 #  - H1 conformity
 #  - Both the geometric and FE reffes have to be the same (and linear - this can be relaxed in the future)
 #  - nvertices == nnodes (no periodicity, etc...)
@@ -277,13 +277,13 @@ end
 function _generate_cell_reffe_clagrangian(z,grid)
   ctype_to_reffe = get_reffes(grid)
   cell_to_ctype = get_cell_type(grid)
-  function fun(reffe)
+  function fun(@nospecialize reffe)
     p = get_polytope(reffe)
     orders = get_orders(reffe)
     LagrangianRefFE(typeof(z),p,orders)
   end
-  # using map instead of lazy_map seems to kill the compiler.
-  ctype_to_reffe = collect(lazy_map(fun,ctype_to_reffe))
+
+  ctype_to_reffe = collect(map(fun,ctype_to_reffe))
   cell_to_reffe = expand_cell_data(ctype_to_reffe,cell_to_ctype)
 end
 

src/FESpaces/DiscreteModelWithFEMaps.jl

@@ -33,12 +33,10 @@ function GridWithFEMap(model,orders::AbstractArray; kwargs...)
   ct = get_cell_type(model)
   ps = lazy_map(Reindex(_ps), ct)
 
-  f(a,b) = LagrangianRefFE(T,a,b)
-  reffes = lazy_map(f,ps,os)
+  reffes = map( (a,b) -> LagrangianRefFE(T,a,b), ps,os)
   Vₕ = FESpace(model,reffes;conformity=:H1,kwargs...)
 
-  fs(a,b) = LagrangianRefFE(Ts,a,b)
-  s_reffes = lazy_map(f,ps,os)
+  s_reffes  = map( (a,b) -> LagrangianRefFE(Ts,a,b), ps,os)
   Vₕ_scal = FESpace(model,s_reffes;conformity=:H1)
 
   grid = get_grid(model)

src/FESpaces/Pullbacks.jl

@@ -159,7 +159,8 @@ function return_cache(k::NormalSignMap,reffe,facet_own_dofs,cell)
   Dc = num_cell_dims(model)
   topo = get_grid_topology(model)
 
-  cell_facets = get_faces(topo, Dc, Dc-1)
+  # Type anotation needed when get_grid_topology isn't inferrable on model (eg DiscreteModelPortion).
+  cell_facets = get_faces(topo, Dc, Dc-1)::Table{Int32, Vector{Int32}, Vector{Int32}}
   cell_facets_cache = array_cache(cell_facets)
 
   return cell_facets, cell_facets_cache, CachedVector(Float64)

src/Fields/AutoDiff.jl

@@ -1,69 +1,48 @@
 # Number differentiation
 
-function gradient(f::Number)
+function gradient(::V) where V<:Number
   function grad_f(x::Point)
-    zero(return_type(outer,x,f))
+    zero(gradient_type(V,x))
   end
 end
 
-function ∇∇(f::Number)
+function ∇∇(::V) where V<:Number
   function hess_f(x::Point)
-    g = gradient(f)(x)
-    gradient(g)(x)
+    zero(gradient_type(V,x,Val(2)))
   end
 end
 
 # Automatic differentiation of functions
 
-function gradient(f::Function)
-  function _gradient(x)
-    gradient(f,x)
-  end
-end
+# gradient(f) = x -> gradient(f,x)
+gradient(f::Function) = Base.Fix1(gradient,f)
+symmetric_gradient(f::Function) = Base.Fix1(symmetric_gradient,f)
+divergence(f::Function) = Base.Fix1(divergence,f)
+curl(f::Function) = Base.Fix1(curl,f)
+laplacian(f::Function) = Base.Fix1(laplacian,f)
 
-function symmetric_gradient(f::Function)
-  function _symmetric_gradient(x)
-    symmetric_gradient(f,x)
-  end
+function gradient(f::Function, x::Point)
+  fx = return_value(f, x)
+  gradient(f, x, fx)
 end
 
-function divergence(f::Function)
-  function _divergence(x)
-    divergence(f,x)
-  end
+function gradient(f::Function, x::Point{A}, fx::Number) where {A}
+  a = ForwardDiff.gradient(f∘Point, get_array(x))
+  VectorValue{A}(a)
 end
 
-function curl(f::Function)
-  function _curl(x)
-    curl(f,x)
-  end
-end
-
-function laplacian(f::Function)
-  function _laplacian(x)
-    laplacian(f,x)
-  end
-end
-
-# Specialization when x::Point
-
-function gradient(f::Function,x::Point)
-  gradient(f,x,return_value(f,x))
-end
-
-function gradient(f::Function,x::Point,fx::Number)
-  VectorValue(ForwardDiff.gradient(f∘Point,get_array(x)))
-end
-
-function gradient(f::Function,x::Point,fx::VectorValue)
-  TensorValue(transpose(ForwardDiff.jacobian(get_array∘f∘Point,get_array(x))))
+function gradient(f::Function, x::Point{A}, fx::VectorValue{B,T}) where {A,B,T}
+  a = transpose(ForwardDiff.jacobian(get_array∘f∘Point, get_array(x)))
+  # This type assert is necessary because of type inferrability issue in the forwardDiff call, for composed gradients
+  TensorValue{A,B}(a)::TensorValue{A,B,T,A*B}
 end
 
 # Implementation for all second order tensor values
 # Does not exploit possible symmetries
-function gradient(f::Function,x::Point{A},fx::S) where S<:MultiValue{Tuple{B,C}} where {A,B,C}
-  a = transpose(ForwardDiff.jacobian(get_array∘f∘Point,get_array(x)))
-  ThirdOrderTensorValue(reshape(a, (A,B,C)))
+function gradient(f::Function, x::Point{A}, fx::S) where S<:MultiValue{Tuple{B,C}} where {A,B,C}
+  a = transpose(ForwardDiff.jacobian(get_array∘f∘Point, get_array(x)))
+  a = reshape(a, (A,B,C))
+  ThirdOrderTensorValue{A,B,C}(a)
 end
 
 function gradient(f::Function,x::Point,fx::MultiValue)
@@ -108,7 +87,7 @@ function divergence(f::Function,x::Point,fx::TensorValue{3,3})
     a[1,1]+a[2,2]+a[3,3],
     a[4,1]+a[5,2]+a[6,3],
     a[7,1]+a[8,2]+a[9,3],
-   )
+  )
 end
 
 function divergence(f::Function,x::Point{2},fx::AbstractSymTensorValue{2})
@@ -127,7 +106,7 @@ function divergence(f::Function,x::Point{3},fx::AbstractSymTensorValue{3})
     a[1,1]+a[2,2]+a[3,3],
     a[2,1]+a[4,2]+a[5,3],
     a[3,1]+a[5,2]+a[6,3],
-   )
+  )
 end
 
 function divergence(f::Function,x::Point,fx::MultiValue)
@@ -149,7 +128,7 @@ end
 function laplacian(f::Function,x::Point{A},fx::VectorValue{B,T}) where {A,B,T}
   a = MMatrix{A*B,A,T}(undef)
   ForwardDiff.jacobian!(a, y->ForwardDiff.jacobian(get_array∘f∘Point,y), get_array(x))
-  VectorValue{B,T}( ntuple(k -> sum(i-> a[(i-1)*B+k,i], 1:A),B) )
+  VectorValue{B,T}( ntuple(k -> sum(i-> a[(i-1)*B+k,i], 1:A), B) )
 end
 
 function laplacian(f::Function,x::Point{A},fx::S) where S<:MultiValue{Tuple{B,C},T} where {A,B,C,T}

src/Fields/FieldArrays.jl

@@ -204,7 +204,11 @@ end
 
 function testvalue(::Type{LinearCombinationField{V,F}}) where {V,F}
   fields = testvalue(F)
-  values = zeros(eltype(V), length(fields), 0)
+  # This method breaks the contract of testvalue for V not subtype of Array
+  #
+  # Ideally, we would extend the API of testvalue with a kwarg that enables
+  # selecting the size of the array when the argument is an array.
+  values = zeros(eltype(V), length(fields), tfill(0,Val(ndims(V)-1))...)
   LinearCombinationField(values,fields,1)
 end
 
@@ -258,12 +262,12 @@ end
 
 function Arrays.testitem(f::LinearCombinationFieldVector{V}) where V
   if !iszero(size(f.values,2))
-    values = f.values
+    values = f.values::V
   elseif V <: Diagonal
     values = Diagonal(zeros(eltype(V), 1))
   else
-    values = zeros(eltype(f.values),size(f.values,1),1)
-  end::V
+    values = zeros(eltype(f.values),size(f.values,1),1)::V
+  end
   LinearCombinationField(values,f.fields,1)
 end
 
@@ -856,3 +860,4 @@ function evaluate!(c,f::ConstantFieldArray{T},x::AbstractArray{<:Point}) where T
   end
   return r
 end
+

src/Fields/Fields.jl

@@ -14,7 +14,7 @@ import Gridap.Arrays: testitem
 
 using Gridap.Helpers
 using Gridap.Helpers: @abstractmethod, @notimplemented
-using Gridap.Helpers: @notimplementedif, @unreachable, @check
+using Gridap.Helpers: @notimplementedif, @unreachable, @check, @check_inferred
 using Gridap.Helpers: tfill, first_and_tail
 
 using Gridap.Algebra: mul!
@@ -31,6 +31,7 @@ using Test
 using StaticArrays
 using LinearAlgebra
 using AutoHashEquals: @auto_hash_equals as @ahe
+using Base: Fix1, Fix2
 
 import LinearAlgebra: det, inv, transpose, tr, cross
 import LinearAlgebra: ⋅, dot