From 45113c3b41dcd3bb9c322415c12bb8db57325d14 Mon Sep 17 00:00:00 2001
From: Jishnu Bhattacharya <jishnub.github@gmail.com>
Date: Sun, 28 Apr 2024 08:07:20 +0530
Subject: [PATCH] Function to query the non-zero index of a OneElement (#345)

* Function to query the non-zero index of a OneElement

* Return a CartesianIndex
---
 src/oneelement.jl | 48 +++++++++++++++++++++++++++++++++++++++++++++++
 test/runtests.jl  |  3 +++
 2 files changed, 51 insertions(+)

diff --git a/src/oneelement.jl b/src/oneelement.jl
index 28bc2c4b..b731071f 100644
--- a/src/oneelement.jl
+++ b/src/oneelement.jl
@@ -47,6 +47,54 @@ Base.@propagate_inbounds function Base.getindex(A::OneElement{T,N}, kj::Vararg{I
     ifelse(kj == A.ind, A.val, zero(T))
 end
 
+"""
+    nzind(A::OneElement{T,N}) -> CartesianIndex{N}
+
+Return the index where `A` contains a non-zero value.
+
+!!! note
+    The indices are not guaranteed to lie within the valid index bounds for `A`,
+    and if `FillArrays.nzind(A) ∉ CartesianIndices(A)` then `all(iszero, A)`.
+    On the other hand, if `FillArrays.nzind(A) in CartesianIndices(A)` then
+    `A[FillArrays.nzind(A)] == FillArrays.getindex_value(A)`
+
+# Examples
+```jldoctest
+julia> A = OneElement(2, (1,2), (2,2))
+2×2 OneElement{Int64, 2, Tuple{Int64, Int64}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}:
+ ⋅  2
+ ⋅  ⋅
+
+julia> FillArrays.nzind(A)
+CartesianIndex(1, 2)
+
+julia> A[FillArrays.nzind(A)]
+2
+```
+"""
+nzind(f::OneElement) = CartesianIndex(f.ind)
+
+"""
+    getindex_value(A::OneElement)
+
+Return the only non-zero value stored in `A`.
+
+!!! note
+    If the index at which the value is stored doesn't lie within the valid indices of `A`, then
+    this returns `zero(eltype(A))`.
+
+# Examples
+```jldoctest
+julia> A = OneElement(2, 3)
+3-element OneElement{Int64, 1, Tuple{Int64}, Tuple{Base.OneTo{Int64}}}:
+ ⋅
+ 1
+ ⋅
+
+julia> FillArrays.getindex_value(A)
+1
+```
+"""
 getindex_value(A::OneElement) = all(in.(A.ind, axes(A))) ? A.val : zero(eltype(A))
 
 Base.AbstractArray{T,N}(A::OneElement{<:Any,N}) where {T,N} = OneElement(T(A.val), A.ind, A.axes)
diff --git a/test/runtests.jl b/test/runtests.jl
index e7a5eda6..2fd3c5ca 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -2114,11 +2114,13 @@ end
 
 @testset "OneElement" begin
     A = OneElement(2, (), ())
+    @test FillArrays.nzind(A) == CartesianIndex()
     @test A == Fill(2, ())
     @test A[] === 2
 
     e₁ = OneElement(2, 5)
     @test e₁ == [0,1,0,0,0]
+    @test FillArrays.nzind(e₁) == CartesianIndex(2)
     @test_throws BoundsError e₁[6]
 
     f₁ = AbstractArray{Float64}(e₁)
@@ -2138,6 +2140,7 @@ end
 
     V = OneElement(2, (2,3), (3,4))
     @test V == [0 0 0 0; 0 0 2 0; 0 0 0 0]
+    @test FillArrays.nzind(V) == CartesianIndex(2,3)
 
     Vf = AbstractArray{Float64}(V)
     @test Vf isa OneElement{Float64,2}