snap: refactor GPU snap detection into helpers

src/bonsai/bonsai/tool/raycast.py

@@ -317,6 +317,322 @@ def _ensure_wireframe_batches(obj) -> dict[str, tuple[GPUBatch, int, list]]:
         _wireframe_batch_cache[cache_key] = batches
     return batches
 
+def _get_tris_render_ops(objs_to_raycast):
+    """Build render ops for solid (triangle) objects to raycast.
+
+    Each mesh contributes a single TRIANGLES batch and a slot base that 
+    encodes its index in the global ``_obj_list`` (slot 0 is reserved for 
+    the background). The batch is drawn unlit so the object index can be
+    read back from the framebuffer. 
+
+    Args: 
+        objs_to_raycast: iterable of candidate objects. 
+
+    Returns:
+        list[tuple[GPUBatch, Matrix, int]]: ``(batch, world_matrix, slot_base)``
+        for every mesh with faces. Populates ``_obj_list`` as a side effect.
+    """
+    global _obj_list
+
+    render_ops: list[tuple[GPUBatch, Matrix, int]] = []
+
+    for snap_obj in objs_to_raycast:
+        if snap_obj.type != "MESH":
+            continue
+        if not hasattr(snap_obj.data, "polygons"):
+            continue
+        if len(snap_obj.data.polygons) == 0:
+            continue
+
+        batch, cut = _ensure_triangle_batches(snap_obj)
+        if batch is None:
+            continue
+
+        obj_index = len(_obj_list)
+        _obj_list.append(snap_obj)
+        slot_base = obj_index + 1  # slot 0 = background
+        render_ops.append((batch, snap_obj.matrix_world.copy(), slot_base))
+    return render_ops
+    
+def _create_tris_snaps(context, event, mouse_read_rect, buffers_list, last_buf, xray_mode):
+    """Decode the triangle readback buffer(s) into face snaps.
+ 
+    In xray mode each object is read back as a single pixel under the 
+    cursor (``buffers_list``); otherwise the center pixel of the readback 
+    region (``last_buf``) is decoded. Every hit object is then ray cast for 
+    real to find the exact face, producing one ``Face`` snap per hit. 
+ 
+    Args: 
+        context: Blender context. 
+        event: the event carrying the cursor position.
+        mouse: ``(mx, read_x, my, read_y)`` cursor and readback origin. 
+        buffers_list: per-object single-pixel buffers (xray mode only). 
+        last_buf: full readback region buffer (non-xray mode).
+        xray_mode: whether solid xray rendering is active.
+ 
+    Returns:
+        tuple[list[dict], bpy.types.Object | None]: the face snaps and the
+        closest hit object, or ``([], None)`` when nothing was hit. 
+    """
+           
+    global _obj_list
+
+    w, h, mx, my, read_x, read_y = mouse_read_rect
+    # Decode hits
+    hits: set[int] = set()
+
+    if xray_mode:
+        vals_read: set[int] = set()
+        # Each buffer is a single pixel read back right under the
+        # cursor. When the cursor is outside the region there is
+        # nothing to snap to, matching the previous bounds check.
+        if not (0 <= mx < w and 0 <= my < h):
+            return [], None
+        for buf in buffers_list:
+            pixel_data = buf.to_list()
+            if not pixel_data or not pixel_data[0]:
+                return [], None
+            px = pixel_data[0][0]
+            val = _decode_wireframe_pixel(px[0], px[1], px[2], px[3])
+            if val in vals_read:  # avoid getting all the tris from the same object
+                continue
+            vals_read.add(val)
+            if val > 0:
+                obj_index = val - 1
+                if obj_index < len(_obj_list):
+                    hits.add(obj_index)
+    else:
+        pixel_data = last_buf.to_list()
+        if not pixel_data or not pixel_data[0]:
+            return [], None
+        centre_x = mx - int(read_x)
+        centre_y = my - int(read_y)
+        if 0 <= centre_y < len(pixel_data) and 0 <= centre_x < len(pixel_data[0]):
+            px = pixel_data[centre_y][centre_x]
+            val = _decode_wireframe_pixel(px[0], px[1], px[2], px[3])
+            if val > 0:
+                obj_index = val - 1
+                if obj_index < len(_obj_list):
+                    hits.add(obj_index)
+
+    if not hits:
+        return [], None
+
+    snaps: list[dict] = []
+    closest_obj = None
+    closest_dist = float("inf")
+    ray_origin, _, _ = tool.Raycast.get_viewport_ray_data(context, event)
+    for obj_index in hits:
+        obj = _obj_list[obj_index]
+        hit_obj, hit, face_index = tool.Raycast.cast_rays_to_single_object(context, event, obj)
+        if hit:
+            snap: dict = {
+                "point": hit,
+                "type": "Face",
+                "group": "Object",
+                "object": hit_obj,
+                "face_index": face_index,
+                # "is_cut": cut, # Used later in snap
+                "distance": 9,  # High value so it has low priority
+            }
+            dist = (hit - ray_origin).length
+            if dist < closest_dist:
+                closest_dist = dist
+                closest_obj = obj
+
+            snaps.append(snap)
+
+    return snaps, closest_obj
+    
+def _get_wireframe_render_ops(objs_to_raycast):
+    """Build render ops for wireframe (non-solid) objects.
+
+    Boundary points and lines of each object are assigned sequential slot 
+    IDs across all objects, so every vertex and edge gets a unique encoded
+    ID. Per-object slot ranges are recorded in ``obj_slots`` for decoding.
+
+    Args: 
+    objs_to_raycast: iterable of candidate objects. 
+
+    Returns:
+    tuple[list[tuple[GPUBatch, Matrix, int]], list[tuple]]: 
+    ``(render_ops, obj_slots)`` where ``render_ops`` holds
+    ``(batch, world_matrix, slot_base)`` and each ``obj_slots`` 
+    entry is ``(snap_obj, pts_start, n_pts, lines_start, n_lines)``.
+    """
+       
+    render_ops: list[tuple[GPUBatch, Matrix, int]] = []
+    obj_slots: list[tuple] = []  # [(snap_obj, pts_start, n_pts, lines_start, n_lines), ...]
+
+    slot = 1  # slot 0 = background
+
+    for snap_obj in objs_to_raycast:
+        # avoids creating batches for solid objects
+        if hasattr(snap_obj.data, "polygons") and len(snap_obj.data.polygons) > 0:
+            continue
+
+        batches = _ensure_wireframe_batches(snap_obj)
+        if not batches:
+            continue
+
+        world_mat = snap_obj.matrix_world.copy()
+        pts_start = 0
+        n_pts = 0
+        lines_start = 0
+        n_lines = 0
+
+        pts_data = batches.get("POINTS")
+        if pts_data is not None:
+            batch, n_pts, _ = pts_data
+            pts_start = slot
+            render_ops.append((batch, world_mat, slot))
+            slot += n_pts
+
+        lines_data = batches.get("LINES")
+        if lines_data is not None:
+            batch, n_lines, _ = lines_data
+            lines_start = slot
+            render_ops.append((batch, world_mat, slot))
+            slot += n_lines
+
+        if n_pts > 0 or n_lines > 0:
+            obj_slots.append((snap_obj, pts_start, n_pts, lines_start, n_lines))
+
+    return render_ops, obj_slots
+
+def _create_wireframe_snaps(context, event, mouse_read_rect, obj_slots, last_buf):
+    """Decode the wireframe readback buffer into vertex/edge snaps. 
+
+    Finds the closest non-zero pixel to the cursor, maps its encoded slot ID
+    back to a vertex or edge via ``obj_slots``, then builds the candidate 
+    snaps (Vertex, Edge, Edge Center, plus endpoint Vertex snaps within the 
+    snap threshold).
+
+    Args: 
+        context: Blender context. 
+        event: the event carrying the cursor position.
+        mouse: ``(mx, read_x, my, read_y)`` cursor and readback origin. 
+        obj_slots: ``[(snap_obj, pts_start, n_pts, lines_start, n_lines), ...]``. 
+        last_buf: full readback region buffer.
+
+    Returns:
+        tuple[list[dict], None]: the wireframe snaps, or ``([], None)`` when
+        no non-zero pixel is found near the cursor. 
+    """ 
+    global _wireframe_batch_cache
+
+    w, h, mx, my, read_x, read_y = mouse_read_rect
+    centre = (mx - int(read_x), my - int(read_y))
+    pixel_data = last_buf.to_list()
+    best = _find_closest_wireframe_pixel(pixel_data, *centre)
+    if best is None:
+        return [], None
+    encoded, dx, dy = best
+
+    # Decode and build snap dicts
+
+    rv3d = context.region_data
+    snaps: list[dict] = []
+    snap_threshold = tool.Raycast.calculate_snap_threshold(rv3d.view_distance)
+
+    # Compute view ray for 3D proximity calculations
+    _, ray_target, ray_direction = tool.Raycast.get_viewport_ray_data(context, event)
+    try:
+        loc = tool.Cad.region_2d_to_location_3d_np(context.region, rv3d, (mx, my), ray_direction)
+    except Exception:
+        loc = ray_target
+
+    for snap_obj, pts_start, n_pts, lines_start, n_lines in obj_slots:
+        if n_pts > 0 and pts_start <= encoded < pts_start + n_pts:
+            vi = encoded - pts_start
+            batches = _wireframe_batch_cache.get(id(snap_obj))
+            if batches:
+                pts_data = batches.get("POINTS")
+                if pts_data:
+                    _, _, coords = pts_data
+                    if vi < len(coords):
+                        local_pos = Vector(coords[vi])
+                        world_pos = snap_obj.matrix_world @ local_pos
+                        # Compute proper 3D distance from vertex to view ray
+                        proj = tool.Cad.point_on_edge(world_pos, (ray_target, loc))
+                        distance = (world_pos - proj).length
+                        snaps.append(
+                            {
+                                "object": snap_obj,
+                                "type": "Vertex",
+                                "point": world_pos,
+                                "distance": distance,
+                                "group": "Wireframe",
+                            }
+                        )
+            break
+
+        if n_lines > 0 and lines_start <= encoded < lines_start + n_lines:
+            ei = encoded - lines_start
+            batches = _wireframe_batch_cache.get(id(snap_obj))
+            if batches:
+                lines_data = batches.get("LINES")
+                if lines_data:
+                    _, _, edge_pairs = lines_data
+                    if ei < len(edge_pairs):
+                        c0, c1 = edge_pairs[ei]
+                        mw = snap_obj.matrix_world
+                        v0 = mw @ Vector(c0)
+                        v1 = mw @ Vector(c1)
+
+                        # Compute closest point on edge to view ray
+                        intersection = tool.Cad.intersect_edges_v2((ray_target, loc), (v0, v1))
+                        if intersection[0] is not None and tool.Cad.is_point_on_edge(intersection[1], (v0, v1)):
+                            edge_point = intersection[1].copy()
+                            distance = (intersection[1] - intersection[0]).length
+                        else:
+                            # Fallback to midpoint if lines are parallel
+                            edge_point = (v0 + v1) / 2
+                            proj = tool.Cad.point_on_edge(edge_point, (ray_target, loc))
+                            distance = (edge_point - proj).length
+
+                        snaps.append(
+                            {
+                                "object": snap_obj,
+                                "type": "Edge",
+                                "point": edge_point,
+                                "edge_verts": (v0, v1),
+                                "distance": distance,
+                                "group": "Wireframe",
+                            }
+                        )
+
+                        # Edge Center snap (midpoint)
+                        mid = (v0 + v1) / 2 # TODO Allow divisions by other values
+                        mid_proj = tool.Cad.point_on_edge(mid, (ray_target, loc))
+                        mid_dist = (mid - mid_proj).length
+                        snaps.append(
+                            {
+                                "object": snap_obj,
+                                "type": "Edge Center",
+                                "point": mid,
+                                "distance": mid_dist,
+                                "group": "Wireframe",
+                            }
+                        )
+
+                        # Also include vertex snaps for edge endpoints
+                        for vtx in (v0, v1):
+                            proj = tool.Cad.point_on_edge(vtx, (ray_target, loc))
+                            vtx_dist = (vtx - proj).length
+                            if vtx_dist < snap_threshold:
+                                snaps.append(
+                                    {
+                                        "object": snap_obj,
+                                        "type": "Vertex",
+                                        "point": vtx,
+                                        "distance": vtx_dist,
+                                        "group": "Wireframe",
+                                    }
+                                )
+            break
+
+    return snaps, None
 
 class Raycast(bonsai.core.tool.Raycast):
     offset = 10
