Make: extend getindex: accept non-Int integers and vectors

Author: rfourquet

src/distributions.jl

@@ -13,10 +13,18 @@ struct Make{T,X<:Tuple, XX<:Tuple} <: Distribution{T}
 end
 
 # @inline necessary for one @inferred test on arrays
-@inline Base.getindex(m::Make{T,X}, i::Int) where {T,X} =
+@inline function Base.getindex(m::Make{T,X}, i::Integer) where {T,X}
+    i = Int(i)
     fieldtype(X, i) <: Type ?
         fieldtype(X, i).parameters[1] :
         m.x[i]
+end
+
+@inline Base.getindex(m::Make, idxs::AbstractVector{<:Integer}) =
+    ntuple(i->m[idxs[i]], length(idxs))
+# seems faster than `Tuple(m[i] for i in idxs)`
+
+Base.lastindex(m::Make) = lastindex(m.x)
 
 @generated function Make{T}(X...) where T
     XX = Tuple{(x <: Type ? Nothing : x for x in X)...}

test/runtests.jl

@@ -602,6 +602,19 @@ Base.rand(rng::AbstractRNG,
     @test all(∈(3:6), a)
 end
 
+## Make getindex
+
+@testset "Make getindex" begin
+    m = make(Pair, 1:2, Bool, make(String, 3))
+    @test m isa RandomExtensions.Make3
+    @test m[1] == 1:2
+    @test m[0x2] == Bool
+    @test m[3] isa RandomExtensions.Make2
+    @test m[1:2] == (m[1], m[2])
+    @test m[[2, 3]] == m[2:end] == (m[2], m[3])
+    @test m[big(3)][1:end] == m[3][1:2]
+end
+
 ## @rand
 
 struct Die