Lyft shaders

- The ubershaders actually compile, but uniforms/material ids are broken
 - Convert vertex/fragment adapter components into generic builder
 - Support arbitrary adapted function signatures
 - Make an attempt at cleaning up generation code

src/main/java/com/jozufozu/flywheel/backend/Backend.java

@@ -1,5 +1,6 @@
 package com.jozufozu.flywheel.backend;
 
+import org.jetbrains.annotations.Contract;
 import org.jetbrains.annotations.Nullable;
 import org.slf4j.Logger;
 
@@ -60,6 +61,7 @@ public class Backend {
 		return TYPE != BackendTypes.OFF;
 	}
 
+	@Contract("null -> false")
 	public static boolean canUseInstancing(@Nullable Level level) {
 		return isOn() && isFlywheelLevel(level);
 	}

src/main/java/com/jozufozu/flywheel/backend/instancing/compile/FlwCompiler.java

@@ -10,6 +10,7 @@ import java.util.stream.Collectors;
 import com.google.common.collect.ArrayListMultimap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Multimap;
+import com.jozufozu.flywheel.Flywheel;
 import com.jozufozu.flywheel.api.context.ContextShader;
 import com.jozufozu.flywheel.api.pipeline.Pipeline;
 import com.jozufozu.flywheel.api.struct.StructType;
@@ -28,6 +29,8 @@ import com.jozufozu.flywheel.core.SourceComponent;
 import com.jozufozu.flywheel.core.pipeline.SimplePipeline;
 import com.jozufozu.flywheel.core.source.ShaderLoadingException;
 import com.jozufozu.flywheel.core.source.ShaderSources;
