lemur/src/graphics.rs

use std::{
    sync::{
        atomic::{AtomicU64, Ordering},
        mpsc, Arc, Mutex, MutexGuard,
    },
    thread,
};

use anyhow::{bail, Result};
use itertools::Itertools as _;
use wgpu::{
    Device, Extent3d, ImageCopyTexture, ImageDataLayout, Origin3d, Queue, Texture,
    TextureDescriptor, TextureFormat, TextureUsages, TextureView, TextureViewDescriptor,
};

#[derive(Debug)]
pub struct TextureSink {
    buffers: Arc<BufferPool>,
    sink: mpsc::Sender<u64>,
}

impl TextureSink {
    pub fn new(device: &Device, queue: Arc<Queue>) -> (Self, TextureView) {
        let texture = Self::create_texture(device);
        let view = texture.create_view(&TextureViewDescriptor::default());
        let buffers = Arc::new(BufferPool::new());
        let (sink, source) = mpsc::channel();
        let bufs = buffers.clone();
        thread::spawn(move || {
            let mut local_buf = vec![0; 384 * 224 * 2];
            while let Ok(id) = source.recv() {
                {
                    let Some(bytes) = bufs.read(id) else {
                        continue;
                    };
                    local_buf.copy_from_slice(bytes.as_slice());
                }
                Self::write_texture(&queue, &texture, local_buf.as_slice());
            }
        });
        let sink = Self { buffers, sink };
        (sink, view)
    }

    pub fn queue_render(&mut self, bytes: &[u8]) -> Result<()> {
        let id = {
            let (mut buf, id) = self.buffers.write()?;
            buf.copy_from_slice(bytes);
            id
        };
        self.sink.send(id)?;
        Ok(())
    }

    fn create_texture(device: &Device) -> Texture {
        let desc = TextureDescriptor {
            label: Some("eyes"),
            size: Extent3d {
                width: 384,
                height: 224,
                depth_or_array_layers: 1,
            },
            mip_level_count: 1,
            sample_count: 1,
            dimension: wgpu::TextureDimension::D2,
            format: TextureFormat::Rg8Unorm,
            usage: TextureUsages::COPY_SRC
                | TextureUsages::COPY_DST
                | TextureUsages::TEXTURE_BINDING,
            view_formats: &[TextureFormat::Rg8Unorm],
        };
        device.create_texture(&desc)
    }

    fn write_texture(queue: &Queue, texture: &Texture, bytes: &[u8]) {
        let texture = ImageCopyTexture {
            texture,
            mip_level: 0,
            origin: Origin3d::ZERO,
            aspect: wgpu::TextureAspect::All,
        };
        let size = Extent3d {
            width: 384,
            height: 224,
            depth_or_array_layers: 1,
        };
        let data_layout = ImageDataLayout {
            offset: 0,
            bytes_per_row: Some(384 * 2),
            rows_per_image: Some(224),
        };
        queue.write_texture(texture, bytes, data_layout, size);
    }
}

#[derive(Debug)]
struct BufferPool {
    buffers: [Buffer; 3],
}
impl BufferPool {
    fn new() -> Self {
        Self {
            buffers: std::array::from_fn(|i| Buffer::new(i as u64)),
        }
    }

    fn read(&self, id: u64) -> Option<MutexGuard<'_, Vec<u8>>> {
        let buf = self
            .buffers
            .iter()
            .find(|buf| buf.id.load(Ordering::Acquire) == id)?;
        buf.data.lock().ok()
    }

    fn write(&self) -> Result<(MutexGuard<'_, Vec<u8>>, u64)> {
        let (min, max) = self
            .buffers
            .iter()
            .minmax_by_key(|buf| buf.id.load(Ordering::Acquire))
            .into_option()
            .unwrap();
        let Ok(lock) = min.data.lock() else {
            bail!("lock was poisoned")
        };
        let id = max.id.load(Ordering::Acquire) + 1;
        min.id.store(id, Ordering::Release);
        Ok((lock, id))
    }
}

#[derive(Debug)]
struct Buffer {
    data: Mutex<Vec<u8>>,
    id: AtomicU64,
}
impl Buffer {
    fn new(id: u64) -> Self {
        Self {
            data: Mutex::new(vec![0; 384 * 224 * 2]),
            id: AtomicU64::new(id),
        }
    }
}