feat: Refactor shader uniform binding to support shader arrays [flame_3d] (#3282)

Refactor shader uniform binding to support shader arrays.

This also decouples the whole shader and uniform byte handling code
(that we should definitely test) from the flutter_gpu primitives that
are impossible to mock (base native classes).

This adds tests that ensure the arrays are bound as they should -
however the underlying flutter_gpu code does not seem to work. See [this
PR](https://github.com/flame-engine/flame/pull/3284) for a test of using
this to support an arbitrary number of lights.

Either way, we can merge this as is as this refactors the underlying
structure to support arrays when ready, and make it more testable as
well.

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,7 +13,7 @@ class SpatialMaterial extends Material {
   })  : albedoTexture = albedoTexture ?? Texture.standard,
         super(
           vertexShader: Shader(
-            _library['TextureVertex']!,
+            name: 'TextureVertex',
             slots: [
               UniformSlot.value('VertexInfo', {
                 'model',
@@ -28,7 +27,7 @@ class SpatialMaterial extends Material {
             ],
           ),
           fragmentShader: Shader(
-            _library['TextureFragment']!,
+            name: 'TextureFragment',
             slots: [
               UniformSlot.sampler('albedoTexture'),
               UniformSlot.value('Material', {
@@ -108,9 +107,5 @@ class SpatialMaterial extends Material {
     device.lightingInfo.apply(fragmentShader);
   }
 
-  static final _library = gpu.ShaderLibrary.fromAsset(
-    'packages/flame_3d/assets/shaders/spatial_material.shaderbundle',
-  )!;
-
   static const _maxJoints = 16;
 }

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);
+  }
 
-  final Map<String, UniformInstance> _instances;
+  static final _library = gpu.ShaderLibrary.fromAsset(
+    'packages/flame_3d/assets/shaders/spatial_material.shaderbundle',
+  )!;
+}
+
+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('.');
+  void _setTypedValue<K, T>(String key, T value) {
+    final groups = parseKey(key);
 
-    // 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'),
-      );
+    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)
 
-      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}"');
-      }
-
-      _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,7 +5,7 @@ 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);
 
@@ -13,4 +13,8 @@ abstract class UniformInstance<T> extends Resource<T?> {
   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,7 +4,7 @@ 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);
 
@@ -12,4 +12,12 @@ class UniformSampler extends UniformInstance<Texture> {
   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.

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

@@ -11,11 +11,11 @@ import 'package:ordered_set/comparing.dart';
 /// The `[]` operator can be used to set the raw data of a field. If the data is
 /// different from the last set it will recalculated the [resource].
 /// {@endtemplate}
-class UniformValue extends UniformInstance<ByteBuffer> {
+class UniformValue extends UniformInstance<String, ByteBuffer> {
   /// {@macro uniform_value}
   UniformValue(super.slot);
 
-  final Map<int, ({int hash, List<double> data})> _storage = HashMap();
+  final Map<int, ({int hash, Float32List data})> _storage = HashMap();
 
   @override
   ByteBuffer? get resource {
@@ -40,9 +40,9 @@ class UniformValue extends UniformInstance<ByteBuffer> {
     return super.resource;
   }
 
-  List<double>? operator [](String key) => _storage[slot.indexOf(key)]?.data;
+  Float32List? operator [](String key) => _storage[slot.indexOf(key)]?.data;
 
-  void operator []=(String key, List<double> data) {
+  void operator []=(String key, Float32List data) {
     final index = slot.indexOf(key);
 
     // Ensure that we are only setting new data if the hash has changed.
@@ -58,8 +58,28 @@ class UniformValue extends UniformInstance<ByteBuffer> {
     super.resource = null;
   }
 
+  @override
+  String makeKey(int? idx, String? field) {
+    if (idx != null) {
+      throw StateError('idx is not supported for ${slot.name}');
+    }
+    if (field == null) {
+      throw StateError('field is required for ${slot.name}');
+    }
+
+    return field;
+  }
+
   @override
   void bind(GraphicsDevice device) {
     device.bindUniform(slot.resource!, resource!);
   }
+
+  @override
+  void set(String key, ByteBuffer value) {
+    if (!slot.fields.contains(key)) {
+      throw StateError('Field "$key" is unmapped for "${slot.name}"');
+    }
+    this[key] = value.asFloat32List();
+  }
 }

packages/flame_3d/test/resources/shader/uniform_binding_test.dart

@@ -0,0 +1,113 @@
+import 'dart:typed_data';
+
+import 'package:flame_3d/core.dart';
+import 'package:flame_3d/resources.dart';
+import 'package:flame_3d/src/resources/shader.dart';
+import 'package:flutter_test/flutter_test.dart';
+
+void main() {
+  group('uniform bindings', () {
+    test('can bind a vec3 a slot', () {
+      final slot = UniformSlot.value('Vertex', {'position'});
+      final shader = _createShader([slot]);
+
+      shader.setVector3('Vertex.position', Vector3(7, 8, 9));
+
+      final bytes = shader.instances['Vertex']!.resource as ByteBuffer;
+      final result = Vector3.fromBuffer(bytes, 0);
+      expect(result, Vector3(7, 8, 9));
+    });
+
+    test('can bind multiple vector slots', () {
+      final slot = UniformSlot.value('AmbientLight', {'color', 'position'});
+      final shader = _createShader([slot]);
+
+      shader.setVector3('AmbientLight.position', Vector3(7, 8, 9));
+      shader.setVector4('AmbientLight.color', Vector4(4, 3, 2, 1));
+
+      final bytes = shader.instances['AmbientLight']!.resource as ByteBuffer;
+
+      final color = Vector4.fromBuffer(bytes, 0);
+      expect(color, Vector4(4, 3, 2, 1));
+
+      final position = Vector3.fromBuffer(bytes, color.storage.lengthInBytes);
+      expect(position, Vector3(7, 8, 9));
+    });
+
+    test('can bind a mat4 a slot', () {
+      final slot = UniformSlot.value('Vertex', {'camera'});
+      final shader = _createShader([slot]);
+
+      shader.setMatrix4('Vertex.camera', Matrix4.identity());
+
+      final bytes = shader.instances['Vertex']!.resource as ByteBuffer;
+      final result = Matrix4.fromBuffer(bytes, 0);
+      expect(result, Matrix4.identity());
+    });
+
+    test('can bind a vec3 to an array slot', () {
+      final slot = UniformSlot.array('Light', {'position'});
+      final shader = _createShader([slot]);
+
+      shader.setVector3('Light[0].position', Vector3(7, 8, 9));
+
+      final bytes = shader.instances['Light']!.resource as ByteBuffer;
+      final result = Vector3.fromBuffer(bytes, 0);
+
+      expect(result, Vector3(7, 8, 9));
+    });
+
+    test('can bind multiple slots', () {
+      final slots = [
+        UniformSlot.value('Vertex', {'position'}),
+        UniformSlot.value('Material', {'color', 'metallic'}),
+        UniformSlot.array('Light', {'position', 'color'}),
+      ];
+      final shader = _createShader(slots);
+
+      shader.setVector3('Vertex.position', Vector3(1, 2, 3));
+      shader.setVector4('Material.color', Vector4(4, 3, 2, 1));
+      shader.setFloat('Material.metallic', 0.5);
+      shader.setVector3('Light[0].position', Vector3(11, 12, 13));
+      shader.setVector4('Light[0].color', Vector4(14, 15, 16, 17));
+      shader.setVector3('Light[1].position', Vector3(-1, -2, -3));
+      shader.setVector4('Light[1].color', Vector4(-11, -12, -13, -14));
+
+      final vertex = shader.instances['Vertex']!.resource as ByteBuffer;
+      final vertexResult = Vector3.fromBuffer(vertex, 0);
+      expect(vertexResult, Vector3(1, 2, 3));
+
+      final material = shader.instances['Material']!.resource as ByteBuffer;
+      final color = Vector4.fromBuffer(material, 0);
+      expect(color, Vector4(4, 3, 2, 1));
+      final metallic = Float32List.view(material, color.storage.lengthInBytes);
+      expect(metallic[0], 0.5);
+
+      final light = shader.instances['Light']!.resource as ByteBuffer;
+
+      var cursor = 0;
+
+      final light0Position = Vector3.fromBuffer(light, cursor);
+      expect(light0Position, Vector3(11, 12, 13));
+      cursor += light0Position.storage.lengthInBytes;
+
+      final light0Color = Vector4.fromBuffer(light, cursor);
+      expect(light0Color, Vector4(14, 15, 16, 17));
+      cursor += light0Color.storage.lengthInBytes;
+
+      final light1Position = Vector3.fromBuffer(light, cursor);
+      expect(light1Position, Vector3(-1, -2, -3));
+      cursor += light1Position.storage.lengthInBytes;
+
+      final light1Color = Vector4.fromBuffer(light, cursor);
+      expect(light1Color, Vector4(-11, -12, -13, -14));
+    });
+  });
+}
+
+Shader _createShader(List<UniformSlot> slots) {
+  return Shader(
+    name: '-test-',
+    slots: slots,
+  );
+}