albedo texture works

1 files changed, 197 insertions, 89 deletions

diff --git a/client/src/scene_prepare.rs b/client/src/scene_prepare.rs
index 85db160..f7a5bc8 100644
--- a/client/src/scene_prepare.rs
+++ b/client/src/scene_prepare.rs
@@ -1,31 +1,62 @@
 use crate::download::Downloader;
 use anyhow::{Context, Result};
+use image::{ImageFormat, ImageReader};
 use log::debug;
 use std::{
     collections::{HashMap, HashSet},
+    hash::Hash,
+    io::Cursor,
     sync::Arc,
 };
 use weareshared::{
     Affine3A,
     packets::{ReadWrite, Resource},
-    resources::{Attribute, Part, Prefab},
+    resources::{Part, Prefab},
 };
 use wgpu::{
-    Buffer, BufferUsages, Device,
-    util::{BufferInitDescriptor, DeviceExt},
+    BindGroup, BindGroupDescriptor, BindGroupEntry, BindGroupLayout, BindingResource, Buffer,
+    BufferUsages, Device, Extent3d, Queue, SamplerDescriptor, Texture, TextureDescriptor,
+    TextureDimension, TextureFormat, TextureUsages, TextureViewDescriptor,
+    util::{BufferInitDescriptor, DeviceExt, TextureDataOrder},
 };
 
