cranpose-render-wgpu 0.9.21

use std::hash::{Hash, Hasher};

use cranpose_render_common::{
    graph::DrawCommandId,
    raster_cache::{LayerRasterCacheKey, RasterScale},
};
use cranpose_ui_graphics::{
    BRUSH_KIND_LINEAR, BRUSH_KIND_RADIAL, BRUSH_KIND_SWEEP, BlendMode, BrushRecord,
    GradientStopRecord, Point, RECORD_KIND_RECT, RecordLane, RecordTables, Rect, ShapeRecord,
    TileMode,
};

use crate::{
    capture_hash::capture_hasher,
    geometry::{canonicalized_scaled_rect, snap_delta_for_anchor},
    render::hash_f32_for_cache,
    run_store::device_clip,
    scene::{CompositorScene, DrawOp, DrawOpKind, RunDraw},
};

pub(crate) struct OpaquePrefix {
    pub(crate) key: LayerRasterCacheKey,
    pub(crate) command: DrawCommandId,
    pub(crate) z_index: usize,
    pub(crate) device_rect: (f32, f32, f32, f32),
}

pub(crate) struct PrefixContext<'a> {
    pub(crate) scene: &'a CompositorScene,
    pub(crate) base: wgpu::LoadOp<wgpu::Color>,
    pub(crate) page_offset: [f32; 2],
    pub(crate) page_size: (u32, u32),
    pub(crate) scale: f32,
}

struct Candidate<'a> {
    run: &'a RunDraw,
    command: DrawCommandId,
    record: ShapeRecord,
    brush: Option<&'a BrushRecord>,
    stops: &'a [GradientStopRecord],
    explicit: &'a [f32],
}

struct Edges {
    left: f32,
    top: f32,
    right: f32,
    bottom: f32,
}

impl Edges {
    fn is_area(&self) -> bool {
        self.right > self.left && self.bottom > self.top
    }

    fn whole(self) -> Option<Self> {
        whole(self.left)?;
        whole(self.top)?;
        whole(self.right)?;
        whole(self.bottom)?;
        self.is_area().then_some(self)
    }

    fn rect(self) -> Rect {
        Rect {
            x: self.left,
            y: self.top,
            width: self.right - self.left,
            height: self.bottom - self.top,
        }
    }
}

fn whole(value: f32) -> Option<f32> {
    (value.is_finite() && value.fract() == 0.0).then_some(value)
}

fn hash_rect<H: Hasher>(rect: Rect, state: &mut H) {
    for value in [rect.x, rect.y, rect.width, rect.height] {
        hash_f32_for_cache(value, state);
    }
}

fn plain_rect(record: &ShapeRecord) -> bool {
    record.blend_mode() == BlendMode::SrcOver
        && record.kind() == RECORD_KIND_RECT
        && !record.is_stroked()
        && record.radii == [0.0; 4]
}

fn brush_stops<'a>(
    tables: &'a RecordTables,
    brush: &BrushRecord,
) -> Option<(&'a [GradientStopRecord], &'a [f32])> {
    let start = brush.stop_start as usize;
    let stops = tables.stops.get(start..start + brush.stop_count as usize)?;
    let explicit = if brush.explicit_len == u32::MAX {
        &[][..]
    } else {
        let start = brush.explicit_start as usize;
        tables
            .explicit_stops
            .get(start..start + brush.explicit_len as usize)?
    };
    Some((stops, explicit))
}

fn candidate<'a>(scene: &'a CompositorScene, op: &DrawOp) -> Option<Candidate<'a>> {
    let DrawOpKind::Run(index) = op.kind else {
        return None;
    };
    let run = &scene.runs[index];
    let command = run.command?;
    if run.placement.paints() || run.placement.clip_rounded() {
        return None;
    }
    let tables = run.tables();
    let segment = tables.segments.get(run.segments.start as usize)?;
    let record = tables.shapes.get(segment.start as usize)?;
    if segment.lane != RecordLane::Shapes
        || segment.blend != BlendMode::SrcOver
        || !plain_rect(&record)
    {
        return None;
    }
    let brush = match record.brush {
        0 => None,
        index => Some(tables.brushes.get(index as usize - 1)?),
    };
    let (stops, explicit) = match brush {
        Some(brush) => brush_stops(tables, brush)?,
        None => (&[][..], &[][..]),
    };
    Some(Candidate {
        run,
        command,
        record,
        brush,
        stops,
        explicit,
    })
}

