fix: Lookup types from function nodes and use those when transpiling

Author: robertleeplummerjr

src/backend/web-gl/function-node.js

@@ -660,14 +660,29 @@ module.exports = class WebGLFunctionNode extends FunctionNodeBase {
 					break;
 				default:
 					if (i === 0) {
-						if (declaration.init && declaration.init.name && this.declarations[declaration.init.name]) {
-							type = this.declarations[declaration.init.name];
-							retArr.push(type + ' ');
-						} else if (declaration.init && declaration.init.type && declaration.init.type === 'ArrayExpression') {
-							type = 'vec' + declaration.init.elements.length;
-							retArr.push(type + ' ');
-						} else {
-							retArr.push('float ');
+						if (declaration.init) {
+							if (declaration.init.name && this.declarations[declaration.init.name]) {
+								type = this.declarations[declaration.init.name];
+								retArr.push(type + ' ');
+							} else if (declaration.init.type) {
+								switch (declaration.init.type) {
+									case 'ArrayExpression':
+										type = 'vec' + declaration.init.elements.length;
+										retArr.push(type + ' ');
+										break;
+									case 'CallExpression':
+										const node = this.builder.nodeMap[declaration.init.callee.name];
+										if (node && node.returnType) {
+											type = node.returnType;
+											retArr.push(type + ' ');
+										}
+										break;
+									default:
+										retArr.push('float ');
+								}
+							} else {
+								retArr.push('float ');
+							}
 						}
 					}
 					break;

test/features/add-typed-functions.js

@@ -14,9 +14,76 @@
       .setOutput([1]);
     var result = kernel();
     QUnit.assert.equal(result[0], 3);
+    gpu.destroy();
   }
 
   QUnit.test( 'add typed functions - vec2 - (auto)', function() {
     vec2Test(null);
   });
+  QUnit.test( 'add typed functions - vec2 - (webgl2)', function() {
+    vec2Test('webgl2');
+  });
+  QUnit.test( 'add typed functions - vec2 - (webgl)', function() {
+    vec2Test('webgl');
+  });
+})();
+
+(function() {
+  function vec3Test(mode) {
+    var gpu = new GPU({ mode: mode });
+    function typedFunction() {
+      return [1, 2, 3];
+    }
+    gpu.addFunction(typedFunction, {
+      returnType: 'vec3'
+    });
+    var kernel = gpu.createKernel(function() {
+      var result = typedFunction();
+      return result[0] + result[1] + result[2];
+    })
+      .setOutput([1]);
+    var result = kernel();
+    QUnit.assert.equal(result[0], 6);
+    gpu.destroy();
+  }
+
+  QUnit.test( 'add typed functions - vec3 - (auto)', function() {
+    vec3Test(null);
+  });
+  QUnit.test( 'add typed functions - vec3 - (webgl2)', function() {
+    vec3Test('webgl2');
+  });
+  QUnit.test( 'add typed functions - vec3 - (webgl)', function() {
+    vec3Test('webgl');
+  });
+})();
+
+(function() {
+  function vec4Test(mode) {
+    var gpu = new GPU({ mode: mode });
+    function typedFunction() {
+      return [1, 2, 3, 4];
+    }
+    gpu.addFunction(typedFunction, {
+      returnType: 'vec4'
+    });
+    var kernel = gpu.createKernel(function() {
+      var result = typedFunction();
+      return result[0] + result[1] + result[2] + result[3];
+    })
+      .setOutput([1]);
+    var result = kernel();
+    QUnit.assert.equal(result[0], 10);
+    gpu.destroy();
+  }
+
+  QUnit.test( 'add typed functions - vec4 - (auto)', function() {
+    vec4Test(null);
+  });
+  QUnit.test( 'add typed functions - vec4 - (webgl2)', function() {
+    vec4Test('webgl2');
+  });
+  QUnit.test( 'add typed functions - vec4 - (webgl)', function() {
+    vec4Test('webgl');
+  });
 })();