diff --git a/CHANGELOG.md b/CHANGELOG.md
index f707cc93..b1e80de1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,17 +11,30 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added a new `ConformToSphereModifier` acting as an attractor applying a force toward a point (sphere center) to all particles in range, and making particles conform ("stick") to the sphere surface.
 - Added `vec2` and `vec3` functions that allow construction of vectors from dynamic parts.
 - Added missing `VectorValue::as_uvecX()` and `impl From<UVecX> for VectorValue` for all `X=2,3,4`.
+- Added a new `ShaderWriter` helper, common to the init and update modifier contexts, and which replaces the `InitContext` and `UpdateContext`.
+- Added a simple `impl Display for ModifierContext`.
+- Added `Modifier::apply()` which replaces the now deleted `InitModifier::apply_init()` and `UdpdateModifier::apply_update()`.
+- Added `RenderModifier::as_modifier()` to upcast a `RenderModifier` to a `Modifier`.
+- Added `RenderModifier::boxed_render_clone()` to clone a `RenderModifier` into a boxed self (instead of a `BoxedModifier`, as we can't upcast from `Box<dyn RenderModifier>` to `Box<dyn Modifier>`). Added `impl Clone for Box<dyn RenderModifier>` based on this.
+- Added `EffectAsset::add_modifier()` to add a pre-boxed `Modifier` (so, a `BoxedModifier`) to the init or update context, and added `EffectAsset::add_render_modifier()` to add a pre-boxed `RenderModifier` to the render context.
 
 ### Changed
 
 - `ExprHandle` is now `#[repr(transparent)]`, which guarantees that `Option<ExprHandle>` has the same size as `ExprHandle` itself (4 bytes).
 - `EffectProperties::set_if_changed()` now returns the `Mut` variable it takes as input, to allow subsequent calls.
 - `VectorValue::new_uvecX()` now take a `UVecX` instead of individual components, like for all other scalar types.
+- Merged the `InitModifier` and `UpdateModifier` traits into the `Modifier` subtrait; see other changelog entries for details. This helps manage modifiers in a unified way, and generally simplifies writing and maintain modifiers compatible with both the init and update contexts.
+- `EffectAsset::init()` and `EffectAsset::update()` now take a `Modifier`-bound type, and validate its `ModifierContext` is compatible (and panics if not).
+- `EffectAsset::render()` now panics if the modifier is not compatible with the `ModifierContext::Render`. Note that this indicates a malformed render modifier, because all objects implementing `RenderModifier` must include `ModifierContext::Render` in their `Modifier::context()`.
 
 ### Removed
 
 - Removed the `screen_space_size` field from the `SetSizeModifier`. Use the new `ScreenSpaceSizeModifier` to use a screen-space size.
 - Removed the built-in `ForceFieldSource` and associated `ForceFieldModifier`. Use the new `ConformToSphereModifer` instead. The behavior might change a bit as the conforming code is not strictly identical; use the `force_field.rs` example with the `examples_world_inspector` feature to tweak the parameters in real time and observe how they work and change the effect.
+- Removed `InitContext` and `UpdateContext`; they're replaced with `ShaderWriter`.
+- Removed `InitModifer` and `UpdateModifer`. Modifiers for the init and update contexts now only need to implement the base `Modifier` trait.
+- Removed downcast methods from `Modifier` (`as_init()`, `as_init_mut()`, `as_update()`, `as_update_mut()`).
+- Removed the various helper macros `impl_mod_xxx!()` to implement modifier traits; simply implement the trait by hand instead.
 
 ### Fixed
 
diff --git a/src/asset.rs b/src/asset.rs
index f7a3e0bb..f8178693 100644
--- a/src/asset.rs
+++ b/src/asset.rs
@@ -4,12 +4,13 @@ use bevy::{
     utils::{default, thiserror::Error, BoxedFuture, HashSet},
 };
 use serde::{Deserialize, Serialize};
+use std::ops::Deref;
 
 use crate::{
     graph::Value,
-    modifier::{InitModifier, RenderModifier, UpdateModifier},
-    BoxedModifier, ExprHandle, Module, ParticleLayout, Property, PropertyLayout, SimulationSpace,
-    Spawner,
+    modifier::{Modifier, RenderModifier},
+    BoxedModifier, ExprHandle, ModifierContext, Module, ParticleLayout, Property, PropertyLayout,
+    SimulationSpace, Spawner,
 };
 
 /// Type of motion integration applied to the particles of a system.
@@ -212,7 +213,7 @@ pub struct EffectAsset {
     /// Render modifiers defining the effect.
     #[reflect(ignore)]
     // TODO - Can't manage to implement FromReflect for BoxedModifier in a nice way yet
-    render_modifiers: Vec<BoxedModifier>,
+    render_modifiers: Vec<Box<dyn RenderModifier>>,
     /// Properties of the effect.
     ///
     /// Properties must have a unique name. Manually adding two or more
@@ -374,49 +375,101 @@ impl EffectAsset {
 
     /// Add an initialization modifier to the effect.
     ///
-    /// [`with_property()`]: crate::EffectAsset::with_property
-    /// [`add_property()`]: crate::EffectAsset::add_property
+    /// # Panics
+    ///
+    /// Panics if the modifier doesn't support the init context (that is,
+    /// `modifier.context()` returns a flag which doesn't include
+    /// [`ModifierContext::Init`]).
     #[inline]
     pub fn init<M>(mut self, modifier: M) -> Self
     where
-        M: InitModifier + Send + Sync + 'static,
+        M: Modifier + Send + Sync + 'static,
     {
+        assert!(modifier.context().contains(ModifierContext::Init));
         self.init_modifiers.push(Box::new(modifier));
         self
     }
 
     /// Add an update modifier to the effect.
     ///
-    /// [`with_property()`]: crate::EffectAsset::with_property
-    /// [`add_property()`]: crate::EffectAsset::add_property
+    /// # Panics
+    ///
+    /// Panics if the modifier doesn't support the update context (that is,
+    /// `modifier.context()` returns a flag which doesn't include
+    /// [`ModifierContext::Update`]).
     #[inline]
     pub fn update<M>(mut self, modifier: M) -> Self
     where
-        M: UpdateModifier + Send + Sync + 'static,
+        M: Modifier + Send + Sync + 'static,
     {
+        assert!(modifier.context().contains(ModifierContext::Update));
         self.update_modifiers.push(Box::new(modifier));
         self
     }
 
+    /// Add a [`BoxedModifier`] to the specific context.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the context is [`ModifierContext::Render`]; use
+    /// [`add_render_modifier()`] instead.
+    ///
+    /// Panics if the input `context` contains more than one context (the
+    /// bitfield contains more than 1 bit set) or no context at all (zero bit
+    /// set).
+    ///
+    /// Panics if the modifier doesn't support the context specified (that is,
+    /// `modifier.context()` returns a flag which doesn't include `context`).
+    ///
+    /// [`add_render_modifier()`]: crate::EffectAsset::add_render_modifier
+    pub fn add_modifier(mut self, context: ModifierContext, modifier: Box<dyn Modifier>) -> Self {
+        assert!(context == ModifierContext::Init || context == ModifierContext::Update);
+        assert!(modifier.context().contains(context));
+        if context == ModifierContext::Init {
+            self.init_modifiers.push(modifier);
+        } else {
+            self.update_modifiers.push(modifier);
+        }
+        self
+    }
+
     /// Add a render modifier to the effect.
     ///
-    /// [`with_property()`]: crate::EffectAsset::with_property
-    /// [`add_property()`]: crate::EffectAsset::add_property
+    /// # Panics
+    ///
+    /// Panics if the modifier doesn't support the render context (that is,
+    /// `modifier.context()` returns a flag which doesn't include
+    /// [`ModifierContext::Render`]).
     #[inline]
     pub fn render<M>(mut self, modifier: M) -> Self
     where
         M: RenderModifier + Send + Sync + 'static,
     {
+        assert!(modifier.context().contains(ModifierContext::Render));
         self.render_modifiers.push(Box::new(modifier));
         self
     }
 
+    /// Add a [`RenderModifier`] to the render context.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the modifier doesn't support the render context (that is,
+    /// `modifier.context()` returns a flag which doesn't include
+    /// [`ModifierContext::Render`]).
+    pub fn add_render_modifier(mut self, modifier: Box<dyn RenderModifier>) -> Self {
+        assert!(modifier.context().contains(ModifierContext::Render));
+        self.render_modifiers.push(modifier);
+        self
+    }
+
     /// Get a list of all the modifiers of this effect.
-    pub fn modifiers(&self) -> impl Iterator<Item = &BoxedModifier> {
+    pub fn modifiers(&self) -> impl Iterator<Item = &dyn Modifier> {
         self.init_modifiers
             .iter()
-            .chain(self.update_modifiers.iter())
-            .chain(self.render_modifiers.iter())
+            .map(|m| m.deref())
+            .chain(self.update_modifiers.iter().map(|m| m.deref()))
+            .chain(self.render_modifiers.iter().map(|m| m.as_modifier()))
     }
 
     /// Get a list of all the init modifiers of this effect.
@@ -425,8 +478,14 @@ impl EffectAsset {
     /// executing in the [`ModifierContext::Init`] context.
     ///
     /// [`ModifierContext::Init`]: crate::ModifierContext::Init
-    pub fn init_modifiers(&self) -> impl Iterator<Item = &dyn InitModifier> {
-        self.init_modifiers.iter().filter_map(|m| m.as_init())
+    pub fn init_modifiers(&self) -> impl Iterator<Item = &dyn Modifier> {
+        self.init_modifiers.iter().filter_map(|m| {
+            if m.context().contains(ModifierContext::Init) {
+                Some(m.deref())
+            } else {
+                None
+            }
+        })
     }
 
     /// Get a list of all the update modifiers of this effect.
@@ -435,8 +494,14 @@ impl EffectAsset {
     /// executing in the [`ModifierContext::Update`] context.
     ///
     /// [`ModifierContext::Update`]: crate::ModifierContext::Update
-    pub fn update_modifiers(&self) -> impl Iterator<Item = &dyn UpdateModifier> {
-        self.update_modifiers.iter().filter_map(|m| m.as_update())
+    pub fn update_modifiers(&self) -> impl Iterator<Item = &dyn Modifier> {
+        self.update_modifiers.iter().filter_map(|m| {
+            if m.context().contains(ModifierContext::Update) {
+                Some(m.deref())
+            } else {
+                None
+            }
+        })
     }
 
     /// Get a list of all the render modifiers of this effect.
@@ -561,6 +626,31 @@ mod tests {
         assert_eq!(layout.offset("unknown"), None);
     }
 
+    #[test]
+    fn add_modifiers() {
+        let mut m = Module::default();
+        let expr = m.lit(3.);
+
+        for modifier_context in [ModifierContext::Init, ModifierContext::Update] {
+            let effect = EffectAsset::default().add_modifier(
+                modifier_context,
+                Box::new(SetAttributeModifier::new(Attribute::POSITION, expr)),
+            );
+            assert_eq!(effect.modifiers().count(), 1);
+            let m = effect.modifiers().next().unwrap();
+            assert!(m.context().contains(modifier_context));
+        }
+
+        {
+            let effect = EffectAsset::default().add_render_modifier(Box::new(SetColorModifier {
+                color: CpuValue::Single(Vec4::ONE),
+            }));
+            assert_eq!(effect.modifiers().count(), 1);
+            let m = effect.modifiers().next().unwrap();
+            assert!(m.context().contains(ModifierContext::Render));
+        }
+    }
+
     #[test]
     fn test_apply_modifiers() {
         let mut module = Module::default();
@@ -600,17 +690,16 @@ mod tests {
 
         let property_layout = PropertyLayout::default();
         let particle_layout = ParticleLayout::default();
-        let mut init_context = InitContext::new(&property_layout, &particle_layout);
+        let mut init_context =
+            ShaderWriter::new(ModifierContext::Init, &property_layout, &particle_layout);
         assert!(init_pos_sphere
-            .apply_init(&mut module, &mut init_context)
+            .apply(&mut module, &mut init_context)
             .is_ok());
         assert!(init_vel_sphere
-            .apply_init(&mut module, &mut init_context)
-            .is_ok());
-        assert!(init_age.apply_init(&mut module, &mut init_context).is_ok());
-        assert!(init_lifetime
-            .apply_init(&mut module, &mut init_context)
+            .apply(&mut module, &mut init_context)
             .is_ok());
+        assert!(init_age.apply(&mut module, &mut init_context).is_ok());
+        assert!(init_lifetime.apply(&mut module, &mut init_context).is_ok());
         // assert_eq!(effect., init_context.init_code);
 
         let mut module = Module::default();
@@ -618,15 +707,12 @@ mod tests {
         let drag_mod = LinearDragModifier::constant(&mut module, 3.5);
         let property_layout = PropertyLayout::default();
         let particle_layout = ParticleLayout::default();
-        let mut update_context = UpdateContext::new(&property_layout, &particle_layout);
-        assert!(accel_mod
-            .apply_update(&mut module, &mut update_context)
-            .is_ok());
-        assert!(drag_mod
-            .apply_update(&mut module, &mut update_context)
-            .is_ok());
+        let mut update_context =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
+        assert!(accel_mod.apply(&mut module, &mut update_context).is_ok());
+        assert!(drag_mod.apply(&mut module, &mut update_context).is_ok());
         assert!(ConformToSphereModifier::new(origin, one, one, one, one)
-            .apply_update(&mut module, &mut update_context)
+            .apply(&mut module, &mut update_context)
             .is_ok());
         // assert_eq!(effect.update_layout, update_layout);
 
diff --git a/src/graph/expr.rs b/src/graph/expr.rs
index e17eed7d..98bed84a 100644
--- a/src/graph/expr.rs
+++ b/src/graph/expr.rs
@@ -496,18 +496,27 @@ pub enum ExprError {
     /// [`EffectAsset`]: crate::EffectAsset
     #[error("Invalid expression handle: {0:?}")]
     InvalidExprHandleError(String),
+
+    /// Invalid modifier context.
+    ///
+    /// The operation was expecting a given [`ModifierContext`], but instead
+    /// another [`ModifierContext`] was available.
+    #[error("Invalid modifier context {0}, expected {1} instead.")]
+    InvalidModifierContext(ModifierContext, ModifierContext),
 }
 
 /// Evaluation context for transforming expressions into WGSL code.
 ///
 /// The evaluation context references a [`Module`] storing all [`Expr`] in use,
-/// as well as a [`PropertyLayout`] defining existing properties and their
-/// layout in memory. These together define the context within which expressions
-/// are evaluated.
+/// as well as a [`ParticleLayout`] defining the existing attributes of each
+/// particle and their layout in memory, and a [`PropertyLayout`] defining
+/// existing properties and their layout in memory. These together define the
+/// context within which expressions are evaluated.
 ///
 /// A same expression can be valid in one context and invalid in another. The
-/// most common example are [`PropertyExpr`] which are only valid if the
-/// property is actually defined in the evaluation context.
+/// most obvious example is a [`PropertyExpr`] which is only valid if the
+/// property is actually defined in the property layout of the evaluation
+/// context.
 pub trait EvalContext {
     /// Get the modifier context of the evaluation.
     fn modifier_context(&self) -> ModifierContext;
@@ -532,6 +541,10 @@ pub trait EvalContext {
     /// Each time this function is called, a new unique name is generated. The
     /// name is guaranteed to be unique within the current evaluation context
     /// only. Do not use for global top-level identifiers.
+    ///
+    /// The variable name is not registered automatically in the [`Module`]. If
+    /// you call `make_local_var()` but doesn't use the returned name, it won't
+    /// appear in the shader.
     fn make_local_var(&mut self) -> String;
 
     /// Push an intermediate statement during an evaluation.
@@ -548,9 +561,9 @@ pub trait EvalContext {
     /// [`Module`]. The function takes a list of arguments `args`, which are
     /// copied verbatim into the shader code without any validation. The body of
     /// the function is generated by invoking the given closure once with the
-    /// input `module` and a temporary `EvalContext` local to the function. The
-    /// closure must return the generated shader code of the function body. Any
-    /// statement pushed to the temporary function context with
+    /// input `module` and a temporary [`EvalContext`] local to the function.
+    /// The closure must return the generated shader code of the function
+    /// body. Any statement pushed to the temporary function context with
     /// [`EvalContext::push_stmt()`] is emitted inside the function body before
     /// the returned code. The function can subsequently be called from the
     /// parent context by generating code to call `func_name`, with the correct
@@ -763,7 +776,7 @@ impl Expr {
     /// let mut module = Module::default();
     /// # let pl = PropertyLayout::empty();
     /// # let pal = ParticleLayout::default();
-    /// # let mut context = InitContext::new(&pl, &pal);
+    /// # let mut context = ShaderWriter::new(ModifierContext::Update, &pl, &pal);
     /// let handle = module.lit(1.);
     /// let expr = module.get(handle).unwrap();
     /// assert_eq!(Ok("1.".to_string()), expr.eval(&module, &mut context));
@@ -3148,7 +3161,7 @@ impl std::ops::Rem<WriterExpr> for WriterExpr {
 
 #[cfg(test)]
 mod tests {
-    use crate::{prelude::Property, InitContext, MatrixType, ScalarValue, VectorType};
+    use crate::{prelude::Property, MatrixType, ScalarValue, ShaderWriter, VectorType};
 
     use super::*;
     use bevy::{prelude::*, utils::HashSet};
@@ -3181,7 +3194,8 @@ mod tests {
     fn local_var() {
         let property_layout = PropertyLayout::default();
         let particle_layout = ParticleLayout::default();
-        let mut ctx = InitContext::new(&property_layout, &particle_layout);
+        let mut ctx =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
         let mut h = HashSet::new();
         for _ in 0..100 {
             let v = ctx.make_local_var();
@@ -3193,7 +3207,8 @@ mod tests {
     fn make_fn() {
         let property_layout = PropertyLayout::default();
         let particle_layout = ParticleLayout::default();
-        let mut ctx = InitContext::new(&property_layout, &particle_layout);
+        let mut ctx =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
         let mut module = Module::default();
 
         // Make a function
@@ -3232,7 +3247,8 @@ mod tests {
             PropertyLayout::new(&[Property::new("my_prop", ScalarValue::Float(3.))]);
         let particle_layout = ParticleLayout::default();
         let m = w.finish();
-        let mut context = InitContext::new(&property_layout, &particle_layout);
+        let mut context =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
 
         // Evaluate the expression
         let x = m.try_get(x).unwrap();
@@ -3271,7 +3287,8 @@ mod tests {
 
         let property_layout = PropertyLayout::default();
         let particle_layout = ParticleLayout::default();
-        let mut ctx = InitContext::new(&property_layout, &particle_layout);
+        let mut ctx =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
 
         for (expr, op) in [
             (add, "+"),
@@ -3307,7 +3324,8 @@ mod tests {
 
             let property_layout = PropertyLayout::default();
             let particle_layout = ParticleLayout::default();
-            let mut ctx = InitContext::new(&property_layout, &particle_layout);
+            let mut ctx =
+                ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
 
             let expr = ctx.eval(&m, value);
             assert!(expr.is_ok());
@@ -3328,13 +3346,14 @@ mod tests {
             {
                 let property_layout = PropertyLayout::default();
                 let particle_layout = ParticleLayout::default();
-                let mut ctx = InitContext::new(&property_layout, &particle_layout);
+                let mut ctx =
+                    ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
 
                 let expr = ctx.eval(&m, value);
                 assert!(expr.is_ok());
                 let expr = expr.unwrap();
                 assert_eq!(expr, "var0");
-                assert_eq!(ctx.init_code, format!("let var0 = {}rand();\n", prefix));
+                assert_eq!(ctx.main_code, format!("let var0 = {}rand();\n", prefix));
             }
 
             // Vector form
@@ -3345,14 +3364,15 @@ mod tests {
 
                 let property_layout = PropertyLayout::default();
                 let particle_layout = ParticleLayout::default();
-                let mut ctx = InitContext::new(&property_layout, &particle_layout);
+                let mut ctx =
+                    ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
 
                 let expr = ctx.eval(&m, vec);
                 assert!(expr.is_ok());
                 let expr = expr.unwrap();
                 assert_eq!(expr, "var0");
                 assert_eq!(
-                    ctx.init_code,
+                    ctx.main_code,
                     format!("let var0 = {}rand{}();\n", prefix, count)
                 );
             }
@@ -3399,7 +3419,8 @@ mod tests {
 
         let property_layout = PropertyLayout::default();
         let particle_layout = ParticleLayout::default();
-        let mut ctx = InitContext::new(&property_layout, &particle_layout);
+        let mut ctx =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
 
         for (expr, op, inner) in [
             (
@@ -3463,7 +3484,8 @@ mod tests {
 
         let property_layout = PropertyLayout::default();
         let particle_layout = ParticleLayout::default();
-        let mut ctx = InitContext::new(&property_layout, &particle_layout);
+        let mut ctx =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
 
         for (expr, op) in [
             (cross, "cross"),
@@ -3500,7 +3522,8 @@ mod tests {
 
         let property_layout = PropertyLayout::default();
         let particle_layout = ParticleLayout::default();
-        let mut ctx = InitContext::new(&property_layout, &particle_layout);
+        let mut ctx =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
 
         for (expr, op, third) in [(mix, "mix", t), (smoothstep, "smoothstep", x)] {
             let expr = ctx.eval(&m, expr);
@@ -3535,7 +3558,8 @@ mod tests {
 
         let property_layout = PropertyLayout::default();
         let particle_layout = ParticleLayout::default();
-        let mut ctx = InitContext::new(&property_layout, &particle_layout);
+        let mut ctx =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
 
         for (expr, cast, target) in [
             (x, cx, ValueType::Vector(VectorType::VEC3I)),
@@ -3560,7 +3584,8 @@ mod tests {
 
         let property_layout = PropertyLayout::default();
         let particle_layout = ParticleLayout::default();
-        let mut ctx = InitContext::new(&property_layout, &particle_layout);
+        let mut ctx =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
 
         let res = ctx.eval(&m, x);
         assert!(res.is_ok());
@@ -3569,7 +3594,9 @@ mod tests {
 
         // 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 mut ctx =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout)
+                .with_attribute_pointer();
 
         let res = ctx.eval(&m, x);
         assert!(res.is_ok());
@@ -3659,28 +3686,31 @@ mod tests {
         {
             let property_layout = PropertyLayout::default();
             let particle_layout = ParticleLayout::default();
-            let mut ctx = InitContext::new(&property_layout, &particle_layout);
+            let mut ctx =
+                ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
             let value = ctx.eval(&m, a).unwrap();
             assert_eq!(value, "(var0) + (var0)");
-            assert_eq!(ctx.init_code, "let var0 = frand();\n");
+            assert_eq!(ctx.main_code, "let var0 = frand();\n");
         }
 
         {
             let property_layout = PropertyLayout::default();
             let particle_layout = ParticleLayout::default();
-            let mut ctx = InitContext::new(&property_layout, &particle_layout);
+            let mut ctx =
+                ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
             let value = ctx.eval(&m, b).unwrap();
             assert_eq!(value, "mix(var0, var0, var0)");
-            assert_eq!(ctx.init_code, "let var0 = frand();\n");
+            assert_eq!(ctx.main_code, "let var0 = frand();\n");
         }
 
         {
             let property_layout = PropertyLayout::default();
             let particle_layout = ParticleLayout::default();
-            let mut ctx = InitContext::new(&property_layout, &particle_layout);
+            let mut ctx =
+                ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
             let value = ctx.eval(&m, c).unwrap();
             assert_eq!(value, "abs((var0) + (var0))");
-            assert_eq!(ctx.init_code, "let var0 = frand();\n");
+            assert_eq!(ctx.main_code, "let var0 = frand();\n");
         }
     }
 
diff --git a/src/graph/node.rs b/src/graph/node.rs
index f66c6bb0..d3947c66 100644
--- a/src/graph/node.rs
+++ b/src/graph/node.rs
@@ -773,7 +773,7 @@ equal to one."
 mod tests {
     use bevy::prelude::*;
 
-    use crate::{EvalContext, InitContext, ParticleLayout, PropertyLayout};
+    use crate::{EvalContext, ModifierContext, ParticleLayout, PropertyLayout, ShaderWriter};
 
     use super::*;
 
@@ -796,7 +796,8 @@ mod tests {
 
         let property_layout = PropertyLayout::default();
         let particle_layout = ParticleLayout::default();
-        let mut context = InitContext::new(&property_layout, &particle_layout);
+        let mut context =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
         let str = context.eval(&module, out).unwrap();
         assert_eq!(str, "(3.) + (2.)".to_string());
     }
@@ -819,7 +820,8 @@ mod tests {
         let out = outputs[0];
         let property_layout = PropertyLayout::default();
         let particle_layout = ParticleLayout::default();
-        let mut context = InitContext::new(&property_layout, &particle_layout);
+        let mut context =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
         let str = context.eval(&module, out).unwrap();
         assert_eq!(str, "(3.) - (2.)".to_string());
     }
@@ -842,7 +844,8 @@ mod tests {
         let out = outputs[0];
         let property_layout = PropertyLayout::default();
         let particle_layout = ParticleLayout::default();
-        let mut context = InitContext::new(&property_layout, &particle_layout);
+        let mut context =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
         let str = context.eval(&module, out).unwrap();
         assert_eq!(str, "(3.) * (2.)".to_string());
     }
@@ -865,7 +868,8 @@ mod tests {
         let out = outputs[0];
         let property_layout = PropertyLayout::default();
         let particle_layout = ParticleLayout::default();
-        let mut context = InitContext::new(&property_layout, &particle_layout);
+        let mut context =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
         let str = context.eval(&module, out).unwrap();
         assert_eq!(str, "(3.) / (2.)".to_string());
     }
@@ -885,7 +889,8 @@ mod tests {
         let out = outputs[0];
         let property_layout = PropertyLayout::default();
         let particle_layout = ParticleLayout::default();
-        let mut context = InitContext::new(&property_layout, &particle_layout);
+        let mut context =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
         let str = context.eval(&module, out).unwrap();
         assert_eq!(str, format!("particle.{}", Attribute::POSITION.name()));
     }
@@ -904,7 +909,8 @@ mod tests {
         assert_eq!(outputs.len(), 2);
         let property_layout = PropertyLayout::default();
         let particle_layout = ParticleLayout::default();
-        let mut context = InitContext::new(&property_layout, &particle_layout);
+        let mut context =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
         let str0 = context.eval(&module, outputs[0]).unwrap();
         let str1 = context.eval(&module, outputs[1]).unwrap();
         assert_eq!(str0, format!("sim_params.{}", BuiltInOperator::Time.name()));
@@ -928,7 +934,8 @@ mod tests {
         assert_eq!(outputs.len(), 1);
         let property_layout = PropertyLayout::default();
         let particle_layout = ParticleLayout::default();
-        let mut context = InitContext::new(&property_layout, &particle_layout);
+        let mut context =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
         let str = context.eval(&module, outputs[0]).unwrap();
         assert_eq!(str, "normalize(vec3<f32>(1.,1.,1.))".to_string());
     }
diff --git a/src/lib.rs b/src/lib.rs
index b0f8a186..d75deb13 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -847,9 +847,10 @@ impl EffectShaderSource {
 
         // Generate the shader code for the initializing shader
         let (init_code, init_extra, init_sim_space_transform_code) = {
-            let mut init_context = InitContext::new(&property_layout, &particle_layout);
+            let mut init_context =
+                ShaderWriter::new(ModifierContext::Init, &property_layout, &particle_layout);
             for m in asset.init_modifiers() {
-                if let Err(err) = m.apply_init(&mut module, &mut init_context) {
+                if let Err(err) = m.apply(&mut module, &mut init_context) {
                     error!("Failed to compile effect, error in init context: {:?}", err);
                     return Err(ShaderGenerateError::Expr(err));
                 }
@@ -862,17 +863,18 @@ impl EffectShaderSource {
                 }
             };
             (
-                init_context.init_code,
-                init_context.init_extra,
+                init_context.main_code,
+                init_context.extra_code,
                 sim_space_transform_code,
             )
         };
 
         // Generate the shader code for the update shader
         let (mut update_code, update_extra) = {
-            let mut update_context = UpdateContext::new(&property_layout, &particle_layout);
+            let mut update_context =
+                ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
             for m in asset.update_modifiers() {
-                if let Err(err) = m.apply_update(&mut module, &mut update_context) {
+                if let Err(err) = m.apply(&mut module, &mut update_context) {
                     error!(
                         "Failed to compile effect, error in udpate context: {:?}",
                         err
@@ -880,7 +882,7 @@ impl EffectShaderSource {
                     return Err(ShaderGenerateError::Expr(err));
                 }
             }
-            (update_context.update_code, update_context.update_extra)
+            (update_context.main_code, update_context.extra_code)
         };
 
         // Insert Euler motion integration if needed.
@@ -1619,25 +1621,29 @@ else { return c1; }
         let property_layout = PropertyLayout::default();
         {
             // Local is always available
-            let ctx = InitContext::new(&property_layout, &particle_layout);
+            let ctx =
+                ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
             assert!(SimulationSpace::Local.eval(&ctx).is_ok());
             assert!(SimulationSpace::Global.eval(&ctx).is_err());
 
             // Global requires storing the particle's position
             let particle_layout = ParticleLayout::new().append(Attribute::POSITION).build();
-            let ctx = InitContext::new(&property_layout, &particle_layout);
+            let ctx =
+                ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
             assert!(SimulationSpace::Local.eval(&ctx).is_ok());
             assert!(SimulationSpace::Global.eval(&ctx).is_ok());
         }
         {
             // Local is always available
-            let ctx = UpdateContext::new(&property_layout, &particle_layout);
+            let ctx =
+                ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
             assert!(SimulationSpace::Local.eval(&ctx).is_ok());
             assert!(SimulationSpace::Global.eval(&ctx).is_err());
 
             // Global requires storing the particle's position
             let particle_layout = ParticleLayout::new().append(Attribute::POSITION).build();
-            let ctx = UpdateContext::new(&property_layout, &particle_layout);
+            let ctx =
+                ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
             assert!(SimulationSpace::Local.eval(&ctx).is_ok());
             assert!(SimulationSpace::Global.eval(&ctx).is_ok());
         }
diff --git a/src/modifier/accel.rs b/src/modifier/accel.rs
index ce3f153f..4c7b4601 100644
--- a/src/modifier/accel.rs
+++ b/src/modifier/accel.rs
@@ -13,8 +13,7 @@ use serde::{Deserialize, Serialize};
 use crate::{
     calc_func_id,
     graph::{BuiltInExpr, EvalContext, ExprError},
-    impl_mod_update, Attribute, BoxedModifier, ExprHandle, Modifier, ModifierContext, Module,
-    UpdateContext, UpdateModifier,
+    Attribute, BoxedModifier, ExprHandle, Modifier, ModifierContext, Module, ShaderWriter,
 };
 
 /// A modifier to apply a uniform acceleration to all particles each frame, to
@@ -67,14 +66,6 @@ impl Modifier for AccelModifier {
         ModifierContext::Update
     }
 
-    fn as_update(&self) -> Option<&dyn UpdateModifier> {
-        Some(self)
-    }
-
-    fn as_update_mut(&mut self) -> Option<&mut dyn UpdateModifier> {
-        Some(self)
-    }
-
     fn attributes(&self) -> &[Attribute] {
         &[Attribute::VELOCITY]
     }
@@ -82,20 +73,13 @@ impl Modifier for AccelModifier {
     fn boxed_clone(&self) -> BoxedModifier {
         Box::new(*self)
     }
-}
 
-#[typetag::serde]
-impl UpdateModifier for AccelModifier {
-    fn apply_update(
-        &self,
-        module: &mut Module,
-        context: &mut UpdateContext,
-    ) -> Result<(), ExprError> {
+    fn apply(&self, module: &mut Module, context: &mut ShaderWriter) -> Result<(), ExprError> {
         let attr = module.attr(Attribute::VELOCITY);
         let attr = context.eval(module, attr)?;
         let expr = context.eval(module, self.accel)?;
         let dt = BuiltInExpr::new(crate::graph::BuiltInOperator::DeltaTime).eval(context)?;
-        context.update_code += &format!("{} += ({}) * {};", attr, expr, dt);
+        context.main_code += &format!("{} += ({}) * {};", attr, expr, dt);
         Ok(())
     }
 }
@@ -162,18 +146,21 @@ impl RadialAccelModifier {
     }
 }
 
-impl_mod_update!(
-    RadialAccelModifier,
-    &[Attribute::POSITION, Attribute::VELOCITY]
-);
-
 #[typetag::serde]
-impl UpdateModifier for RadialAccelModifier {
-    fn apply_update(
-        &self,
-        module: &mut Module,
-        context: &mut UpdateContext,
-    ) -> Result<(), ExprError> {
+impl Modifier for RadialAccelModifier {
+    fn context(&self) -> ModifierContext {
+        ModifierContext::Update
+    }
+
+    fn attributes(&self) -> &[Attribute] {
+        &[Attribute::POSITION, Attribute::VELOCITY]
+    }
+
+    fn boxed_clone(&self) -> BoxedModifier {
+        Box::new(self.clone())
+    }
+
+    fn apply(&self, module: &mut Module, context: &mut ShaderWriter) -> Result<(), ExprError> {
         let func_id = calc_func_id(self);
         let func_name = format!("radial_accel_{0:016X}", func_id);
 
@@ -197,7 +184,7 @@ impl UpdateModifier for RadialAccelModifier {
             },
         )?;
 
-        context.update_code += &format!("{}(&particle);\n", func_name);
+        context.main_code += &format!("{}(&particle);\n", func_name);
 
         Ok(())
     }
@@ -277,18 +264,21 @@ impl TangentAccelModifier {
     }
 }
 
-impl_mod_update!(
-    TangentAccelModifier,
-    &[Attribute::POSITION, Attribute::VELOCITY]
-);
-
 #[typetag::serde]
-impl UpdateModifier for TangentAccelModifier {
-    fn apply_update(
-        &self,
-        module: &mut Module,
-        context: &mut UpdateContext,
-    ) -> Result<(), ExprError> {
+impl Modifier for TangentAccelModifier {
+    fn context(&self) -> ModifierContext {
+        ModifierContext::Update
+    }
+
+    fn attributes(&self) -> &[Attribute] {
+        &[Attribute::POSITION, Attribute::VELOCITY]
+    }
+
+    fn boxed_clone(&self) -> BoxedModifier {
+        Box::new(self.clone())
+    }
+
+    fn apply(&self, module: &mut Module, context: &mut ShaderWriter) -> Result<(), ExprError> {
         let func_id = calc_func_id(self);
         let func_name = format!("tangent_accel_{0:016X}", func_id);
 
@@ -296,7 +286,7 @@ impl UpdateModifier for TangentAccelModifier {
         let axis = context.eval(module, self.axis)?;
         let accel = context.eval(module, self.accel)?;
 
-        context.update_extra += &format!(
+        context.extra_code += &format!(
             r##"fn {}(particle: ptr<function, Particle>) {{
     let radial = normalize((*particle).{} - {});
     let tangent = normalize(cross({}, radial));
@@ -311,7 +301,7 @@ impl UpdateModifier for TangentAccelModifier {
             accel,
         );
 
-        context.update_code += &format!("{}(&particle);\n", func_name);
+        context.main_code += &format!("{}(&particle);\n", func_name);
 
         Ok(())
     }
@@ -331,10 +321,11 @@ mod tests {
 
         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());
+        let mut context =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
+        assert!(modifier.apply(&mut module, &mut context).is_ok());
 
-        assert!(context.update_code.contains(&accel.to_wgsl_string()));
+        assert!(context.main_code.contains(&accel.to_wgsl_string()));
     }
 
     #[test]
@@ -346,19 +337,21 @@ mod tests {
         let origin = Vec3::new(-1.2, 5.3, -8.5);
         let accel = 6.;
         let modifier = RadialAccelModifier::constant(&mut module, origin, accel);
-        let mut context = UpdateContext::new(&property_layout, &particle_layout);
-        assert!(modifier.apply_update(&mut module, &mut context).is_ok());
+        let mut context =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
+        assert!(modifier.apply(&mut module, &mut context).is_ok());
         // TODO: less weak check...
-        assert!(context.update_extra.contains(&accel.to_wgsl_string()));
+        assert!(context.extra_code.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());
+        let mut context =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
+        assert!(modifier.apply(&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"));
+        assert!(context.extra_code.contains(Attribute::POSITION.name()));
+        assert!(context.extra_code.contains("my_prop"));
     }
 
     #[test]
@@ -371,10 +364,11 @@ mod tests {
 
         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());
+        let mut context =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
+        assert!(modifier.apply(&mut module, &mut context).is_ok());
 
         // TODO: less weak check...
-        assert!(context.update_extra.contains(&accel.to_wgsl_string()));
+        assert!(context.extra_code.contains(&accel.to_wgsl_string()));
     }
 }
diff --git a/src/modifier/attr.rs b/src/modifier/attr.rs
index e36de4a2..10b355e8 100644
--- a/src/modifier/attr.rs
+++ b/src/modifier/attr.rs
@@ -10,8 +10,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::{
     graph::{EvalContext, ExprError},
-    Attribute, BoxedModifier, ExprHandle, InitContext, InitModifier, Modifier, ModifierContext,
-    Module, UpdateContext, UpdateModifier,
+    Attribute, BoxedModifier, ExprHandle, Modifier, ModifierContext, Module, ShaderWriter,
 };
 
 /// A modifier to assign a value to a particle attribute.
@@ -99,22 +98,6 @@ impl Modifier for SetAttributeModifier {
         ModifierContext::Init | ModifierContext::Update
     }
 
-    fn as_init(&self) -> Option<&dyn InitModifier> {
-        Some(self)
-    }
-
-    fn as_init_mut(&mut self) -> Option<&mut dyn InitModifier> {
-        Some(self)
-    }
-
-    fn as_update(&self) -> Option<&dyn UpdateModifier> {
-        Some(self)
-    }
-
-    fn as_update_mut(&mut self) -> Option<&mut dyn UpdateModifier> {
-        Some(self)
-    }
-
     fn attributes(&self) -> &[Attribute] {
         std::slice::from_ref(&self.attribute)
     }
@@ -122,26 +105,10 @@ impl Modifier for SetAttributeModifier {
     fn boxed_clone(&self) -> BoxedModifier {
         Box::new(*self)
     }
-}
-
-#[typetag::serde]
-impl InitModifier for SetAttributeModifier {
-    fn apply_init(&self, module: &mut Module, context: &mut InitContext) -> Result<(), ExprError> {
-        let code = self.eval(module, context)?;
-        context.init_code += &code;
-        Ok(())
-    }
-}
 
-#[typetag::serde]
-impl UpdateModifier for SetAttributeModifier {
-    fn apply_update(
-        &self,
-        module: &mut Module,
-        context: &mut UpdateContext,
-    ) -> Result<(), ExprError> {
+    fn apply(&self, module: &mut Module, context: &mut ShaderWriter) -> Result<(), ExprError> {
         let code = self.eval(module, context)?;
-        context.update_code += &code;
+        context.main_code += &code;
         Ok(())
     }
 }
diff --git a/src/modifier/force.rs b/src/modifier/force.rs
index a987fe1f..a0492a08 100644
--- a/src/modifier/force.rs
+++ b/src/modifier/force.rs
@@ -11,7 +11,7 @@ use serde::{Deserialize, Serialize};
 use crate::{
     calc_func_id,
     graph::{BuiltInOperator, EvalContext, ExprError},
-    impl_mod_update, Attribute, ExprHandle, Module, UpdateContext, UpdateModifier,
+    Attribute, BoxedModifier, ExprHandle, Modifier, ModifierContext, Module, ShaderWriter,
 };
 
 /// A modifier to apply a force to the particle which makes it conform ("stick")
@@ -159,18 +159,21 @@ impl ConformToSphereModifier {
     }
 }
 
-impl_mod_update!(
-    ConformToSphereModifier,
-    &[Attribute::POSITION, Attribute::VELOCITY]
-);
-
 #[typetag::serde]
-impl UpdateModifier for ConformToSphereModifier {
-    fn apply_update(
-        &self,
-        module: &mut Module,
-        context: &mut UpdateContext,
-    ) -> Result<(), ExprError> {
+impl Modifier for ConformToSphereModifier {
+    fn context(&self) -> ModifierContext {
+        ModifierContext::Update
+    }
+
+    fn attributes(&self) -> &[Attribute] {
+        &[Attribute::POSITION, Attribute::VELOCITY]
+    }
+
+    fn boxed_clone(&self) -> BoxedModifier {
+        Box::new(self.clone())
+    }
+
+    fn apply(&self, module: &mut Module, context: &mut ShaderWriter) -> Result<(), ExprError> {
         let func_id = calc_func_id(self);
         let func_name = format!("force_field_{0:016X}", func_id);
 
@@ -230,7 +233,7 @@ impl UpdateModifier for ConformToSphereModifier {
             },
         )?;
 
-        context.update_code += &format!("{}(&particle);\n", func_name);
+        context.main_code += &format!("{}(&particle);\n", func_name);
 
         Ok(())
     }
@@ -266,15 +269,21 @@ impl LinearDragModifier {
     }
 }
 
-impl_mod_update!(LinearDragModifier, &[Attribute::VELOCITY]);
-
 #[typetag::serde]
-impl UpdateModifier for LinearDragModifier {
-    fn apply_update(
-        &self,
-        module: &mut Module,
-        context: &mut UpdateContext,
-    ) -> Result<(), ExprError> {
+impl Modifier for LinearDragModifier {
+    fn context(&self) -> ModifierContext {
+        ModifierContext::Update
+    }
+
+    fn attributes(&self) -> &[Attribute] {
+        &[Attribute::VELOCITY]
+    }
+
+    fn boxed_clone(&self) -> BoxedModifier {
+        Box::new(self.clone())
+    }
+
+    fn apply(&self, module: &mut Module, context: &mut ShaderWriter) -> Result<(), ExprError> {
         let m = module;
         let attr = m.attr(Attribute::VELOCITY);
         let dt = m.builtin(BuiltInOperator::DeltaTime);
@@ -285,14 +294,14 @@ impl UpdateModifier for LinearDragModifier {
         let expr = m.max(zero, one_minus_drag_dt);
         let attr = context.eval(m, attr)?;
         let expr = context.eval(m, expr)?;
-        context.update_code += &format!("{} *= {};", attr, expr);
+        context.main_code += &format!("{} *= {};", attr, expr);
         Ok(())
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use crate::{ParticleLayout, PropertyLayout, UpdateContext};
+    use crate::{ParticleLayout, PropertyLayout, ShaderWriter};
 
     use super::*;
 
@@ -303,9 +312,10 @@ mod tests {
 
         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());
+        let mut context =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
+        assert!(modifier.apply(&mut module, &mut context).is_ok());
 
-        assert!(context.update_code.contains("3.5")); // TODO - less weak check
+        assert!(context.main_code.contains("3.5")); // TODO - less weak check
     }
 }
diff --git a/src/modifier/kill.rs b/src/modifier/kill.rs
index c27f80eb..acc97372 100644
--- a/src/modifier/kill.rs
+++ b/src/modifier/kill.rs
@@ -8,7 +8,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::{
     graph::{EvalContext, ExprError},
-    impl_mod_update, Attribute, ExprHandle, Module, UpdateContext, UpdateModifier,
+    Attribute, BoxedModifier, ExprHandle, Modifier, ModifierContext, Module, ShaderWriter,
 };
 
 /// A modifier killing all particles that enter or exit a sphere.
@@ -60,15 +60,21 @@ impl KillSphereModifier {
     }
 }
 
-impl_mod_update!(KillSphereModifier, &[Attribute::POSITION]);
-
 #[typetag::serde]
-impl UpdateModifier for KillSphereModifier {
-    fn apply_update(
-        &self,
-        module: &mut Module,
-        context: &mut UpdateContext,
-    ) -> Result<(), ExprError> {
+impl Modifier for KillSphereModifier {
+    fn context(&self) -> ModifierContext {
+        ModifierContext::Update
+    }
+
+    fn attributes(&self) -> &[Attribute] {
+        &[Attribute::POSITION]
+    }
+
+    fn boxed_clone(&self) -> BoxedModifier {
+        Box::new(self.clone())
+    }
+
+    fn apply(&self, module: &mut Module, context: &mut ShaderWriter) -> Result<(), ExprError> {
         let pos = module.attr(Attribute::POSITION);
         let diff = module.sub(pos, self.center);
         let sqr_dist = module.dot(diff, diff);
@@ -79,7 +85,7 @@ impl UpdateModifier for KillSphereModifier {
         };
         let expr = context.eval(module, cmp)?;
 
-        context.update_code += &format!(
+        context.main_code += &format!(
             r#"if ({}) {{
     is_alive = false;
 }}
@@ -135,15 +141,21 @@ impl KillAabbModifier {
     }
 }
 
-impl_mod_update!(KillAabbModifier, &[Attribute::POSITION]);
-
 #[typetag::serde]
-impl UpdateModifier for KillAabbModifier {
-    fn apply_update(
-        &self,
-        module: &mut Module,
-        context: &mut UpdateContext,
-    ) -> Result<(), ExprError> {
+impl Modifier for KillAabbModifier {
+    fn context(&self) -> ModifierContext {
+        ModifierContext::Update
+    }
+
+    fn attributes(&self) -> &[Attribute] {
+        &[Attribute::POSITION]
+    }
+
+    fn boxed_clone(&self) -> BoxedModifier {
+        Box::new(self.clone())
+    }
+
+    fn apply(&self, module: &mut Module, context: &mut ShaderWriter) -> Result<(), ExprError> {
         let pos = module.attr(Attribute::POSITION);
         let diff = module.sub(pos, self.center);
         let dist = module.abs(diff);
@@ -159,7 +171,7 @@ impl UpdateModifier for KillAabbModifier {
         };
         let expr = context.eval(module, reduce)?;
 
-        context.update_code += &format!(
+        context.main_code += &format!(
             r#"if ({}) {{
     is_alive = false;
 }}
@@ -186,9 +198,11 @@ mod tests {
 
         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());
+        let mut context =
+            ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
+        assert!(modifier.apply(&mut module, &mut context).is_ok());
 
-        assert!(context.update_code.contains("is_alive = false")); // TODO - less weak check
+        assert!(context.main_code.contains("is_alive = false")); // TODO - less
+                                                                 // weak check
     }
 }
diff --git a/src/modifier/mod.rs b/src/modifier/mod.rs
index 0099b90a..9a3bb163 100644
--- a/src/modifier/mod.rs
+++ b/src/modifier/mod.rs
@@ -100,32 +100,39 @@ bitflags! {
     }
 }
 
+impl std::fmt::Display for ModifierContext {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let mut s = String::new();
+        if self.contains(ModifierContext::Init) {
+            s = "Init".to_string();
+        }
+        if self.contains(ModifierContext::Update) {
+            if s.is_empty() {
+                s = "Update".to_string();
+            } else {
+                s += " | Update";
+            }
+        }
+        if self.contains(ModifierContext::Render) {
+            if s.is_empty() {
+                s = "Render".to_string();
+            } else {
+                s += " | Render";
+            }
+        }
+        if s.is_empty() {
+            s = "None".to_string();
+        }
+        write!(f, "{}", s)
+    }
+}
+
 /// Trait describing a modifier customizing an effect pipeline.
 #[typetag::serde]
 pub trait Modifier: Reflect + Send + Sync + 'static {
     /// Get the context this modifier applies to.
     fn context(&self) -> ModifierContext;
 
-    /// Try to cast this modifier to an [`InitModifier`].
-    fn as_init(&self) -> Option<&dyn InitModifier> {
-        None
-    }
-
-    /// Try to cast this modifier to an [`InitModifier`].
-    fn as_init_mut(&mut self) -> Option<&mut dyn InitModifier> {
-        None
-    }
-
-    /// Try to cast this modifier to an [`UpdateModifier`].
-    fn as_update(&self) -> Option<&dyn UpdateModifier> {
-        None
-    }
-
-    /// Try to cast this modifier to an [`UpdateModifier`].
-    fn as_update_mut(&mut self) -> Option<&mut dyn UpdateModifier> {
-        None
-    }
-
     /// Try to cast this modifier to a [`RenderModifier`].
     fn as_render(&self) -> Option<&dyn RenderModifier> {
         None
@@ -142,6 +149,9 @@ pub trait Modifier: Reflect + Send + Sync + 'static {
 
     /// Clone self.
     fn boxed_clone(&self) -> BoxedModifier;
+
+    /// Apply the modifier to generate code.
+    fn apply(&self, module: &mut Module, context: &mut ShaderWriter) -> Result<(), ExprError>;
 }
 
 /// Boxed version of [`Modifier`].
@@ -153,18 +163,31 @@ impl Clone for BoxedModifier {
     }
 }
 
-/// Particle initializing shader code generation context.
+/// Shader code writer.
+///
+/// Writer utility to generate shader code. The writer works in a defined
+/// context, for a given [`ModifierContext`] and a particular effect setup
+/// ([`ParticleLayout`] and [`PropertyLayout`]).
 #[derive(Debug, PartialEq)]
-pub struct InitContext<'a> {
-    /// Main particle initializing code, which needs to assign the fields of the
-    /// `particle` struct instance.
-    pub init_code: String,
-    /// Extra functions emitted at top level, which `init_code` can call.
-    pub init_extra: String,
+pub struct ShaderWriter<'a> {
+    /// Main shader compute code emitted.
+    ///
+    /// This is the WGSL code emitted into the target [`ModifierContext`]. The
+    /// context dictates what variables are available (this is currently
+    /// implicit and requires knownledge of the target context; there's little
+    /// validation that the emitted code is valid).
+    pub main_code: String,
+    /// Extra functions emitted at shader top level.
+    ///
+    /// This contains optional WGSL code emitted at shader top level. This
+    /// generally contains functions called from `main_code`.
+    pub extra_code: String,
     /// Layout of properties for the current effect.
     pub property_layout: &'a PropertyLayout,
     /// Layout of attributes of a particle for the current effect.
     pub particle_layout: &'a ParticleLayout,
+    /// Modifier context the writer is being used from.
+    modifier_context: ModifierContext,
     /// Counter for unique variable names.
     var_counter: u32,
     /// Cache of evaluated expressions.
@@ -173,14 +196,19 @@ pub struct InitContext<'a> {
     is_attribute_pointer: bool,
 }
 
-impl<'a> InitContext<'a> {
+impl<'a> ShaderWriter<'a> {
     /// Create a new init context.
-    pub fn new(property_layout: &'a PropertyLayout, particle_layout: &'a ParticleLayout) -> Self {
+    pub fn new(
+        modifier_context: ModifierContext,
+        property_layout: &'a PropertyLayout,
+        particle_layout: &'a ParticleLayout,
+    ) -> Self {
         Self {
-            init_code: String::new(),
-            init_extra: String::new(),
+            main_code: String::new(),
+            extra_code: String::new(),
             property_layout,
             particle_layout,
+            modifier_context,
             var_counter: 0,
             expr_cache: Default::default(),
             is_attribute_pointer: false,
@@ -194,9 +222,9 @@ impl<'a> InitContext<'a> {
     }
 }
 
-impl<'a> EvalContext for InitContext<'a> {
+impl<'a> EvalContext for ShaderWriter<'a> {
     fn modifier_context(&self) -> ModifierContext {
-        ModifierContext::Init
+        self.modifier_context
     }
 
     fn property_layout(&self) -> &PropertyLayout {
@@ -227,8 +255,8 @@ impl<'a> EvalContext for InitContext<'a> {
     }
 
     fn push_stmt(&mut self, stmt: &str) {
-        self.init_code += stmt;
-        self.init_code += "\n";
+        self.main_code += stmt;
+        self.main_code += "\n";
     }
 
     fn make_fn(
@@ -240,20 +268,24 @@ impl<'a> EvalContext for InitContext<'a> {
     ) -> Result<(), ExprError> {
         // Generate a temporary context for the function content itself
         // FIXME - Dynamic with_attribute_pointer()!
-        let mut ctx =
-            InitContext::new(self.property_layout, self.particle_layout).with_attribute_pointer();
+        let mut ctx = ShaderWriter::new(
+            self.modifier_context,
+            self.property_layout,
+            self.particle_layout,
+        )
+        .with_attribute_pointer();
 
         // Evaluate the function content
         let body = f(module, &mut ctx)?;
 
         // Append any extra
-        self.init_extra += &ctx.init_extra;
+        self.extra_code += &ctx.extra_code;
 
         // Append the function itself
-        self.init_extra += &format!(
+        self.extra_code += &format!(
             r##"fn {0}({1}) {{
 {2}{3}}}"##,
-            func_name, args, ctx.init_code, body
+            func_name, args, ctx.main_code, body
         );
 
         Ok(())
@@ -264,137 +296,6 @@ impl<'a> EvalContext for InitContext<'a> {
     }
 }
 
-/// Trait to customize the initializing of newly spawned particles.
-#[typetag::serde]
-pub trait InitModifier: Modifier {
-    /// Append the initializing code.
-    fn apply_init(&self, module: &mut Module, context: &mut InitContext) -> Result<(), ExprError>;
-}
-
-/// Particle update shader code generation context.
-#[derive(Debug, PartialEq)]
-pub struct UpdateContext<'a> {
-    /// Main particle update code, which needs to update the fields of the
-    /// `particle` struct instance.
-    pub update_code: String,
-    /// Extra functions emitted at top level, which `update_code` can call.
-    pub update_extra: String,
-    /// Layout of properties for the current effect.
-    pub property_layout: &'a PropertyLayout,
-    /// Layout of attributes of a particle for the current effect.
-    pub particle_layout: &'a ParticleLayout,
-    /// Counter for unique variable names.
-    var_counter: u32,
-    /// Cache of evaluated expressions.
-    expr_cache: HashMap<ExprHandle, String>,
-    /// Is the attriubute struct a pointer?
-    is_attribute_pointer: bool,
-}
-
-impl<'a> UpdateContext<'a> {
-    /// Create a new update context.
-    pub fn new(property_layout: &'a PropertyLayout, particle_layout: &'a ParticleLayout) -> Self {
-        Self {
-            update_code: String::new(),
-            update_extra: String::new(),
-            property_layout,
-            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 UpdateContext<'a> {
-    fn modifier_context(&self) -> ModifierContext {
-        ModifierContext::Update
-    }
-
-    fn property_layout(&self) -> &PropertyLayout {
-        self.property_layout
-    }
-
-    fn particle_layout(&self) -> &ParticleLayout {
-        self.particle_layout
-    }
-
-    fn eval(&mut self, module: &Module, handle: ExprHandle) -> Result<String, ExprError> {
-        // On cache hit, don't re-evaluate the expression to prevent any duplicate
-        // side-effect.
-        if let Some(s) = self.expr_cache.get(&handle) {
-            Ok(s.clone())
-        } else {
-            module.try_get(handle)?.eval(module, self).map(|s| {
-                self.expr_cache.insert(handle, s.clone());
-                s
-            })
-        }
-    }
-
-    fn make_local_var(&mut self) -> String {
-        let index = self.var_counter;
-        self.var_counter += 1;
-        format!("var{}", index)
-    }
-
-    fn push_stmt(&mut self, stmt: &str) {
-        self.update_code += stmt;
-        self.update_code += "\n";
-    }
-
-    fn make_fn(
-        &mut self,
-        func_name: &str,
-        args: &str,
-        module: &mut Module,
-        f: &mut dyn FnMut(&mut Module, &mut dyn EvalContext) -> Result<String, ExprError>,
-    ) -> Result<(), ExprError> {
-        // Generate a temporary context for the function content itself
-        // 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)?;
-
-        // Append any extra
-        self.update_extra += &ctx.update_extra;
-
-        // Append the function itself
-        self.update_extra += &format!(
-            r##"fn {0}({1}) {{
-            {2};
-        }}
-        "##,
-            func_name, args, body
-        );
-
-        Ok(())
-    }
-
-    fn is_attribute_pointer(&self) -> bool {
-        self.is_attribute_pointer
-    }
-}
-
-/// Trait to customize the updating of existing particles each frame.
-#[typetag::serde]
-pub trait UpdateModifier: Modifier {
-    /// Append the update code.
-    fn apply_update(
-        &self,
-        module: &mut Module,
-        context: &mut UpdateContext,
-    ) -> Result<(), ExprError>;
-}
-
 /// Particle rendering shader code generation context.
 #[derive(Debug, PartialEq)]
 pub struct RenderContext<'a> {
@@ -564,107 +465,20 @@ impl<'a> EvalContext for RenderContext<'a> {
 pub trait RenderModifier: Modifier {
     /// Apply the rendering code.
     fn apply_render(&self, module: &mut Module, context: &mut RenderContext);
-}
-
-/// Macro to implement the [`Modifier`] trait for an init modifier.
-// macro_rules! impl_mod_init {
-//     ($t:ty, $attrs:expr) => {
-//         #[typetag::serde]
-//         impl $crate::Modifier for $t {
-//             fn context(&self) -> $crate::ModifierContext {
-//                 $crate::ModifierContext::Init
-//             }
-
-//             fn as_init(&self) -> Option<&dyn $crate::InitModifier> {
-//                 Some(self)
-//             }
-
-//             fn as_init_mut(&mut self) -> Option<&mut dyn $crate::InitModifier> {
-//                 Some(self)
-//             }
-
-//             fn attributes(&self) -> &[$crate::Attribute] {
-//                 $attrs
-//             }
-
-//             fn boxed_clone(&self) -> $crate::BoxedModifier {
-//                 Box::new(*self)
-//             }
-//         }
-//     };
-// }
-
-// pub(crate) use impl_mod_init;
-
-/// Macro to implement the [`Modifier`] trait for an update modifier.
-macro_rules! impl_mod_update {
-    ($t:ty, $attrs:expr) => {
-        #[typetag::serde]
-        impl $crate::Modifier for $t {
-            fn context(&self) -> $crate::ModifierContext {
-                $crate::ModifierContext::Update
-            }
-
-            fn as_update(&self) -> Option<&dyn $crate::UpdateModifier> {
-                Some(self)
-            }
-
-            fn as_update_mut(&mut self) -> Option<&mut dyn $crate::UpdateModifier> {
-                Some(self)
-            }
 
-            fn attributes(&self) -> &[$crate::Attribute] {
-                $attrs
-            }
+    /// Clone into boxed self.
+    fn boxed_render_clone(&self) -> Box<dyn RenderModifier>;
 
-            fn boxed_clone(&self) -> $crate::BoxedModifier {
-                Box::new(self.clone())
-            }
-        }
-    };
+    /// Upcast to [`Modifier`] trait.
+    fn as_modifier(&self) -> &dyn Modifier;
 }
 
-pub(crate) use impl_mod_update;
-
-/// Macro to implement the [`Modifier`] trait for a modifier which can be used
-/// both in init and update contexts.
-macro_rules! impl_mod_init_update {
-    ($t:ty, $attrs:expr) => {
-        #[typetag::serde]
-        impl $crate::Modifier for $t {
-            fn context(&self) -> $crate::ModifierContext {
-                $crate::ModifierContext::Init | $crate::ModifierContext::Update
-            }
-
-            fn as_init(&self) -> Option<&dyn $crate::InitModifier> {
-                Some(self)
-            }
-
-            fn as_init_mut(&mut self) -> Option<&mut dyn $crate::InitModifier> {
-                Some(self)
-            }
-
-            fn as_update(&self) -> Option<&dyn $crate::UpdateModifier> {
-                Some(self)
-            }
-
-            fn as_update_mut(&mut self) -> Option<&mut dyn $crate::UpdateModifier> {
-                Some(self)
-            }
-
-            fn attributes(&self) -> &[$crate::Attribute] {
-                $attrs
-            }
-
-            fn boxed_clone(&self) -> $crate::BoxedModifier {
-                Box::new(self.clone())
-            }
-        }
-    };
+impl Clone for Box<dyn RenderModifier> {
+    fn clone(&self) -> Self {
+        self.boxed_render_clone()
+    }
 }
 
-pub(crate) use impl_mod_init_update;
-
 /// Macro to implement the [`Modifier`] trait for a render modifier.
 macro_rules! impl_mod_render {
     ($t:ty, $attrs:expr) => {
@@ -689,6 +503,17 @@ macro_rules! impl_mod_render {
             fn boxed_clone(&self) -> $crate::BoxedModifier {
                 Box::new(self.clone())
             }
+
+            fn apply(
+                &self,
+                _module: &mut Module,
+                context: &mut ShaderWriter,
+            ) -> Result<(), ExprError> {
+                Err(ExprError::InvalidModifierContext(
+                    context.modifier_context(),
+                    ModifierContext::Render,
+                ))
+            }
         }
     };
 }
@@ -715,6 +540,30 @@ mod tests {
         }
     }
 
+    #[test]
+    fn modifier_context_display() {
+        assert_eq!("None", format!("{}", ModifierContext::empty()));
+        assert_eq!("Init", format!("{}", ModifierContext::Init));
+        assert_eq!("Update", format!("{}", ModifierContext::Update));
+        assert_eq!("Render", format!("{}", ModifierContext::Render));
+        assert_eq!(
+            "Init | Update",
+            format!("{}", ModifierContext::Init | ModifierContext::Update)
+        );
+        assert_eq!(
+            "Update | Render",
+            format!("{}", ModifierContext::Update | ModifierContext::Render)
+        );
+        assert_eq!(
+            "Init | Render",
+            format!("{}", ModifierContext::Init | ModifierContext::Render)
+        );
+        assert_eq!(
+            "Init | Update | Render",
+            format!("{}", ModifierContext::all())
+        );
+    }
+
     #[test]
     fn reflect() {
         let m = make_test_modifier();
@@ -758,7 +607,7 @@ mod tests {
         let center = module.lit(Vec3::ZERO);
         let axis = module.lit(Vec3::Y);
         let radius = module.lit(1.);
-        let modifiers: &[&dyn InitModifier] = &[
+        let modifiers: &[&dyn Modifier] = &[
             &SetPositionCircleModifier {
                 center,
                 axis,
@@ -792,12 +641,14 @@ mod tests {
             },
         ];
         for &modifier in modifiers.iter() {
+            assert!(modifier.context().contains(ModifierContext::Init));
             let property_layout = PropertyLayout::default();
             let particle_layout = ParticleLayout::default();
-            let mut context = InitContext::new(&property_layout, &particle_layout);
-            assert!(modifier.apply_init(&mut module, &mut context).is_ok());
-            let init_code = context.init_code;
-            let init_extra = context.init_extra;
+            let mut context =
+                ShaderWriter::new(ModifierContext::Init, &property_layout, &particle_layout);
+            assert!(modifier.apply(&mut module, &mut context).is_ok());
+            let main_code = context.main_code;
+            let extra_code = context.extra_code;
 
             let mut particle_layout = ParticleLayout::new();
             for &attr in modifier.attributes() {
@@ -817,13 +668,13 @@ struct Particle {{
     {attributes_code}
 }};
 
-{init_extra}
+{extra_code}
 
 @compute @workgroup_size(64)
 fn main() {{
     var particle = Particle();
     var transform: mat4x4<f32> = mat4x4<f32>();
-{init_code}
+{main_code}
 }}"##
             );
             // println!("code: {:?}", code);
@@ -851,7 +702,7 @@ fn main() {{
         let y_axis = writer.lit(Vec3::Y).expr();
         let one = writer.lit(1.).expr();
         let radius = one;
-        let modifiers: &[&dyn UpdateModifier] = &[
+        let modifiers: &[&dyn Modifier] = &[
             &AccelModifier::new(origin),
             &RadialAccelModifier::new(origin, one),
             &TangentAccelModifier::new(origin, y_axis, one),
@@ -892,12 +743,14 @@ fn main() {{
         ];
         let mut module = writer.finish();
         for &modifier in modifiers.iter() {
+            assert!(modifier.context().contains(ModifierContext::Update));
             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());
-            let update_code = context.update_code;
-            let update_extra = context.update_extra;
+            let mut context =
+                ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
+            assert!(modifier.apply(&mut module, &mut context).is_ok());
+            let update_code = context.main_code;
+            let update_extra = context.extra_code;
 
             let mut particle_layout = ParticleLayout::new();
             for &attr in modifier.attributes() {
@@ -1070,9 +923,10 @@ fn fragment(in: VertexOutput) -> @location(0) vec4<f32> {{
         let property_layout = PropertyLayout::default();
         let particle_layout = ParticleLayout::default();
         let x = module.builtin(BuiltInOperator::Rand(ScalarType::Float.into()));
-        let init: &mut dyn EvalContext = &mut InitContext::new(&property_layout, &particle_layout);
+        let init: &mut dyn EvalContext =
+            &mut ShaderWriter::new(ModifierContext::Init, &property_layout, &particle_layout);
         let update: &mut dyn EvalContext =
-            &mut UpdateContext::new(&property_layout, &particle_layout);
+            &mut ShaderWriter::new(ModifierContext::Update, &property_layout, &particle_layout);
         let render: &mut dyn EvalContext =
             &mut RenderContext::new(&property_layout, &particle_layout);
         for ctx in [init, update, render] {
diff --git a/src/modifier/output.rs b/src/modifier/output.rs
index 6d9b3eb6..5ddca56c 100644
--- a/src/modifier/output.rs
+++ b/src/modifier/output.rs
@@ -5,8 +5,9 @@ use serde::{Deserialize, Serialize};
 use std::hash::Hash;
 
 use crate::{
-    impl_mod_render, Attribute, BoxedModifier, CpuValue, EvalContext, ExprHandle, Gradient,
-    Modifier, ModifierContext, Module, RenderContext, RenderModifier, ShaderCode, ToWgslString,
+    impl_mod_render, Attribute, BoxedModifier, CpuValue, EvalContext, ExprError, ExprHandle,
+    Gradient, Modifier, ModifierContext, Module, RenderContext, RenderModifier, ShaderCode,
+    ShaderWriter, ToWgslString,
 };
 
 /// Mapping of the sample read from a texture image to the base particle color.
@@ -81,6 +82,14 @@ impl RenderModifier for ParticleTextureModifier {
         context.set_particle_texture(self.texture.clone());
         context.image_sample_mapping_code = self.sample_mapping.to_wgsl_string();
     }
+
+    fn boxed_render_clone(&self) -> Box<dyn RenderModifier> {
+        Box::new(self.clone())
+    }
+
+    fn as_modifier(&self) -> &dyn Modifier {
+        self
+    }
 }
 
 /// A modifier to set the rendering color of all particles.
@@ -105,6 +114,14 @@ impl RenderModifier for SetColorModifier {
     fn apply_render(&self, _module: &mut Module, context: &mut RenderContext) {
         context.vertex_code += &format!("color = {0};\n", self.color.to_wgsl_string());
     }
+
+    fn boxed_render_clone(&self) -> Box<dyn RenderModifier> {
+        Box::new(*self)
+    }
+
+    fn as_modifier(&self) -> &dyn Modifier {
+        self
+    }
 }
 
 /// A modifier modulating each particle's color over its lifetime with a
@@ -147,6 +164,14 @@ impl RenderModifier for ColorOverLifetimeModifier {
             Attribute::LIFETIME.name()
         );
     }
+
+    fn boxed_render_clone(&self) -> Box<dyn RenderModifier> {
+        Box::new(self.clone())
+    }
+
+    fn as_modifier(&self) -> &dyn Modifier {
+        self
+    }
 }
 
 /// A modifier to set the size of all particles.
@@ -174,6 +199,14 @@ impl RenderModifier for SetSizeModifier {
     fn apply_render(&self, _module: &mut Module, context: &mut RenderContext) {
         context.vertex_code += &format!("size = {0};\n", self.size.to_wgsl_string());
     }
+
+    fn boxed_render_clone(&self) -> Box<dyn RenderModifier> {
+        Box::new(*self)
+    }
+
+    fn as_modifier(&self) -> &dyn Modifier {
+        self
+    }
 }
 
 /// A modifier modulating each particle's size over its lifetime with a gradient
@@ -220,6 +253,14 @@ impl RenderModifier for SizeOverLifetimeModifier {
             Attribute::LIFETIME.name()
         );
     }
+
+    fn boxed_render_clone(&self) -> Box<dyn RenderModifier> {
+        Box::new(self.clone())
+    }
+
+    fn as_modifier(&self) -> &dyn Modifier {
+        self
+    }
 }
 
 /// Mode of orientation of a particle's local frame.
@@ -346,6 +387,10 @@ impl Modifier for OrientModifier {
     fn boxed_clone(&self) -> BoxedModifier {
         Box::new(*self)
     }
+
+    fn apply(&self, _module: &mut Module, _context: &mut ShaderWriter) -> Result<(), ExprError> {
+        Err(ExprError::TypeError("Wrong modifier context".to_string()))
+    }
 }
 
 #[typetag::serde]
@@ -405,6 +450,14 @@ axis_z = cross(axis_x, axis_y);
             }
         }
     }
+
+    fn boxed_render_clone(&self) -> Box<dyn RenderModifier> {
+        Box::new(*self)
+    }
+
+    fn as_modifier(&self) -> &dyn Modifier {
+        self
+    }
 }
 
 /// A modifier to render particles using flipbook animation.
@@ -502,6 +555,14 @@ impl RenderModifier for FlipbookModifier {
     fn apply_render(&self, _module: &mut Module, context: &mut RenderContext) {
         context.sprite_grid_size = Some(self.sprite_grid_size);
     }
+
+    fn boxed_render_clone(&self) -> Box<dyn RenderModifier> {
+        Box::new(*self)
+    }
+
+    fn as_modifier(&self) -> &dyn Modifier {
+        self
+    }
 }
 
 /// A modifier to interpret the size of all particles in screen-space pixels.
@@ -551,6 +612,14 @@ impl RenderModifier for ScreenSpaceSizeModifier {
             size = (size * w_cs * 2.0) / min(screen_size_pixels.x * projection_scale.x, screen_size_pixels.y * projection_scale.y);\n",
             Attribute::POSITION.name());
     }
+
+    fn boxed_render_clone(&self) -> Box<dyn RenderModifier> {
+        Box::new(*self)
+    }
+
+    fn as_modifier(&self) -> &dyn Modifier {
+        self
+    }
 }
 
 #[cfg(test)]
diff --git a/src/modifier/position.rs b/src/modifier/position.rs
index af2f6189..7c048e91 100644
--- a/src/modifier/position.rs
+++ b/src/modifier/position.rs
@@ -8,8 +8,8 @@ use bevy::prelude::*;
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    calc_func_id, graph::ExprError, impl_mod_init_update, modifier::ShapeDimension, Attribute,
-    EvalContext, ExprHandle, InitContext, InitModifier, Module, UpdateContext, UpdateModifier,
+    calc_func_id, graph::ExprError, modifier::ShapeDimension, Attribute, BoxedModifier,
+    EvalContext, ExprHandle, Modifier, ModifierContext, Module, ShaderWriter,
 };
 
 /// A modifier to set the position of particles on or inside a circle/disc,
@@ -48,8 +48,6 @@ pub struct SetPositionCircleModifier {
     pub dimension: ShapeDimension,
 }
 
-impl_mod_init_update!(SetPositionCircleModifier, &[Attribute::POSITION]);
-
 impl SetPositionCircleModifier {
     fn eval(
         &self,
@@ -111,23 +109,22 @@ impl SetPositionCircleModifier {
 }
 
 #[typetag::serde]
-impl InitModifier for SetPositionCircleModifier {
-    fn apply_init(&self, module: &mut Module, context: &mut InitContext) -> Result<(), ExprError> {
-        let code = self.eval(module, context)?;
-        context.init_code += &code;
-        Ok(())
+impl Modifier for SetPositionCircleModifier {
+    fn context(&self) -> ModifierContext {
+        ModifierContext::Init | ModifierContext::Update
     }
-}
 
-#[typetag::serde]
-impl UpdateModifier for SetPositionCircleModifier {
-    fn apply_update(
-        &self,
-        module: &mut Module,
-        context: &mut UpdateContext,
-    ) -> Result<(), ExprError> {
+    fn attributes(&self) -> &[Attribute] {
+        &[Attribute::POSITION]
+    }
+
+    fn boxed_clone(&self) -> BoxedModifier {
+        Box::new(self.clone())
+    }
+
+    fn apply(&self, module: &mut Module, context: &mut ShaderWriter) -> Result<(), ExprError> {
         let code = self.eval(module, context)?;
-        context.update_code += &code;
+        context.main_code += &code;
         Ok(())
     }
 }
@@ -152,8 +149,6 @@ pub struct SetPositionSphereModifier {
     pub dimension: ShapeDimension,
 }
 
-impl_mod_init_update!(SetPositionSphereModifier, &[Attribute::POSITION]);
-
 impl SetPositionSphereModifier {
     fn eval(
         &self,
@@ -217,23 +212,22 @@ impl SetPositionSphereModifier {
 }
 
 #[typetag::serde]
-impl InitModifier for SetPositionSphereModifier {
-    fn apply_init(&self, module: &mut Module, context: &mut InitContext) -> Result<(), ExprError> {
-        let code = self.eval(module, context)?;
-        context.init_code += &code;
-        Ok(())
+impl Modifier for SetPositionSphereModifier {
+    fn context(&self) -> ModifierContext {
+        ModifierContext::Init | ModifierContext::Update
     }
-}
 
-#[typetag::serde]
-impl UpdateModifier for SetPositionSphereModifier {
-    fn apply_update(
-        &self,
-        module: &mut Module,
-        context: &mut UpdateContext,
-    ) -> Result<(), ExprError> {
+    fn attributes(&self) -> &[Attribute] {
+        &[Attribute::POSITION]
+    }
+
+    fn boxed_clone(&self) -> BoxedModifier {
+        Box::new(self.clone())
+    }
+
+    fn apply(&self, module: &mut Module, context: &mut ShaderWriter) -> Result<(), ExprError> {
         let code = self.eval(module, context)?;
-        context.update_code += &code;
+        context.main_code += &code;
         Ok(())
     }
 }
@@ -271,8 +265,6 @@ pub struct SetPositionCone3dModifier {
     pub dimension: ShapeDimension,
 }
 
-impl_mod_init_update!(SetPositionCone3dModifier, &[Attribute::POSITION]);
-
 impl SetPositionCone3dModifier {
     fn eval(
         &self,
@@ -335,23 +327,22 @@ impl SetPositionCone3dModifier {
 }
 
 #[typetag::serde]
-impl InitModifier for SetPositionCone3dModifier {
-    fn apply_init(&self, module: &mut Module, context: &mut InitContext) -> Result<(), ExprError> {
-        let code = self.eval(module, context)?;
-        context.init_code += &code;
-        Ok(())
+impl Modifier for SetPositionCone3dModifier {
+    fn context(&self) -> ModifierContext {
+        ModifierContext::Init | ModifierContext::Update
     }
-}
 
-#[typetag::serde]
-impl UpdateModifier for SetPositionCone3dModifier {
-    fn apply_update(
-        &self,
-        module: &mut Module,
-        context: &mut UpdateContext,
-    ) -> Result<(), ExprError> {
+    fn attributes(&self) -> &[Attribute] {
+        &[Attribute::POSITION]
+    }
+
+    fn boxed_clone(&self) -> BoxedModifier {
+        Box::new(self.clone())
+    }
+
+    fn apply(&self, module: &mut Module, context: &mut ShaderWriter) -> Result<(), ExprError> {
         let code = self.eval(module, context)?;
-        context.update_code += &code;
+        context.main_code += &code;
         Ok(())
     }
 }
diff --git a/src/modifier/velocity.rs b/src/modifier/velocity.rs
index e6c8c2c6..65ef2cf0 100644
--- a/src/modifier/velocity.rs
+++ b/src/modifier/velocity.rs
@@ -13,8 +13,8 @@ use bevy::prelude::*;
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    calc_func_id, graph::ExprError, impl_mod_init_update, Attribute, EvalContext, ExprHandle,
-    InitContext, InitModifier, Module, UpdateContext, UpdateModifier,
+    calc_func_id, graph::ExprError, Attribute, BoxedModifier, EvalContext, ExprHandle, Modifier,
+    ModifierContext, Module, ShaderWriter,
 };
 
 /// A modifier to set the velocity of particles radially on a circle.
@@ -41,11 +41,6 @@ pub struct SetVelocityCircleModifier {
     pub speed: ExprHandle,
 }
 
-impl_mod_init_update!(
-    SetVelocityCircleModifier,
-    &[Attribute::POSITION, Attribute::VELOCITY]
-);
-
 impl SetVelocityCircleModifier {
     fn eval(
         &self,
@@ -86,23 +81,22 @@ impl SetVelocityCircleModifier {
 }
 
 #[typetag::serde]
-impl InitModifier for SetVelocityCircleModifier {
-    fn apply_init(&self, module: &mut Module, context: &mut InitContext) -> Result<(), ExprError> {
-        let code = self.eval(module, context)?;
-        context.init_code += &code;
-        Ok(())
+impl Modifier for SetVelocityCircleModifier {
+    fn context(&self) -> ModifierContext {
+        ModifierContext::Init | ModifierContext::Update
     }
-}
 
-#[typetag::serde]
-impl UpdateModifier for SetVelocityCircleModifier {
-    fn apply_update(
-        &self,
-        module: &mut Module,
-        context: &mut UpdateContext,
-    ) -> Result<(), ExprError> {
+    fn attributes(&self) -> &[Attribute] {
+        &[Attribute::POSITION, Attribute::VELOCITY]
+    }
+
+    fn boxed_clone(&self) -> BoxedModifier {
+        Box::new(self.clone())
+    }
+
+    fn apply(&self, module: &mut Module, context: &mut ShaderWriter) -> Result<(), ExprError> {
         let code = self.eval(module, context)?;
-        context.update_code += &code;
+        context.main_code += &code;
         Ok(())
     }
 }
@@ -127,11 +121,6 @@ pub struct SetVelocitySphereModifier {
     pub speed: ExprHandle,
 }
 
-impl_mod_init_update!(
-    SetVelocitySphereModifier,
-    &[Attribute::POSITION, Attribute::VELOCITY]
-);
-
 impl SetVelocitySphereModifier {
     fn eval(
         &self,
@@ -152,23 +141,22 @@ impl SetVelocitySphereModifier {
 }
 
 #[typetag::serde]
-impl InitModifier for SetVelocitySphereModifier {
-    fn apply_init(&self, module: &mut Module, context: &mut InitContext) -> Result<(), ExprError> {
-        let code = self.eval(module, context)?;
-        context.init_code += &code;
-        Ok(())
+impl Modifier for SetVelocitySphereModifier {
+    fn context(&self) -> ModifierContext {
+        ModifierContext::Init | ModifierContext::Update
     }
-}
 
-#[typetag::serde]
-impl UpdateModifier for SetVelocitySphereModifier {
-    fn apply_update(
-        &self,
-        module: &mut Module,
-        context: &mut UpdateContext,
-    ) -> Result<(), ExprError> {
+    fn attributes(&self) -> &[Attribute] {
+        &[Attribute::POSITION, Attribute::VELOCITY]
+    }
+
+    fn boxed_clone(&self) -> BoxedModifier {
+        Box::new(self.clone())
+    }
+
+    fn apply(&self, module: &mut Module, context: &mut ShaderWriter) -> Result<(), ExprError> {
         let code = self.eval(module, context)?;
-        context.update_code += &code;
+        context.main_code += &code;
         Ok(())
     }
 }
@@ -198,11 +186,6 @@ pub struct SetVelocityTangentModifier {
     pub speed: ExprHandle,
 }
 
-impl_mod_init_update!(
-    SetVelocityTangentModifier,
-    &[Attribute::POSITION, Attribute::VELOCITY]
-);
-
 impl SetVelocityTangentModifier {
     fn eval(
         &self,
@@ -243,23 +226,22 @@ impl SetVelocityTangentModifier {
 }
 
 #[typetag::serde]
-impl InitModifier for SetVelocityTangentModifier {
-    fn apply_init(&self, module: &mut Module, context: &mut InitContext) -> Result<(), ExprError> {
-        let code = self.eval(module, context)?;
-        context.init_code += &code;
-        Ok(())
+impl Modifier for SetVelocityTangentModifier {
+    fn context(&self) -> ModifierContext {
+        ModifierContext::Init | ModifierContext::Update
     }
-}
 
-#[typetag::serde]
-impl UpdateModifier for SetVelocityTangentModifier {
-    fn apply_update(
-        &self,
-        module: &mut Module,
-        context: &mut UpdateContext,
-    ) -> Result<(), ExprError> {
+    fn attributes(&self) -> &[Attribute] {
+        &[Attribute::POSITION, Attribute::VELOCITY]
+    }
+
+    fn boxed_clone(&self) -> BoxedModifier {
+        Box::new(self.clone())
+    }
+
+    fn apply(&self, module: &mut Module, context: &mut ShaderWriter) -> Result<(), ExprError> {
         let code = self.eval(module, context)?;
-        context.update_code += &code;
+        context.main_code += &code;
         Ok(())
     }
 }