fn is_opaque(candidate: &Candidate<'_>) -> bool {
    match candidate.brush {
        None => candidate.record.color[3] == 1.0,
        Some(brush) => {
            matches!(
                brush.kind,
                BRUSH_KIND_LINEAR | BRUSH_KIND_RADIAL | BRUSH_KIND_SWEEP
            ) && brush.tile_mode == TileMode::Clamp as u32
                && !candidate.stops.is_empty()
                && candidate.stops.iter().all(|stop| stop.color[3] == 1.0)
        }
    }
}

fn device_edges(rect: [f32; 4], offset: Point, scale: f32, canonicalize: bool) -> Edges {
    let [left, top, right, bottom] =
        crate::run_geometry::device_shape_edges(rect, 0.0, offset, scale, canonicalize);
    Edges {
        left,
        top,
        right,
        bottom,
    }
}

fn clip_contains(clip: Rect, edges: &Edges, scale: f32, canonicalize: bool) -> bool {
    let device = if canonicalize {
        canonicalized_scaled_rect(clip, scale)
    } else {
        Rect {
            x: clip.x * scale,
            y: clip.y * scale,
            width: clip.width * scale,
            height: clip.height * scale,
        }
    };
    device.x <= edges.left
        && device.y <= edges.top
        && device.x + device.width >= edges.right
        && device.y + device.height >= edges.bottom
}

fn clamp_to_page(edges: Edges, context: &PrefixContext<'_>) -> Option<Edges> {
    let [page_x, page_y] = context.page_offset;
    whole(page_x)?;
    whole(page_y)?;
    let clamped = Edges {
        left: edges.left.max(page_x),
        top: edges.top.max(page_y),
        right: edges.right.min(page_x + context.page_size.0 as f32),
        bottom: edges.bottom.min(page_y + context.page_size.1 as f32),
    };
    clamped.is_area().then_some(clamped)
}

fn prefix_hash(
    candidate: &Candidate<'_>,
    snap: Point,
    context: &PrefixContext<'_>,
    edges: &Edges,
) -> u64 {
    let mut hasher = capture_hasher();
    bytemuck::bytes_of(&candidate.record).hash(&mut hasher);
    if let Some(brush) = candidate.brush {
        bytemuck::bytes_of(brush).hash(&mut hasher);
        bytemuck::cast_slice::<_, u8>(candidate.stops).hash(&mut hasher);
        for position in candidate.explicit {
            hash_f32_for_cache(*position, &mut hasher);
        }
    }
    let placement = &candidate.run.placement;
    for value in [
        placement.offset.x,
        placement.offset.y,
        snap.x,
        snap.y,
        context.scale,
    ] {
        hash_f32_for_cache(value, &mut hasher);
    }
    match placement.clip {
        Some(clip) => {
            1u8.hash(&mut hasher);
            hash_rect(clip, &mut hasher);
        }
        None => 0u8.hash(&mut hasher),
    }
    match context.base {
        wgpu::LoadOp::Clear(color) => {
            1u8.hash(&mut hasher);
            for channel in [color.r, color.g, color.b, color.a] {
                channel.to_bits().hash(&mut hasher);
            }
        }
        _ => 0u8.hash(&mut hasher),
    }
    let [page_x, page_y] = context.page_offset;
    for value in [
        edges.left,
        edges.top,
        edges.right,
        edges.bottom,
        page_x,
        page_y,
    ] {
        hash_f32_for_cache(value, &mut hasher);
    }
    hasher.finish()
}