@@ -670,57 +986,9 @@ class Raycast(bonsai.core.tool.Raycast):
         obj_slots: list[tuple] = []  # [(snap_obj, pts_start, n_pts, lines_start, n_lines), ...]
 
         if tris:
-            for snap_obj in objs_to_raycast:
+            render_ops = _get_tris_render_ops(objs_to_raycast)
-                if snap_obj.type != "MESH":
-                    continue
-                if not hasattr(snap_obj.data, "polygons"):
-                    continue
-                if len(snap_obj.data.polygons) == 0:
-                    continue
-
-                batch_info = _ensure_triangle_batches(snap_obj)
-                if batch_info is None:
-                    continue
-
-                batch, _ = batch_info
-                obj_index = len(_obj_list)
-                _obj_list.append(snap_obj)
-                slot_base = obj_index + 1  # slot 0 = background
-                render_ops.append((batch, snap_obj.matrix_world.copy(), slot_base))
         else:
-            slot = 1  # slot 0 = background
+            render_ops, obj_slots = _get_wireframe_render_ops(objs_to_raycast)
-
-            for snap_obj in objs_to_raycast:
-                # avoids creating batches for solid objects
-                if hasattr(snap_obj.data, "polygons") and len(snap_obj.data.polygons) > 0:
-                    continue
-
-                batches = _ensure_wireframe_batches(snap_obj)
-                if not batches:
-                    continue
-
-                world_mat = snap_obj.matrix_world.copy()
-                pts_start = 0
-                n_pts = 0
-                lines_start = 0
-                n_lines = 0
-
-                pts_data = batches.get("POINTS")
-                if pts_data is not None:
-                    batch, n_pts, _ = pts_data
-                    pts_start = slot
-                    render_ops.append((batch, world_mat, slot))
-                    slot += n_pts
-
-                lines_data = batches.get("LINES")
-                if lines_data is not None:
-                    batch, n_lines, _ = lines_data
-                    lines_start = slot
-                    render_ops.append((batch, world_mat, slot))
-                    slot += n_lines
-
-                if n_pts > 0 or n_lines > 0:
-                    obj_slots.append((snap_obj, pts_start, n_pts, lines_start, n_lines))
 
         if not render_ops:
             return [], None
