phosphor-crt 0.1.0

const EPS: f32 = 1e-2;

@group(0) @binding(0)
var<uniform> viewport: ViewportUniform;


struct ViewportUniform {
    transform: mat3x3<f32>,  // the transform from scene coordinates to pixel coordinates
    resolution: vec2<u32>,  // the size of the viewport
    beam_radius: f32,  // the size of the beam (in pixels)
    intensity: f32,  // the global intensity factor
    gamma: f32,  // gamma correction exponent
    decay_per_sample: f32,
    num_samples: u32,
};

struct VertexOutput {
    @builtin(position) position: vec4<f32>,

    /*
    uv coordinates in a quad:
    (unit = 1 beam radius)
         v=+1 ▲──────────────────────────────────┐
              │                                  │
         v=0  │     #======================#     │
              │     .                      .     │
         v=-1 └─────.──────────────────────.─────►
              u=-1  .                      .     u=len+1
                    u=0                    u=len
    */
    @location(1) uv: vec2<f32>,
    @location(2) len: f32, // length of the segment, shared by all 4 vertices

    @location(3) decay: f32,
};

fn segment_index(vertex_index: u32, instance_index: u32) -> u32{
    let draw_odd_lines = u32(vertex_index >= 4);
    return 2 * instance_index + draw_odd_lines;
}

@vertex
fn vs_yt_main(
    @location(0) current_value: f32,
    @location(1) next_value: f32,
    @builtin(vertex_index) vertex_index: u32,
    @builtin(instance_index) instance_index: u32
) -> VertexOutput {
    // start and end points in scene coordinates
    let start_scene = vec2f(
        f32(segment_index(vertex_index, instance_index)),
        current_value
    );
    let end_scene = vec2f(
        f32(segment_index(vertex_index, instance_index) + 1),
        next_value
    );

    return vs_common(
        start_scene,
        end_scene,
        vertex_index,
        instance_index,
    );
}

@vertex
fn vs_xy_main(
    @location(0) current: vec2f,
    @location(1) next: vec2f,
    @builtin(vertex_index) vertex_index: u32,
    @builtin(instance_index) instance_index: u32
) -> VertexOutput {
    return vs_common(
        current,
        next,
        vertex_index,
        instance_index,
    );
}


fn vs_common(
    start_scene: vec2f,
    end_scene: vec2f,
    vertex_index: u32,
    instance_index: u32
) -> VertexOutput {
    let corner_index = vertex_index & 3;

    // map to pixel coordinates (where the beam has a 1:1 aspect ratio)
    let start = (viewport.transform * vec3f(start_scene, 1.)).xy;
    let end = (viewport.transform * vec3f(end_scene, 1.)).xy;

    var dir = (end - start) / viewport.beam_radius;
    var len = length(dir);  // segment length in beam radii
    if len > EPS {
        dir = dir / len;
    } else {
        // If the segment is too short, just draw a square
        dir = vec2(1.0, 0.0);
        len = 0.;
    }

    var normal = vec2(-dir.y, dir.x);

    var current: vec2f;
    var tangential_offset: f32;
    var u: f32;
    if corner_index >= 2u {
        current = end;
        tangential_offset = 1.;
        u = len + 1.;
    } else {
        current = start;
        tangential_offset = -1.;
        u = -1.;
    }
    let lateral_offset = f32(i32(corner_index & 1) * 2 - 1);

    var vertex_pos_pixel = current + (tangential_offset * dir + lateral_offset * normal) * viewport.beam_radius;
    let vertex_pos_clip = vec2f(
        vertex_pos_pixel.x / f32(viewport.resolution.x) * 2. - 1.,
        vertex_pos_pixel.y / f32(viewport.resolution.y) * 2. - 1.
    );

    var out: VertexOutput;
    out.position = vec4(
        vertex_pos_clip,
        0.,
        1.
    );
    out.uv = vec2(u, lateral_offset);
    out.len = len;
    // Calculate decay: earlier samples (lower index) should be dimmer
    // current_sample_position goes from ≈0 (first sample) to ≈num_samples-1 (last sample)
    var current_sample_position = f32(segment_index(vertex_index, instance_index));
    if len > EPS {
        current_sample_position += u / len;
    }
    out.decay = viewport.decay_per_sample * (f32(viewport.num_samples) - 1.0 - current_sample_position);

    return out;
}

const HALF_SQRT_PI: f32 =  0.88622692545275801365;

fn gaussian(x: f32) -> f32 {
    return smoothstep(1.6, 0., abs(x));
}

fn erf(x: f32) -> f32 {
    return smoothstep(-2., 2., x) * 2. - 1.;
}


@fragment
fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
    let len = in.len;
    var value: f32 = 1.;

    let sigma = 0.25;  // sigma for gaussian beam (relative to beam_radius)

    if len/sigma < EPS {
        // If the beam segment is too short, just calculate intensity at the position.
        value *= gaussian(length(in.uv) / sigma);
    } else {
        let x = in.uv.x;
        let y = in.uv.y;
        value *= HALF_SQRT_PI * sigma / len * (erf(x / sigma) - erf((x - len) / sigma));
        value *= gaussian(y / sigma);
    }
    value *= exp(-in.decay);

    return vec4(value, 0., 0., 1.);
}