path: root/wgpu/src/triangle.rs

//! Draw meshes of triangles.
mod msaa;

use crate::buffer::r#static::Buffer;
use crate::layer::mesh::{self, Mesh};
use crate::settings;
use crate::Transformation;

use iced_graphics::Size;

#[cfg(feature = "tracing")]
use tracing::info_span;

#[derive(Debug)]
pub struct Pipeline {
    blit: Option<msaa::Blit>,
    solid: solid::Pipeline,

    /// Gradients are currently not supported on WASM targets due to their need of storage buffers.
    #[cfg(not(target_arch = "wasm32"))]
    gradient: gradient::Pipeline,

    layers: Vec<Layer>,
    prepare_layer: usize,
}

#[derive(Debug)]
struct Layer {
    index_buffer: Buffer<u32>,
    index_strides: Vec<u32>,
    solid: solid::Layer,

    #[cfg(not(target_arch = "wasm32"))]
    gradient: gradient::Layer,
}

impl Layer {
    fn new(
        device: &wgpu::Device,
        solid: &solid::Pipeline,
        #[cfg(not(target_arch = "wasm32"))] gradient: &gradient::Pipeline,
    ) -> Self {
        Self {
            index_buffer: Buffer::new(
                device,
                "iced_wgpu::triangle index buffer",
                wgpu::BufferUsages::INDEX | wgpu::BufferUsages::COPY_DST,
            ),
            index_strides: Vec::new(),
            solid: solid::Layer::new(device, &solid.constants_layout),

            #[cfg(not(target_arch = "wasm32"))]
            gradient: gradient::Layer::new(device, &gradient.constants_layout),
        }
    }

