From 2b1dbdbdef8e9a5cd8df3be7469c07c55c358fbc Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Thu, 5 Sep 2024 16:31:39 +0200
Subject: [PATCH] reflect on MPI.Datatype

---
 src/operators.jl    | 15 +++++++++++++++
 test/test_reduce.jl |  3 ++-
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/src/operators.jl b/src/operators.jl
index 63cfe178d..ddbb20c11 100644
--- a/src/operators.jl
+++ b/src/operators.jl
@@ -91,6 +91,21 @@ function (w::OpWrapper{F,T})(_a::Ptr{Cvoid}, _b::Ptr{Cvoid}, _len::Ptr{Cint}, t:
 end
 
 
+function (w::OpWrapper{F,Any})(_a::Ptr{Cvoid}, _b::Ptr{Cvoid}, _len::Ptr{Cint}, t::Ptr{MPI_Datatype}) where {F}
+    len = unsafe_load(_len)
+    T = to_type(Datatype(unsafe_load(t))) # Ptr might actually point to a Julia object so we could unsafe_pointer_to_objref?
+    @assert isconcretetype(T)
+    function copy(::Type{T}) where T
+        a = Ptr{T}(_a)
+        b = Ptr{T}(_b)
+        for i = 1:len
+            unsafe_store!(b, w.f(unsafe_load(a,i), unsafe_load(b,i)), i)
+        end
+    end
+    copy(T)
+    return nothing
+end
+
 function Op(f, T=Any; iscommutative=false)
     @static if MPI_LIBRARY == "MicrosoftMPI" && Sys.WORD_SIZE == 32
         error("User-defined reduction operators are not supported on 32-bit Windows.\nSee https://github.com/JuliaParallel/MPI.jl/issues/246 for more details.")
diff --git a/test/test_reduce.jl b/test/test_reduce.jl
index 6ea6cd967..1e72f58a0 100644
--- a/test/test_reduce.jl
+++ b/test/test_reduce.jl
@@ -62,7 +62,8 @@ end
 function my_reduce(x, y)
     2x+y-x
 end
-MPI.@Op(my_reduce, Int)
+# MPI.@Op(my_reduce, Int)
+MPI.@Op(my_reduce, Any)
 
 if can_do_closures
     operators = [MPI.SUM, +, my_reduce, (x,y) -> 2x+y-x]