feat: Refactor shader uniform binding to support shader arrays

Author: luanpotter

packages/flame_3d/lib/src/resources/material/material.dart

@@ -15,8 +15,8 @@ abstract class Material extends Resource<gpu.RenderPipeline> {
         _fragmentShader = fragmentShader,
         super(
           gpu.gpuContext.createRenderPipeline(
-            vertexShader.resource,
-            fragmentShader.resource,
+            vertexShader.compile().resource,
+            fragmentShader.compile().resource,
           ),
         );
 
@@ -25,8 +25,8 @@ abstract class Material extends Resource<gpu.RenderPipeline> {
     var resource = super.resource;
     if (_recreateResource) {
       resource = super.resource = gpu.gpuContext.createRenderPipeline(
-        _vertexShader.resource,
-        _fragmentShader.resource,
+        _vertexShader.compile().resource,
+        _fragmentShader.compile().resource,
       );
       _recreateResource = false;
     }

packages/flame_3d/lib/src/resources/material/spatial_material.dart

@@ -3,7 +3,6 @@ import 'dart:ui';
 import 'package:flame_3d/game.dart';
 import 'package:flame_3d/graphics.dart';
 import 'package:flame_3d/resources.dart';
-import 'package:flutter_gpu/gpu.dart' as gpu;
 
 class SpatialMaterial extends Material {
   SpatialMaterial({
@@ -14,13 +13,13 @@ class SpatialMaterial extends Material {
   })  : albedoTexture = albedoTexture ?? Texture.standard,
         super(
           vertexShader: Shader(
-            _library['TextureVertex']!,
+            name: 'TextureVertex',
             slots: [
               UniformSlot.value('VertexInfo', {'model', 'view', 'projection'}),
             ],
           ),
           fragmentShader: Shader(
-            _library['TextureFragment']!,
+            name: 'TextureFragment',
             slots: [
               UniformSlot.sampler('albedoTexture'),
               UniformSlot.value('Material', {
@@ -85,8 +84,4 @@ class SpatialMaterial extends Material {
   void _applyLights(GraphicsDevice device) {
     device.lightingInfo.apply(fragmentShader);
   }
-
-  static final _library = gpu.ShaderLibrary.fromAsset(
-    'packages/flame_3d/assets/shaders/spatial_material.shaderbundle',
-  )!;
 }

packages/flame_3d/lib/src/resources/shader.dart

@@ -1,4 +1,5 @@
 export 'shader/shader.dart';
+export 'shader/uniform_array.dart';
 export 'shader/uniform_instance.dart';
 export 'shader/uniform_sampler.dart';
 export 'shader/uniform_slot.dart';

packages/flame_3d/lib/src/resources/shader/shader.dart

@@ -1,4 +1,3 @@
-import 'dart:collection';
 import 'dart:typed_data';
 import 'dart:ui';
 
@@ -10,24 +9,48 @@ import 'package:flutter_gpu/gpu.dart' as gpu;
 /// {@template shader}
 ///
 /// {@endtemplate}
-class Shader extends Resource<gpu.Shader> {
+class ShaderResource extends Resource<gpu.Shader> {
+  final Shader shader;
+
   /// {@macro shader}
-  Shader(
+  ShaderResource(
     super.resource, {
+    required String name,
     List<UniformSlot> slots = const [],
-  })  : _slots = slots,
-        _instances = {} {
+  }) : shader = Shader(name: name, slots: slots) {
     for (final slot in slots) {
       slot.resource = resource.getUniformSlot(slot.name);
     }
   }
 
-  final List<UniformSlot> _slots;
+  factory ShaderResource.create({
+    required String name,
+    required List<UniformSlot> slots,
+  }) {
+    final shader = _library[name];
+    if (shader == null) {
+      throw StateError('Shader "$name" not found in library');
+    }
+    return ShaderResource(shader, name: name, slots: slots);
+  }
+
+  static final _library = gpu.ShaderLibrary.fromAsset(
+    'packages/flame_3d/assets/shaders/spatial_material.shaderbundle',
+  )!;
+}
 
-  final Map<String, UniformInstance> _instances;
+class Shader {
+  final String name;
+  final List<UniformSlot> slots;
+  final Map<String, UniformInstance> instances = {};
+
+  Shader({
+    required this.name,
+    required this.slots,
+  });
 
   /// Set a [Texture] at the given [key] on the buffer.
-  void setTexture(String key, Texture texture) => _setSampler(key, texture);
+  void setTexture(String key, Texture texture) => _setTypedValue(key, texture);
 
   /// Set a [Vector2] at the given [key] on the buffer.
   void setVector2(String key, Vector2 vector) => _setValue(key, vector.storage);
@@ -45,7 +68,7 @@ class Shader extends Resource<gpu.Shader> {
 
   /// Set a [double] at the given [key] on the buffer.
   void setFloat(String key, double value) {
-    _setValue(key, [value]);
+    _setValue(key, _encodeFloat32(value));
   }
 
   /// Set a [Matrix2] at the given [key] on the buffer.
@@ -60,50 +83,63 @@ class Shader extends Resource<gpu.Shader> {
   void setColor(String key, Color color) => _setValue(key, color.storage);
 
   void bind(GraphicsDevice device) {
-    for (final slot in _slots) {
-      _instances[slot.name]?.bind(device);
+    for (final slot in slots) {
+      instances[slot.name]?.bind(device);
     }
   }
 
   /// Set the [data] to the [UniformSlot] identified by [key].
-  void _setValue(String key, List<double> data) {
-    final (uniform, field) = _getInstance<UniformValue>(key);
-    uniform[field!] = data;
+  void _setValue(String key, Float32List data) {
+    _setTypedValue(key, data.buffer);
   }
 
-  void _setSampler(String key, Texture data) {
-    final (uniform, _) = _getInstance<UniformSampler>(key);
-    uniform.resource = data;
+  List<String?> parseKey(String key) {
+    // examples: albedoTexture, Light[2].position, or Foo.bar
+    final regex = RegExp(r'^(\w+)(?:\[(\d+)\])?(?:\.(\w+))?$');
+    return regex.firstMatch(key)?.groups([1, 2, 3]) ?? [];
   }
 
   /// Get the slot for the [key], it only calculates it once for every unique
   /// [key].
-  (T, String?) _getInstance<T extends UniformInstance>(String key) {
-    final keys = key.split('.');
-
-    // Check if we already have a uniform instance created.
-    if (!_instances.containsKey(keys.first)) {
-      // If the slot or it's property isn't mapped in the uniform it will be
-      // enforced.
-      final slot = _slots.firstWhere(
-        (e) => e.name == keys.first,
-        orElse: () => throw StateError('Uniform "$key" is unmapped'),
-      );
+  void _setTypedValue<K, T>(String key, T value) {
+    final groups = parseKey(key);
 
-      final instance = slot.create();
-      if (instance is UniformValue &&
-          keys.length > 1 &&
-          !slot.fields.contains(keys[1])) {
-        throw StateError('Field "${keys[1]}" is unmapped for "${keys.first}"');
-      }
+    final object = groups[0]; // e.g. Light, albedoTexture
+    final idx = _maybeParseInt(groups[1]); // e.g. 2 (optional)
+    final field = groups[2]; // e.g. position (optional)
 
-      _instances[slot.name] = instance;
+    if (object == null) {
+      throw StateError('Uniform "$key" is missing an object');
     }
 
-    return (_instances[keys.first], keys.elementAtOrNull(1)) as (T, String?);
+    final instance = instances.putIfAbsent(object, () {
+      final slot = slots.firstWhere(
+        (e) => e.name == object,
+        orElse: () => throw StateError('Uniform "$object" is unmapped'),
+      );
+      return slot.create();
+    }) as UniformInstance<K, T>;
+
+    final k = instance.makeKey(idx, field);
+    instance.set(k, value);
   }
 
   static Float32List _encodeUint32(int value, Endian endian) {
     return (ByteData(16)..setUint32(0, value, endian)).buffer.asFloat32List();
   }
+
+  static Float32List _encodeFloat32(double value) {
+    return Float32List.fromList([value]);
+  }
+
+  static int? _maybeParseInt(String? value) {
+    if (value == null) {
+      return null;
+    }
+    return int.parse(value);
+  }
+
+  ShaderResource compile() {
+    return ShaderResource.create(name: name, slots: slots);
+  }
 }

packages/flame_3d/lib/src/resources/shader/uniform_array.dart

@@ -0,0 +1,89 @@
+import 'dart:collection';
+import 'dart:typed_data';
+
+import 'package:flame_3d/graphics.dart';
+import 'package:flame_3d/resources.dart';
+
+typedef UniformArrayKey = ({
+  int idx,
+  String field,
+});
+
+/// {@template uniform_value}
+/// Instance of a uniform array. Represented by a [ByteBuffer].
+/// {@endtemplate}
+class UniformArray extends UniformInstance<UniformArrayKey, ByteBuffer> {
+  /// {@macro uniform_value}
+  UniformArray(super.slot);
+
+  final List<Map<int, ({int hash, List<double> data})>> _storage = [];
+
+  @override
+  ByteBuffer? get resource {
+    if (super.resource == null) {
+      final data = <double>[];
+      for (final element in _storage) {
+        var previousIndex = -1;
+        for (final entry in element.entries) {
+          if (previousIndex + 1 != entry.key) {
+            final field = slot.fields.indexed
+                .firstWhere((e) => e.$1 == previousIndex + 1);
+            throw StateError(
+              'Uniform ${slot.name}.${field.$2} was not set',
+            );
+          }
+          previousIndex = entry.key;
+          data.addAll(entry.value.data);
+        }
+      }
+      super.resource = Float32List.fromList(data).buffer;
+    }
+
+    return super.resource;
+  }
+
+  Map<int, ({int hash, List<double> data})> _get(int idx) {
+    while (idx >= _storage.length) {
+      _storage.add(HashMap());
+    }
+    return _storage[idx];
+  }
+
+  List<double>? get(int idx, String key) => _get(idx)[slot.indexOf(key)]?.data;
+
+  @override
+  void set(UniformArrayKey key, ByteBuffer buffer) {
+    final storage = _get(key.idx);
+    final index = slot.indexOf(key.field);
+
+    // Ensure that we are only setting new data if the hash has changed.
+    final data = buffer.asFloat32List();
+    final hash = Object.hashAll(data);
+    if (storage[index]?.hash == hash) {
+      return;
+    }
+
+    // Store the storage at the given slot index.
+    storage[index] = (data: data, hash: hash);
+
+    // Clear the cache.
+    super.resource = null;
+  }
+
+  @override
+  UniformArrayKey makeKey(int? idx, String? field) {
+    if (idx == null) {
+      throw StateError('idx is required for ${slot.name}');
+    }
+    if (field == null) {
+      throw StateError('field is required for ${slot.name}');
+    }
+
+    return (idx: idx, field: field);
+  }
+
+  @override
+  void bind(GraphicsDevice device) {
+    device.bindUniform(slot.resource!, resource!);
+  }
+}

packages/flame_3d/lib/src/resources/shader/uniform_instance.dart

@@ -5,12 +5,16 @@ import 'package:flame_3d/resources.dart';
 /// An instance of a [UniformSlot] that can cache the [resource] that will be
 /// bound to a [Shader].
 /// {@endtemplate}
-abstract class UniformInstance<T> extends Resource<T?> {
+abstract class UniformInstance<K, T> extends Resource<T?> {
   /// {@macro uniform_instance}
   UniformInstance(this.slot) : super(null);
 
   /// The slot this instance belongs too.
   final UniformSlot slot;
 
   void bind(GraphicsDevice device);
+
+  void set(K key, T value);
+
+  K makeKey(int? idx, String? field);
 }

packages/flame_3d/lib/src/resources/shader/uniform_sampler.dart

@@ -4,12 +4,20 @@ import 'package:flame_3d/resources.dart';
 /// {@template uniform_sampler}
 /// Instance of a uniform sampler. Represented by a [Texture].
 /// {@endtemplate}
-class UniformSampler extends UniformInstance<Texture> {
+class UniformSampler extends UniformInstance<void, Texture> {
   /// {@macro uniform_sampler}
   UniformSampler(super.slot);
 
   @override
   void bind(GraphicsDevice device) {
     device.bindTexture(slot.resource!, resource!);
   }
+
+  @override
+  void set(void key, Texture value) {
+    resource = value;
+  }
+
+  @override
+  void makeKey(int? idx, String? field) {}
 }

packages/flame_3d/lib/src/resources/shader/uniform_slot.dart

@@ -21,6 +21,14 @@ class UniformSlot extends Resource<gpu.UniformSlot?> {
   UniformSlot.value(String name, Set<String> fields)
       : this._(name, fields, UniformValue.new);
 
+  /// {@macro uniform_slot}
+  ///
+  /// Used for array uniforms in shaders.
+  ///
+  /// The [fields] should be defined in order as they appear in the struct.
+  UniformSlot.array(String name, Set<String> fields)
+      : this._(name, fields, UniformArray.new);
+
   /// {@macro uniform_slot}
   ///
   /// Used for sampler uniforms in shaders.