StarPilot/openpilot/system/ui/lib/shader_polygon.py

import pyray as rl
import numpy as np
from dataclasses import dataclass
from typing import Any, Optional, cast
from openpilot.system.ui.lib.application import gui_app, GL_VERSION

MAX_GRADIENT_COLORS = 20  # includes stops as well


@dataclass
class Gradient:
  start: tuple[float, float]
  end: tuple[float, float]
  colors: list[rl.Color]
  stops: list[float]

  def __post_init__(self):
    if len(self.colors) > MAX_GRADIENT_COLORS:
      self.colors = self.colors[:MAX_GRADIENT_COLORS]
      print(f"Warning: Gradient colors truncated to {MAX_GRADIENT_COLORS} entries")

    if len(self.stops) > MAX_GRADIENT_COLORS:
      self.stops = self.stops[:MAX_GRADIENT_COLORS]
      print(f"Warning: Gradient stops truncated to {MAX_GRADIENT_COLORS} entries")

    if not len(self.stops):
      color_count = min(len(self.colors), MAX_GRADIENT_COLORS)
      self.stops = [i / max(1, color_count - 1) for i in range(color_count)]


FRAGMENT_SHADER = GL_VERSION + """
in vec2 fragTexCoord;
out vec4 finalColor;

uniform vec4 fillColor;

// Gradient line defined in *screen pixels*
uniform int useGradient;
uniform vec2 gradientStart;  // e.g. vec2(0, 0)
uniform vec2 gradientEnd;    // e.g. vec2(0, screenHeight)
uniform vec4 gradientColors[20];
uniform float gradientStops[20];
uniform int gradientColorCount;

vec4 getGradientColor(vec2 p) {
  // Compute t from screen-space position
  vec2 d = gradientStart - gradientEnd;
  float len2 = max(dot(d, d), 1e-6);
  float t = clamp(dot(p - gradientEnd, d) / len2, 0.0, 1.0);

  // Clamp to range
  float t0 = gradientStops[0];
  float tn = gradientStops[gradientColorCount-1];
  if (t <= t0) return gradientColors[0];
  if (t >= tn) return gradientColors[gradientColorCount-1];

  for (int i = 0; i < gradientColorCount - 1; i++) {
    float a = gradientStops[i];
    float b = gradientStops[i+1];
    if (t >= a && t <= b) {
      float k = (t - a) / max(b - a, 1e-6);
      return mix(gradientColors[i], gradientColors[i+1], k);
    }
  }

  return gradientColors[gradientColorCount-1];
}

void main() {
  // TODO: do proper antialiasing
  finalColor = useGradient == 1 ? getGradientColor(gl_FragCoord.xy) : fillColor;
}
"""

# Default vertex shader
VERTEX_SHADER = GL_VERSION + """
in vec3 vertexPosition;
in vec2 vertexTexCoord;
out vec2 fragTexCoord;
uniform mat4 mvp;

void main() {
  fragTexCoord = vertexTexCoord;
  gl_Position = mvp * vec4(vertexPosition, 1.0);
}
"""

UNIFORM_INT = rl.ShaderUniformDataType.SHADER_UNIFORM_INT
UNIFORM_FLOAT = rl.ShaderUniformDataType.SHADER_UNIFORM_FLOAT
UNIFORM_VEC2 = rl.ShaderUniformDataType.SHADER_UNIFORM_VEC2
UNIFORM_VEC4 = rl.ShaderUniformDataType.SHADER_UNIFORM_VEC4


class ShaderState:
  _instance: Any = None

  @classmethod
  def get_instance(cls):
    if cls._instance is None:
      cls._instance = cls()
    return cls._instance

  def __init__(self):
    if ShaderState._instance is not None:
      raise Exception("This class is a singleton. Use get_instance() instead.")

    self.initialized = False
    self.shader = None

    # Shader uniform locations
    self.locations = {
      'fillColor': None,
      'useGradient': None,
      'gradientStart': None,
      'gradientEnd': None,
      'gradientColors': None,
      'gradientStops': None,
      'gradientColorCount': None,
      'mvp': None,
    }

    # Pre-allocated FFI objects
    self.fill_color_ptr = rl.ffi.new("float[]", [0.0, 0.0, 0.0, 0.0])
    self.use_gradient_ptr = rl.ffi.new("int[]", [0])
    self.color_count_ptr = rl.ffi.new("int[]", [0])
    self.gradient_colors_ptr = rl.ffi.new("float[]", MAX_GRADIENT_COLORS * 4)
    self.gradient_stops_ptr = rl.ffi.new("float[]", MAX_GRADIENT_COLORS)
    # Last uploaded value per uniform: a program keeps its uniforms, so equal values need no upload.
    self.uniform_values: dict[str, tuple] = {}

  def initialize(self):
    if self.initialized:
      return

    self.shader = rl.load_shader_from_memory(VERTEX_SHADER, FRAGMENT_SHADER)
    self.uniform_values.clear()

    # Cache all uniform locations
    for uniform in self.locations.keys():
      self.locations[uniform] = rl.get_shader_location(self.shader, uniform)

    # Orthographic MVP (origin top-left)
    proj = rl.matrix_ortho(0, gui_app.width, gui_app.height, 0, -1, 1)
    rl.set_shader_value_matrix(self.shader, self.locations['mvp'], proj)

    self.initialized = True

  def cleanup(self):
    if not self.initialized:
      return
    if self.shader:
      rl.unload_shader(self.shader)
      self.shader = None
    self.uniform_values.clear()

    self.initialized = False

  def set_uniform(self, name: str, values: tuple, pointer, kind: int, count: int | None = None) -> None:
    """Upload only a changed value. ``values`` is compared, not the pointer, so in-place color edits still upload."""
    if self.uniform_values.get(name) == values:
      return
    if count is None:
      rl.set_shader_value(self.shader, self.locations[name], pointer, kind)
    else:
      rl.set_shader_value_v(self.shader, self.locations[name], pointer, kind, count)
    self.uniform_values[name] = values


