Understanding do notation, type-stability for function with other functions as arguments


#1

If I do

function apply_fn2a(fns::Tuple,a)
    Tuple(fn(a) for fn in fns)
end

using StatsBase
const x = rand(1:5,1000)
@code_warntype apply_fn2a((countmap, sum, mean), x)

then the output type is simply a generic Tuple. However in the examples below Julia can correctly infer a specific type of tuple as output. Is there a way to make the code above be able to infer the output tuple type?

For example

function apply_fn2a(fn::Function,a)
    fn(a)
end

Initially I thought about how type-stability work with apply_fn2a? The output type of apply_fn2a will depend on fn. But given the do notation is a language feature and from my testing, I can see that Julia can infer the type of output apply_fn2a if fn is output-type-stable (if that’s the right term), which is very nice and clever. E.g. @code_warntype apply_fn2a(sum, rand(1000)) and

@code_warntype apply_fn2a(rand(1000)) do x
    x[end] - x[1]
end

show that Julia can infer the types.

But what if I want to write a function that lets a user supply multiple functions and each function has different output. I found that I can just use a tuple of functions and Julia will still correctly infer the types

function apply_fn2a(fn::Tuple{Function, Function},a)
    (fn[1](a), fn[2](a))
end

const x = rand(1:5, 100)
using StatsBase
@code_warntype apply_fn2a((countmap, sum), x)

#2

Note that Tuple{...} constructs a tuple type, while tuple and (...) construct a tuple. You probably want something like

function apply_fn2a(fns::Tuple,a)
    tuple((fn(a) for fn in fns)...)
end

which is inferred properly.