main.rs

use {
    bytemuck::{cast_slice, Pod, Zeroable},
    glam::{Vec2, Vec3},
    inline_spirv::inline_spirv,
    screen_13::prelude::*,
    std::sync::Arc,
};

// Copied in helper class to have self-contained example
pub struct BufferUploader<'d> {
    device: &'d Arc<Device>,
    queue_index: usize,
    buffers: Vec<(Arc<Buffer>, Arc<Buffer>)>,
}

impl<'d> BufferUploader<'d> {
    pub fn new(device: &'d Arc<Device>, queue_index: usize) -> Self {
        Self {
            device,
            queue_index,
            buffers: Vec::with_capacity(8),
        }
    }

    pub fn buffer_from_slice(
        &mut self,
        usage: vk::BufferUsageFlags,
        slice: &[u8],
    ) -> Result<Arc<Buffer>, DriverError> {
        let dst_buf = Arc::new(Buffer::create(
            self.device,
            BufferInfo::new(
                slice.len() as u64,
                usage | vk::BufferUsageFlags::TRANSFER_DST,
            ),
        )?);
        let src_buf = Arc::new(Buffer::create_from_slice(
            self.device,
            vk::BufferUsageFlags::TRANSFER_SRC,
            slice,
        )?);
        self.buffers.push((src_buf, dst_buf.clone()));
        Ok(dst_buf)
    }

    pub fn upload(self) -> Result<(), DriverError> {
        let mut pool = LazyPool::new(self.device);
        let mut graph = RenderGraph::new();
        for (src_buf, dst_buf) in self.buffers.iter() {
            let src_node = graph.bind_node(src_buf);
            let dst_node = graph.bind_node(dst_buf);
            graph.copy_buffer(src_node, dst_node);
        }
        graph.resolve().submit(&mut pool, self.queue_index)?;
        Ok(())
    }
}


#[repr(C)]
#[derive(Clone, Copy, Debug, Default, Zeroable, Pod)]
pub struct TileVertex {
    pub position: Vec3,
    pub normal: Vec3,
    pub uvs: Vec2,
}

#[repr(C)]
#[derive(Clone, Copy, Debug, Default, Zeroable, Pod)]
pub struct ChunkVertex {
    pub tile_vertex_index: u32,
    pub texture_id: u32,
}

#[repr(C)]
#[derive(Clone, Copy, Debug, Default, Zeroable, Pod)]
pub struct ChunkInstanceData {
    pub chunk_buffer_id: u32,
}

#[repr(C)]
#[derive(Clone, Copy, Debug, Default, Zeroable, Pod)]
struct DrawIndirectCommand {
    vertex_count: u32,
    instance_count: u32,
    first_vertex: u32,
    first_instance: u32,
}

fn main() -> Result<(), DisplayError> {
    tracing_subscriber::fmt::init(); // TODO: if debug true RUST_LOG should not be required

    let event_loop = EventLoop::new()
        .debug(true)
        .window(|window| window.with_inner_size(LogicalSize::new(512, 512)))
        .build()?;

    let mut uploader = BufferUploader::new(&event_loop.device, 0);
    let tile_buffer = create_tile_buffer(&mut uploader)?;
    let chunk_buffers = create_chunk_buffers(&mut uploader)?;
    let instance_buffer = create_instance_buffer(&mut uploader)?;
    let indirect_buffer = create_indirect_buffer(&mut uploader)?;
    uploader.upload().expect("buffers to upload successfully");

    let images = create_images(&event_loop.device)?;
    let pipeline = create_graphic_pipeline(&event_loop.device)?;

    event_loop.run(|frame| {
        let indirect_buffer_node = frame.render_graph.bind_node(&indirect_buffer);
        let instance_buffer_node = frame.render_graph.bind_node(&instance_buffer);

        let mut pass = frame
            .render_graph
            .begin_pass("Test")
            .bind_pipeline(&pipeline)
            .access_node(indirect_buffer_node, AccessType::IndirectBuffer)
            .access_node(instance_buffer_node, AccessType::VertexBuffer);

        for (idx, image) in images.iter().enumerate() {
            let image_node = pass.bind_node(image);
            pass = pass.read_descriptor((0, 0, [idx as u32]), image_node);
        }

        {
            let tile_buffer_node = pass.bind_node(&tile_buffer);
            pass = pass.read_descriptor((0, 1), tile_buffer_node);
        }

        for (idx, chunk_buffer) in chunk_buffers.iter().enumerate() {
            let chunk_buffer_node = pass.bind_node(chunk_buffer);
            pass = pass.read_descriptor((0, 2, [idx as u32]), chunk_buffer_node);
        }

        pass.clear_color(0, frame.swapchain_image)
            .store_color(0, frame.swapchain_image)
            .record_subpass(move |subpass, _| {
                subpass
                    .bind_vertex_buffer(instance_buffer_node)
                    .draw_indirect(indirect_buffer_node, 0, 64, 16);
            });
    })
}

