Adjust code so that push! etc. don't trigger scalar indexing. Add tests for CuArrays.

Author: darsnack

Project.toml

@@ -11,6 +11,7 @@ julia = "1"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 
 [targets]
-test = ["Test"]
+test = ["CUDA", "Test"]

src/CircularArrayBuffers.jl

@@ -5,7 +5,7 @@ using Adapt
 export CircularArrayBuffer, CircularVectorBuffer, capacity, isfull
 
 """
-    CircularArrayBuffer{T}(sz::Integer...) -> CircularArrayBuffer{T, N}
+    CircularArrayBuffer{T}(sz::Integer...) -> CircularArrayBuffer{T, N, Array{T, N}}
 
 `CircularArrayBuffer` uses a `N`-dimension `Array` of size `sz` to serve as a buffer for
 `N-1`-dimension `Array`s of the same size.
@@ -33,12 +33,22 @@ end
 Adapt.adapt_structure(to, cb::CircularArrayBuffer) =
     CircularArrayBuffer(adapt(to, cb.buffer), cb.first, cb.nframes, cb.step_size)
 
+function Base.show(io::IO, ::MIME"text/plain", cb::CircularArrayBuffer{T}) where T
+    print(io, ndims(cb) == 1 ? "CircularVectorBuffer(" : "CircularArrayBuffer(")
+    Base.showarg(io, cb.buffer, false)
+    print(io, ") with eltype $T:\n")
+    Base.print_array(io, adapt(Array, cb))
+    return nothing
+end
+
 Base.IndexStyle(::CircularArrayBuffer) = IndexLinear()
 
 Base.size(cb::CircularArrayBuffer{T,N}, i::Integer) where {T,N} = i == N ? cb.nframes : size(cb.buffer, i)
 Base.size(cb::CircularArrayBuffer{T,N}) where {T,N} = ntuple(i -> size(cb, i), N)
 Base.getindex(cb::CircularArrayBuffer{T,N}, i::Int) where {T,N} = getindex(cb.buffer, _buffer_index(cb, i))
+Base.getindex(cb::CircularArrayBuffer{T,N}, I...) where {T,N} = getindex(cb.buffer, Base.front(I)..., _buffer_frame(cb, Base.last(I)))
 Base.setindex!(cb::CircularArrayBuffer{T,N}, v, i::Int) where {T,N} = setindex!(cb.buffer, v, _buffer_index(cb, i))
+Base.setindex!(cb::CircularArrayBuffer{T,N}, v, I...) where {T,N} = setindex!(cb.buffer, v, Base.front(I)..., _buffer_frame(cb, Base.last(I)))
 
 capacity(cb::CircularArrayBuffer{T,N}) where {T,N} = size(cb.buffer, N)
 isfull(cb::CircularArrayBuffer) = cb.nframes == capacity(cb)
@@ -52,6 +62,7 @@ Base.isempty(cb::CircularArrayBuffer) = cb.nframes == 0
         ind
     end
 end
+@inline _buffer_index(cb::CircularArrayBuffer, I::AbstractVector{<:Integer}) = map(Base.Fix1(_buffer_index, cb), I)
 
 @inline function _buffer_frame(cb::CircularArrayBuffer, i::Int)
     n = capacity(cb)
@@ -63,7 +74,7 @@ end
     end
 end
 
-_buffer_frame(cb::CircularArrayBuffer, I::Vector{Int}) = map(i -> _buffer_frame(cb, i), I)
+_buffer_frame(cb::CircularArrayBuffer, I::AbstractVector{<:Integer}) = map(i -> _buffer_frame(cb, i), I)
 
 function Base.empty!(cb::CircularArrayBuffer)
     cb.nframes = 0
@@ -77,13 +88,14 @@ function Base.push!(cb::CircularArrayBuffer{T,N}, data) where {T,N}
         cb.nframes += 1
     end
     if N == 1
+        i = _buffer_frame(cb, cb.nframes)
         if ndims(data) == 0
-            cb[cb.nframes] = data[]
+            cb.buffer[i:i] .= data[]
         else
-            cb[cb.nframes] = data
+            cb.buffer[i:i] .= data
         end
     else
-        cb[ntuple(_ -> (:), N - 1)..., cb.nframes] .= data
+        cb.buffer[ntuple(_ -> (:), N - 1)..., _buffer_frame(cb, cb.nframes)] .= data
     end
     cb
 end

test/runtests.jl

@@ -1,7 +1,10 @@
 using CircularArrayBuffers
 using Test
+using Adapt
+using CUDA
+CUDA.allowscalar(false)
 
-@testset "CircularArrayBuffers.jl" begin
+@testset "CircularArrayBuffers (Array)" begin
     A = ones(2, 2)
     C = ones(Float32, 2, 2)
 
@@ -165,3 +168,176 @@ using Test
         ]
     end
 end
+
+@testset "CircularArrayBuffers (CuArray)" begin
+    A = CUDA.ones(2, 2)
+    Ac = adapt(Array, A)
+    C = CUDA.ones(Float32, 2, 2)
+
+    @testset "Adapt" begin
+        X = CircularArrayBuffer(rand(2, 3))
+        Xc = adapt(CuArray, X)
+        @test Xc isa CircularArrayBuffer{Float64, 2, <:CuArray}
+        @test adapt(Array, Xc) == X
+    end
+
+    # https://github.com/JuliaReinforcementLearning/ReinforcementLearning.jl/issues/551
+    @testset "1D with 0d data" begin
+        b = adapt(CuArray, CircularArrayBuffer{Int}(3))
+        CUDA.@allowscalar push!(b, CUDA.zeros(Int, ()))
+        @test length(b) == 1
+        @test CUDA.@allowscalar b[1] == 0
+    end
+
+    @testset "1D Int" begin
+        b = adapt(CuArray, CircularArrayBuffer{Int}(3))
+
+        @test eltype(b) == Int
+        @test capacity(b) == 3
+        @test isfull(b) == false
+        @test isempty(b) == true
+        @test length(b) == 0
+        @test size(b) == (0,)
+        # element must has the exact same length with the element of buffer
+        @test_throws Exception push!(b, [1, 2])
+
+        for x in 1:3
+            push!(b, x)
+        end
+
+        @test capacity(b) == 3
+        @test isfull(b) == true
+        @test length(b) == 3
+        @test size(b) == (3,)
+        # scalar indexing is not allowed
+        @test_throws ErrorException b[1]
+        @test_throws ErrorException b[end]
+        @test CUDA.@allowscalar b[1:end] == cu([1, 2, 3])
+
+        for x in 4:5
+            push!(b, x)
+        end
+
+        @test capacity(b) == 3
+        @test length(b) == 3
+        @test size(b) == (3,)
+        @test CUDA.@allowscalar b[1:end] == [3, 4, 5]
+
+        empty!(b)
+        @test isfull(b) == false
+        @test isempty(b) == true
+        @test length(b) == 0
+        @test size(b) == (0,)
+
+        push!(b, 6)
+        @test isfull(b) == false
+        @test isempty(b) == false
+        @test length(b) == 1
+        @test size(b) == (1,)
+        @test CUDA.@allowscalar b[1] == 6
+
+        push!(b, 7)
+        push!(b, 8)
+        @test isfull(b) == true
+        @test isempty(b) == false
+        @test length(b) == 3
+        @test size(b) == (3,)
+        @test CUDA.@allowscalar b[1:3] == cu([6, 7, 8])
+
+        push!(b, 9)
+        @test isfull(b) == true
+        @test isempty(b) == false
+        @test length(b) == 3
+        @test size(b) == (3,)
+        @test CUDA.@allowscalar b[1:3] == cu([7, 8, 9])
+
+        x = CUDA.@allowscalar pop!(b)
+        @test x == 9
+        @test length(b) == 2
+        @test CUDA.@allowscalar b[1:2] == cu([7, 8])
+
+        x = CUDA.@allowscalar popfirst!(b)
+        @test x == 7
+        @test length(b) == 1
+        @test CUDA.@allowscalar b[1] == 8
+
+        x = CUDA.@allowscalar pop!(b)
+        @test x == 8
+        @test length(b) == 0
+
+        @test_throws ArgumentError pop!(b)
+        @test_throws ArgumentError popfirst!(b)
+    end
+
+    @testset "2D Float64" begin
+        b = adapt(CuArray, CircularArrayBuffer{Float64}(2, 2, 3))
+
+        @test eltype(b) == Float64
+        @test capacity(b) == 3
+        @test isfull(b) == false
+        @test length(b) == 0
+        @test size(b) == (2, 2, 0)
+
+        for x in 1:3
+            push!(b, x * A)
+        end
+
+        @test capacity(b) == 3
+        @test isfull(b) == true
+        @test length(b) == 2 * 2 * 3
+        @test size(b) == (2, 2, 3)
+        for i in 1:3
+            @test b[:, :, i] == i * A
+        end
+        @test b[:, :, end] == 3 * A
+
+        for x in 4:5
+            push!(b, x * CUDA.ones(Float64, 2, 2))
+        end
+
+        @test capacity(b) == 3
+        @test length(b) == 2 * 2 * 3
+        @test size(b) == (2, 2, 3)
+        @test b[:, :, 1] == 3 * A
+        @test b[:, :, end] == 5 * A
+
+        # doing b == ... triggers scalar indexing
+        @test CUDA.@allowscalar b == cu(reshape([c for x in 3:5 for c in x * Ac], 2, 2, 3))
+
+        push!(b, 6 * CUDA.ones(Float32, 2, 2))
+        push!(b, 7 * CUDA.ones(Int, 2, 2))
+        @test CUDA.@allowscalar b == cu(reshape([c for x in 5:7 for c in x * Ac], 2, 2, 3))
+
+        x = pop!(b)
+        @test x == 7 * CUDA.ones(Float64, 2, 2)
+        @test CUDA.@allowscalar b == cu(reshape([c for x in 5:6 for c in x * Ac], 2, 2, 2))
+    end
+
+    @testset "append!" begin
+        b = adapt(CuArray, CircularArrayBuffer{Int}(2, 3))
+        append!(b, CUDA.zeros(2))
+        append!(b, 1:4)
+        @test CUDA.@allowscalar b == cu([
+            0 1 3
+            0 2 4
+        ])
+
+
+        b = adapt(CuArray, CircularArrayBuffer{Int}(2, 3))
+        for i in 1:5
+            push!(b, CUDA.fill(i, 2))
+        end
+        empty!(b)
+        append!(b, 1:4)
+        @test CUDA.@allowscalar b == cu([
+            1 3
+            2 4
+        ])
+
+        append!(b, 5:8)
+        @test CUDA.@allowscalar b == cu([
+            3 5 7
+            4 6 8
+        ])
+    end
+end