/// The candidate's device edges and snap, when its placement puts it on
/// whole pixels and no clip cuts into it.
fn placed_edges(candidate: &Candidate<'_>, scale: f32) -> Option<(Edges, Point)> {
    let placement = &candidate.run.placement;
    let snap = placement
        .snap_anchor
        .map(|anchor| snap_delta_for_anchor(anchor, scale))
        .unwrap_or_default();
    let canonicalize = placement.snap_anchor.is_some();
    let offset = Point {
        x: placement.offset.x + snap.x,
        y: placement.offset.y + snap.y,
    };
    let edges = device_edges(candidate.record.rect, offset, scale, canonicalize).whole()?;
    if placement
        .clip
        .is_some_and(|clip| !clip_contains(clip, &edges, scale, canonicalize))
    {
        return None;
    }
    Some((edges, snap))
}

/// The color a flush's first op paints its whole page with, and that op's
/// z, when the op starts with one solid opaque rectangle covering the page:
/// the pass clears to the color instead of shading and blending every pixel
/// of the page, which a tiled GPU does for free.
pub(crate) fn page_fill_color(
    context: &PrefixContext<'_>,
    ops: &[DrawOp],
) -> Option<(wgpu::Color, usize)> {
    let op = ops.first()?;
    let candidate = candidate(context.scene, op)?;
    if candidate.brush.is_some() || !is_opaque(&candidate) {
        return None;
    }
    let (edges, _) = placed_edges(&candidate, context.scale)?;
    let [page_x, page_y] = context.page_offset;
    let covers = edges.left <= page_x
        && edges.top <= page_y
        && edges.right >= page_x + context.page_size.0 as f32
        && edges.bottom >= page_y + context.page_size.1 as f32;
    let [r, g, b, a] = candidate.record.color.map(f64::from);
    covers.then_some((wgpu::Color { r, g, b, a }, op.z_index))
}

pub(crate) fn capture_solid_rect(
    scene: &CompositorScene,
    op: &DrawOp,
    capture: Rect,
    scale: f32,
) -> Option<([f32; 4], Rect, Option<Rect>)> {
    let candidate = candidate(scene, op)?;
    let run = candidate.run;
    if run.segments.end - run.segments.start != 1
        || run
            .tables()
            .segments
            .get(run.segments.start as usize)?
            .count
            != 1
        || candidate.brush.is_some()
        || !is_opaque(&candidate)
    {
        return None;
    }
    let placement = &run.placement;
    let snap = placement
        .snap_anchor
        .map(|anchor| snap_delta_for_anchor(anchor, scale))
        .unwrap_or_default();
    let painted = device_edges(
        candidate.record.rect,
        Point::new(placement.offset.x + snap.x, placement.offset.y + snap.y),
        scale,
        placement.snap_anchor.is_some(),
    )
    .rect()
    .intersect(capture)?;
    let clip = device_clip(placement, scale).map(|[x, y, width, height]| {
        let device = Rect {
            x,
            y,
            width,
            height,
        };
        device.intersect(capture).unwrap_or(device)
    });
    Some((candidate.record.color, painted, clip))
}

pub(crate) fn opaque_prefix(context: &PrefixContext<'_>, ops: &[DrawOp]) -> Option<OpaquePrefix> {
    let op = ops.first()?;
    let candidate = candidate(context.scene, op)?;
    // A solid fill draws for less than copying it back from a cached texture
    // costs, and the cache would hold a page-sized texture for it: on a
    // Pixel 9 Pro caching the benchmarks' solid backgrounds cost 13 to 27 MB
    // of PSS and saved no GPU clock. Only gradients, whose per-pixel stop
    // lookups the copy saves, are cached.
    if candidate.brush.is_none() || !is_opaque(&candidate) {
        return None;
    }
    let scale = context.scale;
    let (edges, snap) = placed_edges(&candidate, scale)?;
    let edges = clamp_to_page(edges, context)?;
    let hash = prefix_hash(&candidate, snap, context, &edges);
    let width = edges.right - edges.left;
    let height = edges.bottom - edges.top;
    Some(OpaquePrefix {
        key: LayerRasterCacheKey::prefix_snapshot(
            hash,
            1,
            Rect {
                x: edges.left,
                y: edges.top,
                width,
                height,
            },
            (width as u32, height as u32),
            RasterScale::from_scale(scale),
        ),
        command: candidate.command,
        z_index: op.z_index,
        device_rect: (edges.left, edges.top, width, height),
    })
}