    fn prepare(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        solid: &solid::Pipeline,
        #[cfg(not(target_arch = "wasm32"))] gradient: &gradient::Pipeline,
        meshes: &[Mesh<'_>],
        transformation: Transformation,
    ) {
        // Count the total amount of vertices & indices we need to handle
        let count = mesh::attribute_count_of(meshes);

        // Then we ensure the current attribute buffers are big enough, resizing if necessary.
        // We are not currently using the return value of these functions as we have no system in
        // place to calculate mesh diff, or to know whether or not that would be more performant for
        // the majority of use cases. Therefore we will write GPU data every frame (for now).
        let _ = self.index_buffer.resize(device, count.indices);
        let _ = self.solid.vertices.resize(device, count.solid_vertices);

        #[cfg(not(target_arch = "wasm32"))]
        let _ = self
            .gradient
            .vertices
            .resize(device, count.gradient_vertices);

        // Prepare dynamic buffers & data store for writing
        self.index_buffer.clear();
        self.index_strides.clear();
        self.solid.vertices.clear();
        self.solid.uniforms.clear();

        #[cfg(not(target_arch = "wasm32"))]
        {
            self.gradient.uniforms.clear();
            self.gradient.vertices.clear();
            self.gradient.storage.clear();
        }

        let mut solid_vertex_offset = 0;
        let mut index_offset = 0;

        #[cfg(not(target_arch = "wasm32"))]
        let mut gradient_vertex_offset = 0;

        for mesh in meshes {
            let origin = mesh.origin();
            let indices = mesh.indices();

            let transform =
                transformation * Transformation::translate(origin.x, origin.y);

            let new_index_offset =
                self.index_buffer.write(queue, index_offset, indices);

            index_offset += new_index_offset;
            self.index_strides.push(indices.len() as u32);

            //push uniform data to CPU buffers
            match mesh {
                Mesh::Solid { buffers, .. } => {
                    self.solid.uniforms.push(&solid::Uniforms::new(transform));

                    let written_bytes = self.solid.vertices.write(
                        queue,
                        solid_vertex_offset,
                        &buffers.vertices,
                    );

                    solid_vertex_offset += written_bytes;
                }
                #[cfg(not(target_arch = "wasm32"))]
                Mesh::Gradient {
                    buffers, gradient, ..
                } => {
                    let written_bytes = self.gradient.vertices.write(
                        queue,
                        gradient_vertex_offset,
                        &buffers.vertices,
                    );

                    gradient_vertex_offset += written_bytes;

                    match gradient {
                        iced_graphics::Gradient::Linear(linear) => {
                            use glam::{IVec4, Vec4};

                            let start_offset = self.gradient.color_stop_offset;
                            let end_offset = (linear.color_stops.len() as i32)
                                + start_offset
                                - 1;

                            self.gradient.uniforms.push(&gradient::Uniforms {
                                transform: transform.into(),
                                direction: Vec4::new(
                                    linear.start.x,
                                    linear.start.y,
                                    linear.end.x,
                                    linear.end.y,
                                ),
                                stop_range: IVec4::new(
                                    start_offset,
                                    end_offset,
                                    0,
                                    0,
                                ),
                            });

                            self.gradient.color_stop_offset = end_offset + 1;

                            let stops: Vec<gradient::ColorStop> = linear
                                .color_stops
                                .iter()
                                .map(|stop| {
                                    let [r, g, b, a] = stop.color.into_linear();

                                    gradient::ColorStop {
                                        offset: stop.offset,
                                        color: Vec4::new(r, g, b, a),
                                    }
                                })
                                .collect();

                            self.gradient
                                .color_stops_pending_write
                                .color_stops
                                .extend(stops);
                        }
                    }
                }
                #[cfg(target_arch = "wasm32")]
                Mesh::Gradient { .. } => {}
            }
        }

        // Write uniform data to GPU
        if count.solid_vertices > 0 {
            let uniforms_resized = self.solid.uniforms.resize(device);

            if uniforms_resized {
                self.solid.constants = solid::Layer::bind_group(
                    device,
                    self.solid.uniforms.raw(),
                    &solid.constants_layout,
                )
            }

            self.solid.uniforms.write(queue);
        }

        #[cfg(not(target_arch = "wasm32"))]
        if count.gradient_vertices > 0 {
            // First write the pending color stops to the CPU buffer
            self.gradient
                .storage
                .push(&self.gradient.color_stops_pending_write);

            // Resize buffers if needed
            let uniforms_resized = self.gradient.uniforms.resize(device);
            let storage_resized = self.gradient.storage.resize(device);

            if uniforms_resized || storage_resized {
                self.gradient.constants = gradient::Layer::bind_group(
                    device,
                    self.gradient.uniforms.raw(),
                    self.gradient.storage.raw(),
                    &gradient.constants_layout,
                );
            }

            // Write to GPU
            self.gradient.uniforms.write(queue);
            self.gradient.storage.write(queue);

            // Cleanup
            self.gradient.color_stop_offset = 0;
            self.gradient.color_stops_pending_write.color_stops.clear();
        }
    }

    fn render<'a>(
        &'a self,
        solid: &'a solid::Pipeline,
        #[cfg(not(target_arch = "wasm32"))] gradient: &'a gradient::Pipeline,
        meshes: &[Mesh<'_>],
        scale_factor: f32,
        render_pass: &mut wgpu::RenderPass<'a>,
    ) {
        let mut num_solids = 0;
        #[cfg(not(target_arch = "wasm32"))]
        let mut num_gradients = 0;
        let mut last_is_solid = None;

        for (index, mesh) in meshes.iter().enumerate() {
            let clip_bounds = (mesh.clip_bounds() * scale_factor).snap();

            render_pass.set_scissor_rect(
                clip_bounds.x,
                clip_bounds.y,
                clip_bounds.width,
                clip_bounds.height,
            );

            match mesh {
                Mesh::Solid { .. } => {
                    if !last_is_solid.unwrap_or(false) {
                        render_pass.set_pipeline(&solid.pipeline);

                        last_is_solid = Some(true);
                    }

                    render_pass.set_bind_group(
                        0,
                        &self.solid.constants,
                        &[self.solid.uniforms.offset_at_index(num_solids)],
                    );

                    render_pass.set_vertex_buffer(
                        0,
                        self.solid.vertices.slice_from_index(num_solids),
                    );

                    num_solids += 1;
                }
                #[cfg(not(target_arch = "wasm32"))]
                Mesh::Gradient { .. } => {
                    if last_is_solid.unwrap_or(true) {
                        render_pass.set_pipeline(&gradient.pipeline);

                        last_is_solid = Some(false);
                    }

                    render_pass.set_bind_group(
                        0,
                        &self.gradient.constants,
                        &[self
                            .gradient
                            .uniforms
                            .offset_at_index(num_gradients)],
                    );

                    render_pass.set_vertex_buffer(
                        0,
                        self.gradient.vertices.slice_from_index(num_gradients),
                    );

                    num_gradients += 1;
                }
                #[cfg(target_arch = "wasm32")]
                Mesh::Gradient { .. } => {}
            };

            render_pass.set_index_buffer(
                self.index_buffer.slice_from_index(index),
                wgpu::IndexFormat::Uint32,
            );

            render_pass.draw_indexed(0..self.index_strides[index], 0, 0..1);
        }
    }
}

impl Pipeline {
    pub fn new(
        device: &wgpu::Device,
        format: wgpu::TextureFormat,
        antialiasing: Option<settings::Antialiasing>,
    ) -> Pipeline {
        Pipeline {
            blit: antialiasing.map(|a| msaa::Blit::new(device, format, a)),
            solid: solid::Pipeline::new(device, format, antialiasing),

            #[cfg(not(target_arch = "wasm32"))]
            gradient: gradient::Pipeline::new(device, format, antialiasing),

            layers: Vec::new(),
            prepare_layer: 0,
        }
    }