def _configure_shader_color(state: ShaderState, color: Optional[rl.Color],  # noqa: UP045  # rl.Color is a function, so `rl.Color | None` fails
                            gradient: Gradient | None, origin_rect: rl.Rectangle):
  assert (color is not None) != (gradient is not None), "Either color or gradient must be provided"

  use_gradient = 1 if (gradient is not None and len(gradient.colors) >= 1) else 0
  state.use_gradient_ptr[0] = use_gradient
  state.set_uniform('useGradient', (use_gradient,), state.use_gradient_ptr, UNIFORM_INT)

  if use_gradient:
    gradient = cast(Gradient, gradient)
    colors = tuple((c.r, c.g, c.b, c.a) for c in gradient.colors)
    state.color_count_ptr[0] = len(colors)
    for i, (r, g, b, a) in enumerate(colors):
      base = i * 4
      state.gradient_colors_ptr[base:base + 4] = [r / 255.0, g / 255.0, b / 255.0, a / 255.0]
    state.set_uniform('gradientColors', colors, state.gradient_colors_ptr, UNIFORM_VEC4, len(colors))

    stops = tuple(0.0 if s < 0.0 else 1.0 if s > 1.0 else s for s in map(float, gradient.stops))
    for i, s in enumerate(stops):
      state.gradient_stops_ptr[i] = s
    state.set_uniform('gradientStops', stops, state.gradient_stops_ptr, UNIFORM_FLOAT, len(stops))
    state.set_uniform('gradientColorCount', (len(colors),), state.color_count_ptr, UNIFORM_INT)

    # Map normalized start/end to screen pixels
    start_vec = rl.Vector2(origin_rect.x + gradient.start[0] * origin_rect.width, origin_rect.y + gradient.start[1] * origin_rect.height)
    end_vec = rl.Vector2(origin_rect.x + gradient.end[0] * origin_rect.width, origin_rect.y + gradient.end[1] * origin_rect.height)
    state.set_uniform('gradientStart', (start_vec.x, start_vec.y), start_vec, UNIFORM_VEC2)
    state.set_uniform('gradientEnd', (end_vec.x, end_vec.y), end_vec, UNIFORM_VEC2)
  else:
    color = color or rl.WHITE
    state.fill_color_ptr[0:4] = [color.r / 255.0, color.g / 255.0, color.b / 255.0, color.a / 255.0]
    state.set_uniform('fillColor', (color.r, color.g, color.b, color.a), state.fill_color_ptr, UNIFORM_VEC4)


def triangulate(pts: np.ndarray) -> np.ndarray:
  """Only supports simple polygons with two chains (ribbon)."""

  # interleave points to produce a triangle strip
  if len(pts) % 2 != 0:
    pts = pts[:-1]

  half = len(pts) // 2
  tri_strip = np.empty(pts.shape, dtype=np.float32)
  tri_strip[0::2] = pts[:half]
  tri_strip[1::2] = pts[half:][::-1]
  return tri_strip


def draw_polygon(origin_rect: rl.Rectangle, points: np.ndarray,
                 color: Optional[rl.Color] = None, gradient: Gradient | None = None):  # noqa: UP045  # rl.Color is a function, so `rl.Color | None` fails

  """
  Draw a ribbon polygon (two chains) with a triangle strip and gradient.
  - Input must be [L0..Lk-1, Rk-1..R0], even count, no crossings/holes.
  """
  if len(points) < 3:
    return

  assert (color is not None) != (gradient is not None), "Either color or gradient must be provided"
  pts = np.asarray(points)
  assert pts.ndim == 2 and pts.shape[1] == 2, "points must be (N,2)"

  # Triangulate via interleaving
  tri_strip = triangulate(pts)
  vertices = rl.ffi.from_buffer("Vector2 *", tri_strip)

  if gradient is None:
    rl.draw_triangle_strip(vertices, len(tri_strip), color)
    return

  state = ShaderState.get_instance()
  state.initialize()
  _configure_shader_color(state, color, gradient, origin_rect)

  # Draw strip, color here doesn't matter
  rl.begin_shader_mode(state.shader)
  rl.draw_triangle_strip(vertices, len(tri_strip), rl.WHITE)
  rl.end_shader_mode()


def cleanup_shader_resources():
  state = ShaderState.get_instance()
  state.cleanup()