Lighter tracking

- Deduplicate section tracking logic between Lit and SmoothLit
- Now there is one SectionPropertyImpl which the 2 storages add
  listeners to

common/src/api/java/dev/engine_room/flywheel/api/visual/LitVisual.java

@@ -1,40 +1,13 @@
 package dev.engine_room.flywheel.api.visual;
 
-import java.util.function.LongConsumer;
-
-import org.jetbrains.annotations.ApiStatus;
-
-import net.minecraft.core.SectionPos;
-
 /**
  * A visual that listens to light updates.
  *
  * <p>If your visual moves around in the level at all, you should use {@link TickableVisual} or {@link DynamicVisual},
  * and poll for light yourself along with listening for updates. When your visual moves to a different section, call
- * {@link Notifier#notifySectionsChanged}.</p>
+ * {@link SectionProperty#lightSections}.</p>
  */
-public interface LitVisual extends Visual {
-	/**
-	 * Set the notifier object.
-	 *
-	 * <p>This method is only called once right after the visual
-	 * is created and before {@link #collectLightSections}.</p>
-	 *
-	 * @param notifier The notifier.
-	 */
-	void setLightSectionNotifier(Notifier notifier);
-
-	/**
-	 * Collect the sections that this visual is contained in.
-	 *
-	 * <p>This method is called upon visual creation, and the frame after
-	 * {@link Notifier#notifySectionsChanged} is called.</p>
-	 *
-	 * @param consumer The consumer to provide the sections to.
-	 * @see SectionPos#asLong
-	 */
-	void collectLightSections(LongConsumer consumer);
-
+public non-sealed interface LitVisual extends SectionTrackedVisual {
 	/**
 	 * Called when a section this visual is contained in receives a light update.
 	 *
@@ -48,18 +21,4 @@ public interface LitVisual extends Visual {
 	 * <p>This method not is invoked automatically after visual creation.</p>
 	 */
 	void updateLight(float partialTick);
-
-	/**
-	 * A notifier object that can be used to indicate to the impl
-	 * that the sections a visual is contained in have changed.
-	 */
-	@ApiStatus.NonExtendable
-	interface Notifier {
-		/**
-		 * Invoke this to indicate to the impl that your visual has moved to a different set of sections.
-		 * <br>
-		 * The next frame, the impl will call {@link LitVisual#collectLightSections} again.
-		 */
-		void notifySectionsChanged();
-	}
 }

common/src/api/java/dev/engine_room/flywheel/api/visual/SectionTrackedVisual.java

@@ -0,0 +1,26 @@
+package dev.engine_room.flywheel.api.visual;
+
+import org.jetbrains.annotations.ApiStatus;
+
+import it.unimi.dsi.fastutil.longs.LongSet;
+
+public sealed interface SectionTrackedVisual extends Visual permits SmoothLitVisual, LitVisual {
+	/**
+	 * Set the section property object.
+	 *
+	 * <p>This method is only called once, upon visual creation.
+	 * <br>If the property is assigned to in this method, the
+	 * visual will immediately be tracked in the given sections.
+	 *
+	 * @param property The property.
+	 */
+	void setSectionProperty(SectionProperty property);
+
+	@ApiStatus.NonExtendable
+	interface SectionProperty {
+		/**
+		 * Assign the set of sections this visual wants to track itself in.
+		 */
+		void lightSections(LongSet sections);
+	}
+}

common/src/api/java/dev/engine_room/flywheel/api/visual/SmoothLitVisual.java

@@ -1,11 +1,7 @@
 package dev.engine_room.flywheel.api.visual;
 
-import org.jetbrains.annotations.ApiStatus;
-
-import it.unimi.dsi.fastutil.longs.LongSet;
-
 /**
- * An interface allowing visuals to request light data on the GPU for a set of sections.
+ * A marker interface allowing visuals to request light data on the GPU for a set of sections.
  *
  * <p> Sections passed into {@link SectionProperty#lightSections} will have their light data handed to the
  * backend and queryable by {@code flw_light*} functions in shaders.
@@ -13,21 +9,6 @@ import it.unimi.dsi.fastutil.longs.LongSet;
  * Note that the queryable light data is shared across all visuals, so even if one specific visual does not
  * request a given section, the data will be available if another visual does.
  */
-public interface SmoothLitVisual extends Visual {
-	/**
-	 * Set the section property object.
-	 *
-	 * <p>This method is only called once, upon visual creation.
-	 *
-	 * @param property The property.
-	 */
-	void setSectionProperty(SectionProperty property);
+public non-sealed interface SmoothLitVisual extends SectionTrackedVisual {
 
-	@ApiStatus.NonExtendable
-	interface SectionProperty {
-		/**
-		 * Assign the set of sections this visual wants to have light data for.
-		 */
-		void lightSections(LongSet sections);
-	}
 }

common/src/lib/java/dev/engine_room/flywheel/lib/visual/AbstractBlockEntityVisual.java

@@ -1,7 +1,5 @@
 package dev.engine_room.flywheel.lib.visual;
 
-import java.util.function.LongConsumer;
-
 import org.jetbrains.annotations.Nullable;
 import org.joml.FrustumIntersection;
 
@@ -13,6 +11,7 @@ import dev.engine_room.flywheel.api.visualization.VisualManager;
 import dev.engine_room.flywheel.api.visualization.VisualizationContext;
 import dev.engine_room.flywheel.lib.instance.FlatLit;
 import dev.engine_room.flywheel.lib.math.MoreMath;
+import it.unimi.dsi.fastutil.longs.LongSet;
 import net.minecraft.client.renderer.LevelRenderer;
 import net.minecraft.core.BlockPos;
 import net.minecraft.core.SectionPos;
@@ -40,7 +39,7 @@ public abstract class AbstractBlockEntityVisual<T extends BlockEntity> extends A
 	protected final BlockPos visualPos;
 	protected final BlockState blockState;
 	@Nullable
-	protected LitVisual.Notifier notifier;
+	protected SectionProperty lightSections;
 
 	public AbstractBlockEntityVisual(VisualizationContext ctx, T blockEntity, float partialTick) {
 		super(ctx, blockEntity.getLevel(), partialTick);
@@ -51,13 +50,9 @@ public abstract class AbstractBlockEntityVisual<T extends BlockEntity> extends A
 	}
 
 	@Override
-	public void setLightSectionNotifier(Notifier notifier) {
-		this.notifier = notifier;
-	}
-
-	@Override
-	public void collectLightSections(LongConsumer consumer) {
-		consumer.accept(SectionPos.asLong(pos));
+	public void setSectionProperty(SectionProperty property) {
+		this.lightSections = property;
+		lightSections.lightSections(LongSet.of(SectionPos.asLong(pos)));
 	}
 
 	/**

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

@@ -31,7 +31,7 @@ public class LitVisualStorage {
 	private final Map<LitVisual, LongSet> visuals2Sections = new WeakHashMap<>();
 	private final Long2ObjectMap<List<Updater>> sections2Visuals = new Long2ObjectOpenHashMap<>();
 
-	private final Queue<LitVisual> movedVisuals = new ConcurrentLinkedQueue<>();
+	private final Queue<MovedVisual> movedVisuals = new ConcurrentLinkedQueue<>();
 	private final LongSet sectionsUpdatedThisFrame = new LongOpenHashSet();
 
 	private long updateId = INITIAL_UPDATE_ID;
@@ -65,11 +65,11 @@ public class LitVisualStorage {
 	}
 
 	private void processMoved() {
-		LitVisual visual;
-		while ((visual = movedVisuals.poll()) != null) {
+		MovedVisual moved;
+		while ((moved = movedVisuals.poll()) != null) {
 			// If the visual isn't there when we try to remove it that means it was deleted before we got to it.
-			if (remove(visual)) {
-				add(visual);
+			if (remove(moved.visual)) {
+				updateTracking(moved.tracker, moved.visual);
 			}
 		}
 	}
@@ -90,24 +90,27 @@ public class LitVisualStorage {
 		return visuals2Sections.isEmpty();
 	}
 
-	public void setNotifierAndAdd(LitVisual visual) {
-		visual.setLightSectionNotifier(new LitVisualNotifierImpl(visual));
-		add(visual);
+	public void add(SectionPropertyImpl tracker, LitVisual visual) {
+		var moved = new MovedVisual(tracker, visual);
+		tracker.addListener(() -> movedVisuals.add(moved));
+
+		updateTracking(tracker, visual);
 	}
 
-	private void add(LitVisual visual) {
-		LongSet sections = new LongArraySet();
+	public void updateTracking(SectionPropertyImpl tracker, LitVisual visual) {
+		if (tracker.sections.isEmpty()) {
+			// Add the visual to the map even if sections is empty, this way we can distinguish from deleted visuals
+			visuals2Sections.put(visual, LongSet.of());
 
-		visual.collectLightSections(sections::add);
-
-		// Add the visual to the map even if sections is empty, this way we can distinguish from deleted visuals
-		visuals2Sections.put(visual, sections);
-
-		// Don't bother creating an updater if the visual isn't in any sections.
-		if (sections.isEmpty()) {
+			// Don't bother creating an updater if the visual isn't in any sections.
 			return;
 		}
 
+		// Create a copy of the array, so we know what section to remove the visual from later.
+		var sections = new LongArraySet(tracker.sections);
+
+		visuals2Sections.put(visual, sections);
+
 		var updater = createUpdater(visual, sections.size());
 
 		for (long section : sections) {
@@ -198,16 +201,6 @@ public class LitVisualStorage {
 		}
 	}
 
-	private final class LitVisualNotifierImpl implements LitVisual.Notifier {
-		private final LitVisual litVisual;
-
-		private LitVisualNotifierImpl(LitVisual litVisual) {
-			this.litVisual = litVisual;
-		}
-
-		@Override
-		public void notifySectionsChanged() {
-			movedVisuals.add(litVisual);
-        }
-    }
+	private record MovedVisual(SectionPropertyImpl tracker, LitVisual visual) {
+	}
 }

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

@@ -0,0 +1,26 @@
+package dev.engine_room.flywheel.impl.visualization.storage;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import dev.engine_room.flywheel.api.visual.SmoothLitVisual;
+import it.unimi.dsi.fastutil.longs.LongArraySet;
+import it.unimi.dsi.fastutil.longs.LongSet;
+
+public class SectionPropertyImpl implements SmoothLitVisual.SectionProperty {
+	public final LongSet sections = new LongArraySet();
+
+	private final List<Runnable> listeners = new ArrayList<>(2);
+
+	@Override
+	public void lightSections(LongSet sections) {
+		this.sections.clear();
+		this.sections.addAll(sections);
+
+		listeners.forEach(Runnable::run);
+	}
+
+	public void addListener(Runnable listener) {
+		listeners.add(listener);
+	}
+}

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

@@ -5,13 +5,12 @@ import java.util.Map;
 import org.jetbrains.annotations.Nullable;
 
 import dev.engine_room.flywheel.api.visual.SmoothLitVisual;
-import it.unimi.dsi.fastutil.longs.LongArraySet;
 import it.unimi.dsi.fastutil.longs.LongOpenHashSet;
 import it.unimi.dsi.fastutil.longs.LongSet;
 import it.unimi.dsi.fastutil.objects.Reference2ObjectOpenHashMap;
 
 public class SmoothLitVisualStorage {
-	private final Map<SmoothLitVisual, SectionProperty> visuals = new Reference2ObjectOpenHashMap<>();
+	private final Map<SmoothLitVisual, SectionPropertyImpl> visuals = new Reference2ObjectOpenHashMap<>();
 
 	@Nullable
 	private LongSet cachedSections;
@@ -20,9 +19,13 @@ public class SmoothLitVisualStorage {
 		return cachedSections == null;
 	}
 
+	public void markDirty() {
+		cachedSections = null;
+	}
+
 	public LongSet sections() {
 		cachedSections = new LongOpenHashSet();
-		for (SectionProperty value : visuals.values()) {
+		for (var value : visuals.values()) {
 			cachedSections.addAll(value.sections);
 		}
 		return cachedSections;
@@ -32,25 +35,18 @@ public class SmoothLitVisualStorage {
 		visuals.remove(smoothLit);
 	}
 
-	public void add(SmoothLitVisual smoothLit) {
-		var sections = new SectionProperty();
-		visuals.put(smoothLit, sections);
-		smoothLit.setSectionProperty(sections);
+	public void add(SectionPropertyImpl tracker, SmoothLitVisual smoothLit) {
+		visuals.put(smoothLit, tracker);
+
+		tracker.addListener(this::markDirty);
+
+		if (!tracker.sections.isEmpty()) {
+			markDirty();
+		}
 	}
 
 	public void clear() {
 		visuals.clear();
 	}
 
-	private final class SectionProperty implements SmoothLitVisual.SectionProperty {
-		private final LongSet sections = new LongArraySet();
-
-		@Override
-		public void lightSections(LongSet sections) {
-			this.sections.clear();
-			this.sections.addAll(sections);
-
-			SmoothLitVisualStorage.this.cachedSections = null;
-		}
-	}
 }

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

@@ -11,6 +11,7 @@ import org.jetbrains.annotations.Nullable;
 import dev.engine_room.flywheel.api.task.Plan;
 import dev.engine_room.flywheel.api.visual.DynamicVisual;
 import dev.engine_room.flywheel.api.visual.LitVisual;
+import dev.engine_room.flywheel.api.visual.SectionTrackedVisual;
 import dev.engine_room.flywheel.api.visual.SmoothLitVisual;
 import dev.engine_room.flywheel.api.visual.TickableVisual;
 import dev.engine_room.flywheel.api.visual.Visual;
@@ -159,12 +160,19 @@ public abstract class Storage<T> {
 			}
 		}
 
-		if (visual instanceof LitVisual lit) {
-			litVisuals.setNotifierAndAdd(lit);
-		}
+		if (visual instanceof SectionTrackedVisual tracked) {
+			SectionPropertyImpl sectionProperty = new SectionPropertyImpl();
 
-		if (visual instanceof SmoothLitVisual smoothLit) {
-			smoothLitVisuals.add(smoothLit);
+			// Give the visual a chance to fill in the property.
+			tracked.setSectionProperty(sectionProperty);
+
+			if (visual instanceof LitVisual lit) {
+				litVisuals.add(sectionProperty, lit);
+			}
+
+			if (visual instanceof SmoothLitVisual smoothLit) {
+				smoothLitVisuals.add(sectionProperty, smoothLit);
+			}
 		}
 	}