Handle attribute pointer in functions (#277)

Ensure that the generated code within a function, which generally
accesses the particle attribute struct via a pointer, is emitted with a
pointer indirection.

This is a partial fix, as `make_fn()` could be called with a non-pointer
attribute struct, which would break with the reverse error. However for
now it's only called in the update context with a pointer.

Fixes #275

src/graph/expr.rs

@@ -560,6 +560,15 @@ pub trait EvalContext {
         module: &mut Module,
         f: &mut dyn FnMut(&mut Module, &mut dyn EvalContext) -> Result<String, ExprError>,
     ) -> Result<(), ExprError>;
+
+    /// Check if the particle attribute struct is a pointer?
+    ///
+    /// In some context the attribute struct (named 'particle' in WGSL code) is
+    /// a pointer instead of being a struct instance. This happens in particular
+    /// when defining a function for a modifier, and passing the attribute
+    /// struct to be modified. In that case the generated code needs to emit
+    /// a pointer indirection code to access the fields of the struct.
+    fn is_attribute_pointer(&self) -> bool;
 }
 
 /// Language expression producing a value.
@@ -923,8 +932,12 @@ impl AttributeExpr {
     }
 
     /// Evaluate the expression in the given context.
-    pub fn eval(&self, _context: &dyn EvalContext) -> Result<String, ExprError> {
-        Ok(self.to_wgsl_string())
+    pub fn eval(&self, context: &dyn EvalContext) -> Result<String, ExprError> {
+        if context.is_attribute_pointer() {
+            Ok(format!("(*particle).{}", self.attr.name()))
+        } else {
+            Ok(format!("particle.{}", self.attr.name()))
+        }
     }
 }
 
@@ -3491,6 +3504,30 @@ mod tests {
         }
     }
 
+    #[test]
+    fn attribute_pointer() {
+        let mut m = Module::default();
+        let x = m.attr(Attribute::POSITION);
+
+        let property_layout = PropertyLayout::default();
+        let particle_layout = ParticleLayout::default();
+        let mut ctx = InitContext::new(&property_layout, &particle_layout);
+
+        let res = ctx.eval(&m, x);
+        assert!(res.is_ok());
+        let xx = res.ok().unwrap();
+        assert_eq!(xx, format!("particle.{}", Attribute::POSITION.name()));
+
+        // Use a different context; it's invalid to reuse a mutated context, as the
+        // expression cache will have been generated with the wrong context.
+        let mut ctx = InitContext::new(&property_layout, &particle_layout).with_attribute_pointer();
+
+        let res = ctx.eval(&m, x);
+        assert!(res.is_ok());
+        let xx = res.ok().unwrap();
+        assert_eq!(xx, format!("(*particle).{}", Attribute::POSITION.name()));
+    }
+
     #[test]
     #[should_panic]
     fn invalid_cast_vector_to_scalar() {

src/modifier/accel.rs

@@ -177,21 +177,25 @@ impl UpdateModifier for RadialAccelModifier {
         let func_id = calc_func_id(self);
         let func_name = format!("radial_accel_{0:016X}", func_id);
 
-        let origin = context.eval(module, self.origin)?;
-        let accel = context.eval(module, self.accel)?;
-
-        context.update_extra += &format!(
-            r##"fn {}(particle: ptr<function, Particle>) {{
-    let radial = normalize((*particle).{} - {});
-    (*particle).{} += radial * (({}) * sim_params.delta_time);
-}}
-"##,
-            func_name,
-            Attribute::POSITION.name(),
-            origin,
-            Attribute::VELOCITY.name(),
-            accel,
-        );
+        context.make_fn(
+            &func_name,
+            "particle: ptr<function, Particle>",
+            module,
+            &mut |m: &mut Module, ctx: &mut dyn EvalContext| -> Result<String, ExprError> {
+                let origin = ctx.eval(m, self.origin)?;
+                let accel = ctx.eval(m, self.accel)?;
+
+                Ok(format!(
+                    r##"let radial = normalize((*particle).{} - {});
+            (*particle).{} += radial * (({}) * sim_params.delta_time);
+        "##,
+                    Attribute::POSITION.name(),
+                    origin,
+                    Attribute::VELOCITY.name(),
+                    accel,
+                ))
+            },
+        )?;
 
         context.update_code += &format!("{}(&particle);\n", func_name);
 