fn create_images(device: &Arc<Device>) -> Result<Vec<Arc<Image>>, DisplayError> {
    let mut textures = Vec::with_capacity(64);

    let (b, a) = (0.0, 1.0);
    let mut graph = RenderGraph::new();
    for y in 0..8 {
        for x in 0..8 {
            let texture = Arc::new(Image::create(
                device,
                ImageInfo::new_2d(
                    vk::Format::R8G8B8A8_UNORM,
                    100,
                    100,
                    vk::ImageUsageFlags::SAMPLED | vk::ImageUsageFlags::TRANSFER_DST,
                ),
            )?);
            let texture_node = graph.bind_node(&texture);
            let r = y as f32 / 7.0;
            let g = x as f32 / 7.0;
            graph.clear_color_image_value(texture_node, [r, g, b, a]);
            textures.push(texture);
        }
    }

    let mut pool = LazyPool::new(device);

    graph.resolve().submit(&mut pool, 0)?;

    Ok(textures)
}

fn create_tile_buffer(uploader: &mut BufferUploader) -> Result<Arc<Buffer>, DriverError> {
    let tile_data = [
        TileVertex {
            position: Vec3::new(0.0, 0.0, 0.0),
            normal: Vec3::new(0.0, 0.0, 1.0),
            uvs: Vec2::new(0.5, 0.5),
        },
        TileVertex {
            position: Vec3::new(0.0, 1.0, 0.0),
            normal: Vec3::new(0.0, 0.0, 1.0),
            uvs: Vec2::new(0.5, 0.5),
        },
        TileVertex {
            position: Vec3::new(1.0, 0.0, 0.0),
            normal: Vec3::new(0.0, 0.0, 1.0),
            uvs: Vec2::new(0.5, 0.5),
        },
        TileVertex {
            position: Vec3::new(0.0, 1.0, 0.0),
            normal: Vec3::new(0.0, 0.0, 1.0),
            uvs: Vec2::new(0.5, 0.5),
        },
        TileVertex {
            position: Vec3::new(1.0, 1.0, 0.0),
            normal: Vec3::new(0.0, 0.0, 1.0),
            uvs: Vec2::new(0.5, 0.5),
        },
        TileVertex {
            position: Vec3::new(1.0, 0.0, 0.0),
            normal: Vec3::new(0.0, 0.0, 1.0),
            uvs: Vec2::new(0.5, 0.5),
        },
    ];
    uploader.buffer_from_slice(vk::BufferUsageFlags::STORAGE_BUFFER, cast_slice(&tile_data))
}

fn create_chunk_buffers(uploader: &mut BufferUploader) -> Result<Vec<Arc<Buffer>>, DriverError> {
    let mut chunk_buffers = Vec::with_capacity(64);
    for texture_id in 0..64 {
        chunk_buffers.push(uploader.buffer_from_slice(
            vk::BufferUsageFlags::STORAGE_BUFFER,
            cast_slice(&[
                ChunkVertex {
                    tile_vertex_index: 0,
                    texture_id,
                },
                ChunkVertex {
                    tile_vertex_index: 1,
                    texture_id,
                },
                ChunkVertex {
                    tile_vertex_index: 2,
                    texture_id,
                },
                ChunkVertex {
                    tile_vertex_index: 3,
                    texture_id,
                },
                ChunkVertex {
                    tile_vertex_index: 4,
                    texture_id,
                },
                ChunkVertex {
                    tile_vertex_index: 5,
                    texture_id,
                },
            ]),
        )?);
    }
    Ok(chunk_buffers)
}