    pub fn prepare(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        meshes: &[Mesh<'_>],
        transformation: Transformation,
    ) {
        #[cfg(feature = "tracing")]
        let _ = info_span!("Wgpu::Triangle", "PREPARE").entered();

        if self.layers.len() <= self.prepare_layer {
            self.layers.push(Layer::new(
                device,
                &self.solid,
                #[cfg(not(target_arch = "wasm32"))]
                &self.gradient,
            ));
        }

        let layer = &mut self.layers[self.prepare_layer];
        layer.prepare(
            device,
            queue,
            &self.solid,
            #[cfg(not(target_arch = "wasm32"))]
            &self.gradient,
            meshes,
            transformation,
        );

        self.prepare_layer += 1;
    }

    pub fn render(
        &mut self,
        device: &wgpu::Device,
        encoder: &mut wgpu::CommandEncoder,
        target: &wgpu::TextureView,
        layer: usize,
        target_size: Size<u32>,
        meshes: &[Mesh<'_>],
        scale_factor: f32,
    ) {
        #[cfg(feature = "tracing")]
        let _ = info_span!("Wgpu::Triangle", "DRAW").entered();

        // Configure render pass
        {
            let (attachment, resolve_target, load) = if let Some(blit) =
                &mut self.blit
            {
                let (attachment, resolve_target) =
                    blit.targets(device, target_size.width, target_size.height);

                (
                    attachment,
                    Some(resolve_target),
                    wgpu::LoadOp::Clear(wgpu::Color::TRANSPARENT),
                )
            } else {
                (target, None, wgpu::LoadOp::Load)
            };

            #[cfg(feature = "tracing")]
            let _ = info_span!("Wgpu::Triangle", "BEGIN_RENDER_PASS").enter();

            let mut render_pass =
                encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                    label: Some("iced_wgpu::triangle render pass"),
                    color_attachments: &[Some(
                        wgpu::RenderPassColorAttachment {
                            view: attachment,
                            resolve_target,
                            ops: wgpu::Operations { load, store: true },
                        },
                    )],
                    depth_stencil_attachment: None,
                });

            let layer = &mut self.layers[layer];

            layer.render(
                &self.solid,
                #[cfg(not(target_arch = "wasm32"))]
                &self.gradient,
                meshes,
                scale_factor,
                &mut render_pass,
            );
        }

        if let Some(blit) = &mut self.blit {
            blit.draw(encoder, target);
        }
    }