+import com.jozufozu.flywheel.core.source.generate.FnSignature;
+import com.jozufozu.flywheel.core.source.generate.GlslExpr;
 import com.jozufozu.flywheel.util.StringUtil;
 
 public class FlwCompiler {
@@ -36,11 +39,10 @@ public class FlwCompiler {
 
 	final long compileStart = System.nanoTime();
 	private final ShaderSources sources;
-	private final VertexMaterialComponent vertexMaterialComponent;
+	private final MaterialAdapterComponent vertexMaterialComponent;
-	private final FragmentMaterialComponent fragmentMaterialComponent;
+	private final MaterialAdapterComponent fragmentMaterialComponent;
 	private final List<PipelineContext> pipelineContexts;
 
-
 	final ShaderCompiler shaderCompiler;
 	final Multimap<Set<UniformProvider>, PipelineContext> uniformProviderGroups = ArrayListMultimap.create();
 	final Map<PipelineContext, GlProgram> pipelinePrograms = new HashMap<>();
@@ -50,8 +52,26 @@ public class FlwCompiler {
 	public FlwCompiler(ShaderSources sources) {
 		this.shaderCompiler = new ShaderCompiler(errors::add);
 		this.sources = sources;
-		this.vertexMaterialComponent = new VertexMaterialComponent(sources, ComponentRegistry.materials.vertexSources());
+		this.vertexMaterialComponent = MaterialAdapterComponent.builder(Flywheel.rl("vertex_material_adapter"))
-		this.fragmentMaterialComponent = new FragmentMaterialComponent(sources, ComponentRegistry.materials.fragmentSources());
+				.materialSources(ComponentRegistry.materials.vertexSources())
+				.adapt(FnSignature.ofVoid("flw_materialVertex"))
+				.switchOn(GlslExpr.variable("flw_materialVertexID"))
+				.build(sources);
+		this.fragmentMaterialComponent = MaterialAdapterComponent.builder(Flywheel.rl("fragment_material_adapter"))
+				.materialSources(ComponentRegistry.materials.fragmentSources())
+				.adapt(FnSignature.ofVoid("flw_materialFragment"))
+				.adapt(FnSignature.create()
+						.returnType("bool")
+						.name("flw_discardPredicate")
+						.arg("vec4", "color")
+						.build(), GlslExpr.literal(false))
+				.adapt(FnSignature.create()
+						.returnType("vec4")
+						.name("flw_fogFilter")
+						.arg("vec4", "color")
+						.build(), GlslExpr.variable("color"))
+				.switchOn(GlslExpr.variable("flw_materialFragmentID"))
+				.build(sources);
 
 		this.pipelineContexts = buildPipelineSet();
 

src/main/java/com/jozufozu/flywheel/backend/instancing/compile/FragmentMaterialComponent.java

@@ -1,29 +0,0 @@
-package com.jozufozu.flywheel.backend.instancing.compile;
-
-import java.util.List;
-
-import com.jozufozu.flywheel.Flywheel;
-import com.jozufozu.flywheel.core.source.FileResolution;
-import com.jozufozu.flywheel.core.source.ShaderSources;
-import com.jozufozu.flywheel.core.source.generate.GlslExpr;
-
-import net.minecraft.resources.ResourceLocation;
-
-public class FragmentMaterialComponent extends MaterialAdapterComponent {
-
-	private static final String flw_materialFragment = "flw_materialFragment";
-	private static final String flw_discardPredicate = "flw_discardPredicate";
-	private static final String flw_fogFilter = "flw_fogFilter";
-	private static final List<String> adaptedFunctions = List.of(flw_materialFragment, flw_discardPredicate, flw_fogFilter);
-
-	private static final GlslExpr flw_materialFragmentID = GlslExpr.variable(flw_materialFragment + "ID");
-
-	public FragmentMaterialComponent(ShaderSources sources, List<FileResolution> sourceMaterials) {
-		super(sources, sourceMaterials, flw_materialFragmentID, adaptedFunctions);
-	}
-
-	@Override
-	public ResourceLocation name() {
-		return Flywheel.rl("fragment_material_adapter");
-	}
-}

src/main/java/com/jozufozu/flywheel/backend/instancing/compile/MaterialAdapterComponent.java

@@ -1,32 +1,161 @@
 package com.jozufozu.flywheel.backend.instancing.compile;
 
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.List;
 
+import javax.annotation.Nonnull;
+
 import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
 
 import com.google.common.collect.ImmutableList;
 import com.jozufozu.flywheel.core.SourceComponent;
 import com.jozufozu.flywheel.core.source.FileResolution;
 import com.jozufozu.flywheel.core.source.ShaderSources;
+import com.jozufozu.flywheel.core.source.generate.FnSignature;
+import com.jozufozu.flywheel.core.source.generate.GlslBlock;
 import com.jozufozu.flywheel.core.source.generate.GlslBuilder;
 import com.jozufozu.flywheel.core.source.generate.GlslExpr;
+import com.jozufozu.flywheel.core.source.generate.GlslSwitch;
 import com.jozufozu.flywheel.util.ResourceUtil;
 
 import net.minecraft.resources.ResourceLocation;
 
-public abstract class MaterialAdapterComponent implements SourceComponent {
+public class MaterialAdapterComponent implements SourceComponent {
 
 	// TODO: material id handling in pipeline shader
+	private final ResourceLocation name;
 	private final GlslExpr switchArg;
-	private final List<String> adaptedFunctions;
+	private final List<AdaptedFn> functionsToAdapt;
-	private final List<RenamedFunctionsSourceComponent> transformedMaterials;
+	private final List<RenamedFunctionsSourceComponent> adaptedComponents;
 
-	// TODO: Create builder and remove Fragment* and Vertex* classes
+	public MaterialAdapterComponent(ResourceLocation name, GlslExpr switchArg, List<AdaptedFn> functionsToAdapt, List<RenamedFunctionsSourceComponent> adaptedComponents) {
-	public MaterialAdapterComponent(ShaderSources sources, List<FileResolution> sourceMaterials, GlslExpr switchArg, List<String> adaptedFunctions) {
+		this.name = name;
 		this.switchArg = switchArg;
-		this.adaptedFunctions = adaptedFunctions;
+		this.functionsToAdapt = functionsToAdapt;
+		this.adaptedComponents = adaptedComponents;
+	}
+
+	public static Builder builder(ResourceLocation name) {
+		return new Builder(name);
+	}
+
+	@Override
+	public ResourceLocation name() {
+		return name;
+	}
+
+	@Override
+	public Collection<? extends SourceComponent> included() {
+		return adaptedComponents;
+	}
+
+	@Override
+	public String source() {
+		var builder = new GlslBuilder();
+
+		for (var adaptedFunction : functionsToAdapt) {
+			builder.function()
+					.signature(adaptedFunction.signature())
+					.body(body -> generateAdapter(body, adaptedFunction));
+		}
+
+		return builder.build();
+	}
+
+	private void generateAdapter(GlslBlock body, AdaptedFn adaptedFunction) {
+		var sw = GlslSwitch.on(switchArg);
+		var fnSignature = adaptedFunction.signature();
+		var fnName = fnSignature.name();
+		var isVoid = fnSignature.isVoid();
+		var fnArgs = fnSignature.createArgExpressions();
+
+		for (int i = 0; i < adaptedComponents.size(); i++) {
+			var component = adaptedComponents.get(i);
+
+			if (!component.replaces(fnName)) {
+				continue;
+			}
+
+			var adaptedCall = GlslExpr.call(component.remapFnName(fnName), fnArgs);
+
+			var block = GlslBlock.create();
+			if (isVoid) {
+				block.eval(adaptedCall)
+						.breakStmt();
+			} else {
+				block.ret(adaptedCall);
+			}
+
+			sw.intCase(i, block);
+		}
+
+		if (!isVoid) {
+			var defaultReturn = adaptedFunction.defaultReturn;
+			if (defaultReturn == null) {
+				throw new IllegalStateException("Function " + fnName + " is not void, but no default return value was provided");
+			}
+			sw.defaultCase(GlslBlock.create()
+					.ret(defaultReturn));
+		}
+
+		body.add(sw);
+	}
+
+	@NotNull
+	private static HashMap<String, String> createAdapterMap(List<AdaptedFn> adaptedFunctions, ResourceLocation loc) {
+		HashMap<String, String> out = new HashMap<>();
+
+		var suffix = '_' + ResourceUtil.toSafeString(loc);
+
+		for (var adapted : adaptedFunctions) {
+			var fnName = adapted.signature()
+					.name();
+			out.put(fnName, fnName + suffix);
+		}
+
+		return out;
+	}
+
+	private record AdaptedFn(FnSignature signature, @Nullable GlslExpr defaultReturn) {
+	}
+
+	public static class Builder {
+		private final ResourceLocation name;
+		private final List<FileResolution> sourceMaterials = new ArrayList<>();
+		private final List<AdaptedFn> adaptedFunctions = new ArrayList<>();
+		private GlslExpr switchArg;
+
+		public Builder(ResourceLocation name) {
+			this.name = name;
+		}
+
+		public Builder materialSources(List<FileResolution> sources) {
+			this.sourceMaterials.addAll(sources);
+			return this;
+		}
+
+		public Builder adapt(FnSignature function) {
+			adaptedFunctions.add(new AdaptedFn(function, null));
+			return this;
+		}
+
+		public Builder adapt(FnSignature function, @Nonnull GlslExpr defaultReturn) {
+			adaptedFunctions.add(new AdaptedFn(function, defaultReturn));
+			return this;
+		}
+
+		public Builder switchOn(GlslExpr expr) {
+			this.switchArg = expr;
+			return this;
+		}
+
+		public MaterialAdapterComponent build(ShaderSources sources) {
+			if (switchArg == null) {
+				throw new NullPointerException("Switch argument must be set");
+			}
 
 			var transformed = ImmutableList.<RenamedFunctionsSourceComponent>builder();
 
@@ -37,51 +166,7 @@ public abstract class MaterialAdapterComponent implements SourceComponent {
 				transformed.add(new RenamedFunctionsSourceComponent(sourceFile, createAdapterMap(adaptedFunctions, loc)));
 			}
 
-		this.transformedMaterials = transformed.build();
+			return new MaterialAdapterComponent(name, switchArg, adaptedFunctions, transformed.build());
 		}
-
-	@Override
-	public Collection<? extends SourceComponent> included() {
-		return transformedMaterials;
-	}
-
-	@Override
-	public String source() {
-		var builder = new GlslBuilder();
-
-		for (String adaptedFunction : adaptedFunctions) {
-			// TODO: support different function signatures
-			builder.function()
-					.returnType("void")
-					.name(adaptedFunction)
-					.body(body -> generateAdapter(body, adaptedFunction));
-		}
-
-		return builder.build();
-	}
-
-	private void generateAdapter(GlslBuilder.BlockBuilder body, String adaptedFunction) {
-		var sw = new GlslBuilder.SwitchBuilder(switchArg);
-		for (int i = 0; i < transformedMaterials.size(); i++) {
-			var variant = transformedMaterials.get(i)
-					.replacement(adaptedFunction);
-
-			sw.case_(i, b -> b.eval(GlslExpr.call(variant))
-					.break_());
-		}
-		body.add(sw.build());
-	}
-
-	@NotNull
-	private static HashMap<String, String> createAdapterMap(List<String> adaptedFunctions, ResourceLocation loc) {
-		HashMap<String, String> out = new HashMap<>();
-
-		var suffix = '_' + ResourceUtil.toSafeString(loc);
-
-		for (String fnName : adaptedFunctions) {
-			out.put(fnName, fnName + suffix);
-		}
-
-		return out;
 	}
 }

src/main/java/com/jozufozu/flywheel/backend/instancing/compile/RenamedFunctionsSourceComponent.java

@@ -11,24 +11,29 @@ import net.minecraft.resources.ResourceLocation;
 public final class RenamedFunctionsSourceComponent implements SourceComponent {
 	private final SourceComponent source;
 	private final Map<String, String> replacements;
+	private final String sourceString;
 
 	public RenamedFunctionsSourceComponent(SourceComponent source, String find, String replace) {
-		this.source = source;
+		this(source, Map.of(find, replace));
-		this.replacements = Map.of(find, replace);
 	}
 
 	public RenamedFunctionsSourceComponent(SourceComponent source, Map<String, String> replacements) {
 		this.source = source;
 		this.replacements = replacements;
+		this.sourceString = source.source();
 	}
 
-	public String replacement(String name) {
+	public String remapFnName(String name) {
 		return replacements.getOrDefault(name, name);
 	}
 
+	public boolean replaces(String name) {
+		return replacements.containsKey(name) && sourceString.contains(name);
+	}
+
 	@Override
 	public String source() {
-		var source = this.source.source();
+		var source = sourceString;
 
 		for (var entry : replacements.entrySet()) {
 			source = source.replace(entry.getKey(), entry.getValue());

src/main/java/com/jozufozu/flywheel/backend/instancing/compile/VertexMaterialComponent.java

@@ -1,27 +0,0 @@
-package com.jozufozu.flywheel.backend.instancing.compile;
-
-import java.util.List;
-
-import com.jozufozu.flywheel.Flywheel;
-import com.jozufozu.flywheel.core.source.FileResolution;
-import com.jozufozu.flywheel.core.source.ShaderSources;
-import com.jozufozu.flywheel.core.source.generate.GlslExpr;
-
-import net.minecraft.resources.ResourceLocation;
-
-public class VertexMaterialComponent extends MaterialAdapterComponent {
-
-	private static final String flw_materialVertex = "flw_materialVertex";
-	private static final List<String> adaptedFunctions = List.of(flw_materialVertex);
-	private static final GlslExpr flw_materialVertexID = GlslExpr.variable(flw_materialVertex + "ID");
-
-
-	public VertexMaterialComponent(ShaderSources sources, List<FileResolution> sourceMaterials) {
-		super(sources, sourceMaterials, flw_materialVertexID, adaptedFunctions);
-	}
-
-	@Override
-	public ResourceLocation name() {
-		return Flywheel.rl("vertex_material_adapter");
-	}
-}

src/main/java/com/jozufozu/flywheel/backend/instancing/indirect/IndirectComponent.java

@@ -12,6 +12,8 @@ import com.jozufozu.flywheel.core.SourceComponent;
 import com.jozufozu.flywheel.core.layout.LayoutItem;
 import com.jozufozu.flywheel.core.source.ShaderSources;
 import com.jozufozu.flywheel.core.source.SourceFile;
+import com.jozufozu.flywheel.core.source.generate.FnSignature;
+import com.jozufozu.flywheel.core.source.generate.GlslBlock;
 import com.jozufozu.flywheel.core.source.generate.GlslBuilder;
 import com.jozufozu.flywheel.core.source.generate.GlslExpr;
 
@@ -21,6 +23,8 @@ public class IndirectComponent implements SourceComponent {
 
 	private static final String UNPACK_ARG = "p";
 	private static final GlslExpr.Variable UNPACKING_VARIABLE = GlslExpr.variable(UNPACK_ARG);
+	private static final String STRUCT_NAME = "IndirectStruct";
+	private static final String PACKED_STRUCT_NAME = STRUCT_NAME + "_packed";
 
 	private final List<LayoutItem> layoutItems;
 	private final ImmutableList<SourceFile> included;
@@ -46,20 +50,19 @@ public class IndirectComponent implements SourceComponent {
 
 	@Override
 	public String source() {
-		return generateIndirect("IndirectStruct");
+		return generateIndirect();
 	}
 
-	public String generateIndirect(String structName) {
+	public String generateIndirect() {
 		var builder = new GlslBuilder();
-		final var packedStructName = structName + "_packed";
+		builder.define("FlwInstance", STRUCT_NAME);
-		builder.define("FlwInstance", structName);
+		builder.define("FlwPackedInstance", PACKED_STRUCT_NAME);
-		builder.define("FlwPackedInstance", packedStructName);
 
 		var packed = builder.struct();
 		builder.blankLine();
 		var instance = builder.struct();
-		packed.setName(packedStructName);
+		packed.setName(PACKED_STRUCT_NAME);
-		instance.setName(structName);
+		instance.setName(STRUCT_NAME);
 
 		for (var field : layoutItems) {
 			field.addPackedToStruct(packed);
@@ -69,13 +72,20 @@ public class IndirectComponent implements SourceComponent {
 		builder.blankLine();
 
 		builder.function()
-				.returnType(structName)
+				.signature(FnSignature.create()
+						.returnType(STRUCT_NAME)
 						.name("flw_unpackInstance")
-				.argumentIn(packedStructName, UNPACK_ARG)
+						.arg(PACKED_STRUCT_NAME, UNPACK_ARG)
-				.body(b -> b.ret(GlslExpr.call(structName, layoutItems.stream()
+						.build())
-						.map(layoutItem -> layoutItem.unpackField(UNPACKING_VARIABLE))
+				.body(this::generateUnpackingBody);
-						.toList())));
 
 		return builder.build();
 	}
+
+	private void generateUnpackingBody(GlslBlock b) {
+		var unpackedFields = layoutItems.stream()
+				.map(layoutItem -> layoutItem.unpackField(UNPACKING_VARIABLE))
+				.toList();
+		b.ret(GlslExpr.call(STRUCT_NAME, unpackedFields));
+	}
 }

src/main/java/com/jozufozu/flywheel/backend/instancing/instancing/InstancedArraysComponent.java

@@ -8,6 +8,8 @@ import com.jozufozu.flywheel.Flywheel;
 import com.jozufozu.flywheel.api.pipeline.Pipeline;
 import com.jozufozu.flywheel.core.SourceComponent;
 import com.jozufozu.flywheel.core.layout.LayoutItem;
+import com.jozufozu.flywheel.core.source.generate.FnSignature;
+import com.jozufozu.flywheel.core.source.generate.GlslBlock;
 import com.jozufozu.flywheel.core.source.generate.GlslBuilder;
 import com.jozufozu.flywheel.core.source.generate.GlslExpr;
 
@@ -15,6 +17,7 @@ import net.minecraft.resources.ResourceLocation;
 
 public class InstancedArraysComponent implements SourceComponent {
 	private static final String ATTRIBUTE_SUFFIX = "_vertex_in";
+	private static final String STRUCT_NAME = "Instance";
 
 	private final List<LayoutItem> layoutItems;
 	private final int baseIndex;
@@ -32,19 +35,15 @@ public class InstancedArraysComponent implements SourceComponent {
 		return Collections.emptyList();
 	}
 
-	@Override
-	public String source() {
-		return generateInstancedArrays("Instance");
-	}
-
 	@Override
 	public ResourceLocation name() {
 		return Flywheel.rl("generated_instanced_arrays");
 	}
 
-	public String generateInstancedArrays(String structName) {
+	@Override
+	public String source() {
 		var builder = new GlslBuilder();
-		builder.define("FlwInstance", structName);
+		builder.define("FlwInstance", STRUCT_NAME);
 
 		int i = baseIndex;
 		for (var field : layoutItems) {
@@ -61,7 +60,7 @@ public class InstancedArraysComponent implements SourceComponent {
 		builder.blankLine();
 
 		var structBuilder = builder.struct();
-		structBuilder.setName(structName);
+		structBuilder.setName(STRUCT_NAME);
 
 		for (var field : layoutItems) {
 			field.addToStruct(structBuilder);
@@ -71,13 +70,16 @@ public class InstancedArraysComponent implements SourceComponent {
 
 		// unpacking function
 		builder.function()
-				.returnType(structName)
+				.signature(FnSignature.of(STRUCT_NAME, "flw_unpackInstance"))
-				.name("flw_unpackInstance")
+				.body(this::generateUnpackingBody);
-				.body(b -> b.ret(GlslExpr.call(structName, layoutItems.stream()
-						.map(it -> new GlslExpr.Variable(it.name() + ATTRIBUTE_SUFFIX))
-						.toList())));
 
 		return builder.build();
 	}
 
+	private void generateUnpackingBody(GlslBlock b) {
+		var fields = layoutItems.stream()
+				.map(it -> new GlslExpr.Variable(it.name() + ATTRIBUTE_SUFFIX))
+				.toList();
+		b.ret(GlslExpr.call(STRUCT_NAME, fields));
+	}
 }

src/main/java/com/jozufozu/flywheel/core/layout/LayoutItem.java

@@ -1,7 +1,7 @@
 package com.jozufozu.flywheel.core.layout;
 
-import com.jozufozu.flywheel.core.source.generate.GlslBuilder;
 import com.jozufozu.flywheel.core.source.generate.GlslExpr;
+import com.jozufozu.flywheel.core.source.generate.GlslStruct;
 
 public record LayoutItem(InputType type, String name) {
 	public GlslExpr unpackField(GlslExpr.Variable struct) {
@@ -9,11 +9,11 @@ public record LayoutItem(InputType type, String name) {
 				.transform(type()::unpack);
 	}
 
-	public void addToStruct(GlslBuilder.StructBuilder structBuilder) {
+	public void addToStruct(GlslStruct glslStruct) {
-		structBuilder.addField(type().typeName(), name());
+		glslStruct.addField(type().typeName(), name());
 	}
 
-	public void addPackedToStruct(GlslBuilder.StructBuilder packed) {
+	public void addPackedToStruct(GlslStruct packed) {
 		packed.addField(type().packedTypeName(), name());
 	}
 }

src/main/java/com/jozufozu/flywheel/core/source/generate/FnSignature.java

@@ -0,0 +1,80 @@
+package com.jozufozu.flywheel.core.source.generate;
+
+import java.util.Collection;
+import java.util.stream.Collectors;
+
+import com.google.common.collect.ImmutableList;
+import com.jozufozu.flywheel.util.Pair;
+
+public record FnSignature(String returnType, String name, ImmutableList<Pair<String, String>> args) {
+
+	public static Builder create() {
+		return new Builder();
+	}
+
+	public static FnSignature of(String returnType, String name) {
+		return FnSignature.create()
+				.returnType(returnType)
+				.name(name)
+				.build();
+	}
+
+	public static FnSignature ofVoid(String name) {
+		return new FnSignature("void", name, ImmutableList.of());
+	}
+
+	public Collection<? extends GlslExpr> createArgExpressions() {
+		return args.stream()
+				.map(Pair::second)
+				.map(GlslExpr::variable)
+				.collect(Collectors.toList());
+	}
+
+	public boolean isVoid() {
+		return "void".equals(returnType);
+	}
+
+	public String fullDeclaration() {
+		return returnType + ' ' + name + '(' + args.stream()
+				.map(p -> p.first() + ' ' + p.second())
+				.collect(Collectors.joining(", ")) + ')';
+	}
+
+	public String signatureDeclaration() {
+		return returnType + ' ' + name + '(' + args.stream()
+				.map(Pair::first)
+				.collect(Collectors.joining(", ")) + ')';
+	}
+
+	public static class Builder {
+		private String returnType;
+		private String name;
+		private final ImmutableList.Builder<Pair<String, String>> args = ImmutableList.builder();
+
+		public Builder returnType(String returnType) {
+			this.returnType = returnType;
+			return this;
+		}
+
+		public Builder name(String name) {
+			this.name = name;
+			return this;
+		}
+
+		public Builder arg(String type, String name) {
+			args.add(Pair.of(type, name));
+			return this;
+		}
+
+		public FnSignature build() {
+			if (returnType == null) {
+				throw new IllegalStateException("returnType not set");
+			}
+			if (name == null) {
+				throw new IllegalStateException("name not set");
+			}
+			return new FnSignature(returnType, name, args.build());
+		}
+	}
+
+}

src/main/java/com/jozufozu/flywheel/core/source/generate/GlslBlock.java

@@ -0,0 +1,39 @@
+package com.jozufozu.flywheel.core.source.generate;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collectors;
+
+public class GlslBlock {
+	private final List<GlslStmt> body = new ArrayList<>();
+
+	public static GlslBlock create() {
+		return new GlslBlock();
+	}
+
+	public GlslBlock add(GlslStmt stmt) {
+		body.add(stmt);
+		return this;
+	}
+
+	public GlslBlock eval(GlslExpr expr) {
+		return add(GlslStmt.eval(expr));
+	}
+
+	public GlslBlock ret(GlslExpr call) {
+		add(GlslStmt.ret(call));
+		return this;
+	}
+
+	public GlslBlock breakStmt() {
+		add(GlslStmt.BREAK);
+		return this;
+	}
+
+	public String prettyPrint() {
+		return body.stream()
+				.map(GlslStmt::prettyPrint)
+				.collect(Collectors.joining("\n"));
+	}
+
+}

src/main/java/com/jozufozu/flywheel/core/source/generate/GlslBuilder.java

@@ -2,31 +2,28 @@ package com.jozufozu.flywheel.core.source.generate;
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.function.Consumer;
 import java.util.stream.Collectors;
 
-import com.jozufozu.flywheel.util.Pair;
-
 public class GlslBuilder {
-	private final List<GlslRootElement> elements = new ArrayList<>();
+	private final List<Declaration> elements = new ArrayList<>();
 
 	public void define(String name, String value) {
 		add(new Define(name, value));
 	}
 
-	public StructBuilder struct() {
+	public GlslStruct struct() {
-		return add(new StructBuilder());
+		return add(new GlslStruct());
 	}
 
-	public FunctionBuilder function() {
+	public GlslFn function() {
-		return add(new FunctionBuilder());
+		return add(new GlslFn());
 	}
 
-	public VertexInputBuilder vertexInput() {
+	public GlslVertexInput vertexInput() {
-		return add(new VertexInputBuilder());
+		return add(new GlslVertexInput());
 	}
 
-	public <T extends GlslRootElement> T add(T element) {
+	public <T extends Declaration> T add(T element) {
 		elements.add(element);
 		return element;
 	}
@@ -37,15 +34,15 @@ public class GlslBuilder {
 
 	public String build() {
 		return elements.stream()
-				.map(GlslRootElement::prettyPrint)
+				.map(Declaration::prettyPrint)
 				.collect(Collectors.joining("\n"));
 	}
 
-	public interface GlslRootElement {
+	public interface Declaration {
 		String prettyPrint();
 	}
 
-	public enum Separators implements GlslRootElement {
+	public enum Separators implements Declaration {
 		BLANK_LINE(""),
 		;
 
@@ -61,169 +58,11 @@ public class GlslBuilder {
 		}
 	}
 
-	public record Define(String name, String value) implements GlslRootElement {
+	public record Define(String name, String value) implements Declaration {
 		@Override
 		public String prettyPrint() {
 			return "#define " + name + " " + value;
 		}
 	}
 
-	public static class VertexInputBuilder implements GlslRootElement {
-
-		private int binding;
-		private String type;
-		private String name;
-
-		public VertexInputBuilder binding(int binding) {
-			this.binding = binding;
-			return this;
-		}
-
-		public VertexInputBuilder type(String type) {
-			this.type = type;
-			return this;
-		}
-
-		public VertexInputBuilder name(String name) {
-			this.name = name;
-			return this;
-		}
-
-		@Override
-		public String prettyPrint() {
-			return "layout(location = " + binding + ") in " + type + " " + name + ";";
-		}
-	}
-
-	public static class StructBuilder implements GlslRootElement {
-
-		private final List<Pair<String, String>> fields = new ArrayList<>();
-		private String name;
-
-		public void setName(String name) {
-			this.name = name;
-		}
-
-		public void addField(String type, String name) {
-			fields.add(Pair.of(type, name));
-		}
-
-		private String buildFields() {
-			return fields.stream()
-					.map(p -> p.first() + ' ' + p.second() + ';')
-					.collect(Collectors.joining("\n"));
-		}
-
-		public String prettyPrint() {
-			return """
-					struct %s {
-					%s
-					};
-					""".formatted(name, buildFields().indent(4));
-		}
-	}
-
-	public static class FunctionBuilder implements GlslRootElement {
-		private final List<Pair<String, String>> arguments = new ArrayList<>();
-		private final BlockBuilder body = new BlockBuilder();
-		private String returnType;
-		private String name;
-
-		public FunctionBuilder returnType(String returnType) {
-			this.returnType = returnType;
-			return this;
-		}
-
-		public FunctionBuilder name(String name) {
-			this.name = name;
-			return this;
-		}
-
-		public FunctionBuilder argument(String type, String name) {
-			arguments.add(Pair.of(type, name));
-			return this;
-		}
-
-		public FunctionBuilder argumentIn(String type, String name) {
-			arguments.add(Pair.of("in " + type, name));
-			return this;
-		}
-
-		public FunctionBuilder body(Consumer<BlockBuilder> f) {
-			f.accept(body);
-			return this;
-		}
-
-		public String prettyPrint() {
-			return """
-					%s %s(%s) {
-					%s
-					}
-					""".formatted(returnType, name, buildArguments(), body.prettyPrint()
-					.indent(4));
-		}
-
-		private String buildArguments() {
-			return arguments.stream()
-					.map(p -> p.first() + ' ' + p.second())
-					.collect(Collectors.joining(", "));
-		}
-	}
-
-	public static class BlockBuilder implements LangItem {
-		private final List<GlslStmt> body = new ArrayList<>();
-
-		public BlockBuilder add(GlslStmt stmt) {
-			body.add(stmt);
-			return this;
-		}
-
-		public BlockBuilder eval(GlslExpr expr) {
-			return add(GlslStmt.eval(expr));
-		}
-
-		public BlockBuilder switchOn(GlslExpr expr, Consumer<SwitchBuilder> f) {
-			var builder = new SwitchBuilder(expr);
-			f.accept(builder);
-			return add(builder.build());
-		}
-
-		public void ret(GlslExpr call) {
-			add(GlslStmt.ret(call));
-		}
-
-		public void break_() {
-			add(GlslStmt.BREAK);
-		}
-
-		@Override
-		public String prettyPrint() {
-			return body.stream()
-					.map(GlslStmt::prettyPrint)
-					.collect(Collectors.joining("\n"));
-		}
-
-	}
-
-	public static class SwitchBuilder {
-
-		private final GlslExpr on;
-
-		private final List<Pair<GlslExpr, BlockBuilder>> cases = new ArrayList<>();
-
-		public SwitchBuilder(GlslExpr on) {
-			this.on = on;
-		}
-
-		public SwitchBuilder case_(int expr, Consumer<BlockBuilder> f) {
-			var builder = new BlockBuilder();
-			f.accept(builder);
-			cases.add(Pair.of(GlslExpr.literal(expr), builder));
-			return this;
-		}
-
-		public GlslStmt.Switch build() {
-			return new GlslStmt.Switch(on, cases);
-		}
-	}
 }

src/main/java/com/jozufozu/flywheel/core/source/generate/GlslExpr.java

@@ -6,7 +6,7 @@ import java.util.stream.Collectors;
 
 import com.google.common.collect.ImmutableList;
 
-public interface GlslExpr extends LangItem {
+public interface GlslExpr {
 
 	/**
 	 * Create a glsl variable with the given name.
@@ -27,7 +27,11 @@ public interface GlslExpr extends LangItem {
 	}
 
 	static GlslExpr literal(int expr) {
-		return new Literal(expr);
+		return new IntLiteral(expr);
+	}
+
+	static GlslExpr literal(boolean expr) {
+		return new BoolLiteral(expr);
 	}
 
 	/**
@@ -70,6 +74,8 @@ public interface GlslExpr extends LangItem {
 		return f.apply(this);
 	}
 
+	String prettyPrint();
+
 	record Variable(String name) implements GlslExpr {
 		@Override
 		public String prettyPrint() {
@@ -117,10 +123,17 @@ public interface GlslExpr extends LangItem {
 
 	}
 
-	record Literal(int value) implements GlslExpr {
+	record IntLiteral(int value) implements GlslExpr {
 		@Override
 		public String prettyPrint() {
 			return Integer.toString(value);
 		}
 	}
+
+	record BoolLiteral(boolean value) implements GlslExpr {
+		@Override
+		public String prettyPrint() {
+			return Boolean.toString(value);
+		}
+	}
 }

src/main/java/com/jozufozu/flywheel/core/source/generate/GlslFn.java

@@ -0,0 +1,28 @@
+package com.jozufozu.flywheel.core.source.generate;
+
+import java.util.function.Consumer;
+
+import com.jozufozu.flywheel.util.StringUtil;
+
+public class GlslFn implements GlslBuilder.Declaration {
+	private final GlslBlock body = new GlslBlock();
+	private FnSignature signature;
+
+	public GlslFn signature(FnSignature signature) {
+		this.signature = signature;
+		return this;
+	}
+
+	public GlslFn body(Consumer<GlslBlock> f) {
+		f.accept(body);
+		return this;
+	}
+
+	public String prettyPrint() {
+		return """
+				%s {
+				%s
+				}
+				""".formatted(signature.fullDeclaration(), StringUtil.indent(body.prettyPrint(), 4));
+	}
+}

src/main/java/com/jozufozu/flywheel/core/source/generate/GlslStmt.java

@@ -1,11 +1,9 @@
 package com.jozufozu.flywheel.core.source.generate;
 
-import java.util.List;
+public interface GlslStmt {
-import java.util.stream.Collectors;
+	GlslStmt BREAK = () -> "break;";
-
+	GlslStmt CONTINUE = () -> "continue;";
-import com.jozufozu.flywheel.util.Pair;
+	GlslStmt RETURN = () -> "return;";
-
-public interface GlslStmt extends LangItem {
 
 	static GlslStmt eval(GlslExpr expr) {
 		return new Eval(expr);
@@ -15,11 +13,7 @@ public interface GlslStmt extends LangItem {
 		return new Return(value);
 	}
 
-	static GlslStmt BREAK = () -> "break;";
+	String prettyPrint();
-
-	static GlslStmt CONTINUE = () -> "continue;";
-
-	static GlslStmt RETURN = () -> "return;";
 
 	record Eval(GlslExpr expr) implements GlslStmt {
 		@Override
@@ -34,27 +28,4 @@ public interface GlslStmt extends LangItem {
 			return "return " + expr.prettyPrint() + ";";
 		}
 	}
-
-	record Switch(GlslExpr expr, List<Pair<GlslExpr, GlslBuilder.BlockBuilder>> body) implements GlslStmt {
-		@Override
-		public String prettyPrint() {
-			var cases = body.stream()
-					.map(Switch::prettyPrintCase)
-					.collect(Collectors.joining("\n"));
-			return """
-					switch (%s) {
-					%s
-					}""".formatted(expr.prettyPrint(), cases);
-		}
-
-		private static String prettyPrintCase(Pair<GlslExpr, GlslBuilder.BlockBuilder> p) {
-			var variant = p.first()
-					.prettyPrint();
-			var block = p.second()
-					.prettyPrint();
-			return """
-					case %s:
-					%s""".formatted(variant, block.indent(4));
-		}
-	}
 }

src/main/java/com/jozufozu/flywheel/core/source/generate/GlslStruct.java

@@ -0,0 +1,35 @@
+package com.jozufozu.flywheel.core.source.generate;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collectors;
+
+import com.jozufozu.flywheel.util.Pair;
+
+public class GlslStruct implements GlslBuilder.Declaration {
+
+	private final List<Pair<String, String>> fields = new ArrayList<>();
+	private String name;
+
+	public void setName(String name) {
+		this.name = name;
+	}
+
+	public void addField(String type, String name) {
+		fields.add(Pair.of(type, name));
+	}
+
+	private String buildFields() {
+		return fields.stream()
+				.map(p -> p.first() + ' ' + p.second() + ';')
+				.collect(Collectors.joining("\n"));
+	}
+
+	public String prettyPrint() {
+		return """
+				struct %s {
+				%s
+				};
+				""".formatted(name, buildFields().indent(4));
+	}
+}

src/main/java/com/jozufozu/flywheel/core/source/generate/GlslSwitch.java

@@ -0,0 +1,64 @@
+package com.jozufozu.flywheel.core.source.generate;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.stream.Collectors;
+
+import org.jetbrains.annotations.NotNull;
+
+import com.jozufozu.flywheel.util.Pair;
+import com.jozufozu.flywheel.util.StringUtil;
+
+public class GlslSwitch implements GlslStmt {
+
+	private final GlslExpr on;
+
+	private final List<Pair<GlslExpr, GlslBlock>> cases = new ArrayList<>();
+	private GlslBlock defaultCase = null;
+
+	private GlslSwitch(GlslExpr on) {
+		this.on = on;
+	}
+
+	public static GlslSwitch on(GlslExpr on) {
+		return new GlslSwitch(on);
+	}
+
+	public void intCase(int expr, GlslBlock block) {
+		cases.add(Pair.of(GlslExpr.literal(expr), block));
+	}
+
+	public void defaultCase(GlslBlock block) {
+		defaultCase = block;
+	}
+
+	@Override
+	public String prettyPrint() {
+		return """
+				switch (%s) {
+				%s
+				}""".formatted(on.prettyPrint(), getCaseStream());
+	}
+
+	@NotNull
+	private String getCaseStream() {
+		var cases = this.cases.stream()
+				.map(GlslSwitch::prettyPrintCase)
+				.collect(Collectors.joining("\n"));
+		if (defaultCase != null) {
+			cases += "\ndefault:\n" + StringUtil.indent(defaultCase.prettyPrint(), 4);
+		}
+		return cases;
+	}
+
+	private static String prettyPrintCase(Pair<GlslExpr, GlslBlock> p) {
+		var variant = p.first()
+				.prettyPrint();
+		var block = p.second()
+				.prettyPrint();
+		return """
+				case %s:
+				%s""".formatted(variant, StringUtil.indent(block, 4));
+	}
+
+}

src/main/java/com/jozufozu/flywheel/core/source/generate/GlslVertexInput.java

@@ -0,0 +1,28 @@
+package com.jozufozu.flywheel.core.source.generate;
+
+public class GlslVertexInput implements GlslBuilder.Declaration {
+
+	private int binding;
+	private String type;
+	private String name;
+
+	public GlslVertexInput binding(int binding) {
+		this.binding = binding;
+		return this;
+	}
+
+	public GlslVertexInput type(String type) {
+		this.type = type;
+		return this;
+	}
+
+	public GlslVertexInput name(String name) {
+		this.name = name;
+		return this;
+	}
+
+	@Override
+	public String prettyPrint() {
+		return "layout(location = " + binding + ") in " + type + " " + name + ";";
+	}
+}

src/main/java/com/jozufozu/flywheel/core/source/generate/LangItem.java

@@ -1,7 +0,0 @@
-package com.jozufozu.flywheel.core.source.generate;
-
-public interface LangItem {
-
-	String prettyPrint();
-
-}

src/main/java/com/jozufozu/flywheel/util/StringUtil.java

@@ -12,6 +12,7 @@ import java.text.DecimalFormat;
 import java.text.NumberFormat;
 import java.util.Arrays;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import javax.annotation.Nonnull;
 
@@ -72,6 +73,25 @@ public class StringUtil {
 		return value.substring(0, len);
 	}
 
+	/**
+	 * Copy of {@link String#indent(int)} with the trailing newline removed.
+	 */
+	public static String indent(String str, int n) {
+		if (str.isEmpty()) {
+			return "";
+		}
+		Stream<String> stream = str.lines();
+		if (n > 0) {
+			final String spaces = repeatChar(' ', n);
+			stream = stream.map(s -> spaces + s);
+		} else if (n == Integer.MIN_VALUE) {
+			stream = stream.map(String::stripLeading);
+		} else if (n < 0) {
+			throw new IllegalArgumentException("Requested indentation (" + n + ") is unsupported");
+		}
+		return stream.collect(Collectors.joining("\n"));
+	}
+
 	@Nonnull
 	public static String readToString(InputStream is) throws IOException {
 		ByteBuffer bytebuffer = null;