refactor: Dynamic allocation of samplers to texture units

This is the first step of a critical refactor to how Iris manages samplers. Previously, samplers were allocated to a fixed texture unit, placing a hard cap on the number of total samplers that could be made available to shader programs.

Now, with the exception of a few reserved samplers managed externally (the primary texture used for textured content, the lightmap, and the overlay texture), Iris dynamically allocates samplers to texture units.

This means that a sampler only takes up a texture unit if it is used in the current program, opening the door to having a potentially unbounded number of samplers available, as long as each program only samples from a small selection of them.

Why now? This is how vanilla 1.17 manages samplers, so it makes things easier for Iris + Sodium on 1.17. It also allows us to finally move around the vanilla texture units back to how they are in vanilla. Finally, it allows the future complete unlocking of the limit of 8 color textures globally.

Overall, this refactor has been necessary for a long time, and now it has become critical. So I've decided to do it now, instead of piling on more kinda hacky code.

src/main/java/net/coderbot/iris/gl/program/GlUniform1iCall.java

@@ -0,0 +1,19 @@
+package net.coderbot.iris.gl.program;
+
+public class GlUniform1iCall {
+	private final int location;
+	private final int value;
+
+	public GlUniform1iCall(int location, int value) {
+		this.location = location;
+		this.value = value;
+	}
+
+	public int getLocation() {
+		return location;
+	}
+
+	public int getValue() {
+		return value;
+	}
+}

src/main/java/net/coderbot/iris/gl/program/Program.java