@@ -315,7 +319,7 @@ impl UpdateModifier for TangentAccelModifier {
 
 #[cfg(test)]
 mod tests {
-    use crate::{ParticleLayout, PropertyLayout, ToWgslString};
+    use crate::{ParticleLayout, Property, PropertyLayout, ToWgslString};
 
     use super::*;
 
@@ -336,17 +340,25 @@ mod tests {
     #[test]
     fn mod_radial_accel() {
         let mut module = Module::default();
+        let property_layout = PropertyLayout::new(&[Property::new("my_prop", 3.)]);
+        let particle_layout = ParticleLayout::default();
+
         let origin = Vec3::new(-1.2, 5.3, -8.5);
         let accel = 6.;
         let modifier = RadialAccelModifier::constant(&mut module, origin, accel);
-
-        let property_layout = PropertyLayout::default();
-        let particle_layout = ParticleLayout::default();
         let mut context = UpdateContext::new(&property_layout, &particle_layout);
         assert!(modifier.apply_update(&mut module, &mut context).is_ok());
-
         // TODO: less weak check...
         assert!(context.update_extra.contains(&accel.to_wgsl_string()));
+
+        let origin = module.attr(Attribute::POSITION);
+        let accel = module.prop("my_prop");
+        let modifier = RadialAccelModifier::new(origin, accel);
+        let mut context = UpdateContext::new(&property_layout, &particle_layout);
+        assert!(modifier.apply_update(&mut module, &mut context).is_ok());
+        // TODO: less weak check...
+        assert!(context.update_extra.contains(Attribute::POSITION.name()));
+        assert!(context.update_extra.contains("my_prop"));
     }
 
     #[test]

src/modifier/mod.rs

@@ -169,6 +169,8 @@ pub struct InitContext<'a> {
     var_counter: u32,
     /// Cache of evaluated expressions.
     expr_cache: HashMap<ExprHandle, String>,
+    /// Is the attriubute struct a pointer?
+    is_attribute_pointer: bool,
 }
 
 impl<'a> InitContext<'a> {
@@ -181,8 +183,15 @@ impl<'a> InitContext<'a> {
             particle_layout,
             var_counter: 0,
             expr_cache: Default::default(),
+            is_attribute_pointer: false,
         }
     }
+
+    /// Mark the attribute struct as being available through a pointer.
+    pub fn with_attribute_pointer(mut self) -> Self {
+        self.is_attribute_pointer = true;
+        self
+    }
 }
 
 impl<'a> EvalContext for InitContext<'a> {
@@ -230,7 +239,9 @@ impl<'a> EvalContext for InitContext<'a> {
         f: &mut dyn FnMut(&mut Module, &mut dyn EvalContext) -> Result<String, ExprError>,
     ) -> Result<(), ExprError> {
         // Generate a temporary context for the function content itself
-        let mut ctx = InitContext::new(self.property_layout, self.particle_layout);
+        // FIXME - Dynamic with_attribute_pointer()!
+        let mut ctx =
+            InitContext::new(self.property_layout, self.particle_layout).with_attribute_pointer();
 
         // Evaluate the function content
         let body = f(module, &mut ctx)?;
@@ -247,6 +258,10 @@ impl<'a> EvalContext for InitContext<'a> {
 
         Ok(())
     }
+
+    fn is_attribute_pointer(&self) -> bool {
+        self.is_attribute_pointer
+    }
 }
 
 /// Trait to customize the initializing of newly spawned particles.
@@ -351,6 +366,8 @@ pub struct UpdateContext<'a> {
     var_counter: u32,
     /// Cache of evaluated expressions.
     expr_cache: HashMap<ExprHandle, String>,
+    /// Is the attriubute struct a pointer?
+    is_attribute_pointer: bool,
 
     // TEMP
     /// Array of force field components with a maximum number of components
@@ -368,9 +385,16 @@ impl<'a> UpdateContext<'a> {
             particle_layout,
             var_counter: 0,
             expr_cache: Default::default(),
+            is_attribute_pointer: false,
             force_field: [ForceFieldSource::default(); ForceFieldSource::MAX_SOURCES],
         }
     }
+
+    /// Mark the attribute struct as being available through a pointer.
+    pub fn with_attribute_pointer(mut self) -> Self {
+        self.is_attribute_pointer = true;
+        self
+    }
 }
 
 impl<'a> EvalContext for UpdateContext<'a> {
@@ -418,7 +442,9 @@ impl<'a> EvalContext for UpdateContext<'a> {
         f: &mut dyn FnMut(&mut Module, &mut dyn EvalContext) -> Result<String, ExprError>,
     ) -> Result<(), ExprError> {
         // Generate a temporary context for the function content itself
-        let mut ctx = UpdateContext::new(self.property_layout, self.particle_layout);
+        // FIXME - Dynamic with_attribute_pointer()!
+        let mut ctx =
+            UpdateContext::new(self.property_layout, self.particle_layout).with_attribute_pointer();
 
         // Evaluate the function content
         let body = f(module, &mut ctx)?;
@@ -437,6 +463,10 @@ impl<'a> EvalContext for UpdateContext<'a> {
 
         Ok(())
     }
+
+    fn is_attribute_pointer(&self) -> bool {
+        self.is_attribute_pointer
+    }
 }
 
 /// Trait to customize the updating of existing particles each frame.
@@ -483,6 +513,8 @@ pub struct RenderContext<'a> {
     var_counter: u32,
     /// Cache of evaluated expressions.
     expr_cache: HashMap<ExprHandle, String>,
+    /// Is the attriubute struct a pointer?
+    is_attribute_pointer: bool,
 }
 
 impl<'a> RenderContext<'a> {
@@ -502,6 +534,7 @@ impl<'a> RenderContext<'a> {
             screen_space_size: false,
             var_counter: 0,
             expr_cache: Default::default(),
+            is_attribute_pointer: false,
         }
     }
 