    pub fn end_frame(&mut self) {
        self.prepare_layer = 0;
    }
}

fn fragment_target(
    texture_format: wgpu::TextureFormat,
) -> Option<wgpu::ColorTargetState> {
    Some(wgpu::ColorTargetState {
        format: texture_format,
        blend: Some(wgpu::BlendState::ALPHA_BLENDING),
        write_mask: wgpu::ColorWrites::ALL,
    })
}

fn primitive_state() -> wgpu::PrimitiveState {
    wgpu::PrimitiveState {
        topology: wgpu::PrimitiveTopology::TriangleList,
        front_face: wgpu::FrontFace::Cw,
        ..Default::default()
    }
}

fn multisample_state(
    antialiasing: Option<settings::Antialiasing>,
) -> wgpu::MultisampleState {
    wgpu::MultisampleState {
        count: antialiasing.map(|a| a.sample_count()).unwrap_or(1),
        mask: !0,
        alpha_to_coverage_enabled: false,
    }
}

mod solid {
    use crate::buffer::dynamic;
    use crate::buffer::r#static::Buffer;
    use crate::settings;
    use crate::triangle;
    use encase::ShaderType;
    use iced_graphics::primitive;
    use iced_graphics::Transformation;

    #[derive(Debug)]
    pub struct Pipeline {
        pub pipeline: wgpu::RenderPipeline,
        pub constants_layout: wgpu::BindGroupLayout,
    }

    #[derive(Debug)]
    pub struct Layer {
        pub vertices: Buffer<primitive::ColoredVertex2D>,
        pub uniforms: dynamic::Buffer<Uniforms>,
        pub constants: wgpu::BindGroup,
    }

    impl Layer {
        pub fn new(
            device: &wgpu::Device,
            constants_layout: &wgpu::BindGroupLayout,
        ) -> Self {
            let vertices = Buffer::new(
                device,
                "iced_wgpu::triangle::solid vertex buffer",
                wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
            );

            let uniforms = dynamic::Buffer::uniform(
                device,
                "iced_wgpu::triangle::solid uniforms",
            );

            let constants =
                Self::bind_group(device, uniforms.raw(), constants_layout);

            Self {
                vertices,
                uniforms,
                constants,
            }
        }

        pub fn bind_group(
            device: &wgpu::Device,
            buffer: &wgpu::Buffer,
            layout: &wgpu::BindGroupLayout,
        ) -> wgpu::BindGroup {
            device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("iced_wgpu::triangle::solid bind group"),
                layout,
                entries: &[wgpu::BindGroupEntry {
                    binding: 0,
                    resource: wgpu::BindingResource::Buffer(
                        wgpu::BufferBinding {
                            buffer,
                            offset: 0,
                            size: Some(Uniforms::min_size()),
                        },
                    ),
                }],
            })
        }
    }

    #[derive(Debug, Clone, Copy, ShaderType)]
    pub struct Uniforms {
        transform: glam::Mat4,
    }

    impl Uniforms {
        pub fn new(transform: Transformation) -> Self {
            Self {
                transform: transform.into(),
            }
        }
    }

    impl Pipeline {
        /// Creates a new [SolidPipeline] using `solid.wgsl` shader.
        pub fn new(
            device: &wgpu::Device,
            format: wgpu::TextureFormat,
            antialiasing: Option<settings::Antialiasing>,
        ) -> Self {
            let constants_layout = device.create_bind_group_layout(
                &wgpu::BindGroupLayoutDescriptor {
                    label: Some("iced_wgpu::triangle::solid bind group layout"),
                    entries: &[wgpu::BindGroupLayoutEntry {
                        binding: 0,
                        visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
                        ty: wgpu::BindingType::Buffer {
                            ty: wgpu::BufferBindingType::Uniform,
                            has_dynamic_offset: true,
                            min_binding_size: Some(Uniforms::min_size()),
                        },
                        count: None,
                    }],
                },
            );

            let layout = device.create_pipeline_layout(
                &wgpu::PipelineLayoutDescriptor {
                    label: Some("iced_wgpu::triangle::solid pipeline layout"),
                    bind_group_layouts: &[&constants_layout],
                    push_constant_ranges: &[],
                },
            );

            let shader =
                device.create_shader_module(wgpu::ShaderModuleDescriptor {
                    label: Some(
                        "iced_wgpu::triangle::solid create shader module",
                    ),
                    source: wgpu::ShaderSource::Wgsl(
                        std::borrow::Cow::Borrowed(include_str!(
                            "shader/solid.wgsl"
                        )),
                    ),
                });

            let pipeline = device.create_render_pipeline(
                &wgpu::RenderPipelineDescriptor {
                    label: Some("iced_wgpu::triangle::solid pipeline"),
                    layout: Some(&layout),
                    vertex: wgpu::VertexState {
                        module: &shader,
                        entry_point: "vs_main",
                        buffers: &[wgpu::VertexBufferLayout {
                            array_stride: std::mem::size_of::<
                                primitive::ColoredVertex2D,
                            >()
                                as u64,
                            step_mode: wgpu::VertexStepMode::Vertex,
                            attributes: &wgpu::vertex_attr_array!(
                                // Position
                                0 => Float32x2,
                                // Color
                                1 => Float32x4,
                            ),
                        }],
                    },
                    fragment: Some(wgpu::FragmentState {
                        module: &shader,
                        entry_point: "fs_main",
                        targets: &[triangle::fragment_target(format)],
                    }),
                    primitive: triangle::primitive_state(),
                    depth_stencil: None,
                    multisample: triangle::multisample_state(antialiasing),
                    multiview: None,
                },
            );

            Self {
                pipeline,
                constants_layout,
            }
        }
    }
}

