From 4b76cf6328c5a44514b16723d9df4f5113ed4d4b Mon Sep 17 00:00:00 2001
From: Zach Corse <zcorse@nvidia.com>
Date: Tue, 2 Jul 2024 08:38:58 -0700
Subject: [PATCH] Adding __new__() methods for all __del__() methods

---
 CHANGELOG.md                      |  1 +
 warp/context.py                   | 26 +++++++++++++++++++-----
 warp/dlpack.py                    |  8 ++++++++
 warp/fabric.py                    |  7 +++++--
 warp/fem/cache.py                 |  5 +++++
 warp/render/render_opengl.py      |  7 +++++--
 warp/tests/test_bvh.py            |  5 +++++
 warp/tests/test_fabricarray.py    |  5 ++++-
 warp/tests/test_hash_grid.py      |  5 +++++
 warp/tests/test_marching_cubes.py |  5 ++++-
 warp/tests/test_mesh.py           |  5 ++++-
 warp/tests/test_streams.py        | 10 ++++++++++
 warp/tests/test_volume.py         |  5 ++++-
 warp/types.py                     | 33 +++++++++++++++++++++++--------
 14 files changed, 106 insertions(+), 21 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8cb916130..cbd154ad1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,7 @@
   loaded `(error)`.
 - `wp.config.verbose = True` now also prints out a message upon the entry to a `wp.ScopedTimer`.
 - Add additional documentation and examples demonstrating wp.copy(), wp.clone(), and array.assign() differentiability
+- Fix adding `__new__()` methods for all class `__del__()` methods to anticipate when a class instance is created but not instantiated before garbage collection
 
 ## [1.2.1] - 2024-06-14
 
diff --git a/warp/context.py b/warp/context.py
index e86677f8a..8bfab1bbd 100644
--- a/warp/context.py
+++ b/warp/context.py
@@ -1940,10 +1940,13 @@ def __exit__(self, exc_type, exc_value, traceback):
 
 
 class Stream:
-    def __init__(self, device=None, **kwargs):
-        self.cuda_stream = None
-        self.owner = False
+    def __new__(cls, *args, **kwargs):
+        instance = super(Stream, cls).__new__(cls)
+        instance.cuda_stream = None
+        instance.owner = False
+        return instance
 
+    def __init__(self, device=None, **kwargs):
         # event used internally for synchronization (cached to avoid creating temporary events)
         self._cached_event = None
 
@@ -2019,9 +2022,12 @@ class Flags:
         BLOCKING_SYNC = 0x1
         DISABLE_TIMING = 0x2
 
-    def __init__(self, device=None, cuda_event=None, enable_timing=False):
-        self.owner = False
+    def __new__(cls, *args, **kwargs):
+        instance = super(Event, cls).__new__(cls)
+        instance.owner = False
+        return instance
 
+    def __init__(self, device=None, cuda_event=None, enable_timing=False):
         device = get_device(device)
         if not device.is_cuda:
             raise RuntimeError(f"Device {device} is not a CUDA device")
@@ -2323,6 +2329,11 @@ def can_access(self, other):
 
 
 class Graph:
+    def __new__(cls, *args, **kwargs):
+        instance = super(Graph, cls).__new__(cls)
+        instance.exec = None
+        return instance
+
     def __init__(self, device: Device, exec: ctypes.c_void_p):
         self.device = device
         self.exec = exec
@@ -3824,6 +3835,11 @@ class RegisteredGLBuffer:
 
     __fallback_warning_shown = False
 
