shadertools: Resolve dependencies in ShaderInputs (#2073)

modules/engine/src/shader-inputs.ts

@@ -61,17 +61,22 @@ export class ShaderInputs<
    * @param modules
    */
   // @ts-expect-error Fix typings
-  constructor(modules: {[P in keyof ShaderPropsT]: ShaderModuleInputs<ShaderPropsT[P]>}) {
-    // TODO - get all dependencies from modules
-    const allModules = _resolveModules(Object.values(modules));
-    log.log(
-      1,
-      'Creating ShaderInputs with modules',
-      allModules.map(m => m.name)
-    )();
+  constructor(modules: {[P in keyof ShaderPropsT]?: ShaderModuleInputs<ShaderPropsT[P]>}) {
+    // Extract modules with dependencies
+    const resolvedModules = _resolveModules(
+      Object.values(modules).filter(module => module.dependencies),
+      false
+    );
+    for (const resolvedModule of resolvedModules) {
+      // @ts-expect-error Fix typings
+      modules[resolvedModule.name] = resolvedModule;
+    }
+
+    log.log(1, 'Creating ShaderInputs with modules', Object.keys(modules))();
 
     // Store the module definitions and create storage for uniform values and binding values, per module
-    this.modules = modules;
+    // @ts-expect-error Fix typings
+    this.modules = modules as {[P in keyof ShaderPropsT]: ShaderModuleInputs<ShaderPropsT[P]>};
     this.moduleUniforms = {} as Record<keyof ShaderPropsT, Record<string, UniformValue>>;
     this.moduleBindings = {} as Record<keyof ShaderPropsT, Record<string, Texture | Sampler>>;
 

modules/engine/test/shader-inputs.spec.ts

@@ -6,6 +6,7 @@ import test from 'tape-promise/tape';
 import {picking} from '../../shadertools/src/index';
 // import {_ShaderInputs as ShaderInputs} from '@luma.gl/engine';
 import {ShaderInputs} from '../src/shader-inputs';
+import {ShaderModule} from '@luma.gl/shadertools';
 
 test('ShaderInputs#picking', t => {
   const shaderInputsUntyped = new ShaderInputs({picking});
@@ -62,3 +63,33 @@ test('ShaderInputs#picking prop merge', t => {
 
   t.end();
 });
+
+test('ShaderInputs#dependencies', t => {
+  type CustomProps = {color: number[]};
+  const custom: ShaderModule<CustomProps> = {
+    name: 'custom',
+    dependencies: [picking],
+    uniformTypes: {
+      color: 'vec3<f32>'
+    }
+  };
+
+  const shaderInputs = new ShaderInputs<{
+    custom: CustomProps;
+    picking: typeof picking.props;
+  }>({custom});
+  t.deepEqual(Object.keys(shaderInputs.modules), ['custom', 'picking']);
+
+  shaderInputs.setProps({
+    custom: {color: [255, 0, 0]},
+    picking: {highlightedObjectColor: [1, 2, 3]}
+  });
+  t.deepEqual(shaderInputs.moduleUniforms.custom.color, [255, 0, 0], 'custom color updated');
+  t.deepEqual(
+    shaderInputs.moduleUniforms.picking.highlightedObjectColor,
+    [1, 2, 3],
+    'highlight object color updated'
+  );
+
+  t.end();
+});

modules/shadertools/src/lib/shader-assembly/resolve-modules.ts

@@ -6,15 +6,34 @@ import type {ShaderModule} from '../shader-module/shader-module';
 import {ShaderModuleInstance} from '../shader-module/shader-module-instance';
 
 /**
- * Instantiate shader modules and esolve any dependencies
+ * Resolve any dependencies and optionally instantiate modules
  */
 export function resolveModules(
   modules: (ShaderModule | ShaderModuleInstance)[]
-): ShaderModuleInstance[] {
-  const instances = ShaderModuleInstance.instantiateModules(modules);
-  return getShaderDependencies(instances);
+): ShaderModuleInstance[];
+export function resolveModules(
+  modules: (ShaderModule | ShaderModuleInstance)[],
+  instantiate: true
+): ShaderModuleInstance[];
+export function resolveModules(modules: ShaderModule[], instantiate: false): ShaderModule[];
+export function resolveModules(
+  modules: ShaderModuleInstance[],
+  instantiate: false
+): ShaderModuleInstance[];
+export function resolveModules(
+  modules: (ShaderModule | ShaderModuleInstance)[],
+  instantiate: boolean = true
+): (ShaderModule | ShaderModuleInstance)[] {
+  return getShaderDependencies(
+    instantiate === false ? modules : ShaderModuleInstance.instantiateModules(modules)
+  );
 }
 
+type AbstractModule = {
+  name: string;
+  dependencies?: AbstractModule[];
+};
+
 /**
  * Takes a list of shader module names and returns a new list of
  * shader module names that includes all dependencies, sorted so
@@ -27,8 +46,8 @@ export function resolveModules(
  * @param modules - Array of modules (inline modules or module names)
  * @return - Array of modules
  */
-function getShaderDependencies(modules: ShaderModuleInstance[]): ShaderModuleInstance[] {
-  const moduleMap: Record<string, ShaderModuleInstance> = {};
+function getShaderDependencies<T extends AbstractModule>(modules: T[]): T[] {
+  const moduleMap: Record<string, T> = {};
   const moduleDepth: Record<string, number> = {};
   getDependencyGraph({modules, level: 0, moduleMap, moduleDepth});
 
@@ -48,10 +67,10 @@ function getShaderDependencies(modules: ShaderModuleInstance[]): ShaderModuleIns
  * @return - Map of module name to its level
  */
 // Adds another level of dependencies to the result map
-export function getDependencyGraph(options: {
-  modules: ShaderModuleInstance[];
+export function getDependencyGraph<T extends AbstractModule>(options: {
+  modules: T[];
   level: number;
-  moduleMap: Record<string, ShaderModuleInstance>;
+  moduleMap: Record<string, T>;
   moduleDepth: Record<string, number>;
 }) {
   const {modules, level, moduleMap, moduleDepth} = options;