#[cfg(not(target_arch = "wasm32"))]
mod gradient {
    use crate::buffer::dynamic;
    use crate::buffer::r#static::Buffer;
    use crate::settings;
    use crate::triangle;

    use encase::ShaderType;
    use glam::{IVec4, Vec4};
    use iced_graphics::primitive;

    #[derive(Debug)]
    pub struct Pipeline {
        pub pipeline: wgpu::RenderPipeline,
        pub constants_layout: wgpu::BindGroupLayout,
    }

    #[derive(Debug)]
    pub struct Layer {
        pub vertices: Buffer<primitive::Vertex2D>,
        pub uniforms: dynamic::Buffer<Uniforms>,
        pub storage: dynamic::Buffer<Storage>,
        pub constants: wgpu::BindGroup,
        pub color_stop_offset: i32,
        //Need to store these and then write them all at once
        //or else they will be padded to 256 and cause gaps in the storage buffer
        pub color_stops_pending_write: Storage,
    }

    impl Layer {
        pub fn new(
            device: &wgpu::Device,
            constants_layout: &wgpu::BindGroupLayout,
        ) -> Self {
            let vertices = Buffer::new(
                device,
                "iced_wgpu::triangle::gradient vertex buffer",
                wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
            );

            let uniforms = dynamic::Buffer::uniform(
                device,
                "iced_wgpu::triangle::gradient uniforms",
            );

            // Note: with a WASM target storage buffers are not supported. Will need to use UBOs & static
            // sized array (eg like the 32-sized array on OpenGL side right now) to make gradients work
            let storage = dynamic::Buffer::storage(
                device,
                "iced_wgpu::triangle::gradient storage",
            );

            let constants = Self::bind_group(
                device,
                uniforms.raw(),
                storage.raw(),
                constants_layout,
            );

            Self {
                vertices,
                uniforms,
                storage,
                constants,
                color_stop_offset: 0,
                color_stops_pending_write: Storage {
                    color_stops: vec![],
                },
            }
        }

