A quick search of this forum yields a number of questions about variable length Tuples. Maybe you have a recursive function that returns a tuple based on depth or something like that.
I have bumped into this issue several times where I know that a tuple will have some reasonable number of elements, but not precisely how many. It is my intuition that I should be able to tell the compiler “Please allocate N number of slots on the stack, even though I might not use them all, and I’ll only getindex on the one’s I’ve populated.”
I wrote a proof of concept that I think achieves this behavior (see below). However, I don’t understand benchmarks well enough to know if my limited tests are just being smartly compiled away or if I am actually achieving allocation free code.
I’d appreciate your feedback on the implementation, bad edge case tests, etc. I’m also wondering, if this does actually work, if there is a package this might fit within? I don’t really want to maintain a package myself, but I’d be happy to contribute to someone else’s thing if they’re up for it.
module UpToNTuples
using Base: issingletontype
struct UpToNTuple{N, T}
data::NTuple{N, T}
len::UInt8
function UpToNTuple(::Val{N}, ::Type{T}, data::NTuple{N,T}, len::UInt8) where {N, T}
return new{N, T}(data, len)
end
end
## A bunch of constructors to add some flexibility
# handle empty case
UpToNTuple(; _...) = throw(ArgumentError("Must provide eltype for an empty `UpToNTuple` like `UpToNTuple{N,T}` or `UpToNTuple{T}`"))
function UpToNTuple{X}(; kwargs...) where X
X isa Integer && throw(ArgumentError("Must provide eltype for an empty `UpToNTuple` like `UpToNTuple{N,T}` or `UpToNTuple{T}`"))
return _UpToNTuple(Val(10), X, (); kwargs...)
end
UpToNTuple{N,T}(; kwargs...) where {N,T} = UpToNTuple{N, T}(Val(N), T, (); kwargs...)
# handle cases with elements
UpToNTuple(data::T; ks...) where {T} = _UpToNTuple(Val(10), eltype(T), data; ks...)
function UpToNTuple{X}(data::T; ks...) where {X, T}
Eltype = X isa Integer ? eltype : X
N = X isa Integer ? X : 10
_UpToNTuple(Val(N), Eltype, data; ks...)
end
UpToNTuple{N,T}(data; ks...) where {N,T} = _UpToNTuple(Val(N), T, data; ks...)
# special case handling for getting default value for empty tuples
function _UpToNTuple(::Val{N}, ::Type{E}, data::T; kwargs...) where {T, E, N}
len = UInt8(length(data))
N > 10 && throw(ArgumentError("Can't construct an UpToNTuple with more than 10 slots"))
len > N && throw(ArgumentError("Too many arguments for UpToNTuple{$N}"))
default_value = if haskey(kwargs, :default_value)
kwargs[:default_value]
elseif iszero(len)
default_value = empty_value(E)
else
default_value = empty_value(first(data))
end::E
d = ntuple(i -> (i <= len ? data[i] : default_value)::E, N)
return UpToNTuple(Val(N), E, d, len)
end
#### constructor helpers
"""
empty_value(x::T)
Return a default empty value for for type `T`.
"""
empty_value(x::T) where T = isbitstype(T) ? x : empty_value(T)
empty_value(::Type{Symbol}) = Symbol("")
empty_value(::Type{T}) where {T<:AbstractString} = T("")
empty_value(::Type{T}) where {T<:Number} = zero(T)
empty_value(::Type{T}) where T = issingletontype(T) ? T() : error("No default empty value for type $T")
# Basic Tuple API
Base.length(u::UpToNTuple) = u.len
Base.eachindex(u::UpToNTuple) = Base.OneTo(u.len)
function Base.getindex(u::UpToNTuple{<:Any, T}, i::Int) where T
i > length(u) && throw(BoundsError(u, i))
return u.data[i]::T
end
function Base.iterate(u::UpToNTuple, state=nothing)
isnothing(state) && return (u.data[1], 2)
i = state
i > length(u) && return nothing
return (u.data[i], i + 1)
end
function push!!(n::UpToNTuple{N, T}, x::T2) where {N, T, T2 <: T}
i = length(n) + 1
i > N && throw(ArgumentError("Too many elements for UpToNTuple{$N}"))
data = ntuple(N) do j
(j == i ? x : n.data[j])::T
end
return UpToNTuple(Val(N), T, data, UInt8(i))::UpToNTuple{N, T}
end
function Base.show(io::IO, x::UpToNTuple{N,T}) where {N, T}
print(io, "UpToNTuple{", N, ", ", T, "}(")
for i in 1:length(x)
show(io, x[i])
i < length(x) && print(io, ", ")
end
print(io, ")")
end
end
Benchmarks
using Chairmarks
f(a, b, c) = UpToNTuples.UpToNTuple((a, b, c))[1];
# I don't see any allocations in the output
@code_native f(1,2,3)
function rec(::Val{N}) where N
numbers = UpToNTuples.UpToNTuple{N, Int}(N)
return rec(N-1, numbers)::UpToNTuples.UpToNTuple{N, Int}
end
function rec(n, numbers)
n == 0 && return numbers
numbers = UpToNTuples.push!!(numbers, n)
return rand((true, false)) ? rec(n-1, numbers) : numbers
end
f() = map(_-> rec(Val(10)), 1:1000)
@b f()
# 19.500 μs (3 allocs: 86.000 KiB)
Thanks to @jakobnissen , @Mason , and @gbaraldi for your help on the slack thread.
Edit: I learned that Tuples allocate after 10 elements, so I set an upper limit to 10. I also improved the benchmark to a slightly more realistic case.
Edit2: Change to API – remove splatting of elements in favor of accepting a single iterator in the constructor. This follows the NTuple API more closely.