yawn/core/lib.rs

//! The arena and allocator compile natively for `cargo test`; everything touching WebGPU is wasm-only.
#![cfg_attr(not(target_arch = "wasm32"), allow(dead_code))]

#[cfg(target_arch = "wasm32")]
use std::cell::RefCell;
#[cfg(target_arch = "wasm32")]
use std::rc::Rc;

#[cfg(target_arch = "wasm32")]
use serde::Deserialize;
#[cfg(target_arch = "wasm32")]
use wasm_bindgen::prelude::*;

#[cfg(target_arch = "wasm32")]
#[path = "render_graph/compiler.rs"]
mod compiler;
#[cfg(target_arch = "wasm32")]
#[path = "renderer/wgpu.rs"]
mod gpu;
#[cfg(target_arch = "wasm32")]
#[path = "render_graph/gpu_resource.rs"]
mod gpu_resource;
#[cfg(target_arch = "wasm32")]
#[path = "render_graph/render_graph.rs"]
mod graph;
#[cfg(target_arch = "wasm32")]
#[path = "renderer/render.rs"]
mod render;
mod render_data;
#[cfg(target_arch = "wasm32")]
#[path = "render_graph/store.rs"]
mod store;

#[cfg(target_arch = "wasm32")]
use gpu::Wgpu;
#[cfg(target_arch = "wasm32")]
use render::RenderLoop;
#[cfg(target_arch = "wasm32")]
use render_data::RenderData;
#[cfg(target_arch = "wasm32")]
use store::Store;

#[cfg(target_arch = "wasm32")]
#[derive(Deserialize)]
struct RowRequest {
    name: String,
    rows: u32,
    stride: u32,
    format: String,
}

#[cfg(target_arch = "wasm32")]
#[wasm_bindgen]
pub struct Core {
    data: Rc<RefCell<RenderData>>,
    gpu: Rc<RefCell<Option<Wgpu>>>,
    store: Rc<RefCell<Store>>,
    render: Rc<RenderLoop>,
}

#[cfg(target_arch = "wasm32")]
#[wasm_bindgen]
impl Core {
    #[wasm_bindgen(constructor)]
    pub fn new(arena_bytes: u32) -> Result<Self, JsError> {
        Ok(Self {
            data: Rc::new(RefCell::new(
                RenderData::new(arena_bytes).map_err(JsError::new)?,
            )),
            gpu: Rc::new(RefCell::new(None)),
            store: Rc::new(RefCell::new(Store::default())),
            render: Rc::new(RenderLoop::new()),
        })
    }

    pub async fn initialize(&self, canvas: web_sys::OffscreenCanvas) -> Result<(), JsError> {
        let gpu = Wgpu::new(canvas)
            .await
            .map_err(|error| JsError::new(&error))?;
        *self.gpu.borrow_mut() = Some(gpu);
        self.render
            .start(self.gpu.clone(), self.store.clone(), self.data.clone());
        Ok(())
    }

    pub fn rows(&self) -> String {
        serde_json::to_string(&self.data.borrow().descriptors()).unwrap()
    }

    /// Optional WebGPU features the device was created with (for example `indirect-first-instance`).
    pub fn features(&self) -> String {
        let gpu = self.gpu.borrow();
        serde_json::to_string(&gpu.as_ref().map(|gpu| &gpu.features[..]).unwrap_or(&[])).unwrap()
    }

    /// Creates or grows an array. When `commit` is set the rows moved: the caller must copy the old
    /// bytes into the returned descriptor and send `commit_rows` before any other request.
    pub fn create_rows(
        &self,
        name: String,
        rows: u32,
        stride: u32,
        format: String,
    ) -> Result<String, JsError> {
        let (rows, moved) = self
            .data
            .borrow_mut()
            .create_rows(name, rows, stride, format)
            .map_err(JsError::new)?;
        if !moved {
            self.refresh_rows(&rows.name)?;
        }
        Ok(serde_json::json!({ "rows": rows, "commit": moved }).to_string())
    }

    /// Like `create_rows` for several arrays; `commit` lists the arrays that moved.
    pub fn create_rows_batch(&self, source: &str) -> Result<String, JsError> {
        let requests: Vec<RowRequest> =
            serde_json::from_str(source).map_err(|_| JsError::new("ROWS"))?;
        if requests.is_empty() {
            return Err(JsError::new("ROWS"));
        }
        let refresh = requests
            .iter()
            .any(|request| self.store.borrow().uses_rows(&request.name));
        let results = {
            let mut data = self.data.borrow_mut();
            let results = requests
                .into_iter()
                .map(|request| {
                    data.create_rows(request.name, request.rows, request.stride, request.format)
                })
                .collect::<Result<Vec<_>, _>>();
            if results.is_err() {
                data.discard_pending();
            }
            results.map_err(JsError::new)?
        };
        let commit = results
            .iter()
            .filter(|(_, moved)| *moved)
            .map(|(rows, _)| rows.name.clone())
            .collect::<Vec<_>>();
        if refresh && results.iter().any(|(_, moved)| !moved) {
            self.refresh()?;
        }
        let descriptors = results
            .into_iter()
            .map(|(rows, _)| rows)
            .collect::<Vec<_>>();
        Ok(serde_json::json!({ "rows": descriptors, "commit": commit }).to_string())
    }