+    def __new__(cls, *args, **kwargs):
+        instance = super(RegisteredGLBuffer, cls).__new__(cls)
+        instance.resource = None
+        return instance
+
     def __init__(self, gl_buffer_id: int, device: Devicelike = None, flags: int = NONE, fallback_to_copy: bool = True):
         """
         Args:
diff --git a/warp/dlpack.py b/warp/dlpack.py
index e5e725113..34de42645 100644
--- a/warp/dlpack.py
+++ b/warp/dlpack.py
@@ -60,10 +60,18 @@
 class _DLPackTensorHolder:
     """Class responsible for deleting DLManagedTensor memory after ownership is transferred from a capsule."""
 
+    def __new__(cls, *args, **kwargs):
+        instance = super(_DLPackTensorHolder, cls).__new__(cls)
+        instance.mem_ptr = None
+        return instance
+
     def __init__(self, mem_ptr):
         self.mem_ptr = mem_ptr
 
     def __del__(self):
+        if not self.mem_ptr:
+            return
+
         managed_tensor = DLManagedTensor.from_address(self.mem_ptr)
         if managed_tensor.deleter:
             managed_tensor.deleter(self.mem_ptr)
diff --git a/warp/fabric.py b/warp/fabric.py
index 72271fab1..eb00e9491 100644
--- a/warp/fabric.py
+++ b/warp/fabric.py
@@ -105,11 +105,14 @@ class fabricarray(noncontiguous_array_base[T]):
     # (initialized when needed)
     _vars = None
 
+    def __new__(cls, *args, **kwargs):
+        instance = super(fabricarray, cls).__new__(cls)
+        instance.deleter = None
+        return instance
+
     def __init__(self, data=None, attrib=None, dtype=Any, ndim=None):
         super().__init__(ARRAY_TYPE_FABRIC)
 
-        self.deleter = None
-
         if data is not None:
             from .context import runtime
 
diff --git a/warp/fem/cache.py b/warp/fem/cache.py
index 142ee529b..3081a489d 100644
--- a/warp/fem/cache.py
+++ b/warp/fem/cache.py
@@ -198,6 +198,11 @@ class Temporary:
     The temporary may also be explicitly returned to the pool before destruction using :meth:`release`.
     """
 
+    def __new__(cls, *args, **kwargs):
+        instance = super(Temporary, cls).__new__(cls)
+        instance._pool = None
+        return instance
+
     def __init__(self, array: wp.array, pool: Optional["TemporaryStore.Pool"] = None, shape=None, dtype=None):
         self._raw_array = array
         self._array_view = array
