polygons

system/ui/lib/shader_polygon.py

@@ -1,9 +1,13 @@
+import os
 import pyray as rl
 import numpy as np
+from raylib import rl as raw_rl
 from dataclasses import dataclass
 from typing import Any, Optional, cast
 from openpilot.system.ui.lib.application import gui_app, GL_VERSION
 
+USE_VERTEX_GRADIENTS = os.getenv("USE_VERTEX_GRADIENTS", "1") == "1"
+
 MAX_GRADIENT_COLORS = 20  # includes stops as well
 
 
@@ -200,6 +204,56 @@ def triangulate(pts: np.ndarray) -> np.ndarray:
   return tri_strip
 
 
+def _draw_strip_vertex_gradient(tri_strip: np.ndarray, gradient: Gradient, origin_rect: rl.Rectangle) -> None:
+  n = len(tri_strip)
+  if n < 3 or not gradient.colors:
+    return
+
+  start_x = origin_rect.x + gradient.start[0] * origin_rect.width
+  start_y = origin_rect.y + gradient.start[1] * origin_rect.height
+  end_x = origin_rect.x + gradient.end[0] * origin_rect.width
+  end_y = origin_rect.y + gradient.end[1] * origin_rect.height
+
+  dx = start_x - end_x
+  dy = start_y - end_y
+  len2 = max(dx * dx + dy * dy, 1e-6)
+
+  vx = tri_strip[:, 0] - end_x
+  vy = tri_strip[:, 1] - end_y
+  t = np.clip((vx * dx + vy * dy) / len2, 0.0, 1.0)
+
+  colors = gradient.colors
+  stops = gradient.stops if len(gradient.stops) == len(colors) else [i / max(1, len(colors) - 1) for i in range(len(colors))]
+  r_stops = [c.r for c in colors]
+  g_stops = [c.g for c in colors]
+  b_stops = [c.b for c in colors]
+  a_stops = [c.a for c in colors]
+
+  r = np.interp(t, stops, r_stops).astype(np.uint8)
+  g = np.interp(t, stops, g_stops).astype(np.uint8)
+  b = np.interp(t, stops, b_stops).astype(np.uint8)
+  a = np.interp(t, stops, a_stops).astype(np.uint8)
+
+  rl_begin = raw_rl.rlBegin
+  rl_end = raw_rl.rlEnd
+  rl_color = raw_rl.rlColor4ub
+  rl_vertex = raw_rl.rlVertex2f
+
+  rl_begin(4)  # RL_TRIANGLES
+  for i in range(2, n):
+    if i % 2 == 0:
+      i0, i1, i2 = i, i - 2, i - 1
+    else:
+      i0, i1, i2 = i, i - 1, i - 2
+    rl_color(r[i0], g[i0], b[i0], a[i0])
+    rl_vertex(tri_strip[i0, 0], tri_strip[i0, 1])
+    rl_color(r[i1], g[i1], b[i1], a[i1])
+    rl_vertex(tri_strip[i1, 0], tri_strip[i1, 1])
+    rl_color(r[i2], g[i2], b[i2], a[i2])
+    rl_vertex(tri_strip[i2, 0], tri_strip[i2, 1])
+  rl_end()
+
+
 def draw_polygon(origin_rect: rl.Rectangle, points: np.ndarray,
                  color: Optional[rl.Color] = None, gradient: Gradient | None = None):
 
