Too many plans

- Add PlannedVisual
- Make ExampleEffect use PlannedVisual
- Remove One2ManyStorage
- Merge Storage, AbstractStorage, and One2OneStorage
- Storage directly provides update and tick plans
- Move work distribution logic to static methods in PlanUtil
- Rename RunOnAllPlan -> ForEachPlan
- Add ContextConsumer and ContextRunnable interfaces and remove
  ContextAgnosticPlan

src/main/java/com/jozufozu/flywheel/api/visual/Effect.java

@@ -1,11 +1,7 @@
 package com.jozufozu.flywheel.api.visual;
 
-import java.util.Collection;
-
 import com.jozufozu.flywheel.api.visualization.VisualizationContext;
 
-// TODO: add level getter?
-// TODO: return single visual instead of many?
 public interface Effect {
-	Collection<EffectVisual<?>> createVisuals(VisualizationContext ctx);
+	EffectVisual<?> visualize(VisualizationContext ctx);
 }

src/main/java/com/jozufozu/flywheel/api/visual/PlannedVisual.java

@@ -0,0 +1,19 @@
+package com.jozufozu.flywheel.api.visual;
+
+import com.jozufozu.flywheel.api.task.Plan;
+import com.jozufozu.flywheel.lib.task.UnitPlan;
+
+/**
+ * An interface giving {@link Visual}s a way to define complex, parallelized update plans.
+ * <p>
+ * Plans allow for
+ */
+public interface PlannedVisual extends Visual {
+	default Plan<VisualFrameContext> planFrame() {
+		return UnitPlan.of();
+	}
+
+	default Plan<VisualTickContext> planTick() {
+		return UnitPlan.of();
+	}
+}

src/main/java/com/jozufozu/flywheel/backend/engine/batching/TransformCall.java

@@ -12,7 +12,7 @@ import com.jozufozu.flywheel.api.material.Material;
 import com.jozufozu.flywheel.api.material.MaterialVertexTransformer;
 import com.jozufozu.flywheel.api.task.Plan;
 import com.jozufozu.flywheel.api.vertex.MutableVertexList;
-import com.jozufozu.flywheel.lib.task.RunOnAllPlan;
+import com.jozufozu.flywheel.lib.task.ForEachPlan;
 import com.jozufozu.flywheel.lib.vertex.VertexTransformations;
 import com.mojang.blaze3d.vertex.PoseStack;
 import com.mojang.math.Matrix3f;
