can render multiple viewports

3 files changed, 245 insertions, 246 deletions

diff --git a/client/src/renderer.rs b/client/src/renderer.rs
index 1fa6991..bb87b6c 100644
--- a/client/src/renderer.rs
+++ b/client/src/renderer.rs
@@ -87,7 +87,8 @@ impl<'a> Renderer<'a> {
             ScenePipeline::new(&device, surface_configuration.format);
         let scene_prepare = ScenePreparer::new(device.clone(), queue.clone(), texture_bgl);
 
-        let ui_renderer = UiRenderer::new(&device, surface_configuration.format);
+        let ui_renderer =
+            UiRenderer::new(device.clone(), queue.clone(), surface_configuration.format);
 
         let depth = device.create_texture(&TextureDescriptor {
             label: None,
@@ -173,14 +174,8 @@ impl<'a> Renderer<'a> {
             projection,
         );
 
-        self.ui_renderer.draw(
-            &self.device,
-            &self.queue,
-            &mut commands,
-            &target_view,
-            &self.depth,
-            projection,
-        );
+        self.ui_renderer
+            .draw(&mut commands, &target_view, &self.depth, projection);
 
         let i = self.queue.submit(Some(commands.finish()));
         self.device.poll(MaintainBase::WaitForSubmissionIndex(i));
diff --git a/client/src/state.rs b/client/src/state.rs
index 70746d8..00a8c0b 100644
--- a/client/src/state.rs
+++ b/client/src/state.rs
@@ -14,15 +14,11 @@
     You should have received a copy of the GNU Affero General Public License
     along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
-use crate::{
-    camera::Camera, download::Downloader, network::Network, renderer::Renderer, ui::UiSurface,
-};
+use crate::{camera::Camera, download::Downloader, network::Network, renderer::Renderer};
 use anyhow::{Context, Result};
-use egui::ViewportId;
 use glam::{Vec2, Vec3};
 use log::{info, warn};
-use rand::random;
-use std::{net::TcpStream, sync::Arc, time::Instant};
+use std::{net::TcpStream, time::Instant};
 use weareshared::{store::ResourceStore, tree::SceneTree};
 use winit::event::MouseButton;
 
@@ -72,21 +68,17 @@ impl<'a> State<'a> {
         if !down || button != MouseButton::Right {
             return;
         }
-        self.renderer.ui_renderer.surfaces.insert(
-            ViewportId::from_hash_of(random::<u128>()),
-            UiSurface {
-                transform: self.camera.new_ui_affine(),
-                content: Arc::new(|ctx| {
-                    egui::Window::new("Funny window")
-                        .default_open(true)
-                        .show(ctx, |ui| {
-                            ui.label("world space ui actually kinda works now");
-                            ui.label("Does input work?");
-                            ui.button("Yes").clicked();
-                        });
-                }),
-            },
-        );
+        self.renderer
+            .ui_renderer
+            .add_surface(self.camera.new_ui_affine(), |ctx| {
+                egui::Window::new("Funny window")
+                    .default_open(true)
+                    .show(ctx, |ui| {
+                        ui.label("world space ui actually kinda works now");
+                        ui.label("Does input work?");
+                        ui.button("Yes").clicked();
+                    });
+            });
     }
     pub fn update(&mut self) -> Result<()> {
         let now = Instant::now();
diff --git a/client/src/ui.rs b/client/src/ui.rs
index e2ba989..2bf8ed5 100644
--- a/client/src/ui.rs
+++ b/client/src/ui.rs
@@ -15,12 +15,17 @@
     along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
 use egui::{
-    Context, ImageData, TextureId, ViewportBuilder, ViewportId, ViewportInfo,
-    epaint::{Primitive, Vertex},
+    Context, ImageData, TextureId, ViewportId, ViewportInfo,
+    epaint::{ImageDelta, Primitive, Vertex},
 };
 use glam::{Affine3A, Mat2, Mat3, Mat4};
 use log::info;
-use std::{collections::HashMap, num::NonZeroU64, sync::Arc};
+use rand::random;
+use std::{
+    collections::HashMap,
+    num::NonZeroU64,
+    sync::{Arc, RwLock},
+};
 use wgpu::{
     AddressMode, BindGroup, BindGroupDescriptor, BindGroupEntry, BindGroupLayout,
     BindGroupLayoutDescriptor, BindGroupLayoutEntry, BindingResource, BindingType, BlendState,
@@ -39,37 +44,24 @@ use wgpu::{
 };
 
 pub struct UiRenderer {
+    device: Arc<Device>,
+    queue: Arc<Queue>,
     ctx: Context,
-
     pipeline: RenderPipeline,
     bind_group_layout: BindGroupLayout,
-
-    index: Buffer,
-    vertex: Buffer,
-    textures: HashMap<TextureId, (BindGroup, Texture, [u32; 2])>,
-
-    pub surfaces: HashMap<ViewportId, UiSurface>,
+    textures: RwLock<HashMap<TextureId, (BindGroup, Texture, [u32; 2])>>,
+    surfaces: HashMap<ViewportId, UiSurface>,
 }
 
 pub struct UiSurface {
     pub transform: Affine3A,
     pub content: Arc<dyn Fn(&Context) + Send + Sync + 'static>,
+    index: Buffer,
+    vertex: Buffer,
 }
 
 impl UiRenderer {
-    pub fn new(device: &Device, format: TextureFormat) -> Self {
-        let index = device.create_buffer(&BufferDescriptor {
-            label: None,
-            size: size_of::<f32>() as u64 * 1024 * 1024,
-            usage: BufferUsages::INDEX | BufferUsages::COPY_DST,
-            mapped_at_creation: false,
-        });
-        let vertex = device.create_buffer(&BufferDescriptor {
-            label: None,
-            size: size_of::<f32>() as u64 * 1024 * 1024,
-            usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
-            mapped_at_creation: false,
-        });
+    pub fn new(device: Arc<Device>, queue: Arc<Queue>, format: TextureFormat) -> Self {
         let module = device.create_shader_module(include_wgsl!("ui.wgsl"));
         let bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
             entries: &[
@@ -144,197 +136,140 @@ impl UiRenderer {
         Self {
             ctx: Context::default(),
             pipeline,
-            index,
-            vertex,
+            device,
+            queue,
             bind_group_layout,
-            textures: HashMap::new(),
+            textures: HashMap::new().into(),
             surfaces: HashMap::new(),
         }
     }
 
-    pub fn create_texture() {}
-
-    pub fn draw(
+    pub fn add_surface(
         &mut self,
-        device: &Device,
-        queue: &Queue,
-        commands: &mut CommandEncoder,
-        target: &TextureView,
-        depth: &TextureView,
-        projection: Mat4,
+        transform: Affine3A,
+        content: impl Fn(&Context) + Send + Sync + 'static,
     ) {
-        let mut raw_input = egui::RawInput::default();
-        raw_input.viewport_id = self
-            .surfaces
-            .keys()
-            .next()
-            .copied()
-            .unwrap_or(ViewportId::ROOT);
-        raw_input.viewports = self
-            .surfaces
-            .keys()
-            .chain(Some(ViewportId::ROOT).iter())
-            .map(|key| {
-                (*key, ViewportInfo {
-                    native_pixels_per_point: Some(2.),
-                    ..Default::default()
-                })
-            })
-            .collect();
-
-        let full_output = self.ctx.run(raw_input, |ctx| {
-            for (id, surf) in &self.surfaces {
-                ctx.show_viewport_immediate(*id, ViewportBuilder::default(), |ctx, _class| {
-                    (surf.content)(ctx)
-                });
-            }
+        let index = self.device.create_buffer(&BufferDescriptor {
+            label: None,
+            size: size_of::<f32>() as u64 * 1024 * 1024,
+            usage: BufferUsages::INDEX | BufferUsages::COPY_DST,
+            mapped_at_creation: false,
+        });
+        let vertex = self.device.create_buffer(&BufferDescriptor {
+            label: None,
+            size: size_of::<Vertex>() as u64 * 1024 * 1024,
+            usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
+            mapped_at_creation: false,
         });
+        self.surfaces
+            .insert(ViewportId::from_hash_of(random::<u128>()), UiSurface {
+                transform,
+                content: Arc::new(content),
+                index,
+                vertex,
+            });
+    }
 
-        for (texid, delta) in full_output.textures_delta.set {
-            let size = Extent3d {
-                depth_or_array_layers: 1,
-                width: delta.image.width() as u32,
-                height: delta.image.height() as u32,
-            };
-            let pixels = match &delta.image {
-                ImageData::Color(color_image) => color_image.pixels.clone(),
-                ImageData::Font(font_image) => font_image.srgba_pixels(None).collect(),
-            };
+    pub fn apply_texture_delta(&self, texid: TextureId, delta: ImageDelta) {
+        let mut textures = self.textures.write().unwrap();
+        let size = Extent3d {
+            depth_or_array_layers: 1,
+            width: delta.image.width() as u32,
+            height: delta.image.height() as u32,
+        };
+        let pixels = match &delta.image {
+            ImageData::Color(color_image) => color_image.pixels.clone(),
+            ImageData::Font(font_image) => font_image.srgba_pixels(None).collect(),
+        };
 
-            if let Some((_texbg, tex, texsize)) = self.textures.get_mut(&texid) {
-                let pos = delta.pos.unwrap_or([0, 0]);
-                queue.write_texture(
-                    ImageCopyTexture {
-                        texture: &tex,
-                        mip_level: 0,
-                        origin: Origin3d {
-                            x: pos[0] as u32,
-                            y: pos[1] as u32,
-                            z: 0,
-                        },
-                        aspect: TextureAspect::All,
-                    },
-                    bytemuck::cast_slice::<_, u8>(&pixels),
-                    ImageDataLayout {
-                        offset: 0,
-                        bytes_per_row: Some(texsize[0] * 4),
-                        rows_per_image: None,
+        if let Some((_texbg, tex, texsize)) = textures.get_mut(&texid) {
+            let pos = delta.pos.unwrap_or([0, 0]);
+            info!("updating UI texture at {pos:?}");
+            self.queue.write_texture(
+                ImageCopyTexture {
+                    texture: &tex,
+                    mip_level: 0,
+                    origin: Origin3d {
+                        x: pos[0] as u32,
+                        y: pos[1] as u32,
+                        z: 0,
                     },
-                    size,
-                );
-            } else {
-                assert_eq!(
-                    delta.pos, None,
-                    "partial update impossible; texture does not yet exist"
-                );
-                info!(
-                    "uploading new UI texture: width={}, height={}",
-                    delta.image.width(),
-                    delta.image.height()
-                );
+                    aspect: TextureAspect::All,
+                },
+                bytemuck::cast_slice::<_, u8>(&pixels),
+                ImageDataLayout {
+                    offset: 0,
+                    bytes_per_row: Some(texsize[0] * 4),
+                    rows_per_image: None,
+                },
+                size,
+            );
+        } else {
+            assert_eq!(
+                delta.pos, None,
+                "partial update impossible; texture does not yet exist"
+            );
+            info!(
+                "uploading new UI texture: width={}, height={}",
+                delta.image.width(),
+                delta.image.height()
+            );
 
-                let texture = device.create_texture_with_data(
-                    &queue,
-                    &TextureDescriptor {
-                        label: None,
-                        size,
-                        mip_level_count: 1,
-                        sample_count: 1,
-                        dimension: TextureDimension::D2,
-                        format: TextureFormat::Rgba8UnormSrgb,
-                        usage: TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_DST,
-                        view_formats: &[],
-                    },
-                    TextureDataOrder::LayerMajor,
-                    bytemuck::cast_slice::<_, u8>(&pixels),
-                );
-                let textureview = texture.create_view(&TextureViewDescriptor::default());
-                let sampler = device.create_sampler(&SamplerDescriptor {
-                    address_mode_u: AddressMode::ClampToEdge,
-                    address_mode_v: AddressMode::ClampToEdge,
-                    mag_filter: FilterMode::Linear,
-                    min_filter: FilterMode::Linear,
-                    ..Default::default()
-                });
-                let bindgroup = device.create_bind_group(&BindGroupDescriptor {
+            let texture = self.device.create_texture_with_data(
+                &self.queue,
+                &TextureDescriptor {
                     label: None,
-                    layout: &self.bind_group_layout,
-                    entries: &[
-                        BindGroupEntry {
-                            binding: 0,
-                            resource: BindingResource::TextureView(&textureview),
-                        },
-                        BindGroupEntry {
-                            binding: 1,
-                            resource: BindingResource::Sampler(&sampler),
-                        },
-                    ],
-                });
-                self.textures.insert(
-                    texid,
-                    (bindgroup, texture, delta.image.size().map(|e| e as u32)),
-                );
-            }
-        }
-
-        let clipped_primitives = self
-            .ctx
-            .tessellate(full_output.shapes, full_output.pixels_per_point);
-
-        let mut index_count = 0;
-        let mut vertex_count = 0;
-        for p in &clipped_primitives {
-            if let Primitive::Mesh(mesh) = &p.primitive {
-                index_count += mesh.indices.len();
-                vertex_count += mesh.vertices.len();
-            }
+                    size,
+                    mip_level_count: 1,
+                    sample_count: 1,
+                    dimension: TextureDimension::D2,
+                    format: TextureFormat::Rgba8UnormSrgb,
+                    usage: TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_DST,
+                    view_formats: &[],
+                },
+                TextureDataOrder::LayerMajor,
+                bytemuck::cast_slice::<_, u8>(&pixels),
+            );
+            let textureview = texture.create_view(&TextureViewDescriptor::default());
+            let sampler = self.device.create_sampler(&SamplerDescriptor {
+                address_mode_u: AddressMode::ClampToEdge,
+                address_mode_v: AddressMode::ClampToEdge,
+                mag_filter: FilterMode::Linear,
+                min_filter: FilterMode::Linear,
+                ..Default::default()
+            });
+            let bindgroup = self.device.create_bind_group(&BindGroupDescriptor {
+                label: None,
+                layout: &self.bind_group_layout,
+                entries: &[
+                    BindGroupEntry {
+                        binding: 0,
+                        resource: BindingResource::TextureView(&textureview),
+                    },
+                    BindGroupEntry {
+                        binding: 1,
+                        resource: BindingResource::Sampler(&sampler),
+                    },
+                ],
+            });
+            textures.insert(
+                texid,
+                (bindgroup, texture, delta.image.size().map(|e| e as u32)),
+            );
         }
+    }
 
-        if index_count == 0 || vertex_count == 0 {
+    pub fn draw(
+        &mut self,
+        commands: &mut CommandEncoder,
+        target: &TextureView,
+        depth: &TextureView,
+        projection: Mat4,
+    ) {
+        if self.surfaces.is_empty() {
             return;
         }
 
-        // TODO realloc buffers and retry if overflowing
-
-        let mut mapped_index = queue
-            .write_buffer_with(
-                &self.index,
-                0,
-                NonZeroU64::new((size_of::<u32>() * index_count) as u64).unwrap(),
-            )
-            .expect("ui index buffer overflow");
-        let mut mapped_vertex = queue
-            .write_buffer_with(
-                &self.vertex,
-                0,
-                NonZeroU64::new((size_of::<Vertex>() * vertex_count) as u64).unwrap(),
-            )
-            .expect("ui vertex buffer overflow");
-
-        let mut index_offset = 0;
-        let mut vertex_offset = 0;
-        let mut slices = Vec::new();
-        for p in clipped_primitives {
-            if let Primitive::Mesh(mesh) = p.primitive {
-                mapped_index[index_offset * size_of::<u32>()
-                    ..(index_offset + mesh.indices.len()) * size_of::<u32>()]
-                    .copy_from_slice(bytemuck::cast_slice(&mesh.indices));
-                mapped_vertex[vertex_offset * size_of::<Vertex>()
-                    ..(vertex_offset + mesh.vertices.len()) * size_of::<Vertex>()]
-                    .copy_from_slice(bytemuck::cast_slice(&mesh.vertices));
-                slices.push((
-                    index_offset as u32..index_offset as u32 + mesh.indices.len() as u32,
-                    vertex_offset as i32,
-                    mesh.texture_id,
-                ));
-                index_offset += mesh.indices.len();
-                vertex_offset += mesh.vertices.len();
-            }
-        }
-
-        assert_eq!(index_count, index_offset);
-        assert_eq!(vertex_count, vertex_offset);
-
         let mut rpass = commands.begin_render_pass(&RenderPassDescriptor {
             label: None,
             color_attachments: &[Some(RenderPassColorAttachment {
@@ -356,30 +291,107 @@ impl UiRenderer {
             ..Default::default()
         });
 
-        let affine = self
+        rpass.set_pipeline(&self.pipeline);
+
+        let mut raw_input = egui::RawInput::default();
+        raw_input.viewport_id = self.surfaces.keys().next().copied().unwrap();
+        raw_input.viewports = self
             .surfaces
-            .values()
-            .next()
-            .map(|s| s.transform)
-            .unwrap_or(Affine3A::IDENTITY);
+            .keys()
+            .map(|k| {
+                (*k, ViewportInfo {
+                    native_pixels_per_point: Some(2.),
+                    ..Default::default()
+                })
+            })
+            .collect();
 
-        let scale = 0.03;
-        let projection = projection
-            * Mat4::from_mat3a(affine.matrix3)
-            * Mat4::from_translation(affine.translation.into())
-            * Mat4::from_mat3(Mat3::from_mat2(Mat2::from_cols_array(&[
-                scale, 0., 0., -scale,
-            ])));
+        for (_viewport_id, surf) in &self.surfaces {
+            let full_output = self.ctx.run(raw_input.clone(), |ctx| (surf.content)(ctx));
 
-        let projection = projection.to_cols_array().map(|v| v.to_le_bytes());
+            for (texid, delta) in full_output.textures_delta.set {
+                self.apply_texture_delta(texid, delta);
+            }
 
-        rpass.set_pipeline(&self.pipeline);
-        rpass.set_push_constants(ShaderStages::VERTEX, 0, projection.as_flattened());
-        rpass.set_index_buffer(self.index.slice(..), IndexFormat::Uint32);
-        rpass.set_vertex_buffer(0, self.vertex.slice(..));
-        for (index, base_vertex, texid) in slices {
-            rpass.set_bind_group(0, &self.textures.get(&texid).unwrap().0, &[]);
-            rpass.draw_indexed(index, base_vertex, 0..1);
+            let clipped_primitives = self
+                .ctx
+                .tessellate(full_output.shapes, full_output.pixels_per_point);
+
+            let mut index_count = 0;
+            let mut vertex_count = 0;
+            for p in &clipped_primitives {
+                if let Primitive::Mesh(mesh) = &p.primitive {
+                    index_count += mesh.indices.len();
+                    vertex_count += mesh.vertices.len();
+                }
+            }
+
+            if index_count == 0 || vertex_count == 0 {
+                return;
+            }
+
+            // TODO realloc buffers and retry if overflowing
+
+            let mut mapped_index = self
+                .queue
+                .write_buffer_with(
+                    &surf.index,
+                    0,
+                    NonZeroU64::new((size_of::<u32>() * index_count) as u64).unwrap(),
+                )
+                .expect("ui index buffer overflow");
+            let mut mapped_vertex = self
+                .queue
+                .write_buffer_with(
+                    &surf.vertex,
+                    0,
+                    NonZeroU64::new((size_of::<Vertex>() * vertex_count) as u64).unwrap(),
+                )
+                .expect("ui vertex buffer overflow");
+
+            let mut index_offset = 0;
+            let mut vertex_offset = 0;
+            let mut slices = Vec::new();
+            for p in clipped_primitives {
+                if let Primitive::Mesh(mesh) = p.primitive {
+                    mapped_index[index_offset * size_of::<u32>()
+                        ..(index_offset + mesh.indices.len()) * size_of::<u32>()]
+                        .copy_from_slice(bytemuck::cast_slice(&mesh.indices));
+                    mapped_vertex[vertex_offset * size_of::<Vertex>()
+                        ..(vertex_offset + mesh.vertices.len()) * size_of::<Vertex>()]
+                        .copy_from_slice(bytemuck::cast_slice(&mesh.vertices));
+                    slices.push((
+                        index_offset as u32..index_offset as u32 + mesh.indices.len() as u32,
+                        vertex_offset as i32,
+                        mesh.texture_id,
+                    ));
+                    index_offset += mesh.indices.len();
+                    vertex_offset += mesh.vertices.len();
+                }
+            }
+
+            assert_eq!(index_count, index_offset);
+            assert_eq!(vertex_count, vertex_offset);
+
+            let scale = 0.03;
+            let projection = projection
+                * Mat4::from_mat3a(surf.transform.matrix3)
+                * Mat4::from_translation(surf.transform.translation.into())
+                * Mat4::from_mat3(Mat3::from_mat2(Mat2::from_cols_array(&[
+                    scale, 0., 0., -scale,
+                ])));
+
+            let projection = projection.to_cols_array().map(|v| v.to_le_bytes());
+
+            rpass.set_push_constants(ShaderStages::VERTEX, 0, projection.as_flattened());
+            rpass.set_index_buffer(surf.index.slice(..), IndexFormat::Uint32);
+            rpass.set_vertex_buffer(0, surf.vertex.slice(..));
+            for (index, base_vertex, texid) in slices {
+                let tex_guard = self.textures.read().unwrap();
+                let bind_group = &tex_guard.get(&texid).unwrap().0;
+                rpass.set_bind_group(0, bind_group, &[]);
+                rpass.draw_indexed(index, base_vertex, 0..1);
+            }
         }
     }
 }