Rewrite gpu collision for spritelists

Author: einarf

arcade/resources/system/shaders/collision/col_trans_vs.glsl

@@ -1,7 +1,7 @@
 #version 330
 // A simple passthrough shader forwarding data to the geomtry shader
 
-in vec3 in_pos;
+in vec4 in_pos;
 in vec2 in_size;
 
 out vec2 pos;

arcade/sprite_list/collision.py

@@ -134,50 +134,7 @@ def _get_nearby_sprites(
     sprite_count = len(sprite_list)
     if sprite_count == 0:
         return []
-
-    # Update the position and size to check
-    ctx = get_window().ctx
-    sprite_list._write_sprite_buffers_to_gpu()
-
-    ctx.collision_detection_program["check_pos"] = sprite.position
-    ctx.collision_detection_program["check_size"] = sprite.width, sprite.height
-
-    # Ensure the result buffer can fit all the sprites (worst case)
-    buffer = ctx.collision_buffer
-    if buffer.size < sprite_count * 4:
-        buffer.orphan(size=sprite_count * 4)
-
-    # Run the transform shader emitting sprites close to the configured position and size.
-    # This runs in a query so we can measure the number of sprites emitted.
-    with ctx.collision_query:
-        sprite_list.sprite_data.geometry.transform(  # type: ignore
-            ctx.collision_detection_program,
-            buffer,
-            vertices=sprite_count,
-        )
-
-    # Store the number of sprites emitted
-    emit_count = ctx.collision_query.primitives_generated
-    # print(
-    #     emit_count,
-    #     ctx.collision_query.time_elapsed,
-    #     ctx.collision_query.time_elapsed / 1_000_000_000,
-    # )
-
-    # If no sprites emitted we can just return an empty list
-    if emit_count == 0:
-        return []
-
-    # # Debug block for transform data to keep around
-    # print("emit_count", emit_count)
-    # data = buffer.read(size=emit_count * 4)
-    # print("bytes", data)
-    # print("data", struct.unpack(f'{emit_count}i', data))
-
-    # .. otherwise build and return a list of the sprites selected by the transform
-    return [
-        sprite_list[i] for i in struct.unpack(f"{emit_count}i", buffer.read(size=emit_count * 4))
-    ]
+    return sprite_list.get_nearby_sprites_gpu(sprite.position, sprite.size)
 
 
 def check_for_collision_with_list(

arcade/sprite_list/sprite_list.py

@@ -8,6 +8,7 @@
 from __future__ import annotations
 
 import random
+import struct
 from abc import abstractmethod
 from array import array
 from collections import deque
@@ -139,6 +140,23 @@ def draw_hit_boxes(
         """
         ...
 
+    @abstractmethod
+    def get_nearby_sprites_gpu(self, pos: Point2, size: Point2) -> list[SpriteType_co]:
+        """
+        Get a list of sprites that are nearby the given position and size
+        using the gpu. No spatial hashing is needed. This is a very fast method
+        to find nearby sprites in large spritelists but is very expensive
+        if the method is called many times per frame or if the sprite list
+        is small.
+
+        Args:
+            pos: The position to check for nearby sprites.
+            size: The size of the area to check for nearby sprites.
+        Returns:
+            A list of sprites nearby the given position and size.
+        """
+        ...
+
     @abstractmethod
     def _write_sprite_buffers_to_gpu(self) -> None: ...
 
@@ -305,13 +323,13 @@ def _init_deferred(self) -> None:
 
         # NOTE: Instantiate the appropriate spritelist data class here
         # Desktop GL (with geo shader)
-        # self._spritelist_data = SpriteListBufferData(
-        #     self.ctx, capacity=self._buf_capacity, atlas=self._atlas
-        # )
-        # WebGL (without geo shader)
-        self._spritelist_data = SpriteListTextureData(
+        self._spritelist_data = SpriteListBufferData(
             self.ctx, capacity=self._buf_capacity, atlas=self._atlas
         )
+        # WebGL (without geo shader)
+        # self._spritelist_data = SpriteListTextureData(
+        #     self.ctx, capacity=self._buf_capacity, atlas=self._atlas
+        # )
 
         self._initialized = True
 
@@ -981,6 +999,30 @@ def draw_hit_boxes(
 
         arcade.draw_lines(points, color=converted_color, line_width=line_thickness)
 
+    def get_nearby_sprites_gpu(self, pos: Point2, size: Point2) -> list[SpriteType]:
+        """
+        Get a list of sprites that are nearby the given position and size
+        using the gpu. No spatial hashing is needed. This is a very fast method
+        to find nearby sprites in large spritelists but is very expensive
+        if the method is called many times per frame or if the sprite list
+        is small.
+
+        Args:
+            pos: The position to check for nearby sprites.
+            size: The size of the area to check for nearby sprites.
+        Returns:
+            A list of sprites nearby the given position and size.
+        """
+        if not self._initialized:
+            self._init_deferred()
+
+        if len(self.sprite_list) == 0:
+            return []
+
+        self._write_sprite_buffers_to_gpu()
+        indices = self._spritelist_data.get_nearby_sprite_indices(pos, size, len(self.sprite_list))
+        return [self.sprite_list[i] for i in indices]
+
     def _grow_sprite_buffers(self) -> None:
         """Double the internal buffer sizes"""
         # Resize sprite buffers if needed
@@ -1294,6 +1336,19 @@ def render(
         """
         raise NotImplementedError("This method should be implemented in subclasses.")
 
+    def get_nearby_sprite_indices(self, pos: Point2, size: Point2, length: int) -> list[int]:
+        """
+        Get indices of sprites that are nearby the given position and size.
+
+        Args:
+            pos: The position to check for nearby sprites.
+            size: The size of the area to check for nearby sprites.
+            length: The number of sprites in the list.
+        Returns:
+            A list of indices of nearby sprites.
+        """
+        raise NotImplementedError("This method should be implemented in subclasses.")
+
 
 class SpriteListBufferData(SpriteListData):
     """Container for all gpu data used by the SpriteList."""
@@ -1577,6 +1632,39 @@ def render(
             if blend_function is not None:
                 self.ctx.blend_func = prev_blend_func
 
+    def get_nearby_sprite_indices(self, pos: Point2, size: Point2, length: int) -> list[int]:
+        """
+        Get indices of sprites that are nearby the given position and size.
+
+        Args:
+            pos: The position to check for nearby sprites.
+            size: The size of the area to check for nearby sprites.
+            length: The number of sprites in the spritelist.
+        Returns:
+            A list of indices of nearby sprites.
+        """
+        ctx = self.ctx
+        ctx.collision_detection_program["check_pos"] = pos
+        ctx.collision_detection_program["check_size"] = size
+
+        # Ensure the result buffer can fit all the sprites (worst case)
+        buffer = ctx.collision_buffer
+        # NOTE: Right now the limit is 1000 hits
+        # Run the transform shader emitting sprites close to the configured position and size.
+        # This runs in a query so we can measure the number of sprites emitted.
+        with ctx.collision_query:
+            self._geometry.transform(  # type: ignore
+                ctx.collision_detection_program,
+                buffer,
+                vertices=length,
+            )
+
+        # Store the number of sprites emitted
+        emit_count = ctx.collision_query.primitives_generated
+        if emit_count == 0:
+            return []
+        return [i for i in struct.unpack(f"{emit_count}i", buffer.read(size=emit_count * 4))]
+
 
 class SpriteListTextureData(SpriteListData):
     """Container for all gpu data used by the SpriteList without buffers."""
@@ -1753,3 +1841,15 @@ def render(
             self.ctx.disable(self.ctx.BLEND)
             if blend_function is not None:
                 self.ctx.blend_func = prev_blend_func
+
+    # def get_nearby_sprite_indices(self, pos: Point2, size: Point2, length: int) -> list[int]:
+    #     """
+    #     Get indices of sprites that are nearby the given position and size.
+
+    #     Args:
+    #         pos: The position to check for nearby sprites.
+    #         size: The size of the area to check for nearby sprites.
+    #         length: The number of sprites in the spritelist.
+    #     Returns:
+    #         A list of indices of nearby sprites.
+    #     """