@@ -788,183 +1056,11 @@ class Raycast(bonsai.core.tool.Raycast):
         gpu.state.depth_mask_set(True)
         gpu.state.depth_test_set("LESS")
 
+        mouse_read_rect = (w, h, mx, my, read_x, read_y)
         if tris:
-            # Decode hits
+            return _create_tris_snaps(context, event, mouse_read_rect, buffers_list, last_buf, xray_mode)
-            hits: set[int] = set()
-
-            if xray_mode:
-                vals_read: set[int] = set()
-                # Each buffer is a single pixel read back right under the
-                # cursor. When the cursor is outside the region there is
-                # nothing to snap to, matching the previous bounds check.
-                if not (0 <= mx < w and 0 <= my < h):
-                    return [], None
-                for buf in buffers_list:
-                    pixel_data = buf.to_list()
-                    if not pixel_data or not pixel_data[0]:
-                        return [], None
-                    px = pixel_data[0][0]
-                    val = _decode_wireframe_pixel(px[0], px[1], px[2], px[3])
-                    if val in vals_read:  # avoid getting all the tris from the same object
-                        continue
-                    vals_read.add(val)
-                    if val > 0:
-                        obj_index = val - 1
-                        if obj_index < len(_obj_list):
-                            hits.add(obj_index)
-            else:
-                pixel_data = last_buf.to_list()
-                if not pixel_data or not pixel_data[0]:
-                    return [], None
-                centre_x = mx - int(read_x)
-                centre_y = my - int(read_y)
-                if 0 <= centre_y < len(pixel_data) and 0 <= centre_x < len(pixel_data[0]):
-                    px = pixel_data[centre_y][centre_x]
-                    val = _decode_wireframe_pixel(px[0], px[1], px[2], px[3])
-                    if val > 0:
-                        obj_index = val - 1
-                        if obj_index < len(_obj_list):
-                            hits.add(obj_index)
-
-            if not hits:
-                return [], None
-
-            snaps: list[dict] = []
-            closest_obj = None
-            closest_dist = float("inf")
-            ray_origin, _, _ = cls.get_viewport_ray_data(context, event)
-            for obj_index in hits:
-                obj = _obj_list[obj_index]
-                hit_obj, hit, face_index = cls.cast_rays_to_single_object(context, event, obj)
-                if hit:
-                    snap: dict = {
-                        "point": hit,
-                        "type": "Face",
-                        "group": "Object",
-                        "object": hit_obj,
-                        "face_index": face_index,
-                        "distance": 9,  # High value so it has low priority
-                    }
-                    dist = (hit - ray_origin).length
-                    if dist < closest_dist:
-                        closest_dist = dist
-                        closest_obj = obj
-
-                    snaps.append(snap)
-
-            return snaps, closest_obj
-
         else:
