Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[SparseArrayDOKs] Add setindex_maybe_grow! and macro @maybe_grow #1434

Merged
merged 10 commits into from
May 14, 2024
Merged

[SparseArrayDOKs] Add setindex_maybe_grow! and macro @maybe_grow #1434

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
e74f478
Add maybe_grow features
emstoudenmire May 13, 2024
320ac28
Minor edits, add check in @maybe_grow macro
emstoudenmire May 13, 2024
4af4cc2
Improve syntax check for @maybe_grow
emstoudenmire May 13, 2024
41ef815
Improve readability
emstoudenmire May 13, 2024
625fa27
Use MacroTools in @maybe_grow
emstoudenmire May 13, 2024
0520e67
Move expression check to separate function
emstoudenmire May 14, 2024
7da8f09
Use dispatch to check setindex! expression
emstoudenmire May 14, 2024
93a65c6
Make dims field a Ref and remove mutable keyword
emstoudenmire May 14, 2024
c952653
Add and test constructor
emstoudenmire May 14, 2024
aba62fd
Update NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl
emstoudenmire May 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 10 additions & 9 deletions NDTensors/src/lib/SparseArrayDOKs/src/sparsearraydok.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -110,22 +110,23 @@ function Base.resize!(a::SparseArrayDOK{<:Any,N}, new_size::NTuple{N,Integer}) w
return a
end

function setindex_maybe_grow!(
a::SparseArrayDOK{<:Any,N}, value, i1::Int, I::Int...
) where {N}
index = (i1, I...)
if any(index .> size(a))
resize!(a, max.(index, size(a)))
function setindex_maybe_grow!(a::SparseArrayDOK{<:Any,N}, value, I::Vararg{Int,N}) where {N}
if any(I .> size(a))
resize!(a, max.(I, size(a)))
end
a[index...] = value
a[I...] = value
return a
end

macro maybe_grow(ex)
if ex.head != :(=)
emstoudenmire marked this conversation as resolved.
Show resolved Hide resolved
error("@maybe_grow must be used with setindex! syntax (@maybe_grow a[inds] = value)")
end
arr_name = esc(ex.args[1].args[1])
index = esc(ex.args[1].args[2:end])
value = esc(ex.args[2])
quote
SparseArrayDOKs.setindex_maybe_grow!($arr_name, $value, $index...)
e = quote
setindex_maybe_grow!($arr_name, $value, $index...)
end
return e
emstoudenmire marked this conversation as resolved.
Show resolved Hide resolved
end
9 changes: 2 additions & 7 deletions NDTensors/src/lib/SparseArrayDOKs/test/runtests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@

using Test: @test, @testset, @test_broken
using NDTensors.SparseArrayDOKs:
SparseArrayDOK, SparseMatrixDOK, @maybe_grow, setindex_maybe_grow!
SparseArrayDOKs, SparseArrayDOK, SparseMatrixDOK, @maybe_grow
using NDTensors.SparseArrayInterface: storage_indices, nstored
using SparseArrays: SparseMatrixCSC, nnz
@testset "SparseArrayDOK (eltype=$elt)" for elt in
Expand Down Expand Up @@ -95,20 +95,16 @@ using SparseArrays: SparseMatrixCSC, nnz
end
end
end

@testset "Maybe Grow Feature" begin
a = SparseArrayDOK{elt,2}((0, 0))

setindex_maybe_grow!(a, 230, 2, 3)
SparseArrayDOKs.setindex_maybe_grow!(a, 230, 2, 3)
@test size(a) == (2, 3)
@test a[2, 3] == 230

# Test @maybe_grow macro
@maybe_grow a[5, 5] = 550
@test size(a) == (5, 5)
@test a[2, 3] == 230
@test a[5, 5] == 550

# Test that size remains same
# if we set at an index smaller than
# the maximum size:
Expand All @@ -117,7 +113,6 @@ using SparseArrays: SparseMatrixCSC, nnz
@test a[2, 3] == 230
@test a[5, 5] == 550
@test a[3, 4] == 340

# Test vector case
v = SparseArrayDOK{elt,1}((0,))
@maybe_grow v[5] = 50
Expand Down
Loading