From 414043c1b01500e8a6de12afa6f6307bb28e71e9 Mon Sep 17 00:00:00 2001
From: Alexander Plavin <alexander@plav.in>
Date: Sun, 17 Nov 2024 20:09:43 -0500
Subject: [PATCH] support ConstructionBase: constructorof and setproperties

---
 src/StructArrays.jl     |  1 +
 src/constructionbase.jl | 17 +++++++++++++++++
 test/runtests.jl        | 42 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 60 insertions(+)
 create mode 100644 src/constructionbase.jl

diff --git a/src/StructArrays.jl b/src/StructArrays.jl
index a8995f40..632fe03c 100644
--- a/src/StructArrays.jl
+++ b/src/StructArrays.jl
@@ -13,6 +13,7 @@ include("utils.jl")
 include("collect.jl")
 include("sort.jl")
 include("lazy.jl")
+include("constructionbase.jl")
 include("tables.jl")
 
 # Implement refarray and refvalue to deal with pooled arrays and weakrefstrings effectively
diff --git a/src/constructionbase.jl b/src/constructionbase.jl
new file mode 100644
index 00000000..24b4f5b4
--- /dev/null
+++ b/src/constructionbase.jl
@@ -0,0 +1,17 @@
+using ConstructionBase
+
+# (named)tuple eltypes: components/fields are all that's needed for the StructArray() constructor
+ConstructionBase.constructorof(::Type{<:StructArray{<:Union{Tuple, NamedTuple}}}) = StructArray
+
+# other eltypes: need to pass eltype to the constructor in addition to components
+ConstructionBase.constructorof(::Type{<:StructArray{T}}) where {T} = function(comps::CT) where {CT}
+    # the resulting eltype is like T, but potentially with different type parameters, eg Complex{Int} -> Complex{Float64}
+    # probe its constructorof to get the right concrete type
+    ET = Base.promote_op(constructorof(T), map(eltype, fieldtypes(CT))...)
+    StructArray{ET}(comps)
+end
+
+# two methods with the same body, required to avoid ambiguities
+# just redirect setproperties to constructorof
+ConstructionBase.setproperties(x::StructArray, patch::NamedTuple) = constructorof(typeof(x))(setproperties(getproperties(x), patch))
+ConstructionBase.setproperties(x::StructArray, patch::Tuple) = constructorof(typeof(x))(setproperties(getproperties(x), patch))
diff --git a/test/runtests.jl b/test/runtests.jl
index 08282b59..0a57817c 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -5,6 +5,7 @@ using StaticArrays
 using TypedTables: Table
 using DataAPI: refarray, refvalue
 using Adapt: adapt, Adapt
+using ConstructionBase: constructorof, setproperties, getproperties, getfields
 using JLArrays
 using LinearAlgebra
 using Test
@@ -529,6 +530,47 @@ end
     @test_throws ArgumentError StructArray(a=[1, 2], b=[3])
 end
 
+@testset "ConstructionBase" begin
+    # first, check the required invariants
+    @testset for obj in Any[
+        StructArray(([1, 2, 3],)),
+        StructArray((Int[],)),
+        StructArray(([1, 2, 3], 4:6)),
+        StructArray(a=[1, 2, 3]),
+        StructArray(a=[1, 2, 3], b=4:6),
+        StructArray([1+2im, 3+4im, 5+6im]),
+        StructArray(ComplexF64[]),
+    ]
+        # constructorof of getfields returns the same object
+        @test constructorof(typeof(obj))(getfields(obj)...) === obj
+
+        # setproperties with getproperties, or with empty properties, returns the same object
+        @test setproperties(obj, getproperties(obj)) === obj
+        if getproperties(obj) isa Tuple
+            @test setproperties(obj, ()) === obj
+        else
+            @test setproperties(obj, (;)) === obj
+        end
+    end
+
+    # now, check less trivial cases: reconstruction with different property types and names
+    s = StructArray(a=[1, 2, 3])
+    @test constructorof(typeof(s))((a=1:3,)) === StructArray(a=1:3)
+    @test constructorof(typeof(s))((b=1.0:3.0,)) === StructArray(b=1.0:3.0)
+
+    s = StructArray(a=[1, 2, 3], b=4:6)
+    @test setproperties(s, a=10:12)::StructArray === StructArray(a=10:12, b=4:6)
+    @test_throws ArgumentError setproperties(s, ccc=10:12)
+
+    s = StructArray(([1, 2, 3], 4:6))
+    @test setproperties(s, (10:12,))::StructArray === StructArray((10:12, 4:6))
+
+    s = StructArray([1+2im, 3+4im, 5+6im])
+    @test constructorof(typeof(s))((re=10:0.1:10.2, im=[1,2,3]))::StructArray == StructArray{ComplexF64}((10:0.1:10.2, [1,2,3]))
+    @test constructorof(typeof(s))((10:0.1:10.2, [1,2,3]))::StructArray == StructArray{ComplexF64}((10:0.1:10.2, [1,2,3]))
+    @test setproperties(s, re=10:0.1:10.2)::StructArray == StructArray{ComplexF64}((10:0.1:10.2, [2,4,6]))
+end
+
 @testset "complex" begin
     a, b = [1 2; 3 4], [4 5; 6 7]
     t = StructArray{ComplexF64}((a, b))