    /// Publishes relocations the caller finished copying, then rebuilds the GPU side once.
    pub fn commit_rows(&self, source: &str) -> Result<(), JsError> {
        let names: Vec<String> = serde_json::from_str(source).map_err(|_| JsError::new("ROWS"))?;
        self.data
            .borrow_mut()
            .commit(&names)
            .map_err(JsError::new)?;
        if names.iter().any(|name| self.store.borrow().uses_rows(name)) {
            self.refresh()?;
        }
        self.data.borrow_mut().mark_dirty();
        Ok(())
    }

    fn refresh_rows(&self, name: &str) -> Result<(), JsError> {
        if let Some(gpu) = self.gpu.borrow().as_ref() {
            self.store
                .borrow_mut()
                .refresh_rows(name, gpu, &self.data.borrow())
                .map_err(|error| JsError::new(&error))?;
        }
        Ok(())
    }

    fn refresh(&self) -> Result<(), JsError> {
        if let Some(gpu) = self.gpu.borrow().as_ref() {
            self.store
                .borrow_mut()
                .refresh(gpu, &self.data.borrow())
                .map_err(|error| JsError::new(&error))?;
        }
        Ok(())
    }

    pub fn delete_rows(&self, name: String) -> Result<(), JsError> {
        if self.store.borrow().uses_rows(&name) {
            return Err(JsError::new("ROWS_ACTIVE"));
        }
        self.data
            .borrow_mut()
            .delete_rows(&name)
            .map_err(JsError::new)
    }

    /// Reserves rows; `commit` is set when the array moved, exactly like `create_rows`.
    pub fn allocate_object(&self, name: &str, count: u32) -> Result<String, JsError> {
        let (id, grew, moved) = self
            .data
            .borrow_mut()
            .allocate_object(name, count)
            .map_err(JsError::new)?;
        let rows = self
            .data
            .borrow()
            .staged(name)
            .ok_or_else(|| JsError::new("ROWS_UNKNOWN"))?
            .clone();
        if grew && !moved {
            self.refresh_rows(name)?;
        }
        Ok(serde_json::json!({ "id": id, "rows": rows, "commit": moved }).to_string())
    }

    pub fn delete_object(&self, name: &str, id: u32) -> Result<(), JsError> {
        self.data
            .borrow_mut()
            .delete_object(name, id)
            .map_err(JsError::new)?;
        self.data.borrow_mut().mark_dirty();
        Ok(())
    }

    /// Uploads rows `first .. first + count` of `name` to every active GPU buffer bound to it.
    pub fn sync_rows(&self, name: &str, first: u32, count: u32) -> Result<(), JsError> {
        let gpu = self.gpu.borrow();
        let gpu = gpu
            .as_ref()
            .ok_or_else(|| JsError::new("WEBGPU_UNINITIALIZED"))?;
        self.store
            .borrow()
            .sync_rows(name, first, count, gpu, &self.data.borrow())
            .map_err(|error| JsError::new(&error))?;
        self.data.borrow_mut().mark_dirty();
        Ok(())
    }

    pub fn compile_graph(&self, source: &str) -> Result<String, JsError> {
        let graph = compiler::compile(source).map_err(JsError::new)?;
        Ok(self.store.borrow_mut().save(graph))
    }

    pub fn switch_loadout(&self, id: &str) -> Result<(), JsError> {
        let gpu = self.gpu.borrow();
        let gpu = gpu
            .as_ref()
            .ok_or_else(|| JsError::new("WEBGPU_UNINITIALIZED"))?;
        self.store
            .borrow_mut()
            .switch(id, gpu, &self.data.borrow())
            .map_err(|error| JsError::new(&error))?;
        self.data.borrow_mut().loadout_ready();
        Ok(())
    }

    pub fn upload_texture(
        &self,
        name: String,
        mip_level: u32,
        image: web_sys::ImageBitmap,
    ) -> Result<(), JsError> {
        let gpu = self.gpu.borrow();
        let gpu = gpu
            .as_ref()
            .ok_or_else(|| JsError::new("WEBGPU_UNINITIALIZED"))?;
        self.store
            .borrow_mut()
            .upload_texture(name, mip_level, image, gpu)
            .map_err(|error| JsError::new(&error))?;
        self.data.borrow_mut().mark_dirty();
        Ok(())
    }

    pub fn delete_texture(&self, name: &str) {
        self.store.borrow_mut().delete_texture(name);
    }

    pub fn play(&self) {
        self.render.play();
    }

    pub fn pause(&self) {
        self.render.pause();
    }

    pub fn set_fps(&self, fps: u32) -> Result<(), JsError> {
        self.render.set_fps(fps).map_err(JsError::new)
    }

    pub fn set_profiler(&self, enabled: bool) -> bool {
        let supported = self
            .gpu
            .borrow()
            .as_ref()
            .is_some_and(|gpu| gpu.timestamp_queries);
        self.render.set_profiling(enabled && supported);
        if enabled && supported {
            self.data.borrow_mut().mark_dirty();
        }
        supported
    }

    /// Installs the worker-local sink that publishes completed GPU timings into shared memory.
    pub fn on_profile(&self, callback: js_sys::Function) {
        self.render.on_profile(callback);
    }
}