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](#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.

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,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);
+  }
+
+  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.