@@ -38,7 +38,7 @@ public class TransformCall<I extends Instance> {
 		meshVertexCount = mesh.getVertexCount();
 		boundingSphere = mesh.mesh.getBoundingSphere();
 
-		drawPlan = RunOnAllPlan.of(instancer::getAll, (instance, ctx) -> {
+		drawPlan = ForEachPlan.of(instancer::getAll, (instance, ctx) -> {
 			var boundingSphere = new Vector4f(this.boundingSphere);
 
 			boundingSphereTransformer.transform(boundingSphere, instance);

src/main/java/com/jozufozu/flywheel/impl/visualization/manager/BlockEntityVisualManager.java

@@ -9,7 +9,6 @@ import com.jozufozu.flywheel.api.visual.BlockEntityVisual;
 import com.jozufozu.flywheel.api.visual.Visual;
 import com.jozufozu.flywheel.api.visualization.VisualizationContext;
 import com.jozufozu.flywheel.impl.visualization.VisualizationHelper;
-import com.jozufozu.flywheel.impl.visualization.storage.One2OneStorage;
 import com.jozufozu.flywheel.impl.visualization.storage.Storage;
 import com.jozufozu.flywheel.util.FlwUtil;
 
@@ -39,7 +38,7 @@ public class BlockEntityVisualManager extends VisualManager<BlockEntity> {
 		}
 	}
 
-	private static class BlockEntityStorage extends One2OneStorage<BlockEntity> {
+	private static class BlockEntityStorage extends Storage<BlockEntity> {
 		private final Long2ObjectMap<BlockEntityVisual<?>> posLookup = new Long2ObjectOpenHashMap<>();
 
 		public BlockEntityStorage(Engine engine) {

src/main/java/com/jozufozu/flywheel/impl/visualization/manager/EffectVisualManager.java

@@ -1,12 +1,9 @@
 package com.jozufozu.flywheel.impl.visualization.manager;
 
-import java.util.Collection;
-
 import com.jozufozu.flywheel.api.backend.Engine;
 import com.jozufozu.flywheel.api.visual.Effect;
-import com.jozufozu.flywheel.api.visual.Visual;
+import com.jozufozu.flywheel.api.visual.EffectVisual;
 import com.jozufozu.flywheel.api.visualization.VisualizationContext;
-import com.jozufozu.flywheel.impl.visualization.storage.One2ManyStorage;
 import com.jozufozu.flywheel.impl.visualization.storage.Storage;
 
 public class EffectVisualManager extends VisualManager<Effect> {
@@ -21,14 +18,14 @@ public class EffectVisualManager extends VisualManager<Effect> {
 		return storage;
 	}
 
-	private static class EffectStorage extends One2ManyStorage<Effect> {
+	private static class EffectStorage extends Storage<Effect> {
 		public EffectStorage(Engine engine) {
 			super(engine);
 		}
 
 		@Override
-		protected Collection<? extends Visual> createRaw(Effect obj) {
-			return obj.createVisuals(new VisualizationContext(engine, engine.renderOrigin()));
+		protected EffectVisual<?> createRaw(Effect obj) {
+			return obj.visualize(new VisualizationContext(engine, engine.renderOrigin()));
 		}
 
 		@Override

src/main/java/com/jozufozu/flywheel/impl/visualization/manager/EntityVisualManager.java

@@ -6,7 +6,6 @@ import com.jozufozu.flywheel.api.backend.Engine;
 import com.jozufozu.flywheel.api.visual.Visual;
 import com.jozufozu.flywheel.api.visualization.VisualizationContext;
 import com.jozufozu.flywheel.impl.visualization.VisualizationHelper;
-import com.jozufozu.flywheel.impl.visualization.storage.One2OneStorage;
 import com.jozufozu.flywheel.impl.visualization.storage.Storage;
 import com.jozufozu.flywheel.util.FlwUtil;
 
@@ -25,7 +24,7 @@ public class EntityVisualManager extends VisualManager<Entity> {
 		return storage;
 	}
 
-	private static class EntityStorage extends One2OneStorage<Entity> {
+	private static class EntityStorage extends Storage<Entity> {
 		public EntityStorage(Engine engine) {
 			super(engine);
 		}

src/main/java/com/jozufozu/flywheel/impl/visualization/manager/VisualManager.java

@@ -4,8 +4,6 @@ import java.util.Queue;
 import java.util.concurrent.ConcurrentLinkedQueue;
 
 import com.jozufozu.flywheel.api.task.Plan;
-import com.jozufozu.flywheel.api.visual.DynamicVisual;
-import com.jozufozu.flywheel.api.visual.TickableVisual;
 import com.jozufozu.flywheel.api.visual.VisualFrameContext;
 import com.jozufozu.flywheel.api.visual.VisualTickContext;
 import com.jozufozu.flywheel.config.FlwConfig;
@@ -16,7 +14,6 @@ import com.jozufozu.flywheel.impl.visualization.ratelimit.DistanceUpdateLimiterI
 import com.jozufozu.flywheel.impl.visualization.ratelimit.NonLimiter;
 import com.jozufozu.flywheel.impl.visualization.storage.Storage;
 import com.jozufozu.flywheel.impl.visualization.storage.Transaction;
-import com.jozufozu.flywheel.lib.task.RunOnAllPlan;
 import com.jozufozu.flywheel.lib.task.SimplePlan;
 import com.jozufozu.flywheel.util.Unit;
 
@@ -72,7 +69,6 @@ public abstract class VisualManager<T> {
 	}
 
 	public Plan<Unit> createRecreationPlan() {
-		// TODO: parallelize recreation?
 		return SimplePlan.of(getStorage()::recreateAll);
 	}
 
@@ -93,7 +89,7 @@ public abstract class VisualManager<T> {
 					tickLimiter.tick();
 					processQueue();
 				})
-				.thenMap(this::createVisualTickContext, RunOnAllPlan.of(getStorage()::getTickableVisuals, TickableVisual::tick));
+				.thenMap(this::createVisualTickContext, getStorage().getTickPlan());
 	}
 
 	public Plan<FrameContext> createFramePlan() {
@@ -101,7 +97,7 @@ public abstract class VisualManager<T> {
 					frameLimiter.tick();
 					processQueue();
 				})
-				.thenMap(this::createVisualContext, RunOnAllPlan.of(getStorage()::getDynamicVisuals, DynamicVisual::beginFrame));
+				.thenMap(this::createVisualContext, getStorage().getFramePlan());
 	}
 
 	private VisualFrameContext createVisualContext(FrameContext ctx) {

src/main/java/com/jozufozu/flywheel/impl/visualization/storage/AbstractStorage.java

@@ -1,41 +0,0 @@
-package com.jozufozu.flywheel.impl.visualization.storage;
-
-import java.util.ArrayList;
-import java.util.List;
-
-import com.jozufozu.flywheel.api.backend.Engine;
-import com.jozufozu.flywheel.api.visual.DynamicVisual;
-import com.jozufozu.flywheel.api.visual.TickableVisual;
-import com.jozufozu.flywheel.api.visual.Visual;
-
-public abstract class AbstractStorage<T> implements Storage<T> {
-	protected final Engine engine;
-	protected final List<TickableVisual> tickableVisuals = new ArrayList<>();
-	protected final List<DynamicVisual> dynamicVisuals = new ArrayList<>();
-
-	protected AbstractStorage(Engine engine) {
-		this.engine = engine;
-	}
-
-	@Override
-	public List<TickableVisual> getTickableVisuals() {
-		return tickableVisuals;
-	}
-
-	@Override
-	public List<DynamicVisual> getDynamicVisuals() {
-		return dynamicVisuals;
-	}
-
-	protected void setup(Visual visual) {
-		visual.init();
-
-		if (visual instanceof TickableVisual tickable) {
-			tickableVisuals.add(tickable);
-		}
-
-		if (visual instanceof DynamicVisual dynamic) {
-			dynamicVisuals.add(dynamic);
-		}
-	}
-}

src/main/java/com/jozufozu/flywheel/impl/visualization/storage/One2ManyStorage.java

@@ -1,86 +0,0 @@
-package com.jozufozu.flywheel.impl.visualization.storage;
-
-import java.util.Collection;
-import java.util.List;
-
-import com.google.common.collect.HashMultimap;
-import com.google.common.collect.Multimap;
-import com.jozufozu.flywheel.api.backend.Engine;
-import com.jozufozu.flywheel.api.visual.Visual;
-
-public abstract class One2ManyStorage<T> extends AbstractStorage<T> {
-	private final Multimap<T, Visual> allVisuals = HashMultimap.create();
-
-	public One2ManyStorage(Engine engine) {
-		super(engine);
-	}
-
-	@Override
-	public Collection<Visual> getAllVisuals() {
-		return allVisuals.values();
-	}
-
-	@Override
-	public void add(T obj) {
-		Collection<Visual> visuals = allVisuals.get(obj);
-
-		if (visuals.isEmpty()) {
-			create(obj);
-		}
-	}
-
-	@Override
-	public void remove(T obj) {
-		Collection<Visual> visuals = allVisuals.removeAll(obj);
-
-		if (visuals.isEmpty()) {
-			return;
-		}
-
-		tickableVisuals.removeAll(visuals);
-		dynamicVisuals.removeAll(visuals);
-		visuals.forEach(Visual::delete);
-	}
-
-	@Override
-	public void update(T obj) {
-		Collection<Visual> visuals = allVisuals.get(obj);
-
-		if (visuals.isEmpty()) {
-			return;
-		}
-
-		// TODO: shouldReset cannot be checked here because all visuals are created at once
-		visuals.forEach(Visual::update);
-	}
-
-	@Override
-	public void recreateAll() {
-		tickableVisuals.clear();
-		dynamicVisuals.clear();
-		allVisuals.values().forEach(Visual::delete);
-
-		List<T> objects = List.copyOf(allVisuals.keySet());
-		allVisuals.clear();
-		objects.forEach(this::create);
-	}
-
-	@Override
-	public void invalidate() {
-		tickableVisuals.clear();
-		dynamicVisuals.clear();
-		allVisuals.values().forEach(Visual::delete);
-		allVisuals.clear();
-	}
-
-	private void create(T obj) {
-		Collection<? extends Visual> visuals = createRaw(obj);
-
-		if (!visuals.isEmpty()) {
-			visuals.forEach(this::setup);
-			allVisuals.putAll(obj, visuals);
-		}
-	}
-
-	protected abstract Collection<? extends Visual> createRaw(T obj);
-}

src/main/java/com/jozufozu/flywheel/impl/visualization/storage/One2OneStorage.java

@@ -1,101 +0,0 @@
-package com.jozufozu.flywheel.impl.visualization.storage;
-
-import java.util.Collection;
-import java.util.HashMap;
-import java.util.Map;
-
-import org.jetbrains.annotations.Nullable;
-
-import com.jozufozu.flywheel.api.backend.Engine;
-import com.jozufozu.flywheel.api.visual.Visual;
-
-public abstract class One2OneStorage<T> extends AbstractStorage<T> {
-	private final Map<T, Visual> visuals = new HashMap<>();
-
-	public One2OneStorage(Engine engine) {
-		super(engine);
-	}
-
-	@Override
-	public Collection<Visual> getAllVisuals() {
-		return visuals.values();
-	}
-
-	@Override
-	public void add(T obj) {
-		Visual visual = visuals.get(obj);
-
-		if (visual == null) {
-			create(obj);
-		}
-	}
-
-	@Override
-	public void remove(T obj) {
-		Visual visual = visuals.remove(obj);
-
-		if (visual == null) {
-			return;
-		}
-
-		tickableVisuals.remove(visual);
-		dynamicVisuals.remove(visual);
-		visual.delete();
-	}
-
-	@Override
-	public void update(T obj) {
-		Visual visual = visuals.get(obj);
-
-		if (visual == null) {
-			return;
-		}
-
-		// resetting visuals is by default used to handle block state changes.
-		if (visual.shouldReset()) {
-			// delete and re-create the visual.
-			// resetting a visual supersedes updating it.
-			remove(obj);
-			create(obj);
-		} else {
-			visual.update();
-		}
-	}
-
-	@Override
-	public void recreateAll() {
-		tickableVisuals.clear();
-		dynamicVisuals.clear();
-		visuals.replaceAll((obj, visual) -> {
-			visual.delete();
-
-			Visual out = createRaw(obj);
-
-			if (out != null) {
-				setup(out);
-			}
-
-			return out;
-		});
-	}
-
-	@Override
-	public void invalidate() {
-		tickableVisuals.clear();
-		dynamicVisuals.clear();
-		visuals.values().forEach(Visual::delete);
-		visuals.clear();
-	}
-
-	private void create(T obj) {
-		Visual visual = createRaw(obj);
-
-		if (visual != null) {
-			setup(visual);
-			visuals.put(obj, visual);
-		}
-	}
-
-	@Nullable
-	protected abstract Visual createRaw(T obj);
-}

src/main/java/com/jozufozu/flywheel/impl/visualization/storage/Storage.java

@@ -1,33 +1,157 @@
 package com.jozufozu.flywheel.impl.visualization.storage;
 
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.Map;
 
+import org.jetbrains.annotations.Nullable;
+
+import com.jozufozu.flywheel.api.backend.Engine;
+import com.jozufozu.flywheel.api.task.Plan;
 import com.jozufozu.flywheel.api.visual.DynamicVisual;
+import com.jozufozu.flywheel.api.visual.PlannedVisual;
 import com.jozufozu.flywheel.api.visual.TickableVisual;
 import com.jozufozu.flywheel.api.visual.Visual;
+import com.jozufozu.flywheel.api.visual.VisualFrameContext;
+import com.jozufozu.flywheel.api.visual.VisualTickContext;
+import com.jozufozu.flywheel.lib.task.ForEachPlan;
 
-public interface Storage<T> {
-	Collection<Visual> getAllVisuals();
+import it.unimi.dsi.fastutil.objects.Reference2ObjectOpenHashMap;
 
-	List<TickableVisual> getTickableVisuals();
+public abstract class Storage<T> {
+	protected final Engine engine;
+	protected final List<TickableVisual> tickableVisuals = new ArrayList<>();
+	protected final List<DynamicVisual> dynamicVisuals = new ArrayList<>();
+	protected final List<PlannedVisual> plannedVisuals = new ArrayList<>();
+	protected final VisualUpdatePlan<VisualFrameContext> framePlan = new VisualUpdatePlan<>(() -> plannedVisuals.stream()
+			.map(PlannedVisual::planFrame)
+			.toList());
+	protected final VisualUpdatePlan<VisualTickContext> tickPlan = new VisualUpdatePlan<>(() -> plannedVisuals.stream()
+			.map(PlannedVisual::planTick)
+			.toList());
 
-	List<DynamicVisual> getDynamicVisuals();
+	private final Map<T, Visual> visuals = new Reference2ObjectOpenHashMap<>();
+
+	public Storage(Engine engine) {
+		this.engine = engine;
+	}
+
+	public Collection<Visual> getAllVisuals() {
+		return visuals.values();
+	}
+
+	public void add(T obj) {
+		Visual visual = visuals.get(obj);
+
+		if (visual == null) {
+			create(obj);
+		}
+	}
+
+	public void remove(T obj) {
+		Visual visual = visuals.remove(obj);
+
+		if (visual == null) {
+			return;
+		}
+
+		tickableVisuals.remove(visual);
+		dynamicVisuals.remove(visual);
+		if (plannedVisuals.remove(visual)) {
+			framePlan.clear();
+		}
+		visual.delete();
+	}
+
+	public void update(T obj) {
+		Visual visual = visuals.get(obj);
+
+		if (visual == null) {
+			return;
+		}
+
+		// resetting visuals is by default used to handle block state changes.
+		if (visual.shouldReset()) {
+			// delete and re-create the visual.
+			// resetting a visual supersedes updating it.
+			remove(obj);
+			create(obj);
+		} else {
+			visual.update();
+		}
+	}
+
+	public void recreateAll() {
+		tickableVisuals.clear();
+		dynamicVisuals.clear();
+		plannedVisuals.clear();
+		visuals.replaceAll((obj, visual) -> {
+			visual.delete();
+
+			Visual out = createRaw(obj);
+
+			if (out != null) {
+				setup(out);
+			}
+
+			return out;
+		});
+	}
+
+	public void invalidate() {
+		tickableVisuals.clear();
+		dynamicVisuals.clear();
+		plannedVisuals.clear();
+		framePlan.clear();
+		tickPlan.clear();
+		visuals.values()
+				.forEach(Visual::delete);
+		visuals.clear();
+	}
+
+	private void create(T obj) {
+		Visual visual = createRaw(obj);
+
+		if (visual != null) {
+			setup(visual);
+			visuals.put(obj, visual);
+		}
+	}
+
+	@Nullable
+	protected abstract Visual createRaw(T obj);
+
+	public Plan<VisualFrameContext> getFramePlan() {
+		return framePlan.and(ForEachPlan.of(() -> dynamicVisuals, DynamicVisual::beginFrame));
+	}
+
+	public Plan<VisualTickContext> getTickPlan() {
+		return tickPlan.and(ForEachPlan.of(() -> tickableVisuals, TickableVisual::tick));
+	}
+
+	private void setup(Visual visual) {
+		visual.init();
+
+		if (visual instanceof TickableVisual tickable) {
+			tickableVisuals.add(tickable);
+		}
+
+		if (visual instanceof DynamicVisual dynamic) {
+			dynamicVisuals.add(dynamic);
+		}
+
+		if (visual instanceof PlannedVisual planned) {
+			plannedVisuals.add(planned);
+			framePlan.add(planned.planFrame());
+			tickPlan.add(planned.planTick());
+		}
+	}
 
 	/**
 	 * Is the given object currently capable of being added?
 	 *
 	 * @return true if the object is currently capable of being visualized.
 	 */
-	boolean willAccept(T obj);
-
-	void add(T obj);
-
-	void remove(T obj);
-
-	void update(T obj);
-
-	void recreateAll();
-
-	void invalidate();
+	public abstract boolean willAccept(T obj);
 }

src/main/java/com/jozufozu/flywheel/impl/visualization/storage/VisualUpdatePlan.java

@@ -0,0 +1,53 @@
+package com.jozufozu.flywheel.impl.visualization.storage;
+
+import java.util.List;
+import java.util.function.Supplier;
+
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import com.jozufozu.flywheel.api.task.Plan;
+import com.jozufozu.flywheel.api.task.TaskExecutor;
+import com.jozufozu.flywheel.lib.task.NestedPlan;
+import com.jozufozu.flywheel.lib.task.SimplyComposedPlan;
+
+public class VisualUpdatePlan<C> implements SimplyComposedPlan<C> {
+	private final Supplier<List<Plan<C>>> initializer;
+	@Nullable
+	private Plan<C> plan;
+	private boolean needsSimplify = true;
+
+	public VisualUpdatePlan(Supplier<List<Plan<C>>> initializer) {
+		this.initializer = initializer;
+	}
+
+	@Override
+	public void execute(TaskExecutor taskExecutor, C context, Runnable onCompletion) {
+		updatePlans().execute(taskExecutor, context, onCompletion);
+	}
+
+	public void add(Plan<C> plan) {
+		if (this.plan == null) {
+			this.plan = plan;
+		} else {
+			this.plan = this.plan.and(plan);
+		}
+		needsSimplify = true;
+	}
+
+	@NotNull
+	private Plan<C> updatePlans() {
+		if (plan == null) {
+			plan = new NestedPlan<>(initializer.get()).maybeSimplify();
+		} else if (needsSimplify) {
+			plan = plan.maybeSimplify();
+		}
+
+		needsSimplify = false;
+		return plan;
+	}
+
+	public void clear() {
+		plan = null;
+	}
+}

src/main/java/com/jozufozu/flywheel/lib/task/BarrierPlan.java

@@ -4,6 +4,10 @@ import com.jozufozu.flywheel.api.task.Plan;
 import com.jozufozu.flywheel.api.task.TaskExecutor;
 
 public record BarrierPlan<C>(Plan<C> first, Plan<C> second) implements SimplyComposedPlan<C> {
+	public static <C> Plan<C> of(Plan<C> first, Plan<C> second) {
+		return new BarrierPlan<>(first, second);
+	}
+
 	@Override
 	public void execute(TaskExecutor taskExecutor, C context, Runnable onCompletion) {
 		first.execute(taskExecutor, context, () -> second.execute(taskExecutor, context, onCompletion));

src/main/java/com/jozufozu/flywheel/lib/task/ContextAgnosticPlan.java

@@ -1,19 +0,0 @@
-package com.jozufozu.flywheel.lib.task;
-
-import com.jozufozu.flywheel.api.task.Plan;
-import com.jozufozu.flywheel.api.task.TaskExecutor;
-
-public interface ContextAgnosticPlan extends SimplyComposedPlan<Object> {
-	@SuppressWarnings("unchecked")
-	default <C> Plan<C> cast() {
-		// The context is entirely ignored, so we can safely cast to any context.
-		return (Plan<C>) this;
-	}
-
-	@Override
-	default void execute(TaskExecutor taskExecutor, Object ignored, Runnable onCompletion) {
-		execute(taskExecutor, onCompletion);
-	}
-
-	void execute(TaskExecutor taskExecutor, Runnable onCompletion);
-}

src/main/java/com/jozufozu/flywheel/lib/task/ContextConsumer.java

@@ -0,0 +1,13 @@
+package com.jozufozu.flywheel.lib.task;
+
+import com.jozufozu.flywheel.api.task.Plan;
+
+/**
+ * A consumer like interface for use with {@link Plan}s.
+ *
+ * @param <C> The context type.
+ */
+@FunctionalInterface
+public interface ContextConsumer<C> {
+	void accept(C context);
+}

src/main/java/com/jozufozu/flywheel/lib/task/ContextRunnable.java

@@ -0,0 +1,16 @@
+package com.jozufozu.flywheel.lib.task;
+
+/**
+ * A {@link ContextConsumer} that ignores the context object.
+ *
+ * @param <C> The context type.
+ */
+@FunctionalInterface
+public interface ContextRunnable<C> extends ContextConsumer<C> {
+	void run();
+
+	@Override
+	default void accept(C ignored) {
+		run();
+	}
+}

src/main/java/com/jozufozu/flywheel/lib/task/ForEachPlan.java

@@ -0,0 +1,25 @@
+package com.jozufozu.flywheel.lib.task;
+
+import java.util.List;
+import java.util.function.BiConsumer;
+import java.util.function.Consumer;
+import java.util.function.Supplier;
+
+import com.jozufozu.flywheel.api.task.Plan;
+import com.jozufozu.flywheel.api.task.TaskExecutor;
+
+public record ForEachPlan<T, C>(Supplier<List<T>> listSupplier,
+								BiConsumer<T, C> action) implements SimplyComposedPlan<C> {
+	public static <T, C> Plan<C> of(Supplier<List<T>> iterable, BiConsumer<T, C> forEach) {
+		return new ForEachPlan<>(iterable, forEach);
+	}
+
+	public static <T, C> Plan<C> of(Supplier<List<T>> iterable, Consumer<T> forEach) {
+		return of(iterable, (t, c) -> forEach.accept(t));
+	}
+
+	@Override
+	public void execute(TaskExecutor taskExecutor, C context, Runnable onCompletion) {
+		taskExecutor.execute(() -> PlanUtil.distribute(taskExecutor, context, onCompletion, listSupplier.get(), action));
+	}
+}

src/main/java/com/jozufozu/flywheel/lib/task/NestedPlan.java

@@ -54,9 +54,28 @@ public record NestedPlan<C>(List<Plan<C>> parallelPlans) implements SimplyCompos
 					.maybeSimplify();
 		}
 
-		var simplifiedTasks = new ArrayList<Runnable>();
+		var simplifiedTasks = new ArrayList<ContextConsumer<C>>();
 		var simplifiedPlans = new ArrayList<Plan<C>>();
-		flattenTasksAndPlans(simplifiedTasks, simplifiedPlans);
+		var toVisit = new ArrayDeque<>(parallelPlans);
+		while (!toVisit.isEmpty()) {
+			var plan = toVisit.pop()
+					.maybeSimplify();
+
+			if (plan == UnitPlan.of()) {
+				continue;
+			}
+
+			if (plan instanceof SimplePlan<C> simplePlan) {
+				// merge all simple plans into one
+				simplifiedTasks.addAll(simplePlan.parallelTasks());
+			} else if (plan instanceof NestedPlan<C> nestedPlan) {
+				// inline and re-visit nested plans
+				toVisit.addAll(nestedPlan.parallelPlans());
+			} else {
+				// /shrug
+				simplifiedPlans.add(plan);
+			}
+		}
 
 		if (simplifiedTasks.isEmpty() && simplifiedPlans.isEmpty()) {
 			// everything got simplified away
@@ -81,27 +100,4 @@ public record NestedPlan<C>(List<Plan<C>> parallelPlans) implements SimplyCompos
 		simplifiedPlans.add(SimplePlan.of(simplifiedTasks));
 		return new NestedPlan<>(simplifiedPlans);
 	}
-
-	private void flattenTasksAndPlans(List<Runnable> simplifiedTasks, List<Plan<C>> simplifiedPlans) {
-		var toVisit = new ArrayDeque<>(parallelPlans);
-		while (!toVisit.isEmpty()) {
-			var plan = toVisit.pop()
-					.maybeSimplify();
-
-			if (plan == UnitPlan.of()) {
-				continue;
-			}
-
-			if (plan instanceof SimplePlan simplePlan) {
-				// merge all simple plans into one
-				simplifiedTasks.addAll(simplePlan.parallelTasks());
-			} else if (plan instanceof NestedPlan<C> nestedPlan) {
-				// inline and re-visit nested plans
-				toVisit.addAll(nestedPlan.parallelPlans());
-			} else {
-				// /shrug
-				simplifiedPlans.add(plan);
-			}
-		}
-	}
 }

src/main/java/com/jozufozu/flywheel/lib/task/OnMainThreadPlan.java

@@ -3,15 +3,19 @@ package com.jozufozu.flywheel.lib.task;
 import com.jozufozu.flywheel.api.task.Plan;
 import com.jozufozu.flywheel.api.task.TaskExecutor;
 
-public record OnMainThreadPlan(Runnable task) implements ContextAgnosticPlan {
-	public static <C> Plan<C> of(Runnable task) {
-		return new OnMainThreadPlan(task).cast();
+public record OnMainThreadPlan<C>(ContextConsumer<C> task) implements SimplyComposedPlan<C> {
+	public static <C> Plan<C> of(ContextConsumer<C> task) {
+		return new OnMainThreadPlan<>(task);
+	}
+
+	public static <C> Plan<C> of(ContextRunnable<C> task) {
+		return new OnMainThreadPlan<>(task);
 	}
 
 	@Override
-	public void execute(TaskExecutor taskExecutor, Runnable onCompletion) {
+	public void execute(TaskExecutor taskExecutor, C context, Runnable onCompletion) {
 		taskExecutor.scheduleForMainThread(() -> {
-			task.run();
+			task.accept(context);
 			onCompletion.run();
 		});
 	}

src/main/java/com/jozufozu/flywheel/lib/task/PlanUtil.java

@@ -0,0 +1,49 @@
+package com.jozufozu.flywheel.lib.task;
+
+import java.util.List;
+import java.util.function.BiConsumer;
+
+import com.jozufozu.flywheel.api.task.TaskExecutor;
+import com.jozufozu.flywheel.lib.math.MoreMath;
+
+public class PlanUtil {
+	public static <C, T> void distribute(TaskExecutor taskExecutor, C context, Runnable onCompletion, List<T> list, BiConsumer<T, C> action) {
+		final int size = list.size();
+
+		if (size == 0) {
+			onCompletion.run();
+		} else if (size <= getChunkSize(taskExecutor, size)) {
+			processList(context, onCompletion, list, action);
+		} else {
+			dispatchChunks(taskExecutor, context, onCompletion, list, action);
+		}
+	}
+
+	public static int getChunkSize(TaskExecutor taskExecutor, int totalSize) {
+		return MoreMath.ceilingDiv(totalSize, taskExecutor.getThreadCount() * 32);
+	}
+
+	static <C, T> void dispatchChunks(TaskExecutor taskExecutor, C context, Runnable onCompletion, List<T> list, BiConsumer<T, C> action) {
+		final int size = list.size();
+		final int chunkSize = getChunkSize(taskExecutor, size);
+
+		var synchronizer = new Synchronizer(MoreMath.ceilingDiv(size, chunkSize), onCompletion);
+		int remaining = size;
+
+		while (remaining > 0) {
+			int end = remaining;
+			remaining -= chunkSize;
+			int start = Math.max(remaining, 0);
+
+			var subList = list.subList(start, end);
+			taskExecutor.execute(() -> processList(context, synchronizer, subList, action));
+		}
+	}
+
+	static <C, T> void processList(C context, Runnable onCompletion, List<T> list, BiConsumer<T, C> action) {
+		for (var t : list) {
+			action.accept(t, context);
+		}
+		onCompletion.run();
+	}
+}

src/main/java/com/jozufozu/flywheel/lib/task/RunOnAllPlan.java

@@ -1,60 +0,0 @@
-package com.jozufozu.flywheel.lib.task;
-
-import java.util.List;
-import java.util.function.BiConsumer;
-import java.util.function.Supplier;
-
-import com.jozufozu.flywheel.api.task.Plan;
-import com.jozufozu.flywheel.api.task.TaskExecutor;
-import com.jozufozu.flywheel.lib.math.MoreMath;
-
-public record RunOnAllPlan<T, C>(Supplier<List<T>> listSupplier,
-								 BiConsumer<T, C> action) implements SimplyComposedPlan<C> {
-	public static <T, C> Plan<C> of(Supplier<List<T>> iterable, BiConsumer<T, C> forEach) {
-		return new RunOnAllPlan<>(iterable, forEach);
-	}
-
-	@Override
-	public void execute(TaskExecutor taskExecutor, C context, Runnable onCompletion) {
-		taskExecutor.execute(() -> {
-			var list = listSupplier.get();
-			final int size = list.size();
-
-			if (size == 0) {
-				onCompletion.run();
-			} else if (size <= getChunkSize(taskExecutor, size)) {
-				processList(list, context, onCompletion);
-			} else {
-				dispatchChunks(list, taskExecutor, context, onCompletion);
-			}
-		});
-	}
-
-	private void dispatchChunks(List<T> suppliedList, TaskExecutor taskExecutor, C context, Runnable onCompletion) {
-		final int size = suppliedList.size();
-		final int chunkSize = getChunkSize(taskExecutor, size);
-
-		var synchronizer = new Synchronizer(MoreMath.ceilingDiv(size, chunkSize), onCompletion);
-		int remaining = size;
-
-		while (remaining > 0) {
-			int end = remaining;
-			remaining -= chunkSize;
-			int start = Math.max(remaining, 0);
-
-			var subList = suppliedList.subList(start, end);
-			taskExecutor.execute(() -> processList(subList, context, synchronizer));
-		}
-	}
-
-	private static int getChunkSize(TaskExecutor taskExecutor, int totalSize) {
-		return MoreMath.ceilingDiv(totalSize, taskExecutor.getThreadCount() * 32);
-	}
-
-	private void processList(List<T> suppliedList, C context, Runnable onCompletion) {
-		for (T t : suppliedList) {
-			action.accept(t, context);
-		}
-		onCompletion.run();
-	}
-}

src/main/java/com/jozufozu/flywheel/lib/task/SimplePlan.java

@@ -2,36 +2,50 @@ package com.jozufozu.flywheel.lib.task;
 
 import java.util.List;
 
+import com.google.common.collect.ImmutableList;
 import com.jozufozu.flywheel.api.task.Plan;
 import com.jozufozu.flywheel.api.task.TaskExecutor;
 
-public record SimplePlan(List<Runnable> parallelTasks) implements ContextAgnosticPlan {
-	public static <C> Plan<C> of(Runnable... tasks) {
-		return new SimplePlan(List.of(tasks)).cast();
+public record SimplePlan<C>(List<ContextConsumer<C>> parallelTasks) implements SimplyComposedPlan<C> {
+	@SafeVarargs
+	public static <C> SimplePlan<C> of(ContextRunnable<C>... tasks) {
+		return new SimplePlan<>(List.of(tasks));
 	}
 
-	public static <C> Plan<C> of(List<Runnable> tasks) {
-		return new SimplePlan(tasks).cast();
+	@SafeVarargs
+	public static <C> SimplePlan<C> of(ContextConsumer<C>... tasks) {
+		return new SimplePlan<>(List.of(tasks));
+	}
+
+	public static <C> SimplePlan<C> of(List<ContextConsumer<C>> tasks) {
+		return new SimplePlan<>(tasks);
 	}
 
 	@Override
-	public void execute(TaskExecutor taskExecutor, Runnable onCompletion) {
+	public void execute(TaskExecutor taskExecutor, C context, Runnable onCompletion) {
 		if (parallelTasks.isEmpty()) {
 			onCompletion.run();
 			return;
 		}
 
-		var synchronizer = new Synchronizer(parallelTasks.size(), onCompletion);
-		for (var task : parallelTasks) {
-			taskExecutor.execute(() -> {
-				task.run();
-				synchronizer.decrementAndEventuallyRun();
-			});
-		}
+		taskExecutor.execute(() -> {
+			PlanUtil.distribute(taskExecutor, context, onCompletion, parallelTasks, ContextConsumer::accept);
+		});
 	}
 
 	@Override
-	public Plan<Object> maybeSimplify() {
+	public Plan<C> and(Plan<C> plan) {
+		if (plan instanceof SimplePlan<C> simple) {
+			return of(ImmutableList.<ContextConsumer<C>>builder()
+					.addAll(parallelTasks)
+					.addAll(simple.parallelTasks)
+					.build());
+		}
+		return SimplyComposedPlan.super.and(plan);
+	}
+
+	@Override
+	public Plan<C> maybeSimplify() {
 		if (parallelTasks.isEmpty()) {
 			return UnitPlan.of();
 		}

src/main/java/com/jozufozu/flywheel/lib/task/ThreadGroupNotifier.java

@@ -4,7 +4,6 @@ package com.jozufozu.flywheel.lib.task;
  * Thin wrapper around Java's built-in object synchronization primitives.
  */
 public class ThreadGroupNotifier {
-
 	public synchronized void awaitNotification() {
 		try {
 			this.wait();

src/main/java/com/jozufozu/flywheel/vanilla/effect/ExampleEffect.java

@@ -1,32 +1,28 @@
 package com.jozufozu.flywheel.vanilla.effect;
 
 import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Collections;
 import java.util.List;
 
 import org.joml.Vector3f;
 
 import com.jozufozu.flywheel.api.event.ReloadRenderersEvent;
 import com.jozufozu.flywheel.api.event.RenderStage;
-import com.jozufozu.flywheel.api.visual.DynamicVisual;
+import com.jozufozu.flywheel.api.task.Plan;
 import com.jozufozu.flywheel.api.visual.Effect;
 import com.jozufozu.flywheel.api.visual.EffectVisual;
-import com.jozufozu.flywheel.api.visual.TickableVisual;
+import com.jozufozu.flywheel.api.visual.PlannedVisual;
 import com.jozufozu.flywheel.api.visual.VisualFrameContext;
 import com.jozufozu.flywheel.api.visual.VisualTickContext;
 import com.jozufozu.flywheel.api.visualization.VisualizationContext;
 import com.jozufozu.flywheel.impl.visualization.VisualizedRenderDispatcher;
-import com.jozufozu.flywheel.lib.box.ImmutableBox;
-import com.jozufozu.flywheel.lib.box.MutableBox;
 import com.jozufozu.flywheel.lib.instance.InstanceTypes;
 import com.jozufozu.flywheel.lib.instance.TransformedInstance;
 import com.jozufozu.flywheel.lib.model.Models;
+import com.jozufozu.flywheel.lib.task.ForEachPlan;
 import com.jozufozu.flywheel.lib.util.AnimationTickHolder;
-import com.jozufozu.flywheel.lib.visual.AbstractVisual;
 
 import net.minecraft.client.Minecraft;
-import net.minecraft.core.BlockPos;
+import net.minecraft.core.Vec3i;
 import net.minecraft.util.Mth;
 import net.minecraft.world.entity.player.Player;
 import net.minecraft.world.level.Level;
@@ -42,7 +38,7 @@ public class ExampleEffect implements Effect {
 	private static final float SPAWN_RADIUS = 8.0f;
 	private static final float LIMIT_RANGE = 10.0f;
 	private static final float SPEED_LIMIT = 0.1f;
-	private static final float RENDER_SCALE = 1 / 16f;
+	private static final float RENDER_SCALE = 2 / 16f;
 
 	private static final float SIGHT_RANGE = 5;
 
@@ -56,20 +52,10 @@ public class ExampleEffect implements Effect {
 
 	private final Level level;
 	private final Vector3f targetPoint;
-	private final BlockPos blockPos;
-	private final ImmutableBox volume;
-
-	private final List<BoidVisual> effects;
-
-	private final List<Boid> boids;
 
 	public ExampleEffect(Level level, Vector3f targetPoint) {
 		this.level = level;
 		this.targetPoint = targetPoint;
-		this.blockPos = new BlockPos(targetPoint.x, targetPoint.y, targetPoint.z);
-		this.volume = MutableBox.from(this.blockPos);
-		this.effects = new ArrayList<>(VISUAL_COUNT);
-		this.boids = new ArrayList<>(VISUAL_COUNT);
 	}
 
 	public static void tick(TickEvent.ClientTickEvent event) {
@@ -109,22 +95,62 @@ public class ExampleEffect implements Effect {
 	}
 
 	@Override
-	public Collection<EffectVisual<?>> createVisuals(VisualizationContext ctx) {
-		effects.clear();
-		boids.clear();
-		for (int i = 0; i < VISUAL_COUNT; i++) {
-			var x = targetPoint.x + level.random.nextFloat(-SPAWN_RADIUS, SPAWN_RADIUS);
-			var y = targetPoint.y + level.random.nextFloat(-SPAWN_RADIUS, SPAWN_RADIUS);
-			var z = targetPoint.z + level.random.nextFloat(-SPAWN_RADIUS, SPAWN_RADIUS);
-
-			Boid boid = new Boid(x, y, z);
-			boids.add(boid);
-			effects.add(new BoidVisual(ctx, level, boid));
-		}
-		return Collections.unmodifiableList(effects);
+	public EffectVisual<?> visualize(VisualizationContext ctx) {
+		return new ExampleVisual(ctx);
 	}
 
-	public class Boid {
+	public class ExampleVisual implements EffectVisual<ExampleEffect>, PlannedVisual {
+		private final List<BoidVisual> effects;
+		private final List<Boid> boids;
+
+		public ExampleVisual(VisualizationContext ctx) {
+			this.effects = new ArrayList<>(VISUAL_COUNT);
+			this.boids = new ArrayList<>(VISUAL_COUNT);
+
+			for (int i = 0; i < VISUAL_COUNT; i++) {
+				var x = targetPoint.x + level.random.nextFloat(-SPAWN_RADIUS, SPAWN_RADIUS);
+				var y = targetPoint.y + level.random.nextFloat(-SPAWN_RADIUS, SPAWN_RADIUS);
+				var z = targetPoint.z + level.random.nextFloat(-SPAWN_RADIUS, SPAWN_RADIUS);
+
+				Boid boid = new Boid(x, y, z);
+				boids.add(boid);
+				effects.add(new BoidVisual(ctx, boid));
+			}
+		}
+
+		@Override
+		public Plan<VisualTickContext> planTick() {
+			Plan<VisualTickContext> beginTick = ForEachPlan.of(() -> boids, Boid::beginTick);
+			return beginTick.then(ForEachPlan.of(() -> effects, boid -> boid.self.tick(boids)));
+		}
+
+		@Override
+		public Plan<VisualFrameContext> planFrame() {
+			return ForEachPlan.of(() -> effects, BoidVisual::beginFrame);
+		}
+
+		@Override
+		public void init() {
+
+		}
+
+		@Override
+		public void update() {
+
+		}
+
+		@Override
+		public boolean shouldReset() {
+			return false;
+		}
+
+		@Override
+		public void delete() {
+			effects.forEach(BoidVisual::_delete);
+		}
+	}
+
+	public static class Boid {
 		final Vector3f lastPosition;
 		final Vector3f position;
 		final Vector3f lastVelocity = new Vector3f(0);
@@ -145,13 +171,11 @@ public class ExampleEffect implements Effect {
 			lastPosition.set(position);
 		}
 
-		public void tick() {
-			beginTick();
-
+		public void tick(List<Boid> swarm) {
 			int seen = 0;
 			coherence.set(0);
 			alignment.set(0);
-			for (Boid boid : boids) {
+			for (Boid boid : swarm) {
 				if (boid == this) {
 					continue;
 				}
@@ -237,44 +261,29 @@ public class ExampleEffect implements Effect {
 		}
 	}
 
-	public class BoidVisual extends AbstractVisual implements EffectVisual<ExampleEffect>, DynamicVisual, TickableVisual {
+	public static class BoidVisual {
 		private final Boid self;
+		private final Vec3i renderOrigin;
 
-		private TransformedInstance instance;
+		private final TransformedInstance instance;
 
-		public BoidVisual(VisualizationContext ctx, Level level, Boid self) {
-			super(ctx, level);
+		public BoidVisual(VisualizationContext ctx, Boid self) {
+			renderOrigin = ctx.renderOrigin();
 			this.self = self;
-		}
 
-		@Override
-		public void init() {
-			instance = instancerProvider.instancer(InstanceTypes.TRANSFORMED, Models.block(Blocks.SHROOMLIGHT.defaultBlockState()), RenderStage.AFTER_PARTICLES)
+			instance = ctx.instancerProvider()
+					.instancer(InstanceTypes.TRANSFORMED, Models.block(Blocks.SHROOMLIGHT.defaultBlockState()), RenderStage.AFTER_PARTICLES)
 					.createInstance();
 
 			instance.setBlockLight(15)
 					.setSkyLight(15);
-
-			super.init();
 		}
 
-		@Override
-		protected void _delete() {
+		public void _delete() {
 			instance.delete();
 		}
 
-		@Override
-		public ImmutableBox getVolume() {
-			return volume;
-		}
-
-		@Override
-		public void tick(VisualTickContext c) {
-			self.tick();
-		}
-
-		@Override
-		public void beginFrame(VisualFrameContext context) {
+		public void beginFrame() {
 			float partialTicks = AnimationTickHolder.getPartialTicks();
 
 			var x = Mth.lerp(partialTicks, self.lastPosition.x, self.position.x);

src/test/java/com/jozufozu/flywheel/lib/task/PlanCompositionTest.java

@@ -8,9 +8,8 @@ import com.jozufozu.flywheel.util.Unit;
 
 public class PlanCompositionTest {
 
-	public static final Runnable NOOP = () -> {
-	};
-	public static final Plan<Unit> SIMPLE = SimplePlan.of(NOOP);
+	public static final Plan<Unit> SIMPLE = SimplePlan.of(() -> {
+	});
 
 	@Test
 	void nestedPlanAnd() {

src/test/java/com/jozufozu/flywheel/lib/task/PlanExecutionTest.java

@@ -52,7 +52,7 @@ class PlanExecutionTest {
 		var sequence = new IntArrayList(barriers + 1);
 		var expected = new IntArrayList(barriers + 1);
 
-		var plan = SimplePlan.<Unit>of(() -> sequence.add(1));
+		Plan<Unit> plan = SimplePlan.of(() -> sequence.add(1));
 		expected.add(1);
 
 		for (int i = 0; i < barriers; i++) {
@@ -71,18 +71,18 @@ class PlanExecutionTest {
 		var lock = new Object();
 		var sequence = new IntArrayList(8);
 
-		Runnable addOne = () -> {
+		ContextRunnable<Unit> addOne = () -> {
 			synchronized (lock) {
 				sequence.add(1);
 			}
 		};
-		Runnable addTwo = () -> {
+		ContextRunnable<Unit> addTwo = () -> {
 			synchronized (lock) {
 				sequence.add(2);
 			}
 		};
 
-		var plan = SimplePlan.<Unit>of(addOne, addOne, addOne, addOne)
+		var plan = SimplePlan.of(addOne, addOne, addOne, addOne)
 				.then(SimplePlan.of(addTwo, addTwo, addTwo, addTwo));
 
 		runAndWait(plan);
@@ -140,7 +140,7 @@ class PlanExecutionTest {
 	@Test
 	void mainThreadPlan() {
 		var done = new AtomicBoolean(false);
-		var plan = new OnMainThreadPlan(() -> done.set(true));
+		var plan = OnMainThreadPlan.of(() -> done.set(true));
 
 		plan.execute(EXECUTOR, Unit.INSTANCE);
 

src/test/java/com/jozufozu/flywheel/lib/task/PlanSimplificationTest.java

@@ -8,7 +8,7 @@ import com.jozufozu.flywheel.util.Unit;
 
 public class PlanSimplificationTest {
 
-	public static final Runnable NOOP = () -> {
+	public static final ContextRunnable<Unit> NOOP = () -> {
 	};
 	public static final Plan<Unit> SIMPLE = SimplePlan.of(NOOP);
 
@@ -36,7 +36,7 @@ public class PlanSimplificationTest {
 
 		Assertions.assertEquals(oneSimple.maybeSimplify(), SIMPLE);
 
-		var mainThreadNoop = new OnMainThreadPlan(NOOP);
+		var mainThreadNoop = new OnMainThreadPlan<>(NOOP);
 		var oneMainThread = NestedPlan.of(mainThreadNoop);
 
 		Assertions.assertEquals(oneMainThread.maybeSimplify(), mainThreadNoop);
@@ -66,7 +66,8 @@ public class PlanSimplificationTest {
 
 	@Test
 	void complexNesting() {
-		var mainThreadNoop = OnMainThreadPlan.<Unit>of(NOOP);
+		var mainThreadNoop = OnMainThreadPlan.<Unit>of(() -> {
+		});
 
 		var nested = NestedPlan.of(mainThreadNoop, SIMPLE);
 		Assertions.assertEquals(nested.maybeSimplify(), nested); // cannot simplify
@@ -78,7 +79,8 @@ public class PlanSimplificationTest {
 
 	@Test
 	void nestedNoSimple() {
-		var mainThreadNoop = OnMainThreadPlan.<Unit>of(NOOP);
+		var mainThreadNoop = OnMainThreadPlan.<Unit>of(() -> {
+		});
 		var barrier = new BarrierPlan<>(SIMPLE, SIMPLE);
 		var oneMainThread = NestedPlan.of(mainThreadNoop, NestedPlan.of(mainThreadNoop, barrier, barrier));