Remove vertex/frag

Author: tychedelia

crates/processing_pyo3/examples/custom_material.py

@@ -7,8 +7,8 @@ def setup():
     size(800, 600)
     mode_3d()
 
-    frag = Shader.load("shaders/custom_material.wesl")
-    mat = Material(fragment=frag, color=[1.0, 0.2, 0.4, 1.0])
+    shader = Shader.load("shaders/custom_material.wesl")
+    mat = Material(shader, color=[1.0, 0.2, 0.4, 1.0])
 
 def draw():
     camera_position(0.0, 0.0, 200.0)

crates/processing_pyo3/src/material.rs

@@ -37,14 +37,10 @@ fn py_to_material_value(value: &Bound<'_, PyAny>) -> PyResult<material::Material
 #[pymethods]
 impl Material {
     #[new]
-    #[pyo3(signature = (vertex=None, fragment=None, **kwargs))]
-    pub fn new(
-        vertex: Option<&Shader>,
-        fragment: Option<&Shader>,
-        kwargs: Option<&Bound<'_, PyDict>>,
-    ) -> PyResult<Self> {
-        let entity = if vertex.is_some() || fragment.is_some() {
-            material_create_custom(vertex.map(|s| s.entity), fragment.map(|s| s.entity))
+    #[pyo3(signature = (shader=None, **kwargs))]
+    pub fn new(shader: Option<&Shader>, kwargs: Option<&Bound<'_, PyDict>>) -> PyResult<Self> {
+        let entity = if let Some(shader) = shader {
+            material_create_custom(shader.entity)
                 .map_err(|e| PyRuntimeError::new_err(format!("{e}")))?
         } else {
             material_create_pbr().map_err(|e| PyRuntimeError::new_err(format!("{e}")))?

crates/processing_render/src/lib.rs

@@ -1333,13 +1333,10 @@ pub fn shader_destroy(entity: Entity) -> error::Result<()> {
     })
 }
 
-pub fn material_create_custom(
-    vertex: Option<Entity>,
-    fragment: Option<Entity>,
-) -> error::Result<Entity> {
+pub fn material_create_custom(shader: Entity) -> error::Result<Entity> {
     app_mut(|app| {
         app.world_mut()
-            .run_system_cached_with(material::custom::create_custom, (vertex, fragment))
+            .run_system_cached_with(material::custom::create_custom, shader)
             .unwrap()
     })
 }

crates/processing_render/src/material/custom.rs

@@ -50,8 +50,9 @@ use crate::render::material::UntypedMaterial;
 #[derive(Asset, TypePath, Clone)]
 pub struct CustomMaterial {
     pub shader: DynamicShader,
-    pub vertex_shader: Option<Handle<ShaderAsset>>,
-    pub fragment_shader: Option<Handle<ShaderAsset>>,
+    pub shader_handle: Handle<ShaderAsset>,
+    pub has_vertex: bool,
+    pub has_fragment: bool,
 }
 
 #[derive(Component)]
@@ -197,34 +198,35 @@ pub fn destroy_shader(In(entity): In<Entity>, mut commands: Commands) -> Result<
 }
 
 pub fn create_custom(
-    In((vertex_entity, fragment_entity)): In<(Option<Entity>, Option<Entity>)>,
+    In(shader_entity): In<Entity>,
     mut commands: Commands,
     shader_programs: Query<&Shader>,
     mut custom_materials: ResMut<Assets<CustomMaterial>>,
 ) -> Result<Entity> {
-    let vertex_program = vertex_entity
-        .map(|e| shader_programs.get(e))
-        .transpose()
-        .map_err(|_| ProcessingError::ShaderNotFound)?;
-    let fragment_program = fragment_entity
-        .map(|e| shader_programs.get(e))
-        .transpose()
+    let program = shader_programs
+        .get(shader_entity)
         .map_err(|_| ProcessingError::ShaderNotFound)?;
 
-    // Prefer fragment module for reflection, fall back to vertex.
-    let reflection_module = fragment_program
-        .map(|p| &p.module)
-        .or(vertex_program.map(|p| &p.module))
-        .ok_or(ProcessingError::ShaderNotFound)?;
-
-    let mut shader = DynamicShader::new(reflection_module.clone())
+    let has_vertex = program
+        .module
+        .entry_points
+        .iter()
+        .any(|ep| ep.stage == naga::ShaderStage::Vertex);
+    let has_fragment = program
+        .module
+        .entry_points
+        .iter()
+        .any(|ep| ep.stage == naga::ShaderStage::Fragment);
+
+    let mut shader = DynamicShader::new(program.module.clone())
         .map_err(|e| ProcessingError::ShaderCompilationError(e.to_string()))?;
     shader.init();
 
     let material = CustomMaterial {
         shader,
-        vertex_shader: vertex_program.map(|p| p.shader_handle.clone()),
-        fragment_shader: fragment_program.map(|p| p.shader_handle.clone()),
+        shader_handle: program.shader_handle.clone(),
+        has_vertex,
+        has_fragment,
     };
     let handle = custom_materials.add(material);
     Ok(commands.spawn(UntypedMaterial(handle.untyped())).id())
@@ -401,11 +403,11 @@ impl ErasedRenderAsset for CustomMaterial {
             ..Default::default()
         };
         properties.add_draw_function(MainPassOpaqueDrawFunction, draw_function);
-        if let Some(vertex) = &source_asset.vertex_shader {
-            properties.add_shader(MaterialVertexShader, vertex.clone());
+        if source_asset.has_vertex {
+            properties.add_shader(MaterialVertexShader, source_asset.shader_handle.clone());
         }
-        if let Some(fragment) = &source_asset.fragment_shader {
-            properties.add_shader(MaterialFragmentShader, fragment.clone());
+        if source_asset.has_fragment {
+            properties.add_shader(MaterialFragmentShader, source_asset.shader_handle.clone());
         }
 
         Ok(PreparedMaterial {

examples/custom_material.rs

@@ -31,8 +31,8 @@ fn sketch() -> error::Result<()> {
     transform_set_position(graphics, 100.0, 100.0, 300.0)?;
     transform_look_at(graphics, 0.0, 0.0, 0.0)?;
 
-    let frag = shader_load("shaders/custom_material.wesl")?;
-    let mat = material_create_custom(None, Some(frag))?;
+    let shader = shader_load("shaders/custom_material.wesl")?;
+    let mat = material_create_custom(shader)?;
     material_set(
         mat,
         "color",
@@ -61,7 +61,7 @@ fn sketch() -> error::Result<()> {
     }
 
     material_destroy(mat)?;
-    shader_destroy(frag)?;
+    shader_destroy(shader)?;
 
     Ok(())
 }