From ab52c22a64ac4d4f6c69c1b2f7775a108ab96b78 Mon Sep 17 00:00:00 2001
From: Ian Hammond <51729722+hammy4815@users.noreply.github.com>
Date: Wed, 15 Feb 2023 18:55:16 -0500
Subject: [PATCH] memory leak fix; Py_XDCREF amp_func set to NULL (#2394)

* memory leak fix; Py_XDCREF amp_func set to NULL

* added test for amp_func memory fix
---
 python/meep.i               |  3 +++
 python/tests/test_source.py | 32 +++++++++++++++++++++++++++++++-
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/python/meep.i b/python/meep.i
index 78f922919..2a9da2bd8 100644
--- a/python/meep.i
+++ b/python/meep.i
@@ -714,6 +714,7 @@ meep::volume_list *make_volume_list(const meep::volume &v, int c,
 
 %typemap(freearg) double (*)(const meep::vec &) {
   Py_XDECREF(py_callback);
+  py_callback = NULL;
 }
 
 %typecheck(SWIG_TYPECHECK_POINTER) double (*)(const meep::vec &) {
@@ -738,6 +739,7 @@ meep::volume_list *make_volume_list(const meep::volume &v, int c,
 
 %typemap(freearg) std::complex<double> (*)(const meep::vec &) {
     Py_XDECREF(py_amp_func);
+    py_amp_func = NULL;
 }
 
 // Typemap suite for vector3
@@ -1050,6 +1052,7 @@ void _get_gradient(PyObject *grad, double scalegrad,
 
 %typemap(freearg) std::complex<double> (*)(const meep::vec &) {
     Py_XDECREF(py_amp_func);
+    py_amp_func = NULL;
 }
 
 %typecheck(SWIG_TYPECHECK_POINTER) PyObject *min_max_loc {
diff --git a/python/tests/test_source.py b/python/tests/test_source.py
index adecb5a38..9fe6e4601 100644
--- a/python/tests/test_source.py
+++ b/python/tests/test_source.py
@@ -4,7 +4,7 @@
 
 import numpy as np
 from meep.geom import Cylinder, Vector3
-from meep.source import ContinuousSource, EigenModeSource, GaussianSource
+from meep.source import ContinuousSource, EigenModeSource, GaussianSource, Source
 
 import meep as mp
 
@@ -12,6 +12,36 @@
 
 
 class TestEigenModeSource(unittest.TestCase):
+    def test_amp_func_change_sources(self):
+        src = ContinuousSource(5.0)
+        center = Vector3()
+        size = Vector3(0, 1, 0)
+
+        ampfunc = lambda X: 1.0
+        amp_source = [
+            Source(src, component=mp.Ez, center=center, size=size, amp_func=ampfunc)
+        ]
+        sim = mp.Simulation(
+            cell_size=Vector3(1, 1, 0), resolution=5, sources=amp_source
+        )
+        sim.run(until=1)
+
+        sim.restart_fields()
+        sim.clear_dft_monitors()
+        default_lattice = [
+            EigenModeSource(src, size=size, center=center),
+            EigenModeSource(src, size=size, center=center),
+        ]
+        sim.change_sources(default_lattice)
+        sim.run(until=1)
+
+        sim.reset_meep()
+        sim.change_sources(amp_source)
+        sim.run(until=1)
+
+        self.assertTrue(sim.sources[0].amp_func is ampfunc)
+        self.assertTrue(sim.sources[0].amp_func is not None)
+
     def test_eig_lattice_defaults(self):
         src = ContinuousSource(5.0)
         center = Vector3()