make partial output container iterable (zero cost) (#47)

* make partial output container iterable (zero cost) * fix test
tpapp · Nov 9, 2023 · 9ec9c6a · 9ec9c6a · tpapp · Nov 9, 2023
1 parent eb8519a
commit 9ec9c6a
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 3 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "SpectralKit"
 uuid = "5c252ae7-b5b6-46ab-a016-b0e3d78320b7"
 authors = ["Tamas K. Papp <tkpapp@gmail.com>"]
-version = "0.14.0"
+version = "0.14.1"
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"

diff --git a/src/derivatives.jl b/src/derivatives.jl
@@ -284,8 +284,8 @@ end
 
 """
 Container for output of evaluating partial derivatives. Each corresponds to an
-specification in a [`∂Specification`](@ref). Can be indexed with integers, or converted
-to a `Tuple`.
+specification in a [`∂Specification`](@ref). Can be indexed with integers, iterated, or
+converted to a `Tuple`.
 """
 struct ∂Output{N,T}
     values::NTuple{N,T}
@@ -303,6 +303,8 @@ end
 
 @inline Base.getindex(∂output::∂Output, i) = ∂output.values[i]
 
+@inline Base.iterate(∂output::∂Output, i...) = Base.iterate(∂output.values, i...)
+
 function _product(∂specification::∂Specification, sources, indices)
     (; lookups) = ∂specification
     p = map(lookups) do lookup

diff --git a/test/test_derivatives.jl b/test/test_derivatives.jl
@@ -87,3 +87,17 @@ end
         end
     end
 end
+
+@testset "iteration for ∂Output" begin
+    o_src = (1.0, 2.0, 3.0)
+    ∂o = SpectralKit.∂Output(o_src)
+    for (o1, o2) in zip(o_src, ∂o)
+        @test o1 ≡ o2
+    end
+    function f(o)
+        o1, o2, o3 = o
+        o1 + o2 + o3
+    end
+    @test @inferred(f(∂o)) == sum(o_src)
+    @test @ballocated($f($∂o)) == 0
+end