diff --git a/warp/render/render_opengl.py b/warp/render/render_opengl.py
index d96324ec7..cc69c772d 100644
--- a/warp/render/render_opengl.py
+++ b/warp/render/render_opengl.py
@@ -666,12 +666,15 @@ class ShapeInstancer:
         [3D point, 3D normal, UV texture coordinates]
     """
 
+    def __new__(cls):
+        instance = super(ShapeInstancer, cls).__new__(cls)
+        instance.instance_transform_gl_buffer = None
+        instance.vao = None
+
     def __init__(self, shape_shader, device):
         self.shape_shader = shape_shader
         self.device = device
         self.face_count = 0
-        self.vao = None
-        self.instance_transform_gl_buffer = None
         self.instance_color1_buffer = None
         self.instance_color2_buffer = None
         self.color1 = (1.0, 1.0, 1.0)
diff --git a/warp/tests/test_bvh.py b/warp/tests/test_bvh.py
index ea4f6e51b..2473529f6 100644
--- a/warp/tests/test_bvh.py
+++ b/warp/tests/test_bvh.py
@@ -153,6 +153,11 @@ def kernel_fn(bvh: wp.uint64):
 
         wp.Kernel(func=kernel_fn)
 
+    def test_bvh_new_del(self):
+        # test the scenario in which a bvh is created but not initialized before gc
+        instance = wp.Bvh.__new__(wp.Bvh)
+        instance.__del__()
+
 
 add_function_test(TestBvh, "test_bvh_aabb", test_bvh_query_aabb, devices=devices)
 add_function_test(TestBvh, "test_bvh_ray", test_bvh_query_ray, devices=devices)
diff --git a/warp/tests/test_fabricarray.py b/warp/tests/test_fabricarray.py
index 29a4a6cc8..b691ae6ba 100644
--- a/warp/tests/test_fabricarray.py
+++ b/warp/tests/test_fabricarray.py
@@ -931,7 +931,10 @@ def test_fabricarrayarray(test, device):
 
 
 class TestFabricArray(unittest.TestCase):
-    pass
+    def test_fabricarray_new_del(self):
+        # test the scenario in which a fabricarray is created but not initialized before gc
+        instance = wp.fabricarray.__new__(wp.fabricarray)
+        instance.__del__()
 
 
 # fabric arrays
diff --git a/warp/tests/test_hash_grid.py b/warp/tests/test_hash_grid.py
index 53ce98722..4f5a236b0 100644
--- a/warp/tests/test_hash_grid.py
+++ b/warp/tests/test_hash_grid.py
@@ -198,6 +198,11 @@ def kernel_fn(grid: wp.uint64):
 
         wp.Kernel(func=kernel_fn)
 
+    def test_hashgrid_new_del(self):
+        # test the scenario in which a hashgrid is created but not initialized before gc
+        instance = wp.HashGrid.__new__(wp.HashGrid)
+        instance.__del__()
+
 
 add_function_test(TestHashGrid, "test_hashgrid_query", test_hashgrid_query, devices=devices)
 add_function_test(TestHashGrid, "test_hashgrid_inputs", test_hashgrid_inputs, devices=devices)
diff --git a/warp/tests/test_marching_cubes.py b/warp/tests/test_marching_cubes.py
index 5682146c3..a19791db2 100644
--- a/warp/tests/test_marching_cubes.py
+++ b/warp/tests/test_marching_cubes.py
@@ -52,7 +52,10 @@ def test_marching_cubes(test, device):
 
 
 class TestMarchingCubes(unittest.TestCase):
-    pass
+    def test_marching_cubes_new_del(self):
+        # test the scenario in which a MarchingCubes instance is created but not initialized before gc
+        instance = wp.MarchingCubes.__new__(wp.MarchingCubes)
+        instance.__del__()
 
 
 add_function_test(TestMarchingCubes, "test_marching_cubes", test_marching_cubes, devices=devices)
diff --git a/warp/tests/test_mesh.py b/warp/tests/test_mesh.py
index b4891d924..a5f62e11b 100644
--- a/warp/tests/test_mesh.py
+++ b/warp/tests/test_mesh.py
@@ -266,7 +266,10 @@ def test_mesh_exceptions(test, device):
 
 
 class TestMesh(unittest.TestCase):
-    pass
+    def test_mesh_new_del(self):
+        # test the scenario in which a mesh is created but not initialized before gc
+        instance = wp.Mesh.__new__(wp.Mesh)
+        instance.__del__()
 
 
 add_function_test(TestMesh, "test_mesh_read_properties", test_mesh_read_properties, devices=devices)
diff --git a/warp/tests/test_streams.py b/warp/tests/test_streams.py
index 47fd092ba..12a96f0f6 100644
--- a/warp/tests/test_streams.py
+++ b/warp/tests/test_streams.py
@@ -464,6 +464,16 @@ def test_stream_scope_graph_mgpu(self):
             # check results
             assert_np_equal(c0.numpy(), np.full(N, fill_value=2 * num_iters))
 
+    def test_stream_new_del(self):
+        # test the scenario in which a Stream is created but not initialized before gc
+        instance = wp.Stream.__new__(wp.Stream)
+        instance.__del__()
+
+    def test_event_new_del(self):
+        # test the scenario in which an Event is created but not initialized before gc
+        instance = wp.Event.__new__(wp.Event)
+        instance.__del__()
+
 
 add_function_test(TestStreams, "test_stream_set", test_stream_set, devices=devices)
 add_function_test(TestStreams, "test_stream_arg_explicit_sync", test_stream_arg_explicit_sync, devices=devices)
diff --git a/warp/tests/test_volume.py b/warp/tests/test_volume.py
index 6b48d41d0..727d24071 100644
--- a/warp/tests/test_volume.py
+++ b/warp/tests/test_volume.py
@@ -844,7 +844,10 @@ def test_volume_from_numpy(test, device):
 
 
 class TestVolume(unittest.TestCase):
-    pass
+    def test_volume_new_del(self):
+        # test the scenario in which a volume is created but not initialized before gc
+        instance = wp.Volume.__new__(wp.Volume)
+        instance.__del__()
 
 
 add_function_test(
diff --git a/warp/types.py b/warp/types.py
index 0eadfe005..b3c446315 100644
--- a/warp/types.py
+++ b/warp/types.py
@@ -2841,6 +2841,11 @@ def array_type_id(a):
 
 
 class Bvh:
+    def __new__(cls, *args, **kwargs):
+        instance = super(Bvh, cls).__new__(cls)
+        instance.id = None
+        return instance
+
     def __init__(self, lowers, uppers):
         """Class representing a bounding volume hierarchy.
 