@@ -215,12 +269,19 @@ def draw_polygon(origin_rect: rl.Rectangle, points: np.ndarray,
 
   # Triangulate via interleaving
   tri_strip = triangulate(pts)
-  vertices = rl.ffi.from_buffer("Vector2 *", tri_strip)
+  if len(tri_strip) < 3:
+    return
 
   if gradient is None:
+    vertices = rl.ffi.from_buffer("Vector2 *", tri_strip)
     rl.draw_triangle_strip(vertices, len(tri_strip), color or rl.WHITE)
     return
 
+  if USE_VERTEX_GRADIENTS:
+    _draw_strip_vertex_gradient(tri_strip, gradient, origin_rect)
+    return
+
+  vertices = rl.ffi.from_buffer("Vector2 *", tri_strip)
   state = ShaderState.get_instance()
   state.initialize()
 

system/ui/lib/tests/test_shader_polygon.py

@@ -59,7 +59,8 @@ def test_triangulate_preserves_coordinates_in_contiguous_float_buffer(points):
 
 @pytest.mark.parametrize("use_gradient", [False, True])
 @pytest.mark.parametrize("point_count", [4, 5, 66])
-def test_draw_polygon_passes_native_vertices(monkeypatch, use_gradient, point_count):
+def test_draw_polygon_legacy_shader_fallback(monkeypatch, use_gradient, point_count):
+  monkeypatch.setattr(shader_polygon, "USE_VERTEX_GRADIENTS", False)
   points = np.arange(point_count * 4, dtype=np.float64).reshape(-1, 2)[::2]
   rect = rl.Rectangle(0, 0, 100, 100)
   color = rl.Color(10, 20, 30, 40)
@@ -84,3 +85,76 @@ def test_draw_polygon_passes_native_vertices(monkeypatch, use_gradient, point_co
 
   assert calls == (["initialize", ("configure", (state, None, gradient, rect)), ("begin", "shader"), "draw", "end"]
                    if use_gradient else ["draw"])
+
+
+@pytest.mark.parametrize("use_gradient", [False, True])
+@pytest.mark.parametrize("point_count", [4, 5, 66])
+def test_draw_polygon_native_vertex_gradients(monkeypatch, use_gradient, point_count):
+  monkeypatch.setattr(shader_polygon, "USE_VERTEX_GRADIENTS", True)
+  points = np.arange(point_count * 4, dtype=np.float64).reshape(-1, 2)[::2]
+  rect = rl.Rectangle(0, 0, 100, 100)
+  color = rl.Color(10, 20, 30, 40)
+  gradient = Gradient((0, 0), (1, 1), [color, rl.Color(*rl.WHITE)], [0, 1]) if use_gradient else None
+
+  raw_calls = []
+  emitted_vertices = []
+  emitted_colors = []
+  fake_raw = SimpleNamespace(
+    rlBegin=lambda m: raw_calls.append(("rlBegin", m)),
+    rlEnd=lambda: raw_calls.append("rlEnd"),
+    rlColor4ub=lambda r, g, b, a: emitted_colors.append((int(r), int(g), int(b), int(a))),
+    rlVertex2f=lambda x, y: emitted_vertices.append((float(x), float(y))),
+  )
+  monkeypatch.setattr(shader_polygon, "raw_rl", fake_raw)
+
+  draw_calls = []
+  def draw_strip(vertices, count, fill):
+    assert rl.ffi.typeof(vertices) == rl.ffi.typeof("Vector2 *")
+    coordinates = [[vertices[i].x, vertices[i].y] for i in range(count)]
+    np.testing.assert_array_equal(coordinates, triangulate(points))
+    assert fill == color
+    draw_calls.append("draw")
+
+  monkeypatch.setattr(rl, "draw_triangle_strip", draw_strip)
+
+  draw_polygon(rect, points, gradient=gradient, color=None if use_gradient else color)
+
+  if use_gradient:
+    strip = triangulate(points)
+    n = len(strip)
+    assert raw_calls == [("rlBegin", 4), "rlEnd"]
+    expected_vertex_count = (n - 2) * 3
+    assert len(emitted_vertices) == expected_vertex_count
+    assert len(emitted_colors) == expected_vertex_count
+    assert draw_calls == []
+  else:
+    assert draw_calls == ["draw"]
+    assert raw_calls == []
+
+
+def test_vertex_gradient_color_interpolation(monkeypatch):
+  monkeypatch.setattr(shader_polygon, "USE_VERTEX_GRADIENTS", True)
+  points = np.array([
+    [0.0, 100.0],
+    [0.0, 0.0],
+    [10.0, 0.0],
+    [10.0, 100.0],
+  ], dtype=np.float32)
+  rect = rl.Rectangle(0, 0, 10, 100)
+  color_start = rl.Color(255, 0, 0, 255)
+  color_end = rl.Color(0, 255, 0, 128)
+  gradient = Gradient((0.0, 1.0), (0.0, 0.0), [color_start, color_end], [0.0, 1.0])
+
+  emitted_colors = []
+  fake_raw = SimpleNamespace(
+    rlBegin=lambda m: None,
+    rlEnd=lambda: None,
+    rlColor4ub=lambda r, g, b, a: emitted_colors.append((int(r), int(g), int(b), int(a))),
+    rlVertex2f=lambda x, y: None,
+  )
+  monkeypatch.setattr(shader_polygon, "raw_rl", fake_raw)
+
+  draw_polygon(rect, points, gradient=gradient)
+
+  assert (255, 0, 0, 255) in emitted_colors
+  assert (0, 255, 0, 128) in emitted_colors