Generate sub-tuples in an efficient, type-stable way

General Usage Performance
1

Given a NTuple{N, T}, I’d like to generate the collection of all sub-tuples that can be obtained by removing one element from the original tuple. This collection of sub-tuples can be represented as an NTuple{N, NTuple{N-1, T}}, which I’d like to be the inferred return-type of the function.

Do you know of any way to do this in an efficient, type-stable way? The “best” I could come up with is this generated function, but I wondered whether there existed a simpler way to achieve this…

@generated function subtuples(t::NTuple{N, T}) where {N, T}
    expr = Expr(:tuple)
    expr.args = map(1:N) do i
        Expr(:tuple, (:(t[$j]) for j in 1:N if i != j)...)
    end
    expr
end

julia> t = (1,2,3)
(1, 2, 3)

julia> typeof(subtuples(t))
Tuple{Tuple{Int64, Int64}, Tuple{Int64, Int64}, Tuple{Int64, Int64}}

julia> using Test; @inferred subtuples(t)
((2, 3), (1, 3), (1, 2))

julia> using BenchmarkTools; @btime subtuples($t)
  1.540 ns (0 allocations: 0 bytes)
((2, 3), (1, 3), (1, 2))

PS: in my use-cases, N is typically 2 or 3

2 
julia> function subtup(tup::NTuple{N, T}) where {N,T}
           tuple((ntuple(N-1) do idx
               tup[idx + (idx >= i)]
           end for i in 1:N)...)
       end
subtup (generic function with 1 method)

julia> @inferred subtup(t)
((2, 3), (1, 3), (1, 2))

julia> @btime subtup($t)
  1.800 ns (0 allocations: 0 bytes)
((2, 3), (1, 3), (1, 2))

julia> @code_warntype subtup(t)
MethodInstance for subtup(::Tuple{Int64, Int64, Int64})
  from subtup(tup::Tuple{Vararg{T, N}}) where {N, T} @ Main REPL[17]:1
Static Parameters
  N = 3
  T = Int64
Arguments
  #self#::Core.Const(subtup)
  tup::Tuple{Int64, Int64, Int64}
Locals
  #11::var"#11#13"{3, Tuple{Int64, Int64, Int64}}
Body::Tuple{Tuple{Int64, Int64}, Tuple{Int64, Int64}, Tuple{Int64, Int64}}
1 ─ %1  = Main.tuple::Core.Const(tuple)
│   %2  = Main.:(var"#11#13")::Core.Const(var"#11#13")
│   %3  = $(Expr(:static_parameter, 1))::Core.Const(3)
│   %4  = Core.typeof(tup)::Core.Const(Tuple{Int64, Int64, Int64})
│   %5  = Core.apply_type(%2, %3, %4)::Core.Const(var"#11#13"{3, Tuple{Int64, Int64, Int64}})
│         (#11 = %new(%5, tup))
│   %7  = #11::var"#11#13"{3, Tuple{Int64, Int64, Int64}}
│   %8  = (1:$(Expr(:static_parameter, 1)))::Core.Const(1:3)
│   %9  = Base.Generator(%7, %8)::Core.PartialStruct(Base.Generator{UnitRange{Int64}, var"#11#13"{3, Tuple{Int64, Int64, Int64}}}, Any[var"#11#13"{3, Tuple{Int64, Int64, Int64}}, Core.Const(1:3)])
│   %10 = Core._apply_iterate(Base.iterate, %1, %9)::Tuple{Tuple{Int64, Int64}, Tuple{Int64, Int64}, Tuple{Int64, Int64}}
└──       return %10

The trick is using the constant N to let the compiler know that everything will have N-1 as its length :slight_smile:

2 Likes
3

Ah, nice! Thanks!