From f972b4e644c5380bfaa578931b1ee7c4e2234fe6 Mon Sep 17 00:00:00 2001
From: Kipton Barros <kbarros@gmail.com>
Date: Mon, 27 Feb 2023 13:49:39 -0700
Subject: [PATCH] Avoid all use of `deepcopy`

Replace with custom clone functions that recurse datastructures and
perform ordinary `copy` on mutable arrays.
---
 src/ReplicaExchangeMC.jl        |  2 +-
 src/System/Ewald.jl             |  6 ++++++
 src/System/Interactions.jl      | 15 ++++++++++++---
 src/System/System.jl            | 32 +++++++++++++++++++++++++++-----
 test/mc_tests/REMC_H_Ising2D.jl |  2 +-
 test/test_energy_consistency.jl | 23 +++++++++++++----------
 6 files changed, 60 insertions(+), 20 deletions(-)

diff --git a/src/ReplicaExchangeMC.jl b/src/ReplicaExchangeMC.jl
index 3299d6cbe..3967819e9 100755
--- a/src/ReplicaExchangeMC.jl
+++ b/src/ReplicaExchangeMC.jl
@@ -92,7 +92,7 @@ function replica_exchange!(replica::Replica)
     S_curr = running_energy(replica.sampler) / get_temp(replica.sampler)
 
     # Backup current configuration
