graph-explorer-render 0.1.0

use bytemuck::{Pod, Zeroable};
use crate::NodeInstanceData;

// NOTE: no `_pad` field here. `vertex_attr_array!` computes attribute offsets by
// summing format sizes (0, 8, 12, 16, 20) — it does not know about manual padding
// fields. Any padding inserted between fields would shift subsequent attribute
// offsets in the array away from the true field offsets, corrupting node data.
// Keeping the struct tightly packed (36 bytes: pos 8 + radius 4 + opacity 4 +
// shape 4 + color 16) keeps the computed offsets correct. This is a vertex
// buffer, not a uniform, so no 16-byte alignment is required.
#[repr(C)]
#[derive(Copy, Clone, Pod, Zeroable)]
struct NodeInstance {
    pos: [f32; 2],
    radius: f32,
    opacity: f32,
    shape: u32,
    color: [f32; 4],
}

pub struct NodePipeline {
    pipeline: wgpu::RenderPipeline,
    instances: wgpu::Buffer,
    count: u32,
    capacity: u64,
}

impl NodePipeline {
    pub fn new(device: &wgpu::Device, format: wgpu::TextureFormat, camera_bgl: &wgpu::BindGroupLayout) -> Self {
        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("node-shader"),
            source: wgpu::ShaderSource::Wgsl(include_str!("shaders/node.wgsl").into()),
        });
        let layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("node-layout"),
            bind_group_layouts: &[camera_bgl],
            push_constant_ranges: &[],
        });
        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("node-pipeline"),
            layout: Some(&layout),
            vertex: wgpu::VertexState {
                module: &shader,
                entry_point: "vs_main",
                buffers: &[wgpu::VertexBufferLayout {
                    array_stride: std::mem::size_of::<NodeInstance>() as u64,
                    step_mode: wgpu::VertexStepMode::Instance,
                    attributes: &wgpu::vertex_attr_array![0 => Float32x2, 1 => Float32, 2 => Float32, 3 => Uint32, 4 => Float32x4],
                }],
                compilation_options: Default::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &shader,
                entry_point: "fs_main",
                targets: &[Some(wgpu::ColorTargetState {
                    format,
                    blend: Some(wgpu::BlendState::ALPHA_BLENDING),
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: Default::default(),
            }),
            primitive: wgpu::PrimitiveState { topology: wgpu::PrimitiveTopology::TriangleList, ..Default::default() },
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
            multiview: None,
            cache: None,
        });
        let capacity = 64;
        let instances = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("node-instances"),
            size: capacity * std::mem::size_of::<NodeInstance>() as u64,
            usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        });
        Self { pipeline, instances, count: 0, capacity }
    }

    pub fn upload(&mut self, device: &wgpu::Device, queue: &wgpu::Queue, nodes: &[NodeInstanceData]) {
        let data: Vec<NodeInstance> = nodes.iter().map(|n| NodeInstance {
            pos: n.pos, radius: n.radius, opacity: n.opacity, shape: n.shape, color: n.color,
        }).collect();
        let needed = data.len() as u64;
        if needed > self.capacity {
            self.capacity = needed.next_power_of_two();
            self.instances = device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("node-instances"),
                size: self.capacity * std::mem::size_of::<NodeInstance>() as u64,
                usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            });
        }
        queue.write_buffer(&self.instances, 0, bytemuck::cast_slice(&data));
        self.count = data.len() as u32;
    }

    pub fn draw<'a>(&'a self, rpass: &mut wgpu::RenderPass<'a>, camera_bg: &'a wgpu::BindGroup) {
        if self.count == 0 { return; }
        rpass.set_pipeline(&self.pipeline);
        rpass.set_bind_group(0, camera_bg, &[]);
        rpass.set_vertex_buffer(0, self.instances.slice(..));
        rpass.draw(0..6, 0..self.count);
    }
}