Fractal/src/fractal.rs

use std::{mem, sync::Arc};

use context::{prelude::*, ContextObject};

#[derive(Debug, Clone)]
struct ShaderInfo {
    max_iter: i32,
    size: f32,
    time: f32,
    factor: f32,
    width: f32,
    height: f32,
}

#[derive(Clone)]
struct Vertex {
    position: [f32; 2],
    uv: [f32; 2],
}

pub struct Fractal {
    context: Arc<Context>,

    descriptor_set: Arc<DescriptorSet>,
    pipeline: Arc<Pipeline>,
    render_target: RenderTarget,

    vertex_buffer: Arc<Buffer<Vertex>>,
    descriptor_buffer: Arc<Buffer<ShaderInfo>>,
}

impl Fractal {
    pub fn new(context: &Arc<Context>) -> VerboseResult<Arc<Self>> {
        let render_core = context.render_core();

        let render_target = RenderTarget::builder()
            .add_sub_pass(
                SubPass::builder(render_core.width(), render_core.height())
                    .set_prepared_targets(
                        render_core.images()?.single()?,
                        0,
                        [0.0, 0.0, 0.0, 0.0],
                        true,
                    )
                    .build(context.device(), context.queue())?,
            )
            .build(context.device())?;

        let descriptor_layout = DescriptorSetLayout::builder()
            .add_layout_binding(
                0,
                VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
                VK_SHADER_STAGE_FRAGMENT_BIT,
                0,
            )
            .build(context.device().clone())?;

        let pipeline_layout = PipelineLayout::builder()
            .add_descriptor_set_layout(&descriptor_layout)
            .build(context.device().clone())?;

        let pipeline = Pipeline::new_graphics()
            .default_rasterization(VK_CULL_MODE_FRONT_BIT, VK_FRONT_FACE_COUNTER_CLOCKWISE)
            .default_multisample(VK_SAMPLE_COUNT_1_BIT)
            .input_assembly(VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST, false)
            .default_color_blend(vec![VkPipelineColorBlendAttachmentState::default()])
            .set_vertex_shader(
                ShaderModule::from_slice(
                    context.device().clone(),
                    include_bytes!("fractal.vert.spv"),
                    ShaderType::Vertex,
                )?,
                vec![VkVertexInputBindingDescription {
                    binding: 0,
                    inputRate: VK_VERTEX_INPUT_RATE_VERTEX,
                    stride: mem::size_of::<Vertex>() as u32,
                }],
                vec![
                    VkVertexInputAttributeDescription {
                        location: 0,
                        binding: 0,
                        format: VK_FORMAT_R32G32_SFLOAT,
                        offset: 0,
                    },
                    VkVertexInputAttributeDescription {
                        location: 1,
                        binding: 0,
                        format: VK_FORMAT_R32G32_SFLOAT,
                        offset: 8,
                    },
                ],
            )
            .set_fragment_shader(ShaderModule::from_slice(
                context.device().clone(),
                include_bytes!("fractal.frag.spv"),
                ShaderType::Fragment,
            )?)
            .add_viewport(VkViewport {
                x: 0.0,
                y: 0.0,
                width: render_core.width() as f32,
                height: render_core.height() as f32,
                minDepth: 0.0,
                maxDepth: 1.0,
            })
            .add_scissor(VkRect2D {
                offset: VkOffset2D { x: 0, y: 0 },
                extent: VkExtent2D {
                    width: render_core.width(),
                    height: render_core.height(),
                },
            })
            .build(
                context.device().clone(),
                &pipeline_layout,
                render_target.render_pass(),
                0,
            )?;

        let descriptor_set = DescriptorPool::builder()
            .set_layout(descriptor_layout)
            .build(context.device().clone())?
            .prepare_set()
            .allocate()?;

        let vertex_buffer = Buffer::builder()
            .set_data(&[
                Vertex {
                    position: [-1.0, -1.0],
                    uv: [0.0, 0.0],
                },
                Vertex {
                    position: [1.0, -1.0],
                    uv: [1.0, 0.0],
                },
                Vertex {
                    position: [1.0, 1.0],
                    uv: [1.0, 1.0],
                },
                Vertex {
                    position: [1.0, 1.0],
                    uv: [1.0, 1.0],
                },
                Vertex {
                    position: [-1.0, 1.0],
                    uv: [0.0, 1.0],
                },
                Vertex {
                    position: [-1.0, -1.0],
                    uv: [0.0, 0.0],
                },
            ])
            .set_memory_usage(MemoryUsage::CpuOnly)
            .set_usage(VK_BUFFER_USAGE_VERTEX_BUFFER_BIT)
            .build(context.device().clone())?;

        let shader_info = Buffer::builder()
            .set_data(&[ShaderInfo {
                max_iter: 1000,
                size: 0.001,
                time: 0.0,
                factor: 1.0,
                width: render_core.width() as f32,
                height: render_core.height() as f32,
            }])
            .set_memory_usage(MemoryUsage::CpuOnly)
            .set_usage(VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT)
            .build(context.device().clone())?;

        descriptor_set.update(&[DescriptorWrite::uniform_buffers(0, &[&shader_info])])?;

        Ok(Arc::new(Self {
            context: context.clone(),

            descriptor_set,
            pipeline,
            render_target,
            vertex_buffer,
            descriptor_buffer: shader_info,
        }))
    }
}

impl ContextObject for Fractal {
    fn name(&self) -> &str {
        "Fractal"
    }

    fn update(&self) -> VerboseResult<()> {
        let mut mapping = self.descriptor_buffer.map_complete()?;

        mapping[0].time = self.context.time().as_secs_f32();

        Ok(())
    }

    fn event(&self, event: Event) -> VerboseResult<()> {
        match event {
            Event::KeyDown(key) => match key {
                Keycode::Escape => self.context.close()?,

                _ => (),
            },

            _ => (),
        }

        Ok(())
    }
}

impl TScene for Fractal {
    fn update(&self) -> VerboseResult<()> {
        Ok(())
    }

    fn process(
        &self,
        buffer_recorder: &mut CommandBufferRecorder<'_>,
        indices: &TargetMode<usize>,
    ) -> VerboseResult<()> {
        match indices {
            TargetMode::Single(index) => {
                self.render_target
                    .begin(buffer_recorder, VK_SUBPASS_CONTENTS_INLINE, *index);

                buffer_recorder.bind_pipeline(&self.pipeline)?;
                buffer_recorder.bind_descriptor_sets_minimal(&[&self.descriptor_set])?;
                buffer_recorder.bind_vertex_buffer(&self.vertex_buffer);
                buffer_recorder.draw_complete_single_instance(self.vertex_buffer.size() as u32);

                self.render_target.end(buffer_recorder);
            }
            TargetMode::Stereo(_, _) => unimplemented!(),
        }

        Ok(())
    }

    fn resize(&self) -> VerboseResult<()> {
        let mut mapping = self.descriptor_buffer.map_complete()?;

        mapping[0].width = self.context.render_core().width() as f32;
        mapping[0].height = self.context.render_core().height() as f32;

        todo!("resize of FB is still missing");

        // Ok(())
    }
}