How can I get this to infer types?

I’m having problems getting my code to infer a return type correctly. Here’s a MWE:

using Test
struct TwoTo{N} end
@inline t(g::Base.Generator{<:TwoTo{N}, F}) where {N,F} = ntuple(i->g.f(i+1), Val{N-1}())
f(r) = t(t(i for i in r) for j in r)
Test.@inferred f(TwoTo{3}())
# ERROR: return type Tuple{Tuple{Int64,Int64},Tuple{Int64,Int64}} does not match inferred return type Tuple{Tuple{Any,Any},Tuple{Any,Any}}

If I change the definition of t to:

@inline t(g::Base.Generator{<:TwoTo{N}, F}) where {N,F} = ntuple(g.f, Val{N-1}())

then type inference works (but the output is different, of course.)

I can’t figure out any way to wrap g.f in a function in a way that does not destroy type inference. Any help would be appreciated!

I think I’ve figured this out. If I change the code slightly:

using Test
struct TwoTo{N} end
@inline t(g::Base.Generator{<:TwoTo{N}, F}) where {N,F} = ntuple(i->g.f(i+1), Val{N-1}())
@inline u(g::Base.Generator{<:TwoTo{N}, F}) where {N,F} = ntuple(i->g.f(i+1), Val{N-1}())
f2(r) = t(u(i for i in r) for j in r)
Test.@inferred f2(TwoTo{3}())

then it works. (Note that u is an exact copy of t, only with a different name.)
Apparently, type inference gives up when t calls itself recursively even though the function argument is different. I can see how this can be a reasonable heuristic. Now I just need to figure out a good way to get around it!

2 Likes