Construction plans

- Not pretty code:tm:
- Deduplicate items in the storage queue serially
- Add and update visuals in parallel
- Throw everything into a single custom Plan impl because the dataflow
  is pretty complicated

common/src/lib/java/dev/engine_room/flywheel/lib/task/Distribute.java

@@ -69,6 +69,57 @@ public static <C, T> void tasks(TaskExecutor taskExecutor, C context, Runnable o
 		}
 	}
 
+	public interface IndexedConsumer<T, C> {
+		void accept(int i, T t, C context);
+	}
+
+	public static <C, T> void indexed(TaskExecutor taskExecutor, C context, Runnable onCompletion, List<T> list, IndexedConsumer<T, C> action) {
+		final int size = list.size();
+
+		if (size == 0) {
+			onCompletion.run();
+			return;
+		}
+
+		final int sliceSize = sliceSize(taskExecutor, size);
+
+		if (size <= sliceSize) {
+			for (int i = 0; i < list.size(); i++) {
+				T t = list.get(i);
+				action.accept(i, t, context);
+			}
+			onCompletion.run();
+		} else if (sliceSize == 1) {
+			var synchronizer = new Synchronizer(size, onCompletion);
+			for (int i = 0; i < list.size(); i++) {
+				T t = list.get(i);
+				final int index = i;
+				taskExecutor.execute(() -> {
+					action.accept(index, t, context);
+					synchronizer.decrementAndEventuallyRun();
+				});
+			}
+		} else {
+			var synchronizer = new Synchronizer(MoreMath.ceilingDiv(size, sliceSize), onCompletion);
+			int remaining = size;
+
+			while (remaining > 0) {
+				int end = remaining;
+				remaining -= sliceSize;
+				final int start = Math.max(remaining, 0);
+
+				var subList = list.subList(start, end);
+				taskExecutor.execute(() -> {
+					for (int i = 0; i < subList.size(); i++) {
+						T t = subList.get(i);
+						action.accept(start + i, t, context);
+					}
+					synchronizer.decrementAndEventuallyRun();
+				});
+			}
+		}
+	}
+
 	/**
 	 * Distribute the given list of tasks in chunks across the threads of the task executor.
 	 *

common/src/main/java/dev/engine_room/flywheel/impl/visualization/VisualManagerImpl.java

@@ -1,19 +1,33 @@
 package dev.engine_room.flywheel.impl.visualization;
 
-import java.util.Queue;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.ConcurrentLinkedQueue;
 
 import dev.engine_room.flywheel.api.task.Plan;
+import dev.engine_room.flywheel.api.task.TaskExecutor;
 import dev.engine_room.flywheel.api.visual.DynamicVisual;
+import dev.engine_room.flywheel.api.visual.LightUpdatedVisual;
+import dev.engine_room.flywheel.api.visual.SectionTrackedVisual;
+import dev.engine_room.flywheel.api.visual.ShaderLightVisual;
 import dev.engine_room.flywheel.api.visual.TickableVisual;
+import dev.engine_room.flywheel.api.visual.Visual;
 import dev.engine_room.flywheel.api.visualization.VisualManager;
+import dev.engine_room.flywheel.impl.visualization.storage.Action;
+import dev.engine_room.flywheel.impl.visualization.storage.LightUpdatedVisualStorage;
+import dev.engine_room.flywheel.impl.visualization.storage.SectionTracker;
 import dev.engine_room.flywheel.impl.visualization.storage.Storage;
 import dev.engine_room.flywheel.impl.visualization.storage.Transaction;
-import dev.engine_room.flywheel.lib.task.SimplePlan;
+import dev.engine_room.flywheel.lib.task.Distribute;
+import dev.engine_room.flywheel.lib.task.SimplyComposedPlan;
+import dev.engine_room.flywheel.lib.task.Synchronizer;
 import it.unimi.dsi.fastutil.longs.LongSet;
+import it.unimi.dsi.fastutil.objects.Reference2IntMap;
+import it.unimi.dsi.fastutil.objects.Reference2IntOpenHashMap;
+import net.minecraft.util.Unit;
 
 public class VisualManagerImpl<T, S extends Storage<T>> implements VisualManager<T> {
-	private final Queue<Transaction<T>> queue = new ConcurrentLinkedQueue<>();
+	private final ConcurrentLinkedQueue<Transaction<T>> queue = new ConcurrentLinkedQueue<>();
 
 	private final S storage;
 
@@ -42,6 +56,10 @@ public void queueAdd(T obj) {
 
 	@Override
 	public void queueRemove(T obj) {
+		if (!getStorage().willAccept(obj)) {
+			return;
+		}
+
 		queue.add(Transaction.remove(obj));
 	}
 
@@ -54,21 +72,13 @@ public void queueUpdate(T obj) {
 		queue.add(Transaction.update(obj));
 	}
 
-	public void processQueue(float partialTick) {
-		var storage = getStorage();
-		Transaction<T> transaction;
-		while ((transaction = queue.poll()) != null) {
-			transaction.apply(storage, partialTick);
-		}
-	}
-
 	public Plan<DynamicVisual.Context> framePlan() {
-		return SimplePlan.<DynamicVisual.Context>of(context -> processQueue(context.partialTick()))
+		return new ProcessQueuePlan<>(DynamicVisual.Context::partialTick)
 				.then(storage.framePlan());
 	}
 
 	public Plan<TickableVisual.Context> tickPlan() {
-		return SimplePlan.<TickableVisual.Context>of(context -> processQueue(1))
+		return new ProcessQueuePlan<TickableVisual.Context>($ -> 1)
 				.then(storage.tickPlan());
 	}
 
@@ -90,4 +100,120 @@ public LongSet gpuLightSections() {
 	public void invalidate() {
 		getStorage().invalidate();
 	}
+
+	interface PartialTick<C> {
+		float partialTick(C context);
+	}
+
+	private class ProcessQueuePlan<C> implements SimplyComposedPlan<C> {
+		private final PartialTick<C> partialTick;
+
+		private ProcessQueuePlan(PartialTick<C> partialTick) {
+			this.partialTick = partialTick;
+		}
+
+		@Override
+		public void execute(TaskExecutor taskExecutor, C context, Runnable onCompletion) {
+			var added = new ArrayList<T>();
+			var removed = new ArrayList<T>();
+			var updated = new ArrayList<T>();
+
+			sortQueue(added, removed, updated);
+
+			var size = added.size();
+			var visualDst = new Visual[size];
+			var trackerDst = new SectionTracker[size];
+
+			var partialTick = this.partialTick.partialTick(context);
+
+			var applyPatch = new Synchronizer(2, () -> {
+				var storage = getStorage();
+				for (int i = 0; i < size; i++) {
+					var visual = visualDst[i];
+
+					if (visual == null) {
+						continue;
+					}
+
+					var obj = added.get(i);
+					var tracker = trackerDst[i];
+
+					storage.add(obj, visual, tracker);
+				}
+
+				Distribute.tasks(taskExecutor, Unit.INSTANCE, onCompletion, updated, (obj, ignored) -> {
+					storage.update(obj, 1);
+				});
+			});
+
+			taskExecutor.execute(() -> {
+				for (var t : removed) {
+					storage.remove(t);
+				}
+
+				applyPatch.decrementAndEventuallyRun();
+			});
+
+			Distribute.indexed(taskExecutor, Unit.INSTANCE, applyPatch, added, (i, obj, ignored) -> {
+				asyncCreate(i, obj, partialTick, trackerDst, visualDst);
+			});
+		}
+	}
+
+	private void sortQueue(List<T> added, List<T> removed, List<T> updated) {
+		Reference2IntMap<T> dedupe = new Reference2IntOpenHashMap<>();
+
+		Transaction<T> transaction;
+		while ((transaction = queue.poll()) != null) {
+			var i = dedupe.getInt(transaction.obj());
+			dedupe.put(transaction.obj(), transaction.action().setBit(i));
+		}
+
+		for (var entry : dedupe.reference2IntEntrySet()) {
+			var obj = entry.getKey();
+			var action = entry.getIntValue();
+
+			var add = Action.ADD.in(action);
+			var remove = Action.REMOVE.in(action);
+			var update = Action.UPDATE.in(action);
+
+			// Add and remove can happen at the same time
+			if (add) {
+				added.add(obj);
+			}
+			if (remove) {
+				removed.add(obj);
+			}
+
+			// But only update if the object will actually be around for it
+			if (update && add == remove) {
+				updated.add(obj);
+			}
+		}
+	}
+
+	private void asyncCreate(int i, T obj, float partialTick, SectionTracker[] trackerDst, Visual[] visualDst) {
+		var visual = storage.createRaw(obj, partialTick);
+
+		if (visual instanceof SectionTrackedVisual tracked) {
+			var tracker = new SectionTracker();
+
+			// Give the visual a chance to invoke the collector.
+			tracked.setSectionCollector(tracker);
+
+			if (visual instanceof LightUpdatedVisual lightUpdated) {
+				lightUpdated.updateLight(partialTick);
+
+				tracker.addListener(new LightUpdatedVisualStorage.MovedVisual(lightUpdated, tracker, storage.lightUpdatedVisuals()));
+			}
+
+			if (visual instanceof ShaderLightVisual) {
+				tracker.addListener(storage.shaderLightVisuals().sectionListener);
+			}
+
+			trackerDst[i] = tracker;
+		}
+
+		visualDst[i] = visual;
+	}
 }

common/src/main/java/dev/engine_room/flywheel/impl/visualization/storage/Action.java

@@ -1,7 +1,22 @@
 package dev.engine_room.flywheel.impl.visualization.storage;
 
 public enum Action {
-	ADD,
-	REMOVE,
-	UPDATE,
+	ADD(1),
+	REMOVE(2),
+	UPDATE(4),
+	;
+
+	private final int bit;
+
+	Action(int bit) {
+		this.bit = bit;
+	}
+
+	public int setBit(int i) {
+		return i | bit;
+	}
+
+	public boolean in(int i) {
+		return (i & bit) != 0;
+	}
 }

common/src/main/java/dev/engine_room/flywheel/impl/visualization/storage/BlockEntityStorage.java

@@ -3,6 +3,7 @@
 import org.jetbrains.annotations.Nullable;
 
 import dev.engine_room.flywheel.api.visual.BlockEntityVisual;
+import dev.engine_room.flywheel.api.visual.Visual;
 import dev.engine_room.flywheel.api.visualization.VisualizationContext;
 import dev.engine_room.flywheel.lib.visualization.VisualizationHelper;
 import it.unimi.dsi.fastutil.longs.Long2ObjectMap;
@@ -50,18 +51,21 @@ public boolean willAccept(BlockEntity blockEntity) {
 
 	@Override
 	@Nullable
-	protected BlockEntityVisual<?> createRaw(BlockEntity obj, float partialTick) {
+	public BlockEntityVisual<?> createRaw(BlockEntity obj, float partialTick) {
 		var visualizer = VisualizationHelper.getVisualizer(obj);
 		if (visualizer == null) {
 			return null;
 		}
 
-		var visual = visualizer.createVisual(visualizationContext, obj, partialTick);
+		return visualizer.createVisual(visualizationContext, obj, partialTick);
+	}
 
-		BlockPos blockPos = obj.getBlockPos();
-		posLookup.put(blockPos.asLong(), visual);
+	@Override
+	public void add(BlockEntity obj, Visual visual, @Nullable SectionTracker tracker) {
+		super.add(obj, visual, tracker);
 
-		return visual;
+		BlockPos blockPos = obj.getBlockPos();
+		posLookup.put(blockPos.asLong(), (BlockEntityVisual<?>) visual);
 	}
 
 	@Override

common/src/main/java/dev/engine_room/flywheel/impl/visualization/storage/EffectStorage.java

@@ -10,7 +10,7 @@ public EffectStorage(VisualizationContext visualizationContext) {
 	}
 
 	@Override
-	protected EffectVisual<?> createRaw(Effect obj, float partialTick) {
+	public EffectVisual<?> createRaw(Effect obj, float partialTick) {
 		return obj.visualize(visualizationContext, partialTick);
 	}
 

common/src/main/java/dev/engine_room/flywheel/impl/visualization/storage/EntityStorage.java

@@ -12,7 +12,7 @@ public EntityStorage(VisualizationContext visualizationContext) {
 	}
 
 	@Override
-	protected EntityVisual<?> createRaw(Entity obj, float partialTick) {
+	public EntityVisual<?> createRaw(Entity obj, float partialTick) {
 		var visualizer = VisualizationHelper.getVisualizer(obj);
 		if (visualizer == null) {
 			return null;

common/src/main/java/dev/engine_room/flywheel/impl/visualization/storage/LightUpdatedVisualStorage.java

@@ -86,9 +86,6 @@ private long getNextUpdateId() {
 	}
 
 	public void add(LightUpdatedVisual visual, SectionTracker tracker) {
-		var moved = new MovedVisual(visual, tracker);
-		tracker.addListener(() -> movedVisuals.add(moved));
-
 		addInner(visual, tracker);
 	}
 
@@ -194,6 +191,20 @@ record Context(long updateId, float partialTick) {
 		}
 	}
 
-	private record MovedVisual(LightUpdatedVisual visual, SectionTracker tracker) {
+	public static final class MovedVisual implements Runnable {
+		private final LightUpdatedVisual visual;
+		private final SectionTracker tracker;
+		private final LightUpdatedVisualStorage parent;
+
+		public MovedVisual(LightUpdatedVisual visual, SectionTracker tracker, LightUpdatedVisualStorage parent) {
+			this.visual = visual;
+			this.tracker = tracker;
+			this.parent = parent;
+		}
+
+		@Override
+		public void run() {
+			parent.movedVisuals.add(this);
+		}
 	}
 }

common/src/main/java/dev/engine_room/flywheel/impl/visualization/storage/ShaderLightVisualStorage.java

@@ -13,6 +13,8 @@ public class ShaderLightVisualStorage {
 	private final LongSet sections = new LongOpenHashSet();
 	private boolean isDirty;
 
+	public final Runnable sectionListener = this::markDirty;
+
 	public LongSet sections() {
 		if (isDirty) {
 			sections.clear();
@@ -35,7 +37,7 @@ public void markDirty() {
 	public void add(ShaderLightVisual visual, SectionTracker tracker) {
 		trackers.put(visual, tracker);
 
-		tracker.addListener(this::markDirty);
+		tracker.addListener(sectionListener);
 
 		if (!tracker.sections().isEmpty()) {
 			markDirty();