        pub fn bind_group(
            device: &wgpu::Device,
            uniform_buffer: &wgpu::Buffer,
            storage_buffer: &wgpu::Buffer,
            layout: &wgpu::BindGroupLayout,
        ) -> wgpu::BindGroup {
            device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("iced_wgpu::triangle::gradient bind group"),
                layout,
                entries: &[
                    wgpu::BindGroupEntry {
                        binding: 0,
                        resource: wgpu::BindingResource::Buffer(
                            wgpu::BufferBinding {
                                buffer: uniform_buffer,
                                offset: 0,
                                size: Some(Uniforms::min_size()),
                            },
                        ),
                    },
                    wgpu::BindGroupEntry {
                        binding: 1,
                        resource: storage_buffer.as_entire_binding(),
                    },
                ],
            })
        }
    }

    #[derive(Debug, ShaderType)]
    pub struct Uniforms {
        pub transform: glam::Mat4,
        //xy = start, zw = end
        pub direction: Vec4,
        //x = start stop, y = end stop, zw = padding
        pub stop_range: IVec4,
    }

    #[derive(Debug, ShaderType)]
    pub struct ColorStop {
        pub color: Vec4,
        pub offset: f32,
    }

    #[derive(Debug, ShaderType)]
    pub struct Storage {
        #[size(runtime)]
        pub color_stops: Vec<ColorStop>,
    }

    impl Pipeline {
        /// Creates a new [GradientPipeline] using `gradient.wgsl` shader.
        pub(super) fn new(
            device: &wgpu::Device,
            format: wgpu::TextureFormat,
            antialiasing: Option<settings::Antialiasing>,
        ) -> Self {
            let constants_layout = device.create_bind_group_layout(
                &wgpu::BindGroupLayoutDescriptor {
                    label: Some(
                        "iced_wgpu::triangle::gradient bind group layout",
                    ),
                    entries: &[
                        wgpu::BindGroupLayoutEntry {
                            binding: 0,
                            visibility: wgpu::ShaderStages::VERTEX_FRAGMENT,
                            ty: wgpu::BindingType::Buffer {
                                ty: wgpu::BufferBindingType::Uniform,
                                has_dynamic_offset: true,
                                min_binding_size: Some(Uniforms::min_size()),
                            },
                            count: None,
                        },
                        wgpu::BindGroupLayoutEntry {
                            binding: 1,
                            visibility: wgpu::ShaderStages::FRAGMENT,
                            ty: wgpu::BindingType::Buffer {
                                ty: wgpu::BufferBindingType::Storage {
                                    read_only: true,
                                },
                                has_dynamic_offset: false,
                                min_binding_size: Some(Storage::min_size()),
                            },
                            count: None,
                        },
                    ],
                },
            );

            let layout = device.create_pipeline_layout(
                &wgpu::PipelineLayoutDescriptor {
                    label: Some(
                        "iced_wgpu::triangle::gradient pipeline layout",
                    ),
                    bind_group_layouts: &[&constants_layout],
                    push_constant_ranges: &[],
                },
            );

            let shader =
                device.create_shader_module(wgpu::ShaderModuleDescriptor {
                    label: Some(
                        "iced_wgpu::triangle::gradient create shader module",
                    ),
                    source: wgpu::ShaderSource::Wgsl(
                        std::borrow::Cow::Borrowed(include_str!(
                            "shader/gradient.wgsl"
                        )),
                    ),
                });

            let pipeline =
                device.create_render_pipeline(
                    &wgpu::RenderPipelineDescriptor {
                        label: Some("iced_wgpu::triangle::gradient pipeline"),
                        layout: Some(&layout),
                        vertex: wgpu::VertexState {
                            module: &shader,
                            entry_point: "vs_main",
                            buffers: &[wgpu::VertexBufferLayout {
                                array_stride: std::mem::size_of::<
                                    primitive::Vertex2D,
                                >(
                                )
                                    as u64,
                                step_mode: wgpu::VertexStepMode::Vertex,
                                attributes: &wgpu::vertex_attr_array!(
                                    // Position
                                    0 => Float32x2,
                                ),
                            }],
                        },
                        fragment: Some(wgpu::FragmentState {
                            module: &shader,
                            entry_point: "fs_main",
                            targets: &[triangle::fragment_target(format)],
                        }),
                        primitive: triangle::primitive_state(),
                        depth_stencil: None,
                        multisample: triangle::multisample_state(antialiasing),
                        multiview: None,
                    },
                );

            Self {
                pipeline,
                constants_layout,
            }
        }
    }
}