retro3d/shaders/retro3d.frag.msl

#include <metal_stdlib>
using namespace metal;

struct Uniforms {
  float4x4 mvp;
  float4x4 model;
  float4x4 view;
  float4x4 projection;
  float4 ambient;       // rgb, unused
  float4 light_dir;     // xyz normalized, unused
  float4 light_color;   // rgb, intensity
  float4 fog_params;    // near, far, unused, enabled
  float4 fog_color;     // rgb, unused
  float4 tint;          // rgba (base_color_factor from glTF)
  float4 style_params;  // style_id, vertex_snap, affine, dither
  float4 resolution;    // w, h, unused, unused
  float4 material_params; // alpha_mode (0=opaque,1=mask,2=blend), alpha_cutoff, unlit, unused
};

struct FragmentIn {
  float4 position [[position]];
  float3 world_normal;
  float2 uv;
  float4 color;
  float fog_factor;
  float3 noperspective_uv;
};

// Dither pattern for Saturn style (4x4 Bayer matrix)
constant float dither_matrix[16] = {
  0.0/16.0,  8.0/16.0,  2.0/16.0, 10.0/16.0,
  12.0/16.0, 4.0/16.0, 14.0/16.0,  6.0/16.0,
  3.0/16.0, 11.0/16.0,  1.0/16.0,  9.0/16.0,
  15.0/16.0, 7.0/16.0, 13.0/16.0,  5.0/16.0
};

float get_dither(float2 pos) {
  int x = int(pos.x) % 4;
  int y = int(pos.y) % 4;
  return dither_matrix[y * 4 + x];
}

// Quantize color to lower bit depth
float3 quantize_color(float3 color, float bits) {
  float levels = pow(2.0, bits) - 1.0;
  return floor(color * levels + 0.5) / levels;
}

fragment float4 fragment_main(
  FragmentIn in [[stage_in]],
  constant Uniforms &uniforms [[buffer(1)]],
  texture2d<float> tex [[texture(0)]],
  sampler samp [[sampler(0)]]
) {
  float style_id = uniforms.style_params.x;
  float affine = uniforms.style_params.z;
  float dither = uniforms.style_params.w;

  // Material params
  float alpha_mode = uniforms.material_params.x;   // 0=opaque, 1=mask, 2=blend
  float alpha_cutoff = uniforms.material_params.y; // default 0.5
  float unlit = uniforms.material_params.z;

  // Get UV coordinates
  float2 uv;
  if (affine > 0.5) {
    // Affine texturing (PS1 style) - divide by interpolated w
    uv = in.noperspective_uv.xy / in.noperspective_uv.z;
  } else {
    uv = in.uv;
  }

  // Sample texture
  float4 tex_color = tex.sample(samp, uv);

  // glTF spec: final color = vertexColor * baseColorFactor * baseColorTexture
  // tint = base_color_factor from material
  float4 base_color = in.color * uniforms.tint * tex_color;

  // Alpha handling based on alpha_mode
  float alpha = base_color.a;

  // MASK mode: discard fragments below cutoff
  if (alpha_mode > 0.5 && alpha_mode < 1.5) {
    if (alpha < alpha_cutoff) {
      discard_fragment();
    }
    alpha = 1.0; // MASK mode outputs fully opaque or discards
  }

  // OPAQUE mode: ignore alpha entirely
  if (alpha_mode < 0.5) {
    alpha = 1.0;
  }

  // BLEND mode: alpha is used as-is (alpha_mode >= 1.5)

  float3 final_color;

  if (unlit > 0.5) {
    // Unlit material - no lighting calculation
    final_color = base_color.rgb;
  } else {
    // Lighting calculation
    float3 normal = normalize(in.world_normal);
    float3 light_dir = normalize(uniforms.light_dir.xyz);
    float ndotl = max(dot(normal, light_dir), 0.0);

    float3 ambient = uniforms.ambient.rgb;
    float3 diffuse = uniforms.light_color.rgb * uniforms.light_color.w * ndotl;
    float3 lighting = ambient + diffuse;

    // Apply lighting
    final_color = base_color.rgb * lighting;
  }

  // Style-specific processing
  if (style_id < 0.5) {
    // PS1 style: 15-bit color (5 bits per channel)
    final_color = quantize_color(final_color, 5.0);
  } else if (style_id < 1.5) {
    // N64 style: smoother, 16-bit color with bilinear filtering
    // (filtering is handled by sampler, just quantize slightly)
    final_color = quantize_color(final_color, 5.0);
  } else {
    // Saturn style: dithered, flat shaded look
    if (dither > 0.5) {
      float d = get_dither(in.position.xy);
      // Add dither before quantization
      final_color += (d - 0.5) * 0.1;
    }
    final_color = quantize_color(final_color, 5.0);
  }

  // Apply fog
  float3 fog_color = uniforms.fog_color.rgb;
  final_color = mix(fog_color, final_color, in.fog_factor);

  return float4(final_color, alpha);
}

// Unlit fragment shader for sprites/UI
fragment float4 fragment_unlit(
  FragmentIn in [[stage_in]],
  constant Uniforms &uniforms [[buffer(1)]],
  texture2d<float> tex [[texture(0)]],
  sampler samp [[sampler(0)]]
) {
  float affine = uniforms.style_params.z;

  float2 uv;
  if (affine > 0.5) {
    uv = in.noperspective_uv.xy / in.noperspective_uv.z;
  } else {
    uv = in.uv;
  }

  float4 tex_color = tex.sample(samp, uv);
  float4 color = in.color * tex_color * uniforms.tint;

  // Alpha test for sprites
  if (color.a < 0.1) {
    discard_fragment();
  }

  return color;
}