@@ -6,17 +6,20 @@
 
 public final class Program extends GlResource {
 	private final ProgramUniforms uniforms;
+	private final ProgramSamplers samplers;
 
-	Program(int program, ProgramUniforms uniforms) {
+	Program(int program, ProgramUniforms uniforms, ProgramSamplers samplers) {
 		super(program);
 
 		this.uniforms = uniforms;
+		this.samplers = samplers;
 	}
 
 	public void use() {
 		GL20C.glUseProgram(getGlId());
 
 		uniforms.update();
+		samplers.update();
 	}
 
 	public void destroyInternal() {

src/main/java/net/coderbot/iris/gl/program/ProgramBuilder.java

@@ -1,6 +1,8 @@
 package net.coderbot.iris.gl.program;
 
+import com.google.common.collect.ImmutableSet;
 import com.mojang.blaze3d.systems.RenderSystem;
+import net.coderbot.iris.gl.sampler.SamplerHolder;
 import net.coderbot.iris.gl.shader.GlShader;
 import net.coderbot.iris.gl.shader.ProgramCreator;
 import net.coderbot.iris.gl.shader.ShaderConstants;
@@ -10,7 +12,9 @@
 import org.lwjgl.opengl.GL20C;
 import org.lwjgl.opengl.GL21C;
 
-public class ProgramBuilder extends ProgramUniforms.Builder {
+import java.util.function.IntSupplier;
+
+public class ProgramBuilder extends ProgramUniforms.Builder implements SamplerHolder {
 	private static final ShaderConstants EMPTY_CONSTANTS = ShaderConstants.builder().build();
 
 	public static final ShaderConstants MACRO_CONSTANTS = ShaderConstants.builder()
@@ -26,18 +30,21 @@ public class ProgramBuilder extends ProgramUniforms.Builder {
 
 
 	private final int program;
+	private ProgramSamplers.Builder samplers;
 
-	private ProgramBuilder(String name, int program) {
+	private ProgramBuilder(String name, int program, ImmutableSet<Integer> reservedTextureUnits) {
 		super(name, program);
 
 		this.program = program;
+		this.samplers = ProgramSamplers.builder(program, reservedTextureUnits);
 	}
 
 	public void bindAttributeLocation(int index, String name) {
 		GL21C.glBindAttribLocation(program, index, name);
 	}
 
-	public static ProgramBuilder begin(String name, @Nullable String vertexSource, @Nullable String geometrySource, @Nullable String fragmentSource) {
+	public static ProgramBuilder begin(String name, @Nullable String vertexSource, @Nullable String geometrySource,
+									   @Nullable String fragmentSource, ImmutableSet<Integer> reservedTextureUnits) {
 		RenderSystem.assertThread(RenderSystem::isOnRenderThread);
 
 		GlShader vertex;
@@ -70,11 +77,11 @@ public static ProgramBuilder begin(String name, @Nullable String vertexSource, @
 
 		fragment.destroy();
 
-		return new ProgramBuilder(name, programId);
+		return new ProgramBuilder(name, programId, reservedTextureUnits);
 	}
 
 	public Program build() {
-		return new Program(program, super.buildUniforms());
+		return new Program(program, super.buildUniforms(), this.samplers.build());
 	}
 
 	private static GlShader buildShader(ShaderType shaderType, String name, @Nullable String source) {
@@ -84,4 +91,19 @@ private static GlShader buildShader(ShaderType shaderType, String name, @Nullabl
 			throw new RuntimeException("Failed to compile " + shaderType + " shader for program " + name, e);
 		}
 	}
+
+	@Override
+	public void addExternalSampler(int textureUnit, String... names) {
+		samplers.addExternalSampler(textureUnit, names);
+	}
+
+	@Override
+	public boolean hasSampler(String name) {
+		return samplers.hasSampler(name);
+	}
+
+	@Override
+	public boolean addDynamicSampler(IntSupplier sampler, String... names) {
+		return samplers.addDynamicSampler(sampler, names);
+	}
 }

src/main/java/net/coderbot/iris/gl/program/ProgramSamplers.java

@@ -0,0 +1,160 @@
+package net.coderbot.iris.gl.program;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.mojang.blaze3d.systems.RenderSystem;
+import net.coderbot.iris.gl.sampler.SamplerBinding;
+import net.coderbot.iris.gl.sampler.SamplerHolder;
+import net.coderbot.iris.gl.sampler.SamplerLimits;
+import org.lwjgl.opengl.GL20C;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+import java.util.function.IntSupplier;
+
+public class ProgramSamplers {
+	private final ImmutableList<SamplerBinding> samplerBindings;
+	private List<GlUniform1iCall> initializer;
+
+	private ProgramSamplers(ImmutableList<SamplerBinding> samplerBindings, List<GlUniform1iCall> initializer) {
+		this.samplerBindings = samplerBindings;
+		this.initializer = initializer;
+	}
+
+	public void update() {
+		if (initializer != null) {
+			for (GlUniform1iCall call : initializer) {
+				GL20C.glUniform1i(call.getLocation(), call.getValue());
+			}
+
+			initializer = null;
+		}
+
+		for (SamplerBinding samplerBinding : samplerBindings) {
+			samplerBinding.update();
+		}
+
+		RenderSystem.activeTexture(GL20C.GL_TEXTURE0);
+	}
+
+	public static Builder builder(int program, Set<Integer> reservedTextureUnits) {
+		return new Builder(program, reservedTextureUnits);
+	}
+
+	public static final class Builder implements SamplerHolder {
+		private final int program;
+		private final ImmutableSet<Integer> reservedTextureUnits;
+		private final ImmutableList.Builder<SamplerBinding> samplers;
+		private final List<GlUniform1iCall> calls;
+		private int remainingUnits;
+		private int nextUnit;
+
+		private Builder(int program, Set<Integer> reservedTextureUnits) {
+			this.program = program;
+			this.reservedTextureUnits = ImmutableSet.copyOf(reservedTextureUnits);
+			this.samplers = ImmutableList.builder();
+			this.calls = new ArrayList<>();
+
+			int maxTextureUnits = SamplerLimits.get().getMaxTextureUnits();
+
+			for (int unit : reservedTextureUnits) {
+				if (unit >= maxTextureUnits) {
+					throw new IllegalStateException("Cannot mark texture unit " + unit + " as reserved because that " +
+							"texture unit isn't available on this system! Only " + maxTextureUnits +
+							" texture units are available.");
+				}
+			}
+
+			this.remainingUnits = maxTextureUnits - reservedTextureUnits.size();
+			this.nextUnit = 0;
+
+			while (reservedTextureUnits.contains(nextUnit)) {
+				nextUnit += 1;
+			}
+
+			//System.out.println("Begin building samplers. Reserved texture units are " + reservedTextureUnits +
+			//		", next texture unit is " + nextUnit + ", there are " + remainingUnits + " units remaining.");
+		}
+
+		// TODO: "default" sampler behavior
+
+		@Override
+		public void addExternalSampler(int textureUnit, String... names) {
+			if (!reservedTextureUnits.contains(textureUnit)) {
+				throw new IllegalArgumentException("Cannot add an externally-managed sampler for texture unit " +
+						textureUnit + " since it isn't in the set of reserved texture units.");
+			}
+
+			for (String name : names) {
+				int location = GL20C.glGetUniformLocation(program, name);
+
+				if (location == -1) {
+					// There's no active sampler with this particular name in the program.
+					continue;
+				}
+
+				// Set up this sampler uniform to use this particular texture unit.
+				//System.out.println("Binding external sampler " + name + " to texture unit " + textureUnit);
+				calls.add(new GlUniform1iCall(location, textureUnit));
+			}
+		}
+
+		@Override
+		public boolean hasSampler(String name) {
+			return GL20C.glGetUniformLocation(program, name) != -1;
+		}
+
+		/**
+		 * Adds a sampler
+		 * @return false if this sampler is not active, true if at least one of the names referred to an active sampler
+		 */
+		@Override
+		public boolean addDynamicSampler(IntSupplier sampler, String... names) {
+			boolean used = false;
+
+			for (String name : names) {
+				int location = GL20C.glGetUniformLocation(program, name);
+
+				if (location == -1) {
+					// There's no active sampler with this particular name in the program.
+					continue;
+				}
+
+				// Make sure that we aren't out of texture units.
+				if (remainingUnits <= 0) {
+					throw new IllegalStateException("No more available texture units while activating sampler " + name);
+				}
+
+				//System.out.println("Binding dynamic sampler " + name + " to texture unit " + nextUnit);
+
+				// Set up this sampler uniform to use this particular texture unit.
+				calls.add(new GlUniform1iCall(location, nextUnit));
+
+				// And mark this texture unit as used.
+				used = true;
+			}
+
+			if (!used) {
+				return false;
+			}
+
+			samplers.add(new SamplerBinding(nextUnit, sampler));
+
+			remainingUnits -= 1;
+			nextUnit += 1;
+
+			while (remainingUnits > 0 && reservedTextureUnits.contains(nextUnit)) {
+				nextUnit += 1;
+			}
+
+			//System.out.println("The next unit is " + nextUnit + ", there are " + remainingUnits + " units remaining.");
+
+			return true;
+		}
+
+		public ProgramSamplers build() {
+			return new ProgramSamplers(samplers.build(), calls);
+		}
+	}
+}

src/main/java/net/coderbot/iris/gl/sampler/SamplerBinding.java

@@ -0,0 +1,21 @@
+package net.coderbot.iris.gl.sampler;
+
+import com.mojang.blaze3d.systems.RenderSystem;
+import org.lwjgl.opengl.GL20C;
+
+import java.util.function.IntSupplier;
+
+public class SamplerBinding {
+	private final int textureUnit;
+	private final IntSupplier texture;
+
+	public SamplerBinding(int textureUnit, IntSupplier texture) {
+		this.textureUnit = textureUnit;
+		this.texture = texture;
+	}
+
+	public void update() {
+		RenderSystem.activeTexture(GL20C.GL_TEXTURE0 + textureUnit);
+		RenderSystem.bindTexture(texture.getAsInt());
+	}
+}

src/main/java/net/coderbot/iris/gl/sampler/SamplerHolder.java

@@ -0,0 +1,9 @@
+package net.coderbot.iris.gl.sampler;
+
+import java.util.function.IntSupplier;
+
+public interface SamplerHolder {
+	void addExternalSampler(int textureUnit, String... names);
+	boolean hasSampler(String name);
+	boolean addDynamicSampler(IntSupplier sampler, String... names);
+}

src/main/java/net/coderbot/iris/gl/sampler/SamplerLimits.java

@@ -0,0 +1,24 @@
+package net.coderbot.iris.gl.sampler;
+
+import org.lwjgl.opengl.GL20C;
+
+public class SamplerLimits {
+	private final int maxTextureUnits;
+	private static SamplerLimits instance;
+
+	private SamplerLimits() {
+		this.maxTextureUnits = GL20C.glGetInteger(GL20C.GL_MAX_TEXTURE_IMAGE_UNITS);
+	}
+
+	public int getMaxTextureUnits() {
+		return maxTextureUnits;
+	}
+
+	public static SamplerLimits get() {
+		if (instance == null) {
+			instance = new SamplerLimits();
+		}
+
+		return instance;
+	}
+}