From 4d92ef0119edf4f2391f35a0fb7bad92d3bc6b51 Mon Sep 17 00:00:00 2001
From: Carter Anderson <mcanders1@gmail.com>
Date: Sat, 14 Mar 2020 18:32:33 -0700
Subject: [PATCH] zero copy uniform bytes

---
 bevy_derive/src/lib.rs                              |  9 ++++++++-
 src/core/bytes.rs                                   | 13 +------------
 src/render/color.rs                                 |  5 ++++-
 .../resource_providers/uniform_resource_provider.rs |  8 ++++++--
 src/render/shader/uniform.rs                        |  3 +--
 src/render/shader/uniforms/local_to_world.rs        |  9 ++++++++-
 6 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/bevy_derive/src/lib.rs b/bevy_derive/src/lib.rs
index ee4681baa6..781bb8f859 100644
--- a/bevy_derive/src/lib.rs
+++ b/bevy_derive/src/lib.rs
@@ -158,7 +158,6 @@ pub fn derive_uniforms(input: TokenStream) -> TokenStream {
                 }
             }
 
-            // TODO: Fix this so uniform_name_uniform_info lines up with getbytes
             fn get_uniform_bytes(&self, name: &str) -> Option<Vec<u8>> {
                 use bevy::core::bytes::GetBytes;
                 match name {
@@ -167,6 +166,14 @@ pub fn derive_uniforms(input: TokenStream) -> TokenStream {
                 }
             }
 
+            fn get_uniform_bytes_ref(&self, name: &str) -> Option<&[u8]> {
+                use bevy::core::bytes::GetBytes;
+                match name {
+                    #(#uniform_name_strings => self.#active_uniform_field_names.get_bytes_ref(),)*
+                    _ => None,
+                }
+            }
+
             fn get_uniform_texture(&self, name: &str) -> Option<bevy::asset::Handle<bevy::render::texture::Texture>> {
                 use bevy::render::shader::GetTexture;
                 match name {
diff --git a/src/core/bytes.rs b/src/core/bytes.rs
index 1e94fe61b5..36410e41e5 100644
--- a/src/core/bytes.rs
+++ b/src/core/bytes.rs
@@ -6,17 +6,6 @@ pub trait GetBytes {
     fn get_bytes_ref(&self) -> Option<&[u8]>;
 }
 
-// TODO: might need to add zerocopy to this crate to impl AsBytes for external crates
-// impl<T> GetBytes for T where T : AsBytes {
-//     fn get_bytes(&self) -> Vec<u8> {
-//         self.as_bytes().into()
-//     }
-
-//     fn get_bytes_ref(&self) -> Option<&[u8]> {
-//         Some(self.as_bytes())
-//     }
-// }
-
 impl GetBytes for Vec4 {
     fn get_bytes(&self) -> Vec<u8> {
         let vec4_array: [f32; 4] = (*self).into();
@@ -24,7 +13,7 @@ impl GetBytes for Vec4 {
     }
 
     fn get_bytes_ref(&self) -> Option<&[u8]> {
-        None
+        Some(self.as_ref().as_bytes())
     }
 }
 
diff --git a/src/render/color.rs b/src/render/color.rs
index 6aee91f715..82648a6fec 100644
--- a/src/render/color.rs
+++ b/src/render/color.rs
@@ -112,6 +112,9 @@ impl GetBytes for ColorSource {
         }
     }
     fn get_bytes_ref(&self) -> Option<&[u8]> {
-        None
+        match *self {
+            ColorSource::Color(ref color) => color.get_bytes_ref(),
+            ColorSource::Texture(ref texture) => texture.get_bytes_ref(), // Texture is not a uniform
+        }
     }
 }
diff --git a/src/render/render_resource/resource_providers/uniform_resource_provider.rs b/src/render/render_resource/resource_providers/uniform_resource_provider.rs
index 899e04db87..4cdad81399 100644
--- a/src/render/render_resource/resource_providers/uniform_resource_provider.rs
+++ b/src/render/render_resource/resource_providers/uniform_resource_provider.rs
@@ -201,8 +201,12 @@ where
                         }
                         // TODO: check if index has changed. if it has, then entity should be updated
                         // TODO: only mem-map entities if their data has changed
-                        // TODO: try getting bytes ref first
-                        if let Some(uniform_bytes) = uniforms.get_uniform_bytes(&name) {
+                        if let Some(uniform_bytes) = uniforms.get_uniform_bytes_ref(&name) {
+                            mapped[offset..(offset + uniform_bytes.len())]
+                                .copy_from_slice(uniform_bytes);
+                            offset += alignment;
+                        }
+                        else if let Some(uniform_bytes) = uniforms.get_uniform_bytes(&name) {
                             mapped[offset..(offset + uniform_bytes.len())]
                                 .copy_from_slice(uniform_bytes.as_slice());
                             offset += alignment;
diff --git a/src/render/shader/uniform.rs b/src/render/shader/uniform.rs
index 4fa717d39e..12600a0ae3 100644
--- a/src/render/shader/uniform.rs
+++ b/src/render/shader/uniform.rs
@@ -17,8 +17,7 @@ pub trait AsUniforms {
     fn get_uniform_texture(&self, name: &str) -> Option<Handle<Texture>>;
     fn get_shader_defs(&self) -> Option<Vec<String>>;
     fn get_field_bind_type(&self, name: &str) -> Option<FieldBindType>;
-    // TODO: support zero-copy uniforms
-    // fn get_uniform_bytes_ref(&self, name: &str) -> Option<&[u8]>;
+    fn get_uniform_bytes_ref(&self, name: &str) -> Option<&[u8]>;
 }
 
 pub trait ShaderDefSuffixProvider {
diff --git a/src/render/shader/uniforms/local_to_world.rs b/src/render/shader/uniforms/local_to_world.rs
index 53c5c22c01..e16f60621f 100644
--- a/src/render/shader/uniforms/local_to_world.rs
+++ b/src/render/shader/uniforms/local_to_world.rs
@@ -22,7 +22,7 @@ impl AsUniforms for bevy_transform::prelude::LocalToWorld {
 
     fn get_uniform_bytes(&self, name: &str) -> Option<Vec<u8>> {
         match name {
-            "Object" => Some(self.0.to_cols_array_2d().as_bytes().into()),
+            "Object" => Some(self.0.as_ref().as_bytes().into()),
             _ => None,
         }
     }
@@ -39,4 +39,11 @@ impl AsUniforms for bevy_transform::prelude::LocalToWorld {
     fn get_uniform_texture(&self, _name: &str) -> Option<Handle<Texture>> {
         None
     }
+
+    fn get_uniform_bytes_ref(&self, name: &str) -> Option<&[u8]> {
+        match name {
+            "Object" => Some(self.0.as_ref().as_bytes()),
+            _ => None,
+        }
+    }
 }