-            centre = (mx - int(read_x), my - int(read_y))
+            return _create_wireframe_snaps(context, event, mouse_read_rect, obj_slots, last_buf)
-            pixel_data = last_buf.to_list()
-            best = _find_closest_wireframe_pixel(pixel_data, *centre)
-            if best is None:
-                return [], None
-            encoded, dx, dy = best
-
-            # Decode and build snap dicts
-
-            snaps: list[dict] = []
-            snap_threshold = cls.calculate_snap_threshold(rv3d.view_distance)
-
-            # Compute view ray for 3D proximity calculations
-            _, ray_target, ray_direction = cls.get_viewport_ray_data(context, event)
-            try:
-                loc = tool.Cad.region_2d_to_location_3d_np(region, rv3d, (mx, my), ray_direction)
-            except Exception:
-                loc = ray_target
-
-            for snap_obj, pts_start, n_pts, lines_start, n_lines in obj_slots:
-                if n_pts > 0 and pts_start <= encoded < pts_start + n_pts:
-                    vi = encoded - pts_start
-                    batches = _wireframe_batch_cache.get(id(snap_obj))
-                    if batches:
-                        pts_data = batches.get("POINTS")
-                        if pts_data:
-                            _, _, coords = pts_data
-                            if vi < len(coords):
-                                local_pos = Vector(coords[vi])
-                                world_pos = snap_obj.matrix_world @ local_pos
-                                # Compute proper 3D distance from vertex to view ray
-                                proj = tool.Cad.point_on_edge(world_pos, (ray_target, loc))
-                                distance = (world_pos - proj).length
-                                snaps.append(
-                                    {
-                                        "object": snap_obj,
-                                        "type": "Vertex",
-                                        "point": world_pos,
-                                        "distance": distance,
-                                        "group": "Wireframe",
-                                    }
-                                )
-                    break
-
-                if n_lines > 0 and lines_start <= encoded < lines_start + n_lines:
-                    ei = encoded - lines_start
-                    batches = _wireframe_batch_cache.get(id(snap_obj))
-                    if batches:
-                        lines_data = batches.get("LINES")
-                        if lines_data:
-                            _, _, edge_pairs = lines_data
-                            if ei < len(edge_pairs):
-                                c0, c1 = edge_pairs[ei]
-                                mw = snap_obj.matrix_world
-                                v0 = mw @ Vector(c0)
-                                v1 = mw @ Vector(c1)
-
-                                # Compute closest point on edge to view ray
-                                intersection = tool.Cad.intersect_edges_v2((ray_target, loc), (v0, v1))
-                                if intersection[0] is not None and tool.Cad.is_point_on_edge(intersection[1], (v0, v1)):
-                                    edge_point = intersection[1].copy()
-                                    distance = (intersection[1] - intersection[0]).length
-                                else:
-                                    # Fallback to midpoint if lines are parallel
-                                    edge_point = (v0 + v1) / 2
-                                    proj = tool.Cad.point_on_edge(edge_point, (ray_target, loc))
-                                    distance = (edge_point - proj).length
-
-                                snaps.append(
-                                    {
-                                        "object": snap_obj,
-                                        "type": "Edge",
-                                        "point": edge_point,
-                                        "edge_verts": (v0, v1),
-                                        "distance": distance,
-                                        "group": "Wireframe",
-                                    }
-                                )
-
-                                # Edge Center snap (midpoint)
-                                mid = (v0 + v1) / 2
-                                mid_proj = tool.Cad.point_on_edge(mid, (ray_target, loc))
-                                mid_dist = (mid - mid_proj).length
-                                snaps.append(
-                                    {
-                                        "object": snap_obj,
-                                        "type": "Edge Center",
-                                        "point": mid,
-                                        "distance": mid_dist,
-                                        "group": "Wireframe",
-                                    }
-                                )
-
-                                # Also include vertex snaps for edge endpoints
-                                for vtx in (v0, v1):
-                                    proj = tool.Cad.point_on_edge(vtx, (ray_target, loc))
-                                    vtx_dist = (vtx - proj).length
-                                    if vtx_dist < snap_threshold:
-                                        snaps.append(
-                                            {
-                                                "object": snap_obj,
-                                                "type": "Vertex",
-                                                "point": vtx,
-                                                "distance": vtx_dist,
-                                                "group": "Wireframe",
-                                            }
-                                        )
-                    break
-
-            return snaps, None
 
     @classmethod
     def get_gpu_solid_snaps(cls, context, event, objs_to_raycast):