-pub struct ScenePreparer {
-    index_buffers: HashMap<Resource, (Arc<Buffer>, u32)>,
-    index_buffers_needed: HashSet<Resource>,
-    vertex_buffers: HashMap<Resource, Arc<Buffer>>,
-    vertex_buffers_needed: HashSet<Resource>,
-    parts: HashMap<Resource, Arc<RPart>>,
-    parts_needed: HashSet<Resource>,
-    pub prefabs: HashMap<Resource, RPrefab>,
-    pub prefabs_needed: HashSet<Resource>,
+pub struct DemandMap<K, V> {
+    values: HashMap<K, V>,
+    needed: HashSet<K>,
+}
+impl<K: Hash + Eq, V: Clone> DemandMap<K, V> {
+    pub fn new() -> Self {
+        Self {
+            needed: HashSet::new(),
+            values: HashMap::new(),
+        }
+    }
+    pub fn insert(&mut self, key: K, value: V) {
+        self.needed.remove(&key);
+        self.values.insert(key, value);
+    }
+    pub fn try_get(&mut self, key: K) -> Option<V> {
+        if let Some(k) = self.values.get(&key) {
+            Some(k.to_owned())
+        } else {
+            self.needed.insert(key);
+            None
+        }
+    }
+}
 
+pub struct ScenePreparer {
     device: Arc<Device>,
+    queue: Arc<Queue>,
+    texture_bgl: BindGroupLayout,
+
+    textures: DemandMap<Resource, (Arc<Texture>, Arc<BindGroup>)>,
+    placeholder_textures: DemandMap<(), (Arc<Texture>, Arc<BindGroup>)>,
+    index_buffers: DemandMap<Resource, (Arc<Buffer>, u32)>,
+    vertex_buffers: DemandMap<Resource, (Arc<Buffer>, u32)>,
+    placeholder_vertex_buffers: DemandMap<(u32, bool), Arc<Buffer>>,
+    parts: DemandMap<Resource, Arc<RPart>>,
+    pub prefabs: DemandMap<Resource, Arc<RPrefab>>,
 }
 
 pub struct RPrefab(pub Vec<(Affine3A, Arc<RPart>)>);
@@ -35,44 +66,42 @@ pub struct RPart {
     pub position: [Arc<Buffer>; 3],
     pub normal: [Arc<Buffer>; 3],
     pub texcoord: [Arc<Buffer>; 2],
+    pub texture: Arc<BindGroup>,
 }
 
 impl ScenePreparer {
-    pub fn new(device: Arc<Device>) -> Self {
+    pub fn new(device: Arc<Device>, queue: Arc<Queue>, texture_bgl: BindGroupLayout) -> Self {
         Self {
-            index_buffers: HashMap::new(),
-            vertex_buffers: HashMap::new(),
-            vertex_buffers_needed: HashSet::new(),
-            parts: HashMap::new(),
-            parts_needed: HashSet::new(),
-            prefabs: HashMap::new(),
-            prefabs_needed: HashSet::new(),
-            index_buffers_needed: HashSet::new(),
+            texture_bgl,
+            index_buffers: DemandMap::new(),
+            vertex_buffers: DemandMap::new(),
+            parts: DemandMap::new(),
+            prefabs: DemandMap::new(),
+            textures: DemandMap::new(),
+            placeholder_vertex_buffers: DemandMap::new(),
+            placeholder_textures: DemandMap::new(),
             device,
+            queue,
         }
     }
     pub fn update(&mut self, dls: &mut Downloader) -> Result<()> {
-        let mut done = Vec::new();
-        for pres in &self.prefabs_needed {
-            if let Some(buf) = dls.try_get(*pres)? {
+        for pres in self.prefabs.needed.clone() {
+            if let Some(buf) = dls.try_get(pres)? {
                 let prefab = Prefab::read(&mut buf.as_slice()).context("parsing prefab")?;
                 let mut rprefab = RPrefab(Vec::new());
                 for (aff, partres) in &prefab.0 {
-                    if let Some(part) = self.parts.get(partres) {
+                    if let Some(part) = self.parts.try_get(*partres) {
                         rprefab.0.push((*aff, part.clone()));
-                    } else {
-                        self.parts_needed.insert(*partres);
                     }
                 }
                 if rprefab.0.len() == prefab.0.len() {
-                    self.prefabs.insert(*pres, rprefab);
+                    self.prefabs.insert(pres, Arc::new(rprefab));
                     debug!("prefab created ({pres})");
-                    done.push(*pres);
                 }
             }
         }
-        for pres in &self.index_buffers_needed {
-            if let Some(buf) = dls.try_get(*pres)? {
+        for pres in self.index_buffers.needed.clone() {
+            if let Some(buf) = dls.try_get(pres)? {
                 let buf = buf
                     .into_iter()
                     .array_chunks::<2>()
@@ -86,13 +115,12 @@ impl ScenePreparer {
                     usage: BufferUsages::INDEX | BufferUsages::COPY_DST,
                 });
                 self.index_buffers
-                    .insert(*pres, (Arc::new(buffer), buf.len() as u32 / 2));
-                debug!("index buffer created (len={}) {pres}", buf.len() / 6);
-                done.push(*pres);
+                    .insert(pres, (Arc::new(buffer), (buf.len() / 2) as u32));
+                debug!("index buffer created (len={}) {pres}", buf.len() / 2);
             }
         }
-        for pres in &self.vertex_buffers_needed {
-            if let Some(buf) = dls.try_get(*pres)? {
+        for pres in self.vertex_buffers.needed.clone() {
+            if let Some(buf) = dls.try_get(pres)? {
                 let buf = buf
                     .into_iter()
                     .array_chunks::<4>()
@@ -105,94 +133,174 @@ impl ScenePreparer {
                     label: None,
                     usage: BufferUsages::VERTEX | BufferUsages::COPY_DST,
                 });
-                self.vertex_buffers.insert(*pres, Arc::new(buffer));
+                self.vertex_buffers
+                    .insert(pres, (Arc::new(buffer), (buf.len() / 4) as u32));
                 debug!(
                     "vertex attribute buffer created (len={}) {pres}",
                     buf.len() / 4
                 );
-                done.push(*pres);
             }
         }
-        for pres in &self.parts_needed {
-            if let Some(buf) = dls.try_get(*pres)? {
+        for pres in self.textures.needed.clone() {
+            if let Some(buf) = dls.try_get(pres)? {
+                let image = ImageReader::new(Cursor::new(buf)).with_guessed_format()?;
+                let image = image.decode()?;
+                let image = image.to_rgba8();
+                let image_raw = image.to_vec();
+                let tex_bg = create_texture(
+                    &self.device,
+                    &self.queue,
+                    &self.texture_bgl,
+                    &image_raw,
+                    image.width(),
+                    image.height(),
+                );
+                self.textures.insert(pres, tex_bg);
+            }
+        }
+        for pres in self.placeholder_textures.needed.clone() {
+            let tex_bg = create_texture(
+                &self.device,
+                &self.queue,
+                &self.texture_bgl,
+                &[255, 255, 255, 255],
+                1,
+                1,
+            );
+            self.placeholder_textures.insert(pres, tex_bg);
+        }
+        for pres in self.parts.needed.clone() {
+            if let Some(buf) = dls.try_get(pres)? {
                 let part = Part::read(&mut buf.as_slice()).context("parsing part")?;
-                if let (Some(indexres), Some(positionres), Some(normalres), Some(texcoordres)) = (
-                    part.index,
-                    part.va_position,
-                    part.va_normal,
-                    part.va_texcoord,
-                ) {
-                    let Some((index, index_count)) = self.index_buffers.get(&indexres).cloned()
-                    else {
-                        self.index_buffers_needed.insert(indexres);
+                if let (Some(indexres), Some(positionres)) = (part.index, part.va_position) {
+                    let Some((index, index_count)) = self.index_buffers.try_get(indexres) else {
+                        self.index_buffers.needed.insert(indexres);
                         continue;
                     };
                     let mut position = Vec::new();
+                    let mut vertex_count = 0;
                     for vr in positionres {
-                        match vr {
-                            Attribute::Constant(_) => todo!(),
-                            Attribute::Vertex(resource) => {
-                                if let Some(vertex) = self.vertex_buffers.get(&resource).cloned() {
-                                    position.push(vertex);
-                                } else {
-                                    self.vertex_buffers_needed.insert(resource);
-                                };
-                            }
-                            Attribute::Texture(_resource, _ch) => todo!(),
+                        if let Some((vertex, n)) = self.vertex_buffers.try_get(vr) {
+                            vertex_count = n;
+                            position.push(vertex);
                         }
                     }
                     let mut normal = Vec::new();
-                    for vr in normalres {
-                        match vr {
-                            Attribute::Constant(_) => todo!(),
-                            Attribute::Vertex(resource) => {
-                                if let Some(vertex) = self.vertex_buffers.get(&resource).cloned() {
-                                    normal.push(vertex);
-                                } else {
-                                    self.vertex_buffers_needed.insert(resource);
-                                };
+                    if let Some(normalres) = part.va_normal {
+                        for vr in normalres {
+                            if let Some((vertex, _)) = self.vertex_buffers.try_get(vr) {
+                                normal.push(vertex);
+                            }
+                        }
+                    } else {
+                        // TODO generate normals
+                        for _ in 0..3 {
+                            if let Some(buf) = self
+                                .placeholder_vertex_buffers
+                                .try_get((vertex_count, false))
+                            {
+                                normal.push(buf);
                             }
-                            Attribute::Texture(_resource, _ch) => todo!(),
                         }
                     }
                     let mut texcoord = Vec::new();
-                    for vr in texcoordres {
-                        match vr {
-                            Attribute::Constant(_) => todo!(),
-                            Attribute::Vertex(resource) => {
-                                if let Some(vertex) = self.vertex_buffers.get(&resource).cloned() {
-                                    texcoord.push(vertex);
-                                } else {
-                                    self.vertex_buffers_needed.insert(resource);
-                                };
+                    if let Some(texcoordres) = part.va_texcoord {
+                        for vr in texcoordres {
+                            if let Some((vertex, _)) = self.vertex_buffers.try_get(vr) {
+                                texcoord.push(vertex);
                             }
-                            Attribute::Texture(_resource, _ch) => todo!(),
+                        }
+                    } else {
+                        // TODO generate UVs
+                        for _ in 0..3 {
+                            if let Some(buf) = self
+                                .placeholder_vertex_buffers
+                                .try_get((vertex_count, false))
+                            {
+                                texcoord.push(buf);
+                            }
+                        }
+                    }
+                    let mut texture = None;
+                    if let Some(albedores) = part.tex_pbr_albedo {
+                        if let Some((_tex, bg)) = self.textures.try_get(albedores) {
+                            texture = Some(bg)
+                        }
+                    } else {
+                        if let Some((_tex, bg)) = self.placeholder_textures.try_get(()) {
+                            texture = Some(bg)
                         }
                     }
-                    if texcoord.len() == 2 && normal.len() == 3 && position.len() == 3 {
+
+                    if texcoord.len() == 2
+                        && normal.len() == 3
+                        && position.len() == 3
+                        && texture.is_some()
+                    {
                         debug!("part created ({pres})");
                         self.parts.insert(
-                            *pres,
+                            pres,
                             Arc::new(RPart {
                                 index_count,
                                 index,
                                 texcoord: texcoord.try_into().unwrap(),
                                 normal: normal.try_into().unwrap(),
                                 position: position.try_into().unwrap(),
+                                texture: texture.unwrap(),
                             }),
                         );
-                        done.push(*pres);
                     }
                 }
             }
         }
-
-        for d in done {
-            self.parts_needed.remove(&d);
-            self.prefabs_needed.remove(&d);
-            self.index_buffers_needed.remove(&d);
-            self.vertex_buffers_needed.remove(&d);
-        }
         Ok(())
     }
 }
+
+fn create_texture(
+    device: &Device,
+    queue: &Queue,
+    bgl: &BindGroupLayout,
+    data: &[u8],
+    width: u32,
+    height: u32,
+) -> (Arc<Texture>, Arc<BindGroup>) {
+    let texture = device.create_texture_with_data(
+        &queue,
+        &TextureDescriptor {
+            label: None,
+            size: Extent3d {
+                depth_or_array_layers: 1,
+                width,
+                height,
+            },
+            mip_level_count: 1,
+            sample_count: 1,
+            dimension: TextureDimension::D2,
+            format: TextureFormat::Rgba8UnormSrgb,
+            usage: TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_DST,
+            view_formats: &[],
+        },
+        TextureDataOrder::LayerMajor,
+        data,
+    );
+    let textureview = texture.create_view(&TextureViewDescriptor::default());
+    let sampler = device.create_sampler(&SamplerDescriptor {
+        ..Default::default()
+    });
+    let bindgroup = device.create_bind_group(&BindGroupDescriptor {
+        label: None,
+        layout: &bgl,
+        entries: &[
+            BindGroupEntry {
+                binding: 0,
+                resource: BindingResource::TextureView(&textureview),
+            },
+            BindGroupEntry {
+                binding: 1,
+                resource: BindingResource::Sampler(&sampler),
+            },
+        ],
+    });
+    (Arc::new(texture), Arc::new(bindgroup))
+}