fn create_instance_buffer(uploader: &mut BufferUploader) -> Result<Arc<Buffer>, DriverError> {
    let mut instance_data = Vec::with_capacity(64);
    for chunk_buffer_id in 0..64 {
        instance_data.push(ChunkInstanceData { chunk_buffer_id });
    }
    uploader.buffer_from_slice(
        vk::BufferUsageFlags::VERTEX_BUFFER,
        cast_slice(&instance_data),
    )
}

fn create_indirect_buffer(uploader: &mut BufferUploader) -> Result<Arc<Buffer>, DriverError> {
    let mut draw_cmds = Vec::with_capacity(64);
    for first_instance in 0..64 {
        draw_cmds.push(DrawIndirectCommand {
            vertex_count: 6,
            instance_count: 1,
            first_vertex: 0,
            first_instance,
        });
    }
    uploader.buffer_from_slice(
        vk::BufferUsageFlags::INDIRECT_BUFFER,
        cast_slice(&draw_cmds),
    )
}

fn create_graphic_pipeline(device: &Arc<Device>) -> Result<Arc<GraphicPipeline>, DriverError> {
    Ok(Arc::new(GraphicPipeline::create(
        device,
        GraphicPipelineInfo::default(),
        [
            Shader::new_vertex(
                inline_spirv!(
                    r#"
                    #version 460 core
                    #extension GL_EXT_nonuniform_qualifier: require

                    struct TileVertex {
                        float position[3];
                        float normal[3];
                        float uvs[2];
                    };

                    struct ChunkVertex {
                        uint tile_vertex_index;
                        uint texture_id;
                    };

                    layout(location = 0) in uint chunk_buffer_id_ibind0;

                    layout(location = 0) flat out uint texture_id;
                    layout(location = 1) out vec2 uvs;

                    layout(set = 0, binding = 1) readonly buffer TileVertices { TileVertex tile_vertices[]; };
                    layout(set = 0, binding = 2) readonly buffer ChunkBuffer { ChunkVertex chunk_vertices[]; } chunk_buffers[];

                    void main() {
                        ChunkVertex chunk_vertex = chunk_buffers[nonuniformEXT(chunk_buffer_id_ibind0)].chunk_vertices[gl_VertexIndex];
                        TileVertex tile_vertex = tile_vertices[nonuniformEXT(chunk_vertex.tile_vertex_index)];

                        uint x = gl_InstanceIndex % 8;
                        uint y = gl_InstanceIndex / 8;
                        vec3 scale = vec3(vec2(1.0 / 8.0), 1.0);
                        vec3 offset = vec3((float(x) - 4.0) * scale.x, (float(y) - 4.0) * scale.y, 0);
                        vec3 position = vec3(
                            tile_vertex.position[0],
                            tile_vertex.position[1],
                            tile_vertex.position[2]
                        );
                        gl_Position = vec4(position * scale + offset, 1);
                        texture_id = uint(chunk_vertex.texture_id);
                        uvs = vec2(tile_vertex.uvs[0], tile_vertex.uvs[1]);
                    }
                    "#,
                    vert
                )
                .as_slice(),
            ),
            Shader::new_fragment(
                inline_spirv!(
                    r#"
                    #version 460 core
                    #extension GL_EXT_nonuniform_qualifier: require

                    layout(set = 0, binding = 0) uniform sampler2D textures_nnr[];

                    layout(location = 0) out vec4 f_color;

                    layout(location = 0) flat in uint texture_id;
                    layout(location = 1) in vec2 uvs;

                    void main() {
                        f_color = texture( textures_nnr[nonuniformEXT(texture_id)], uvs);
                    }
                    "#,
                    frag
                )
                .as_slice(),
            ),
        ],
    )?))
}