Allow passing shader defines when resolving shader programs

assets/shaders/mesh.glsl

@@ -292,6 +292,9 @@ vec3 ibl(int matIdx, vec3 N, vec3 V) {
 
 void main() {
   int matIdx = draw_data[DrawID].materialIdx;
+#if OVERRIDE_COLOR
+  FragColor = getAlbedo(matIdx);
+#else
   if (getAlbedo(matIdx).a < 0.5) {
     FragColor = vec4(0);
     return;
@@ -320,6 +323,7 @@ void main() {
   finalColor += getEmission(matIdx);
 
   FragColor = vec4(finalColor, getAlbedo(matIdx).a);
+#endif
 }
 
 

src/AssetManager.zig

@@ -110,22 +110,24 @@ const AssetWatcher = struct {
 
         var updatedList = std.SegmentedList(AssetId, 128){};
 
-        var loaded_assets: std.AutoHashMapUnmanaged(AssetId, LoadedAsset) = .{};
         var modified_times: std.AutoHashMapUnmanaged(AssetId, i128) = .{};
 
         {
             self.assetman.rw_lock.lockShared();
             defer self.assetman.rw_lock.unlockShared();
 
-            loaded_assets = try self.assetman.loaded_assets.clone(self.fba.allocator());
             modified_times = try self.assetman.modified_times.clone(self.fba.allocator());
         }
 
         {
-            var iter = loaded_assets.iterator();
+            var iter = modified_times.iterator();
             while (iter.next()) |entry| {
-                const modified_time = modified_times.get(entry.key_ptr.*) orelse @panic("Modified Time Missing");
+                const modified_time = entry.value_ptr.*;
-                if (self.assetman.didUpdate(asset_manifest.getPath(entry.key_ptr.*), modified_time)) {
+                const asset_path = asset_manifest.getPath(entry.key_ptr.*);
+
+                // Might happen for shader permuted assets
+                if (asset_path.len == 0) continue;
+                if (self.assetman.didUpdate(asset_path, modified_time)) {
                     try updatedList.append(self.fba.allocator(), entry.key_ptr.*);
                 }
             }
@@ -183,20 +185,43 @@ fn resolveAsset(self: *AssetManager, handle: AssetId) ?*LoadedAsset {
     return self.loaded_assets.getPtr(handle);
 }
 
-pub fn resolveShader(self: *AssetManager, handle: Handle.Shader) LoadedShader {
+const DefinePair = struct {
+    key: []const u8,
+    value: []const u8,
+};
+
+fn permuteAssetIdDefines(id: AssetId, defines: []const DefinePair) AssetId {
+    var hash = std.hash.Wyhash.init(id);
+    hash.update(&std.mem.toBytes(id));
+    for (defines) |def| {
+        hash.update(def.key);
+        hash.update(def.value);
+    }
+    return hash.final();
+}
+
+pub fn resolveShaderWithDefines(self: *AssetManager, handle: Handle.Shader, defines: []const DefinePair) LoadedShader {
     if (handle.id == 0) return NullShader;
 
-    if (self.resolveAsset(handle.id)) |asset| {
+    const permuted_asset_id = permuteAssetIdDefines(handle.id, defines);
+
+    if (self.resolveAsset(permuted_asset_id)) |asset| {
         return asset.shader;
     }
 
-    return self.loadShader(handle.id);
+    return self.loadShader(handle.id, permuted_asset_id, defines);
 }
 
 pub fn resolveShaderProgram(self: *AssetManager, handle: Handle.ShaderProgram) LoadedShaderProgram {
+    return self.resolveShaderProgramWithDefines(handle, &.{});
+}
+
+pub fn resolveShaderProgramWithDefines(self: *AssetManager, handle: Handle.ShaderProgram, defines: []const DefinePair) LoadedShaderProgram {
     if (handle.id == 0) return NullShaderProgram;
 
-    if (self.resolveAsset(handle.id)) |asset| {
+    const permuted_asset_id = permuteAssetIdDefines(handle.id, defines);
+
+    if (self.resolveAsset(permuted_asset_id)) |asset| {
         switch (asset.*) {
             .shaderProgram => |shader| {
                 return shader;
@@ -205,7 +230,7 @@ pub fn resolveShaderProgram(self: *AssetManager, handle: Handle.ShaderProgram) L
         }
     }
 
-    return self.loadShaderProgram(handle);
+    return self.loadShaderProgram(handle, permuted_asset_id, defines);
 }
 
 pub fn resolveMesh(self: *AssetManager, handle: Handle.Mesh) LoadedMesh {
@@ -281,15 +306,15 @@ pub const ShaderProgramDefinition = struct {
     fragment: []const u8,
 };
 
-pub fn loadShaderProgram(self: *AssetManager, handle: Handle.ShaderProgram) LoadedShaderProgram {
+pub fn loadShaderProgram(self: *AssetManager, handle: Handle.ShaderProgram, permuted_id: AssetId, defines: []const DefinePair) LoadedShaderProgram {
-    return self.loadShaderProgramErr(handle.id) catch |err| {
+    return self.loadShaderProgramErr(handle.id, permuted_id, defines) catch |err| {
         std.log.err("Failed to load shader program {}\n", .{err});
 
         return NullShaderProgram;
     };
 }
 
-fn loadShaderProgramErr(self: *AssetManager, id: AssetId) !LoadedShaderProgram {
+fn loadShaderProgramErr(self: *AssetManager, id: AssetId, permuted_id: AssetId, defines: []const DefinePair) !LoadedShaderProgram {
     const data = try self.loadFile(self.frame_arena, asset_manifest.getPath(id), SHADER_MAX_BYTES);
     const program = formats.ShaderProgram.fromBuffer(data.bytes);
 
@@ -300,11 +325,7 @@ fn loadShaderProgramErr(self: *AssetManager, id: AssetId) !LoadedShaderProgram {
 
     // TODO: !!! this will keep shader source in memory as long as shader program is in memory
     // probably don't want this!
-    const shader = self.resolveShader(program.shader);
+    const shader = self.resolveShaderWithDefines(program.shader, defines);
-
-    // TODO: !!! Will evict shader program if shader source is evicted. Only want this for watch changes, not
-    // normal eviction!
-    try self.addDependencies(id, &.{program.shader.id});
 
     const prog = gl.createProgram();
     errdefer gl.deleteProgram(prog);
@@ -353,16 +374,19 @@ fn loadShaderProgramErr(self: *AssetManager, id: AssetId) !LoadedShaderProgram {
 
     const loaded_shader_program = LoadedShaderProgram{
         .program = prog,
+        .permuted_id = permuted_id,
     };
 
     {
         self.rw_lock.lock();
         defer self.rw_lock.unlock();
 
-        try self.loaded_assets.put(self.allocator, id, .{
+        try self.loaded_assets.put(self.allocator, permuted_id, .{
             .shaderProgram = loaded_shader_program,
         });
         try self.modified_times.put(self.allocator, id, data.modified);
+
+        try self.addDependencies(permuted_id, &.{ id, shader.permuted_id });
     }
 
     return loaded_shader_program;
@@ -370,10 +394,12 @@ fn loadShaderProgramErr(self: *AssetManager, id: AssetId) !LoadedShaderProgram {
 
 const NullShader = LoadedShader{
     .source = "",
+    .permuted_id = 0,
 };
 
 const NullShaderProgram = LoadedShaderProgram{
     .program = 0,
+    .permuted_id = 0,
 };
 
 const NullMesh = LoadedMesh{
@@ -658,10 +684,12 @@ const LoadedAsset = union(enum) {
 
 const LoadedShader = struct {
     source: []const u8,
+    permuted_id: AssetId,
 };
 
 const LoadedShaderProgram = struct {
     program: gl.GLuint,
+    permuted_id: AssetId,
 };
 
 pub const LoadedMesh = struct {
@@ -778,8 +806,8 @@ fn loadFile(self: *AssetManager, allocator: std.mem.Allocator, path: []const u8,
     return .{ .bytes = bytes, .modified = meta.modified() };
 }
 
-fn loadShader(self: *AssetManager, id: AssetId) LoadedShader {
+fn loadShader(self: *AssetManager, id: AssetId, permuted_id: AssetId, defines: []const DefinePair) LoadedShader {
-    return self.loadShaderErr(id) catch |err| {
+    return self.loadShaderErr(id, permuted_id, defines) catch |err| {
         std.log.err("Error: {} when loading shader id {} {s}", .{ err, id, asset_manifest.getPath(id) });
         return NullShader;
     };
@@ -1320,7 +1348,7 @@ test "ShaderTokenizer" {
     }
 }
 
-fn loadShaderErr(self: *AssetManager, id: AssetId) !LoadedShader {
+fn loadShaderErr(self: *AssetManager, id: AssetId, permuted_id: AssetId, defines: []const DefinePair) !LoadedShader {
     const path = asset_manifest.getPath(id);
     const dir = std.fs.path.dirname(path) orelse @panic("No dir");
 
@@ -1330,6 +1358,13 @@ fn loadShaderErr(self: *AssetManager, id: AssetId) !LoadedShader {
     var preprocessed_segments = std.SegmentedList([]const u8, 64){};
     var final_len: usize = 0;
 
+    // Just append defines here, no need to manually preprocess
+    for (defines) |define| {
+        const define_str = try std.fmt.allocPrint(self.frame_arena, "#define {s} {s}\n", .{ define.key, define.value });
+        try preprocessed_segments.append(self.frame_arena, define_str);
+        final_len += define_str.len;
+    }
+
     // Preprocess
     {
         var tokenizer = ShaderTokenizer.init(data.bytes);
@@ -1341,7 +1376,6 @@ fn loadShaderErr(self: *AssetManager, id: AssetId) !LoadedShader {
             switch (token.type) {
                 .Directive => {
                     if (std.mem.eql(u8, token.text, "include")) {
-
                         // Append section of text up to this directive
                         try preprocessed_segments.append(self.frame_arena, data.bytes[last_offset..token.start]);
                         final_len += token.start - last_offset;
@@ -1358,8 +1392,8 @@ fn loadShaderErr(self: *AssetManager, id: AssetId) !LoadedShader {
 
                         const included_asset_id = assets.AssetPath.fromString(included_file_path).hash();
                         if (included_asset_id != 0) {
-                            try included_asset_ids.append(included_asset_id);
+                            const included_shader = self.resolveShaderWithDefines(.{ .id = included_asset_id }, defines);
-                            const included_shader = try self.loadShaderErr(included_asset_id);
+                            try included_asset_ids.append(included_shader.permuted_id);
                             try preprocessed_segments.append(self.frame_arena, included_shader.source);
                             final_len += included_shader.source.len;
                         }
@@ -1391,15 +1425,16 @@ fn loadShaderErr(self: *AssetManager, id: AssetId) !LoadedShader {
         }
     }
 
-    const loaded_shader = LoadedShader{ .source = result_source };
+    const loaded_shader = LoadedShader{ .source = result_source, .permuted_id = permuted_id };
     {
         self.rw_lock.lock();
         defer self.rw_lock.unlock();
 
-        try self.loaded_assets.put(self.allocator, id, .{ .shader = loaded_shader });
+        try self.loaded_assets.put(self.allocator, permuted_id, .{ .shader = loaded_shader });
         try self.modified_times.put(self.allocator, id, data.modified);
 
-        try self.addDependencies(id, included_asset_ids.items);
+        try self.addDependencies(permuted_id, &.{id});
+        try self.addDependencies(permuted_id, included_asset_ids.items);
     }
 
     return loaded_shader;

src/Render.zig

@@ -1034,7 +1034,7 @@ pub fn finish(self: *Render) void {
 
     // Main pass
     {
-        gl.useProgram(self.assetman.resolveShaderProgram(a.ShaderPrograms.shaders.mesh).program);
+        gl.useProgram(self.assetman.resolveShaderProgramWithDefines(a.ShaderPrograms.shaders.mesh, &.{.{ .key = "OVERRIDE_COLOR", .value = "1" }}).program);
         gl.bindVertexArray(self.mesh_vao);
         gl.depthFunc(gl.EQUAL);
         defer gl.depthFunc(gl.LEQUAL);