@@ -535,6 +568,12 @@ impl<'a> RenderContext<'a> {
         let func_name = format!("size_gradient_{0:016X}", func_id);
         func_name
     }
+
+    /// Mark the attribute struct as being available through a pointer.
+    pub fn with_attribute_pointer(mut self) -> Self {
+        self.is_attribute_pointer = true;
+        self
+    }
 }
 
 impl<'a> EvalContext for RenderContext<'a> {
@@ -583,7 +622,9 @@ impl<'a> EvalContext for RenderContext<'a> {
         f: &mut dyn FnMut(&mut Module, &mut dyn EvalContext) -> Result<String, ExprError>,
     ) -> Result<(), ExprError> {
         // Generate a temporary context for the function content itself
-        let mut ctx = RenderContext::new(self.property_layout, self.particle_layout);
+        // FIXME - Dynamic with_attribute_pointer()!
+        let mut ctx =
+            RenderContext::new(self.property_layout, self.particle_layout).with_attribute_pointer();
 
         // Evaluate the function content
         let body = f(module, &mut ctx)?;
@@ -602,6 +643,10 @@ impl<'a> EvalContext for RenderContext<'a> {
 
         Ok(())
     }
+
+    fn is_attribute_pointer(&self) -> bool {
+        self.is_attribute_pointer
+    }
 }
 
 /// Trait to customize the rendering of alive particles each frame.