-    backup_spins = deepcopy(replica.sampler.sys.dipoles)
+    backup_spins = copy(replica.sampler.sys.dipoles)
 
     # Swap trial configuration with partner
     MPI.Sendrecv!(
diff --git a/src/System/Ewald.jl b/src/System/Ewald.jl
index 9748d300c..cf156f90b 100644
--- a/src/System/Ewald.jl
+++ b/src/System/Ewald.jl
@@ -27,6 +27,12 @@ function Ewald(sys::System{N}) where N
     return Ewald(A, ϕ, FA, Fs, Fϕ, plan, ift_plan)
 end
 
+# Clone all mutable state within Ewald. Note that `A`, `FA`, and plans should be
+# immutable data.
+function clone_ewald(ewald::Ewald)
+    (; A, ϕ, FA, Fs, Fϕ, plan, ift_plan) = ewald
+    return Ewald(A, copy(ϕ), FA, copy(Fs), copy(Fϕ), plan, ift_plan)
+end
 
 # Tensor product of 3-vectors
 (⊗)(a::Vec3,b::Vec3) = reshape(kron(a,b), 3, 3)
diff --git a/src/System/Interactions.jl b/src/System/Interactions.jl
index 4d620dada..ec83f2527 100644
--- a/src/System/Interactions.jl
+++ b/src/System/Interactions.jl
@@ -7,6 +7,13 @@ function empty_interactions(nb, N)
     end
 end
 
+# Creates a clone of the lists of exchange interactions, which can be mutably
+# updated.
+function clone_interactions(ints::Interactions)
+    (; aniso, heisen, exchange, biquad) = ints
+    return Interactions(aniso, copy(heisen), copy(exchange), copy(biquad))
+end
+
 function interactions_homog(sys::System{N}) where N
     return sys.interactions_union :: Vector{Interactions}
 end
@@ -33,11 +40,13 @@ function to_inhomogeneous(sys::System{N}) where N
     is_homogeneous(sys) || error("System is already inhomogeneous.")
     ints = interactions_homog(sys)
 
-    ret = deepcopy(sys)
+    ret = clone_system(sys)
     nb = nbasis(ret.crystal)
     ret.interactions_union = Array{Interactions}(undef, ret.latsize..., nb)
-    for cell in all_cells(ret)
-        ret.interactions_union[cell, :] = deepcopy(ints)
+    for b in 1:nbasis(ret.crystal)
+        for cell in all_cells(ret)
+            ret.interactions_union[cell, b] = clone_interactions(ints[b])
+        end
     end
 
     return ret
diff --git a/src/System/System.jl b/src/System/System.jl
index 3dd6e31b9..df76c69a3 100644
--- a/src/System/System.jl
+++ b/src/System/System.jl
@@ -86,11 +86,33 @@ function Base.show(io::IO, ::MIME"text/plain", sys::System{N}) where N
     end
 end
 
+# Per Julia developers, `deepcopy` is memory unsafe, especially in conjunction
+# with C libraries. We were observing very confusing crashes that surfaced in
+# the FFTW library, https://github.com/JuliaLang/julia/issues/48722. To prevent
+# this from happening again, avoid all uses of `deepcopy`, and create our own
+# stack of `clone` functions instead.
+Base.deepcopy(_::System) = error("Use `clone_system` instead of `deepcopy`.")
+
+# Creates a clone of the system where all the mutable internal data is copied.
+# It should be thread-safe to use the original and the copied systems, without
+# any restrictions.
+function clone_system(sys::System{N}) where N
+    (; origin, mode, crystal, latsize, Ns, gs, κs, extfield, interactions_union, ewald, dipoles, coherents, units, rng) = sys
+
+    origin_clone = isnothing(origin) ? nothing : clone_system(origin)
+    ewald_clone  = isnothing(ewald)  ? nothing : clone_ewald(ewald)
+
+    # Dynamically dispatch to the correct `map` function for either homogeneous
+    # (Vector) or inhomogeneous interactions (4D Array)
+    interactions_clone = map(clone_interactions, interactions_union)
+    
+    # Empty buffers are required for thread safety.
+    empty_dipole_buffers = Array{Vec3, 4}[]
+    empty_coherent_buffers = Array{CVec{N}, 4}[]
 
-function clone_spin_state(sys::System{N}) where N
-    System(sys.origin, sys.mode, sys.crystal, sys.latsize, sys.Ns, sys.gs, sys.κs, sys.extfield,
-        sys.interactions_union, sys.ewald, copy(sys.dipoles), copy(sys.coherents),
-        sys.dipole_buffers, sys.coherent_buffers, sys.units, copy(sys.rng))
+    System(origin_clone, mode, crystal, latsize, Ns, copy(gs), copy(κs), copy(extfield),
+           interactions_clone, ewald_clone, copy(dipoles), copy(coherents),
+           empty_dipole_buffers, empty_coherent_buffers, units, copy(rng))
 end
 
 
@@ -362,7 +384,7 @@ function reshape_geometry_aux(sys::System{N}, new_latsize::NTuple{3, Int}, new_c
     # reshapings, `sys.origin` keeps its original meaning. Make a deep copy so
     # that the new system fully owns `origin`, and mutable updates to the
     # previous system won't affect this one.
-    origin = deepcopy(isnothing(sys.origin) ? sys : sys.origin)
+    origin = clone_system(isnothing(sys.origin) ? sys : sys.origin)
 
     # If `new_cell_size == I`, we can effectively restore the unit cell of
     # `origin`, but with `new_latsize`
diff --git a/test/mc_tests/REMC_H_Ising2D.jl b/test/mc_tests/REMC_H_Ising2D.jl
index aec504c67..8f4b34c04 100644
--- a/test/mc_tests/REMC_H_Ising2D.jl
+++ b/test/mc_tests/REMC_H_Ising2D.jl
@@ -46,7 +46,7 @@ replica = Replica(IsingSampler(system, kT, 1), α)
 #  the sampling distribution or the system
 function set_α!(replica::Replica, α::Float64)
     replica.α = α
-    copy_sites = deepcopy(replica.sampler.system.sites)
+    copy_sites = copy(replica.sampler.system.sites)
     replica.sampler.system = create_system(α)
     replica.sampler.system.sites .= copy_sites
     reset_running_energy!(replica.sampler)
diff --git a/test/test_energy_consistency.jl b/test/test_energy_consistency.jl
index 0a95af018..81010c760 100644
--- a/test/test_energy_consistency.jl
+++ b/test/test_energy_consistency.jl
@@ -8,7 +8,7 @@
         add_linear_interactions!(sys, mode)
         add_quadratic_interactions!(sys, mode)
         add_quartic_interactions!(sys, mode)
-        # enable_dipole_dipole!(sys) # workaround segfault
+        enable_dipole_dipole!(sys)
 
         # Random field
         for idx in Sunny.all_sites(sys)
@@ -18,16 +18,19 @@
         rand!(sys.rng, sys.κs)
         # Random spins
         randomize_spins!(sys)
-        # Make inhomogeneous
-        if inhomog
-            sys = to_inhomogeneous(sys)
-            set_vacancy_at!(sys, (1,1,1,1))
-            set_exchange_at!(sys, 0.5, Bond(1, 2, [1,0,0]), (1,1,1,1))
-            set_biquadratic_at!(sys, 0.7, Bond(2, 3, [0,-1,0]), (3,1,1,2))
-            set_anisotropy_at!(sys, 0.4*(𝒮[1]^4+𝒮[2]^4+𝒮[3]^4), (2,2,2,4))
-        end
 
-        return sys
+        if !inhomog
+            return sys
+        else
+            # Add some inhomogeneous interactions
+            sys2 = to_inhomogeneous(sys)
+            @test energy(sys2) ≈ energy(sys)
+            set_vacancy_at!(sys2, (1,1,1,1))
+            set_exchange_at!(sys2, 0.5, Bond(1, 2, [1,0,0]), (1,1,1,1))
+            set_biquadratic_at!(sys2, 0.7, Bond(2, 3, [0,-1,0]), (3,1,1,2))
+            set_anisotropy_at!(sys2, 0.4*(𝒮[1]^4+𝒮[2]^4+𝒮[3]^4), (2,2,2,4))
+            return sys2
+        end
     end