@@ -2853,8 +2858,6 @@ def __init__(self, lowers, uppers):
             uppers (:class:`warp.array`): Array of upper bounds :class:`warp.vec3`
         """
 
-        self.id = 0
-
         if len(lowers) != len(uppers):
             raise RuntimeError("Bvh the same number of lower and upper bounds must be provided")
 
@@ -2916,6 +2919,11 @@ class Mesh:
         "indices": Var("indices", array(dtype=int32)),
     }
 
+    def __new__(cls, *args, **kwargs):
+        instance = super(Mesh, cls).__new__(cls)
+        instance.id = None
+        return instance
+
     def __init__(self, points=None, indices=None, velocities=None, support_winding_number=False):
         """Class representing a triangle mesh.
 
@@ -2930,8 +2938,6 @@ def __init__(self, points=None, indices=None, velocities=None, support_winding_n
             support_winding_number (bool): If true the mesh will build additional datastructures to support `wp.mesh_query_point_sign_winding_number()` queries
         """
 
-        self.id = 0
-
         if points.device != indices.device:
             raise RuntimeError("Mesh points and indices must live on the same device")
 
@@ -3001,6 +3007,11 @@ class Volume:
     #: Enum value to specify trilinear interpolation during sampling
     LINEAR = constant(1)
 
+    def __new__(cls, *args, **kwargs):
+        instance = super(Volume, cls).__new__(cls)
+        instance.id = None
+        return instance
+
     def __init__(self, data: array, copy: bool = True):
         """Class representing a sparse grid.
 
@@ -3009,8 +3020,6 @@ def __init__(self, data: array, copy: bool = True):
             copy (bool): Whether the incoming data will be copied or aliased
         """
 
-        self.id = 0
-
         # keep a runtime reference for orderly destruction
         self.runtime = warp.context.runtime
 
@@ -4487,6 +4496,11 @@ def adj_batched_matmul(
 
 
 class HashGrid:
+    def __new__(cls, *args, **kwargs):
+        instance = super(HashGrid, cls).__new__(cls)
+        instance.id = None
+        return instance
+
     def __init__(self, dim_x, dim_y, dim_z, device=None):
         """Class representing a hash grid object for accelerated point queries.
 
@@ -4500,8 +4514,6 @@ def __init__(self, dim_x, dim_y, dim_z, device=None):
             dim_z (int): Number of cells in z-axis
         """
 
-        self.id = 0
-
         self.runtime = warp.context.runtime
 
         self.device = self.runtime.get_device(device)
@@ -4559,6 +4571,11 @@ def __del__(self):
 
 
 class MarchingCubes:
+    def __new__(cls, *args, **kwargs):
+        instance = super(MarchingCubes, cls).__new__(cls)
+        instance.id = None
+        return instance
+
     def __init__(self, nx: int, ny: int, nz: int, max_verts: int, max_tris: int, device=None):
         """CUDA-based Marching Cubes algorithm to extract a 2D surface mesh from a 3D volume.