gpu/painter/src/lib.rs

// Copyright 2022 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

use std::{borrow::Cow, mem, slice};

use bytemuck::{Pod, Zeroable};
pub use wgpu;
use wgpu::util::DeviceExt;

// mod compact_pixel_segment;

// pub use compact_pixel_segment::CompactPixelSegment;

pub const TILE_WIDTH: usize = 16;
pub const TILE_HEIGHT: usize = 4;

#[derive(Clone, Copy, Debug, Pod, Zeroable)]
#[repr(C)]
pub struct Style {
    pub fill_rule: u32,
    pub color: Color,
    pub blend_mode: u32,
}

#[derive(Clone, Copy, Debug, Pod, Zeroable)]
#[repr(C)]
pub struct Color {
    pub r: f32,
    pub g: f32,
    pub b: f32,
    pub a: f32,
}

#[derive(Copy, Clone)]
#[repr(C)]
pub struct Config {
    pub segments_len: u32,
    pub width: u32,
    pub height: u32,
    pub _padding: u32,
    pub clear_color: Color,
}

impl Config {
    pub fn as_byte_slice(&self) -> &[u8] {
        unsafe { slice::from_raw_parts((self as *const _) as *const u8, mem::size_of::<Self>()) }
    }
}

#[derive(Debug)]
pub struct PaintContext {
    pub pipeline: wgpu::ComputePipeline,
    pub bind_group_layout: wgpu::BindGroupLayout,
}

pub fn init(device: &wgpu::Device) -> PaintContext {
    let module = device.create_shader_module(&wgpu::ShaderModuleDescriptor {
        label: None,
        source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(include_str!("paint.wgsl"))),
    });

    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
        label: None,
        layout: None,
        module: &module,
        entry_point: "paint",
    });

    let bind_group_layout = pipeline.get_bind_group_layout(0);

    PaintContext { pipeline, bind_group_layout }
}

pub fn encode(
    device: &wgpu::Device,
    encoder: &mut wgpu::CommandEncoder,
    context: &PaintContext,
    texture_view: &wgpu::TextureView,
    segments_buffer: &wgpu::Buffer,
    segments_len: u32,
    style_indices_buffer: &wgpu::Buffer,
    styles_buffer: &wgpu::Buffer,
    width: u32,
    height: u32,
    clear_color: Color,
    timestamp: Option<(&wgpu::QuerySet, u32, u32)>,
) {
    let config = Config { segments_len, width, height, clear_color, _padding: 0 };

    let config_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
        label: None,
        contents: config.as_byte_slice(),
        usage: wgpu::BufferUsages::UNIFORM,
    });

    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &context.bind_group_layout,
        entries: &[
            wgpu::BindGroupEntry { binding: 0, resource: config_buffer.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 1, resource: segments_buffer.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 2, resource: style_indices_buffer.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 3, resource: styles_buffer.as_entire_binding() },
            wgpu::BindGroupEntry {
                binding: 4,
                resource: wgpu::BindingResource::TextureView(&texture_view),
            },
        ],
    });

    let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { label: None });
    pass.set_pipeline(&context.pipeline);
    pass.set_bind_group(0, &bind_group, &[]);

    if let Some((timestamp, start_index, _)) = timestamp {
        pass.write_timestamp(&timestamp, start_index);
    }

    pass.dispatch_workgroups(((height as usize + TILE_HEIGHT - 1) / TILE_HEIGHT) as u32, 1, 1);

    if let Some((timestamp, _, end_index)) = timestamp {
        pass.write_timestamp(&timestamp, end_index);
    }
}