@shaderfrog/core 1.5.3 → 2.0.0-beta.1

package/graph/context.d.ts

@@ -1,4 +1,4 @@
-import type { GlslSyntaxError } from '@shaderfrog/glsl-parser';
+import { type GlslSyntaxError } from '@shaderfrog/glsl-parser';
 import { AstNode, Program } from '@shaderfrog/glsl-parser/ast';
 import { Engine, EngineContext } from '../engine';
 import { NodeInput } from './base-node';

package/graph/graph.d.ts

@@ -2,6 +2,7 @@ import { Program } from '@shaderfrog/glsl-parser/ast';
 import { Engine, EngineContext } from '../engine';
 import { NodeErrors } from './context';
 import { ShaderSections } from './shader-sections';
+import { FrogProgram } from '../util/ast';
 import { DataNode } from './data-nodes';
 import { Edge } from './edge';
 import { CodeNode, SourceNode } from './code-nodes';
@@ -15,7 +16,7 @@ export declare const doesLinkThruShader: (graph: Graph, node: GraphNode) => bool
 export declare const nodeName: (node: GraphNode) => string;
 export declare const mangleName: (name: string, node: GraphNode, nextSibling?: GraphNode) => string;
 export declare const mangleVar: (name: string, engine: Engine, node: GraphNode, sibling?: GraphNode) => string;
-export declare const mangleEntireProgram: (engine: Engine, ast: Program, node: GraphNode, sibling?: GraphNode) => void;
+export declare const mangleEntireProgram: (engine: Engine, ast: FrogProgram, node: GraphNode, sibling?: GraphNode) => void;
 export declare const mangleMainFn: (ast: Program, node: GraphNode, sibling?: GraphNode) => void;
 export declare const ensureFromNode: (graph: Graph, inputEdge: Edge) => GraphNode;
 export declare const resetGraphIds: (graph: Graph) => Graph;

package/graph/graph.js

@@ -130,13 +130,13 @@ export var mangleName = function (name, node, nextSibling) {
 };
 export var mangleVar = function (name, engine, node, sibling) { return (engine.preserve.has(name) ? name : mangleName(name, node, sibling)); };
 export var mangleEntireProgram = function (engine, ast, node, sibling) {
-    renameBindings(ast.scopes[0], function (name, n) {
-        return n.doNotDescope ? name : mangleVar(name, engine, node, sibling);
+    ast.scopes[0].bindings = renameBindings(ast.scopes[0].bindings, function (name) {
+        return name === ast.outVar ? name : mangleVar(name, engine, node, sibling);
     });
     mangleMainFn(ast, node, sibling);
 };
 export var mangleMainFn = function (ast, node, sibling) {
-    renameFunctions(ast.scopes[0], function (name) {
+    ast.scopes[0].functions = renameFunctions(ast.scopes[0].functions, function (name) {
         return name === 'main' ? nodeName(node) : mangleName(name, node, sibling);
     });
 };
@@ -365,11 +365,11 @@ export var compileNode = function (engine, graph, edges, engineContext, node, ac
                         ];
                         // @ts-ignore
                         var scope = fc.ast.scopes[0];
-                        renameBindings(scope, function (name, node) {
-                            return node.type !== 'declaration' && name === 'vUv'
-                                ? 'vv'
-                                : name;
-                        });
+                        // renameBindings(scope, (name, node) => {
+                        //   return node.type !== 'declaration' && name === 'vUv'
+                        //     ? 'vv'
+                        //     : name;
+                        // });
                     }
                     // })
                 }
@@ -482,7 +482,8 @@ export var compileSource = function (graph, engine, ctx) { return __awaiter(void
                     return [2 /*return*/, result];
                 }
                 compileResult = compileGraph(ctx, engine, graph);
-                fragmentResult = generate(shaderSectionsToProgram(compileResult.fragment, engine.mergeOptions).program);
+                fragmentResult = generate(shaderSectionsToProgram(compileResult.fragment, engine.mergeOptions)
+                    .program);
                 vertexResult = generate(shaderSectionsToProgram(compileResult.vertex, engine.mergeOptions).program);
                 dataInputs = filterGraphNodes(graph, [compileResult.outputFrag, compileResult.outputVert], { input: isDataInput }).inputs;
                 dataNodes = Object.entries(dataInputs).reduce(function (acc, _a) {

package/graph/graph.test.js

@@ -1,11 +1,52 @@
+var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
+    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
+    return new (P || (P = Promise))(function (resolve, reject) {
+        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
+        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
+        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
+        step((generator = generator.apply(thisArg, _arguments || [])).next());
+    });
+};
+var __generator = (this && this.__generator) || function (thisArg, body) {
+    var _ = { label: 0, sent: function() { if (t[0] & 1) throw t[1]; return t[1]; }, trys: [], ops: [] }, f, y, t, g;
+    return g = { next: verb(0), "throw": verb(1), "return": verb(2) }, typeof Symbol === "function" && (g[Symbol.iterator] = function() { return this; }), g;
+    function verb(n) { return function (v) { return step([n, v]); }; }
+    function step(op) {
+        if (f) throw new TypeError("Generator is already executing.");
+        while (g && (g = 0, op[0] && (_ = 0)), _) try {
+            if (f = 1, y && (t = op[0] & 2 ? y["return"] : op[0] ? y["throw"] || ((t = y["return"]) && t.call(y), 0) : y.next) && !(t = t.call(y, op[1])).done) return t;
+            if (y = 0, t) op = [op[0] & 2, t.value];
+            switch (op[0]) {
+                case 0: case 1: t = op; break;
+                case 4: _.label++; return { value: op[1], done: false };
+                case 5: _.label++; y = op[1]; op = [0]; continue;
+                case 7: op = _.ops.pop(); _.trys.pop(); continue;
+                default:
+                    if (!(t = _.trys, t = t.length > 0 && t[t.length - 1]) && (op[0] === 6 || op[0] === 2)) { _ = 0; continue; }
+                    if (op[0] === 3 && (!t || (op[1] > t[0] && op[1] < t[3]))) { _.label = op[1]; break; }
+                    if (op[0] === 6 && _.label < t[1]) { _.label = t[1]; t = op; break; }
+                    if (t && _.label < t[2]) { _.label = t[2]; _.ops.push(op); break; }
+                    if (t[2]) _.ops.pop();
+                    _.trys.pop(); continue;
+            }
+            op = body.call(thisArg, _);
+        } catch (e) { op = [6, e]; y = 0; } finally { f = t = 0; }
+        if (op[0] & 5) throw op[1]; return { value: op[0] ? op[1] : void 0, done: true };
+    }
+};
+import { expect, describe, it } from 'vitest';
 import util from 'util';
 import { parser } from '@shaderfrog/glsl-parser';
 import { generate } from '@shaderfrog/glsl-parser';
-import { addNode } from './graph-node';
+import { addNode, outputNode, sourceNode } from './graph-node';
 import { shaderSectionsToProgram, mergeShaderSections, findShaderSections, } from './shader-sections';
 import { numberNode } from './data-nodes';
 import { makeEdge } from './edge';
 import { evaluateNode } from './evaluate';
+import { compileSource } from './graph';
+import { texture2DStrategy } from '../strategy';
+import { isError } from './context';
+import { fail } from '../test-util';
 var inspect = function (thing) {
     return console.log(util.inspect(thing, false, null, true));
 };
@@ -65,60 +106,112 @@ var engine = {
     preserve: new Set(),
     parsers: {},
 };
-// it('graph compiler arbitrary helper test', () => {
-//   const graph: Graph = {
-//     nodes: [
-//       outputNode('0', 'Output v', p, 'vertex'),
-//       outputNode('1', 'Output f', p, 'fragment'),
-//       makeSourceNode(
-//         '2',
-//         `uniform sampler2D image1;
-// uniform sampler2D image2;
-// void main() {
-//   vec3 col = texture2D(image1, posTurn - 0.4 * time).rgb + 1.0;
-//   vec3 col = texture2D(image2, negTurn - 0.4 * time).rgb + 2.0;
-// }
-// `,
-//         'fragment'
-//       ),
-//       makeSourceNode(
-//         '3',
-//         `void main() {
-//     return vec4(0.0);
-// }
-// `,
-//         'fragment'
-//       ),
-//       makeSourceNode(
-//         '4',
-//         `void main() {
-//     return vec4(1.0);
-// }
-// `,
-//         'fragment'
-//       ),
-//     ],
-//     edges: [
-//       makeEdge(id(), '2', '1', 'out', 'filler_frogFragOut', 'fragment'),
-//       makeEdge(id(), '3', '2', 'out', 'filler_image1', 'fragment'),
-//       makeEdge(id(), '4', '2', 'out', 'filler_image2', 'fragment'),
-//     ],
-//   };
-//   const engineContext: EngineContext = {
-//     engine: 'three',
-//     nodes: {},
-//     runtime: {},
-//     debuggingNonsense: {},
-//   };
-//   const result = compileGraph(engineContext, engine, graph);
-//   const built = generate(
-//     shaderSectionsToProgram(result.fragment, {
-//       includePrecisions: true,
-//       includeVersion: true,
-//     }).program
-//   );
-//   expect(built).toBe('hi');
-// });
+var makeSourceNode = function (id, source, stage, strategies) {
+    if (strategies === void 0) { strategies = [texture2DStrategy()]; }
+    return sourceNode(id, "Shader ".concat(id), p, {
+        version: 2,
+        preprocess: false,
+        strategies: strategies,
+        uniforms: [],
+    }, source, stage);
+};
+/**
+ * What exactly am I doing here?
+ *
+ * I opened shaderfrog to start looking at the backfilling case and inlining
+ * function calls at the top of functions
+ *
+ * WHie doing that I found jest not to work well anyore and switched to vitest,
+ * which is fine, but with esm by default I can't stub the mangleName() function
+ * call, which means mangling happens as-is in the tests.
+ *
+ * Without changing the mangling strategy, the strategies.test.ts file fails
+ * because the uniform strategy looks for a mangled variable name, but the
+ * program itself isn't mangled.
+ *
+ * One way to fix this is to make fillers not have to care about mangling names,
+ * which would be simpler on the surface.
+ *
+ * Then everything in the tests broke and you found out the reason why was
+ * trying to use scopes to rename things, and most of the ast manipulation steps
+ * don't modify scopes, so you made some of them modify scopes, and now things
+ * are fucked
+ */
+it('compileSource() fragment produces inlined output', function () { return __awaiter(void 0, void 0, void 0, function () {
+    var outV, outF, imageReplacemMe, input1, input2, graph, engineContext, result, imgOut;
+    return __generator(this, function (_a) {
+        switch (_a.label) {
+            case 0:
+                outV = outputNode(id(), 'Output v', p, 'vertex');
+                outF = outputNode(id(), 'Output f', p, 'fragment');
+                imageReplacemMe = makeSourceNode(id(), "uniform sampler2D image1;\nuniform sampler2D image2;\nvoid main() {\n  vec3 col1 = texture2D(image1, posTurn - 0.4 * time).rgb + 1.0;\n  vec3 col2 = texture2D(image2, negTurn - 0.4 * time).rgb + 2.0;\n  gl_FragColor = vec4(col1 + col2, 1.0);\n}\n", 'fragment');
+                input1 = makeSourceNode(id(), "float a = 1.0;\nvoid main() {\n  gl_FragColor = vec4(0.0);\n}\n", 'fragment');
+                input2 = makeSourceNode(id(), "float a = 2.0;\nvoid main() {\n  gl_FragColor = vec4(1.0);\n}\n", 'fragment');
+                graph = {
+                    nodes: [outV, outF, imageReplacemMe, input1, input2],
+                    edges: [
+                        makeEdge(id(), imageReplacemMe.id, outF.id, 'out', 'filler_frogFragOut', 'fragment'),
+                        makeEdge(id(), input1.id, imageReplacemMe.id, 'out', 'filler_image1', 'fragment'),
+                        makeEdge(id(), input2.id, imageReplacemMe.id, 'out', 'filler_image2', 'fragment'),
+                        makeEdge(id(), input2.id, imageReplacemMe.id, 'out', 'filler_image2', 'fragment'),
+                    ],
+                };
+                engineContext = {
+                    engine: 'three',
+                    nodes: {},
+                    runtime: {},
+                    debuggingNonsense: {},
+                };
+                return [4 /*yield*/, compileSource(graph, engine, engineContext)];
+            case 1:
+                result = _a.sent();
+                if (isError(result)) {
+                    fail(result);
+                }
+                expect(result.fragmentResult).toContain("vec4 main_Shader_".concat(input1.id, "() {"));
+                expect(result.fragmentResult).toContain("vec4 main_Shader_".concat(input2.id, "() {"));
+                imgOut = "frogOut_".concat(imageReplacemMe.id);
+                expect(result.fragmentResult).toContain("vec4 ".concat(imgOut, ";"));
+                expect(result.fragmentResult)
+                    .toContain("vec4 main_Shader_".concat(imageReplacemMe.id, "() {\n  vec3 col1 = main_Shader_").concat(input1.id, "().rgb + 1.0;\n  vec3 col2 = main_Shader_").concat(input2.id, "().rgb + 2.0;\n  ").concat(imgOut, " = vec4(col1 + col2, 1.0);\n  return ").concat(imgOut, ";\n}"));
+                return [2 /*return*/];
+        }
+    });
+}); });
+it('compileSource() base case', function () { return __awaiter(void 0, void 0, void 0, function () {
+    var outV, outF, imageReplacemMe, graph, engineContext, result, imgOut;
+    return __generator(this, function (_a) {
+        switch (_a.label) {
+            case 0:
+                outV = outputNode(id(), 'Output v', p, 'vertex');
+                outF = outputNode(id(), 'Output f', p, 'fragment');
+                imageReplacemMe = makeSourceNode(id(), "float a = 1.0;\nvoid main() {\n  gl_FragColor = vec4(1.0);\n}\n", 'fragment');
+                graph = {
+                    nodes: [outV, outF, imageReplacemMe],
+                    edges: [
+                        makeEdge(id(), imageReplacemMe.id, outF.id, 'out', 'filler_frogFragOut', 'fragment'),
+                    ],
+                };
+                engineContext = {
+                    engine: 'three',
+                    nodes: {},
+                    runtime: {},
+                    debuggingNonsense: {},
+                };
+                return [4 /*yield*/, compileSource(graph, engine, engineContext)];
+            case 1:
+                result = _a.sent();
+                if (isError(result)) {
+                    fail(result);
+                }
+                imgOut = "frogOut_".concat(imageReplacemMe.id);
+                expect(result.fragmentResult).toContain("vec4 ".concat(imgOut, ";"));
+                expect(result.fragmentResult)
+                    .toContain("vec4 main_Shader_".concat(imageReplacemMe.id, "() {\n  ").concat(imgOut, " = vec4(1.0);\n  return ").concat(imgOut, ";\n}"));
+                return [2 /*return*/];
+        }
+    });
+}); });
 describe('evaluateNode()', function () {
     it('evaluates binary nodes', function () {
         var finalAdd = addNode(id(), p);

package/graph/parsers.js

@@ -24,7 +24,7 @@ var __spreadArray = (this && this.__spreadArray) || function (to, from, pack) {
     return to.concat(ar || Array.prototype.slice.call(from));
 };
 var _a;
-import { generate, parser } from '@shaderfrog/glsl-parser';
+import { parser } from '@shaderfrog/glsl-parser';
 import { visit, } from '@shaderfrog/glsl-parser/ast';
 import preprocess from '@shaderfrog/glsl-parser/preprocessor';
 import { convert300MainToReturn, from2To3, makeExpression, makeExpressionWithScopes, makeFnBodyStatementWithScopes, makeFnStatement, } from '../util/ast';
@@ -86,14 +86,11 @@ export var coreParsers = (_a = {},
                 ast = parser.parse(preprocessed);
                 if (node.config.version === 2 && node.stage) {
                     from2To3(ast, node.stage);
-                    log('converted ', node, 'to version 3', {
-                        code: generate(ast),
-                    });
                 }
                 // This assumes that expressionOnly nodes don't have a stage and that all
                 // fragment source code shades have main function, which is probably wrong
                 if (node.stage === 'fragment') {
-                    convert300MainToReturn(node.id, ast);
+                    convert300MainToReturn(ast);
                 }
             }
             return ast;
@@ -133,7 +130,7 @@ export var coreParsers = (_a = {},
                     nodeInput(MAGIC_OUTPUT_STMTS, "filler_".concat(MAGIC_OUTPUT_STMTS), 'filler', 'rgba', ['code'], false),
                     function (fillerAst) {
                         var fn = ast.program.find(function (stmt) { return stmt.type === 'function'; });
-                        fn === null || fn === void 0 ? void 0 : fn.body.statements.unshift(makeFnStatement(generateFiller(fillerAst)));
+                        fn === null || fn === void 0 ? void 0 : fn.body.statements.unshift(makeFnStatement(generateFiller(fillerAst))[0]);
                         return ast;
                     },
                 ],

package/graph/shader-sections.js

@@ -62,7 +62,7 @@ export var highestPrecisions = function (nodes) {
         return (__assign(__assign({}, precisions), (_a = {}, _a[stmt.declaration.specifier.specifier.token] = higherPrecision(precisions[stmt.declaration.specifier.specifier.token], stmt.declaration.qualifier.token), _a)));
     }, {})).map(function (_a) {
         var _b = __read(_a, 2), typeName = _b[0], precision = _b[1];
-        return makeStatement("precision ".concat(precision, " ").concat(typeName));
+        return makeStatement("precision ".concat(precision, " ").concat(typeName))[0];
     });
 };
 export var dedupeQualifiedStatements = function (statements, qualifier) {
@@ -74,7 +74,7 @@ export var dedupeQualifiedStatements = function (statements, qualifier) {
         }, {})), _a)));
     }, {})).map(function (_a) {
         var _b = __read(_a, 2), type = _b[0], varNames = _b[1];
-        return makeStatement("".concat(qualifier, " ").concat(type, " ").concat(Object.keys(varNames).join(', ')));
+        return makeStatement("".concat(qualifier, " ").concat(type, " ").concat(Object.keys(varNames).join(', ')))[0];
     });
 };
 /**
@@ -165,7 +165,7 @@ export var dedupeUniforms = function (statements) {
         var _b = __read(_a, 2), type = _b[0], variables = _b[1];
         return makeStatement("uniform ".concat(type, " ").concat(Object.values(variables)
             .map(function (v) { return v.generated; })
-            .join(', ')));
+            .join(', ')))[0];
     });
 };
 export var mergeShaderSections = function (s1, s2) {
@@ -237,11 +237,6 @@ export var findShaderSections = function (ast) {
         else if (node.type === 'declaration_statement' &&
             'specified_type' in node.declaration &&
             ((_m = (_l = (_k = node.declaration) === null || _k === void 0 ? void 0 : _k.specified_type) === null || _l === void 0 ? void 0 : _l.qualifiers) === null || _m === void 0 ? void 0 : _m.find(function (n) { return 'token' in n && n.token === 'in'; }))) {
-            if (generate(node).includes('main_Fireball')) {
-                console.log('findShaderSections\n', generate(ast));
-                console.log("Tracking inStatement \"".concat(generate(node), "\""), node);
-                debugger;
-            }
             return __assign(__assign({}, sections), { inStatements: sections.inStatements.concat(node) });
         }
         else if (node.type === 'declaration_statement' &&

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@shaderfrog/core",
-  "version": "1.5.3",
+  "version": "2.0.0-beta.1",
   "description": "Shaderfrog core",
   "type": "module",
   "files": [
@@ -17,8 +17,8 @@
     "prepare": "npm run build && ./prepublish.sh",
     "postpublish": "./postbuild.sh",
     "build": "./build.sh",
-    "watch-test": "jest --watch",
-    "test": "jest --colors"
+    "watch-test": "vitest --watch",
+    "test": "vitest"
   },
   "repository": {
     "type": "git",
@@ -34,21 +34,20 @@
     "@babel/core": "^7.21.8",
     "@babel/preset-env": "^7.21.5",
     "@babel/preset-typescript": "^7.21.5",
-    "@types/jest": "^29.5.1",
+    "@swc/core": "^1.6.7",
     "@types/lodash.groupby": "^4.6.7",
     "@types/three": "^0.156.0",
-    "babel-jest": "^29.5.0",
     "babylonjs": "^6.2.0",
-    "jest": "^29.5.0",
-    "prettier": "^2.8.8",
+    "prettier": "^3.3.2",
     "three": "^0.156.1",
-    "typescript": "^5.0.4"
+    "typescript": "^5.5.3",
+    "vitest": "^1.6.0"
   },
   "dependencies": {
     "lodash.groupby": "^4.6.0"
   },
   "peerDependencies": {
-    "@shaderfrog/glsl-parser": "^4.0.0",
+    "@shaderfrog/glsl-parser": "^5.0.0-beta.3",
     "babylonjs": ">=4",
     "playcanvas": "^1.65.3",
     "three": ">=0.50"

package/plugins/babylon/importers.js

@@ -22,7 +22,7 @@ var importers = {
             // Babylon has no normalmatrix. They do have a normal attribute. So undo any
             // multiplication by normalMatrix?
             var seen = {};
-            renameBindings(ast.scopes[0], function (name) {
+            ast.scopes[0].bindings = renameBindings(ast.scopes[0].bindings, function (name) {
                 // console.log({ name }, 'seen:', seen[name]);
                 var renamed = name === 'vUv'
                     ? 'vMainUV1'

package/plugins/three/importers.js

@@ -2,7 +2,7 @@ import { renameBindings } from '@shaderfrog/glsl-parser/parser/utils';
 var importers = {
     babylon: {
         convertAst: function (ast, type) {
-            renameBindings(ast.scopes[0], function (name) {
+            ast.scopes[0].bindings = renameBindings(ast.scopes[0].bindings, function (name) {
                 return name === 'vMainUV1' ? 'vUv' : name === 'vNormalW' ? 'vNormal' : name;
             });
         },

package/plugins/three/threngine.test.d.ts

@@ -0,0 +1 @@
+export {};

package/plugins/three/threngine.test.js

@@ -0,0 +1,89 @@
+var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
+    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
+    return new (P || (P = Promise))(function (resolve, reject) {
+        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
+        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
+        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
+        step((generator = generator.apply(thisArg, _arguments || [])).next());
+    });
+};
+var __generator = (this && this.__generator) || function (thisArg, body) {
+    var _ = { label: 0, sent: function() { if (t[0] & 1) throw t[1]; return t[1]; }, trys: [], ops: [] }, f, y, t, g;
+    return g = { next: verb(0), "throw": verb(1), "return": verb(2) }, typeof Symbol === "function" && (g[Symbol.iterator] = function() { return this; }), g;
+    function verb(n) { return function (v) { return step([n, v]); }; }
+    function step(op) {
+        if (f) throw new TypeError("Generator is already executing.");
+        while (g && (g = 0, op[0] && (_ = 0)), _) try {
+            if (f = 1, y && (t = op[0] & 2 ? y["return"] : op[0] ? y["throw"] || ((t = y["return"]) && t.call(y), 0) : y.next) && !(t = t.call(y, op[1])).done) return t;
+            if (y = 0, t) op = [op[0] & 2, t.value];
+            switch (op[0]) {
+                case 0: case 1: t = op; break;
+                case 4: _.label++; return { value: op[1], done: false };
+                case 5: _.label++; y = op[1]; op = [0]; continue;
+                case 7: op = _.ops.pop(); _.trys.pop(); continue;
+                default:
+                    if (!(t = _.trys, t = t.length > 0 && t[t.length - 1]) && (op[0] === 6 || op[0] === 2)) { _ = 0; continue; }
+                    if (op[0] === 3 && (!t || (op[1] > t[0] && op[1] < t[3]))) { _.label = op[1]; break; }
+                    if (op[0] === 6 && _.label < t[1]) { _.label = t[1]; t = op; break; }
+                    if (t && _.label < t[2]) { _.label = t[2]; _.ops.push(op); break; }
+                    if (t[2]) _.ops.pop();
+                    _.trys.pop(); continue;
+            }
+            op = body.call(thisArg, _);
+        } catch (e) { op = [6, e]; y = 0; } finally { f = t = 0; }
+        if (op[0] & 5) throw op[1]; return { value: op[0] ? op[1] : void 0, done: true };
+    }
+};
+import { expect, it } from 'vitest';
+import { outputNode, sourceNode } from '../../graph/graph-node';
+import { makeEdge } from '../../graph/edge';
+import { compileSource } from '../../graph/graph';
+import { texture2DStrategy } from '../../strategy';
+import { isError } from '../../graph/context';
+import { threngine } from './threngine';
+import { makeId } from '../../util/id';
+import { fail } from '../../test-util';
+var p = { x: 0, y: 0 };
+var makeSourceNode = function (id, source, stage, strategies) {
+    if (strategies === void 0) { strategies = [texture2DStrategy()]; }
+    return sourceNode(id, "Shader ".concat(id), p, {
+        version: 2,
+        preprocess: false,
+        strategies: strategies,
+        uniforms: [],
+    }, source, stage);
+};
+it('threngine compileSource() and manipulateAst()', function () { return __awaiter(void 0, void 0, void 0, function () {
+    var outV, outF, vertInput, graph, engineContext, result;
+    return __generator(this, function (_a) {
+        switch (_a.label) {
+            case 0:
+                outV = outputNode(makeId(), 'Output v', p, 'vertex');
+                outF = outputNode(makeId(), 'Output f', p, 'fragment');
+                vertInput = makeSourceNode(makeId(), "uniform vec4 modelViewMatrix;\nattribute vec3 position;\nfloat a = 2.0;\nvoid main() {\n  gl_Position = modelViewMatrix * vec4(position, 1.0);\n}\n", 'vertex');
+                graph = {
+                    nodes: [outV, outF, vertInput],
+                    edges: [
+                        makeEdge(makeId(), vertInput.id, outV.id, 'out', 'filler_gl_Position', 'vertex'),
+                    ],
+                };
+                engineContext = {
+                    engine: 'three',
+                    nodes: {},
+                    runtime: {},
+                    debuggingNonsense: {},
+                };
+                return [4 /*yield*/, compileSource(graph, threngine, engineContext)];
+            case 1:
+                result = _a.sent();
+                if (isError(result)) {
+                    fail(result);
+                }
+                // Threngine has parsers for vertex shaders, make sure that is set properly
+                expect(result.vertexResult).toContain("vec4 main_Shader_".concat(vertInput.id, "() {\n  vec4 frogOut = modelViewMatrix * vec4(position, 1.0);\n  return frogOut;\n}"));
+                // Check that it inlned. For fun.
+                expect(result.vertexResult).toContain("gl_Position = main_Shader_".concat(vertInput.id, "();"));
+                return [2 /*return*/];
+        }
+    });
+}); });

package/strategy/assignemntTo.d.ts

@@ -3,7 +3,8 @@ export interface AssignemntToStrategy extends BaseStrategy {
     type: StrategyType.ASSIGNMENT_TO;
     config: {
         assignTo: string;
+        nth?: number;
     };
 }
-export declare const assignemntToStrategy: (assignTo: string) => AssignemntToStrategy;
+export declare const assignemntToStrategy: (assignTo: string, nth?: number) => AssignemntToStrategy;
 export declare const applyAssignmentToStrategy: ApplyStrategy<AssignemntToStrategy>;

package/strategy/assignemntTo.js

@@ -2,13 +2,16 @@ import { findAssignmentTo } from '../util/ast';
 import { nodeInput } from '../graph/base-node';
 import { StrategyType } from '.';
 // Constructor
-export var assignemntToStrategy = function (assignTo) { return ({
-    type: StrategyType.ASSIGNMENT_TO,
-    config: { assignTo: assignTo },
-}); };
+export var assignemntToStrategy = function (assignTo, nth) {
+    if (nth === void 0) { nth = 1; }
+    return ({
+        type: StrategyType.ASSIGNMENT_TO,
+        config: { assignTo: assignTo, nth: nth },
+    });
+};
 // Apply the strategy
 export var applyAssignmentToStrategy = function (strategy, ast, graphNode, siblingNode) {
-    var assignNode = findAssignmentTo(ast, strategy.config.assignTo);
+    var assignNode = findAssignmentTo(ast, strategy.config.assignTo, strategy.config.nth || 1);
     var name = strategy.config.assignTo;
     return assignNode
         ? [

package/strategy/inject.d.ts

@@ -11,5 +11,5 @@ export interface InjectStrategy extends BaseStrategy {
  * Inject source code into an AST, in a find-and-replace style. This only
  * operates on statements right now
  */
-export declare const injectStrategy: (config: InjectStrategy['config']) => InjectStrategy;
+export declare const injectStrategy: (config: InjectStrategy["config"]) => InjectStrategy;
 export declare const applyInjectStrategy: ApplyStrategy<InjectStrategy>;

package/strategy/inject.js

@@ -1,14 +1,3 @@
-var __assign = (this && this.__assign) || function () {
-    __assign = Object.assign || function(t) {
-        for (var s, i = 1, n = arguments.length; i < n; i++) {
-            s = arguments[i];
-            for (var p in s) if (Object.prototype.hasOwnProperty.call(s, p))
-                t[p] = s[p];
-        }
-        return t;
-    };
-    return __assign.apply(this, arguments);
-};
 var __read = (this && this.__read) || function (o, n) {
     var m = typeof Symbol === "function" && o[Symbol.iterator];
     if (!m) return o;
@@ -37,6 +26,7 @@ var __spreadArray = (this && this.__spreadArray) || function (to, from, pack) {
 import { generate } from '@shaderfrog/glsl-parser';
 import { nodeInput } from '../graph/base-node';
 import { StrategyType } from '.';
+import { transferWhitespace } from '../util/whitespace';
 /**
  * Inject source code into an AST, in a find-and-replace style. This only
  * operates on statements right now
@@ -45,26 +35,6 @@ export var injectStrategy = function (config) { return ({
     type: StrategyType.INJECT,
     config: config,
 }); };
-// Typescript fucked up flat https://stackoverflow.com/a/61420611/743464
-var combineWs = function (a, b) {
-    return [a, b].flat(Infinity);
-};
-// Move whitespace from one node to another. Since whitespace is trailing, if a
-// node is injected after a previous node, move the whitespace from the earlier
-// node to the later one. This keeps comments in the same place
-var transferWhitespace = function (to, from) {
-    return 'semi' in to && 'semi' in from
-        ? [
-            __assign(__assign({}, to), { semi: __assign(__assign({}, to.semi), { whitespace: from.semi.whitespace }) }),
-            __assign(__assign({}, from), { semi: __assign(__assign({}, from.semi), { whitespace: '\n' }) }),
-        ]
-        : 'whitespace' in to && 'semi' in from
-            ? [
-                __assign(__assign({}, to), { whitespace: combineWs(to.whitespace, from.semi.whitespace) }),
-                __assign(__assign({}, from), { semi: __assign(__assign({}, from.semi), { whitespace: '\n' }) }),
-            ]
-            : [to, from];
-};
 export var applyInjectStrategy = function (strategy, ast, graphNode, siblingNode) {
     var program = ast.program;
     var _a = strategy.config, find = _a.find, count = _a.count, insert = _a.insert;

package/strategy/stratgies.test.js

@@ -14,23 +14,12 @@ var __read = (this && this.__read) || function (o, n) {
     }
     return ar;
 };
+import { expect, it } from 'vitest';
 import { parser } from '@shaderfrog/glsl-parser';
 import { generate } from '@shaderfrog/glsl-parser';
 import { applyStrategy, StrategyType } from '.';
-import * as graphModule from '../graph/graph';
 import { makeExpression } from '../util/ast';
 import preprocess from '@shaderfrog/glsl-parser/preprocessor';
-var orig;
-beforeEach(function () {
-    orig = graphModule.mangleName;
-    // Terrible hack. in the real world, strategies are applied after mangling
-    // @ts-ignore
-    graphModule.mangleName = function (name) { return name; };
-});
-afterEach(function () {
-    // @ts-ignore
-    graphModule.mangleName = orig;
-});
 it('named attribute strategy`', function () {
     var source = "\nin vec3 replaceThisAtrribute;\nvoid main() {\n  vec2 y = replaceThisAtrribute;\n}\n";
     var ast = parser.parse(source, { quiet: true });
@@ -92,9 +81,65 @@ it('inject strategy before', function () {
     // Should fill references
     expect(result).toBe("\nuniform float x;\n// Some comment\nvoid main() {\n/* some comment */\nsomeOtherCall(x, y, z);\nsomeOtherCall(x, y, z);\nre(x, y, z);\n// Middle comment\nsomeOtherCall(x, y, z);\nsomeOtherCall(x, y, z);\nre(x, y, z);\n// Final comment\n}");
 });
+var constructor = function () { return ({
+    config: {
+        version: 3,
+        preprocess: false,
+        strategies: [],
+        uniforms: [],
+    },
+    id: '1',
+    name: '1',
+    engine: true,
+    type: '',
+    inputs: [],
+    outputs: [],
+    position: { x: 0, y: 0 },
+    source: '',
+    stage: undefined,
+}); };
+var engine = {
+    name: 'three',
+    displayName: 'Three.js',
+    evaluateNode: function (node) {
+        if (node.type === 'number') {
+            return parseFloat(node.value);
+        }
+        return node.value;
+    },
+    constructors: {
+        physical: constructor,
+        toon: constructor,
+    },
+    mergeOptions: {
+        includePrecisions: true,
+        includeVersion: true,
+    },
+    importers: {},
+    preserve: new Set(),
+    parsers: {},
+};
 it('correctly fills with uniform strategy', function () {
     var _a, _b, _c;
     var ast = parser.parse("\nlayout(std140,column_major) uniform;\nuniform sampler2D image;\nuniform vec4 input, output, other;\nuniform vec4 zenput;\nuniform Light0 { vec4 y; } x;\nvec3 topLevel = vec3(0.0);\nvoid other(in vec3 param) {}\nvoid main() {\n  vec4 computed = texture2D(image, uvPow * 1.0);\n  vec4 x = input;\n  vec4 y = output;\n  vec4 z = zenput;\n}", { quiet: true });
+    // TODO: Experimenting with strategy tests where we mangle in the test to
+    // avoid having to mangle in the strategy, in service of maybe mangling the
+    // AST as part of producing context. But as the test shows -
+    // mangleEntireProgram does NOT modify binding names
+    //
+    // You started updating binding names in the parser but realized that
+    // technically a mangler can produce different results for different nodes
+    // during the rename, since the parser takes in the node to mangle.
+    //
+    // which raised the question about why pass in the node at all to the mangler?
+    // looks like it's for "doNotDescope" hack to avoid renaming a specific
+    // varaible.
+    //
+    // But maybe that could be done here instead? And mangleEntireProgram could be
+    // aware of the output varaibles to ignore? Which means we need to track the
+    // output varialbe names somewhere... do we alredy?
+    var node = { name: 'fake', id: '1' };
+    // mangleEntireProgram(engine, ast, node);
     var fillers = applyStrategy({ type: StrategyType.UNIFORM, config: {} }, ast, {}, {});
     // It should find uniforms with simple types, excluding sampler2D
     expect(fillers.map(function (_a) {
@@ -160,5 +205,5 @@ it('Make sure texture2D finds preprocessed texture() call', function () {
     expect(applyStrategy({ type: StrategyType.TEXTURE_2D, config: {} }, ast, {}, {}).map(function (_a) {
         var _b = __read(_a, 1), name = _b[0].displayName;
         return name;
-    })).toEqual(['normalMapx']);
+    })).toEqual(['normalMap']);
 });

package/strategy/uniform.js

@@ -14,7 +14,6 @@ var __read = (this && this.__read) || function (o, n) {
     }
     return ar;
 };
-import { mangleName } from '../graph/graph';
 import { nodeInput } from '../graph/base-node';
 import { StrategyType } from '.';
 import { generateFiller } from '../util/ast';
@@ -155,7 +154,8 @@ export var applyUniformStrategy = function (strategy, ast, graphNode, siblingNod
             return names.map(function (name) { return [
                 nodeInput(name, "uniform_".concat(name), 'uniform', graphDataType, ['code', 'data'], true),
                 function (filler) {
-                    var mangledName = mangleName(name, graphNode, siblingNode);
+                    //const mangledName = mangleName(name, graphNode, siblingNode);
+                    var mangledName = name;
                     // Remove the declaration line, or the declared uniform
                     if (declarations_1.length === 1) {
                         program.program.splice(program.program.indexOf(node), 1);

package/util/ast.d.ts

@@ -1,18 +1,30 @@
-/**
- * Utility functions to work with ASTs
- */
-import { Filler } from '../graph/parsers';
-import { AstNode, ExpressionStatementNode, FunctionNode, DeclarationNode } from '@shaderfrog/glsl-parser/ast';
+import { AstNode, ExpressionStatementNode, FunctionNode, DeclarationStatementNode, DeclarationNode } from '@shaderfrog/glsl-parser/ast';
 import { Program } from '@shaderfrog/glsl-parser/ast';
 import { ShaderStage } from '../graph/graph-types';
 import { Scope } from '@shaderfrog/glsl-parser/parser/scope';
+import { Filler } from '../graph/parsers';
+export interface FrogProgram extends Program {
+    outVar?: string;
+}
 export declare const findVec4Constructor: (ast: AstNode) => AstNode | undefined;
-export declare const findAssignmentTo: (ast: AstNode | Program, assignTo: string) => DeclarationNode | ExpressionStatementNode;
+export declare const findAssignmentTo: (ast: AstNode | Program, assignTo: string, nth?: number) => DeclarationNode | ExpressionStatementNode;
 export declare const findDeclarationOf: (ast: AstNode | Program, declarationOf: string) => DeclarationNode | undefined;
 export declare const from2To3: (ast: Program, stage: ShaderStage) => void;
-export declare const outDeclaration: (name: string) => Object;
-export declare const makeStatement: (stmt: string) => AstNode;
-export declare const makeFnStatement: (fnStmt: string) => AstNode;
+export declare const makeStatement: (stmt: string) => readonly [DeclarationStatementNode | FunctionNode | import("@shaderfrog/glsl-parser/ast").PreprocessorNode, Scope];
+export declare const makeFnStatement: (fnStmt: string) => readonly [AstNode, Scope];
+/**
+ * Add a new scope into an existing one. Meant to be used for adding net new
+ * lines of coe to an AST, and wanting to add the new scope generated from those
+ * lines.
+ *
+ * DO NOT USE THIS TO MERGE SCOPES! If both the left and right scope contain the
+ * same binding name, this will override the left scope outright, rather than
+ * merge the binding.references.
+ *
+ * One reason I chose not to make a full merge: What happens if both sides
+ * contain a binding.declaration?
+ */
+export declare const addNewScope: (left: Scope, right: Scope) => Scope;
 export declare const makeExpression: (expr: string) => AstNode;
 export declare const makeExpressionWithScopes: (expr: string) => {
     scope: Scope;
@@ -22,9 +34,16 @@ export declare const makeFnBodyStatementWithScopes: (body: string) => {
     scope: Scope;
     statements: AstNode[];
 };
-export declare const findFn: (ast: Program, name: string) => FunctionNode | undefined;
+export declare const findFn: (name: string) => (ast: Program) => FunctionNode | undefined;
+export declare const findMain: (ast: Program) => FunctionNode | undefined;
+export declare const findMainOrThrow: (ast: Program) => FunctionNode;
 export declare const returnGlPosition: (fnName: string, ast: Program) => void;
 export declare const returnGlPositionHardCoded: (fnName: string, ast: Program, returnType: string, hardCodedReturn: string) => void;
 export declare const returnGlPositionVec3Right: (fnName: string, ast: Program) => void;
-export declare const convert300MainToReturn: (suffix: string, ast: Program) => void;
+/**
+ * For either a fragment or vertex AST, convert the main() function that sets
+ * gl_FragColor, gl_Position, or "out vec4 ____" into a main() function that
+ * returns a vec4.
+ */
+export declare const convert300MainToReturn: (ast: FrogProgram) => void;
 export declare const generateFiller: (filler: Filler) => string;

package/util/ast.js

@@ -1,3 +1,14 @@
+var __assign = (this && this.__assign) || function () {
+    __assign = Object.assign || function(t) {
+        for (var s, i = 1, n = arguments.length; i < n; i++) {
+            s = arguments[i];
+            for (var p in s) if (Object.prototype.hasOwnProperty.call(s, p))
+                t[p] = s[p];
+        }
+        return t;
+    };
+    return __assign.apply(this, arguments);
+};
 var __read = (this && this.__read) || function (o, n) {
     var m = typeof Symbol === "function" && o[Symbol.iterator];
     if (!m) return o;
@@ -25,6 +36,8 @@ var __spreadArray = (this && this.__spreadArray) || function (to, from, pack) {
 };
 import { parser, generate } from '@shaderfrog/glsl-parser';
 import { visit, } from '@shaderfrog/glsl-parser/ast';
+import { addFnStmtWithIndent } from './whitespace';
+import { renameBinding, } from '@shaderfrog/glsl-parser/parser/utils';
 var log = function () {
     var _a;
     var args = [];
@@ -49,14 +62,20 @@ export var findVec4Constructor = function (ast) {
     visit(ast, visitors);
     return parent;
 };
-export var findAssignmentTo = function (ast, assignTo) {
+export var findAssignmentTo = function (ast, assignTo, nth) {
+    if (nth === void 0) { nth = 1; }
     var assign;
+    var foundth = 0;
     var visitors = {
         expression_statement: {
             enter: function (path) {
                 var _a, _b;
                 if (((_b = (_a = path.node.expression) === null || _a === void 0 ? void 0 : _a.left) === null || _b === void 0 ? void 0 : _b.identifier) === assignTo) {
-                    assign = path.node;
+                    foundth++;
+                    if (foundth === nth) {
+                        assign = path.node;
+                        path.stop();
+                    }
                 }
                 path.skip();
             },
@@ -66,7 +85,11 @@ export var findAssignmentTo = function (ast, assignTo) {
                 var _a, _b;
                 var foundDecl = (_b = (_a = path.node.declaration) === null || _a === void 0 ? void 0 : _a.declarations) === null || _b === void 0 ? void 0 : _b.find(function (decl) { var _a; return ((_a = decl === null || decl === void 0 ? void 0 : decl.identifier) === null || _a === void 0 ? void 0 : _a.identifier) === assignTo; });
                 if (foundDecl === null || foundDecl === void 0 ? void 0 : foundDecl.initializer) {
-                    assign = foundDecl;
+                    foundth++;
+                    if (foundth === nth) {
+                        assign = foundDecl;
+                        path.stop();
+                    }
                 }
                 path.skip();
             },
@@ -101,7 +124,11 @@ export var from2To3 = function (ast, stage) {
     //   _: '\n',
     // });
     if (stage === 'fragment') {
-        ast.program.unshift(makeStatement("out vec4 ".concat(glOut)));
+        // Add in "out vec4 fragmentColor" to convert gl_FragColor to an out statement
+        var _a = __read(makeStatement("out vec4 ".concat(glOut)), 2), outStmt = _a[0], scope = _a[1];
+        ast.program.unshift(outStmt);
+        // Add the out statement variable to the scope
+        ast.scopes[0] = addNewScope(ast.scopes[0], scope);
     }
     visit(ast, {
         function_call: {
@@ -117,6 +144,7 @@ export var from2To3 = function (ast, stage) {
             enter: function (path) {
                 if (path.node.identifier === 'gl_FragColor') {
                     path.node.identifier = glOut;
+                    ast.scopes[0].bindings[glOut].references.push(path.node);
                 }
             },
         },
@@ -131,36 +159,6 @@ export var from2To3 = function (ast, stage) {
         },
     });
 };
-export var outDeclaration = function (name) { return ({
-    type: 'declaration_statement',
-    declaration: {
-        type: 'declarator_list',
-        specified_type: {
-            type: 'fully_specified_type',
-            qualifiers: [{ type: 'keyword', token: 'out', whitespace: ' ' }],
-            specifier: {
-                type: 'type_specifier',
-                specifier: { type: 'keyword', token: 'vec4', whitespace: ' ' },
-                quantifier: null,
-            },
-        },
-        declarations: [
-            {
-                type: 'declaration',
-                identifier: {
-                    type: 'identifier',
-                    identifier: name,
-                    whitespace: undefined,
-                },
-                quantifier: null,
-                operator: undefined,
-                initializer: undefined,
-            },
-        ],
-        commas: [],
-    },
-    semi: { type: 'literal', literal: ';', whitespace: '\n    ' },
-}); };
 export var makeStatement = function (stmt) {
     // log(`Parsing "${stmt}"`);
     var ast;
@@ -172,20 +170,38 @@ export var makeStatement = function (stmt) {
         throw new Error("Error parsing stmt \"".concat(stmt, "\": ").concat(error === null || error === void 0 ? void 0 : error.message));
     }
     // log(util.inspect(ast, false, null, true));
-    return ast.program[0];
+    return [ast.program[0], ast.scopes[0]];
 };
 export var makeFnStatement = function (fnStmt) {
     var ast;
     try {
-        ast = parser.parse("\n  void main() {\n      ".concat(fnStmt, ";\n    }"), { quiet: true });
+        // Create a statement with no trailing nor leading whitespace
+        ast = parser.parse("void main() {".concat(fnStmt, ";}"), { quiet: true });
     }
     catch (error) {
         console.error({ fnStmt: fnStmt, error: error });
         throw new Error("Error parsing fnStmt \"".concat(fnStmt, "\": ").concat(error === null || error === void 0 ? void 0 : error.message));
     }
     // log(util.inspect(ast, false, null, true));
-    return ast.program[0].body.statements[0];
+    var n = ast.program[0].body.statements[0];
+    n.semi.whitespace = '';
+    return [n, ast.scopes[1]];
 };
+/**
+ * Add a new scope into an existing one. Meant to be used for adding net new
+ * lines of coe to an AST, and wanting to add the new scope generated from those
+ * lines.
+ *
+ * DO NOT USE THIS TO MERGE SCOPES! If both the left and right scope contain the
+ * same binding name, this will override the left scope outright, rather than
+ * merge the binding.references.
+ *
+ * One reason I chose not to make a full merge: What happens if both sides
+ * contain a binding.declaration?
+ */
+export var addNewScope = function (left, right) { return (__assign(__assign({}, left), { 
+    // name, parent comes from left
+    bindings: __assign(__assign({}, left.bindings), right.bindings), types: __assign(__assign({}, left.types), right.types), functions: __assign(__assign({}, left.functions), right.functions) })); };
 export var makeExpression = function (expr) {
     var ast;
     try {
@@ -229,10 +245,21 @@ export var makeFnBodyStatementWithScopes = function (body) {
         statements: ast.program[0].body.statements,
     };
 };
-export var findFn = function (ast, name) {
-    return ast.program.find(function (stmt) {
-        return stmt.type === 'function' && stmt.prototype.header.name.identifier === name;
-    });
+export var findFn = function (name) {
+    return function (ast) {
+        return ast.program.find(function (stmt) {
+            return stmt.type === 'function' &&
+                stmt.prototype.header.name.identifier === name;
+        });
+    };
+};
+export var findMain = findFn('main');
+export var findMainOrThrow = function (ast) {
+    var main = findMain(ast);
+    if (!main) {
+        throw new Error('No main function found!');
+    }
+    return main;
 };
 export var returnGlPosition = function (fnName, ast) {
     return convertVertexMain(fnName, ast, 'vec4', function (assign) { return assign.expression.right; });
@@ -264,9 +291,10 @@ export var returnGlPositionVec3Right = function (fnName, ast) {
         return found;
     });
 };
+var replacedReturn = 'frogOut';
 var convertVertexMain = function (fnName, ast, returnType, generateRight) {
     var mainReturnVar = "frogOut";
-    var main = findFn(ast, fnName);
+    var main = findFn(fnName)(ast);
     if (!main) {
         throw new Error("No ".concat(fnName, " fn found!"));
     }
@@ -282,14 +310,23 @@ var convertVertexMain = function (fnName, ast, returnType, generateRight) {
     if (!assign) {
         throw new Error("No gl position assign found in main fn!");
     }
-    var rtnStmt = makeFnStatement("".concat(returnType, " ").concat(mainReturnVar, " = 1.0"));
+    var rtnStmt = makeFnStatement("".concat(returnType, " ").concat(mainReturnVar, " = 1.0"))[0];
     rtnStmt.declaration.declarations[0].initializer =
         generateRight(assign);
+    rtnStmt.semi.whitespace = '\n';
     main.body.statements.splice(main.body.statements.indexOf(assign), 1, rtnStmt);
-    main.body.statements.push(makeFnStatement("return ".concat(mainReturnVar)));
+    main.body.statements = addFnStmtWithIndent(main, "return ".concat(mainReturnVar));
 };
-export var convert300MainToReturn = function (suffix, ast) {
-    var mainReturnVar = "frogOut_".concat(suffix);
+/**
+ * For either a fragment or vertex AST, convert the main() function that sets
+ * gl_FragColor, gl_Position, or "out vec4 ____" into a main() function that
+ * returns a vec4.
+ */
+export var convert300MainToReturn = function (ast) {
+    // Convert the main function to return a vec4
+    var main = findMainOrThrow(ast);
+    main.prototype.header.returnType.specifier.specifier.token =
+        'vec4';
     // Find the output variable, as in "pc_fragColor" from  "out highp vec4 pc_fragColor;"
     var outName;
     ast.program.find(function (line, index) {
@@ -300,7 +337,8 @@ export var convert300MainToReturn = function (suffix, ast) {
         ((_b = (_a = declaration === null || declaration === void 0 ? void 0 : declaration.specified_type) === null || _a === void 0 ? void 0 : _a.qualifiers) === null || _b === void 0 ? void 0 : _b.find(function (n) { return n.token === 'out'; })) &&
             declaration.specified_type.specifier.specifier.token ===
                 'vec4') {
-            // Remove the out declaration
+            // Remove the out declaration. This does NOT yet remove the declaration
+            // from the scope, that's done below
             ast.program.splice(index, 1);
             outName = declaration.declarations[0].identifier.identifier;
             return true;
@@ -310,27 +348,23 @@ export var convert300MainToReturn = function (suffix, ast) {
         console.error(generate(ast));
         throw new Error('No "out vec4" line found in the fragment shader');
     }
-    ast.program.unshift(makeStatement("vec4 ".concat(mainReturnVar)));
-    visit(ast, {
-        identifier: {
-            enter: function (path) {
-                if (path.node.identifier === outName) {
-                    path.node.identifier = mainReturnVar;
-                    // @ts-ignore
-                    path.node.doNotDescope = true; // hack because this var is in the scope which gets renamed later
-                }
-            },
-        },
-        function: {
-            enter: function (path) {
-                if (path.node.prototype.header.name.identifier === 'main') {
-                    path.node.prototype.header.returnType.specifier
-                        .specifier.token = 'vec4';
-                    path.node.body.statements.push(makeFnStatement("return ".concat(mainReturnVar)));
-                }
-            },
-        },
-    });
+    // Store the variable to avoid descoping it later
+    ast.outVar = outName;
+    // Rename the scope entry of "out vec4 ___" to our return variable, and rename
+    // all references to our new variable
+    ast.scopes[0].bindings[replacedReturn] = renameBinding(ast.scopes[0].bindings[outName], replacedReturn);
+    delete ast.scopes[0].bindings[outName];
+    // Add the declaration of the return variable to the top of the program, and
+    // add it to the AST scope, including the declaration
+    var decl = makeStatement("vec4 ".concat(replacedReturn))[0];
+    ast.program.unshift(decl);
+    ast.scopes[0].bindings[replacedReturn].declaration = decl;
+    ast.scopes[0].bindings[replacedReturn].references.push(decl.declaration.declarations[0]);
+    // Add a return statement to the main() function and add the return variable
+    // to scope
+    var rtn = makeFnStatement("return ".concat(replacedReturn))[0];
+    main.body.statements = addFnStmtWithIndent(main, rtn);
+    ast.scopes[0].bindings[replacedReturn].references.push(rtn.expression);
 };
 export var generateFiller = function (filler) {
     if (!filler) {

package/util/ensure.d.ts

@@ -1 +1 @@
-export declare const ensure: <T>(argument: T, message?: string) => T;
+export declare const ensure: <T>(argument: T | undefined | null, message?: string) => T;

package/util/whitespace.d.ts

@@ -0,0 +1,18 @@
+/**
+ * Utility functions to work with whitespace on nodes.
+ *
+ * This is for manual manipulaiton of code where adding new lines should
+ * attempt to keep the indentation of the original program.
+ *
+ * Another overal option is simply to pretty print, which this library
+ * does not yet support.
+ */
+import { AstNode, FunctionNode, Whitespace } from '@shaderfrog/glsl-parser/ast';
+export declare const combineWs: (a: string | string[], b: string | string[]) => string[];
+export declare const transferWhitespace: (to: AstNode, from: AstNode) => [AstNode, AstNode];
+export declare const getLiteralIndent: (node: {
+    whitespace: Whitespace;
+}) => string;
+export declare const tryAddTrailingWhitespace: <T extends AstNode>(node: T, ws: string) => T;
+export declare const guessFnIndent: (fnBody: FunctionNode) => string;
+export declare const addFnStmtWithIndent: (fnBody: FunctionNode, newNode: string | AstNode) => AstNode[];

package/util/whitespace.js

@@ -0,0 +1,91 @@
+var __assign = (this && this.__assign) || function () {
+    __assign = Object.assign || function(t) {
+        for (var s, i = 1, n = arguments.length; i < n; i++) {
+            s = arguments[i];
+            for (var p in s) if (Object.prototype.hasOwnProperty.call(s, p))
+                t[p] = s[p];
+        }
+        return t;
+    };
+    return __assign.apply(this, arguments);
+};
+var __read = (this && this.__read) || function (o, n) {
+    var m = typeof Symbol === "function" && o[Symbol.iterator];
+    if (!m) return o;
+    var i = m.call(o), r, ar = [], e;
+    try {
+        while ((n === void 0 || n-- > 0) && !(r = i.next()).done) ar.push(r.value);
+    }
+    catch (error) { e = { error: error }; }
+    finally {
+        try {
+            if (r && !r.done && (m = i["return"])) m.call(i);
+        }
+        finally { if (e) throw e.error; }
+    }
+    return ar;
+};
+var __spreadArray = (this && this.__spreadArray) || function (to, from, pack) {
+    if (pack || arguments.length === 2) for (var i = 0, l = from.length, ar; i < l; i++) {
+        if (ar || !(i in from)) {
+            if (!ar) ar = Array.prototype.slice.call(from, 0, i);
+            ar[i] = from[i];
+        }
+    }
+    return to.concat(ar || Array.prototype.slice.call(from));
+};
+import { makeFnStatement } from './ast';
+var log = function () {
+    var _a;
+    var args = [];
+    for (var _i = 0; _i < arguments.length; _i++) {
+        args[_i] = arguments[_i];
+    }
+    return (_a = console.log).call.apply(_a, __spreadArray([console, '\x1b[31m(core.whitespace)\x1b[0m'], __read(args), false));
+};
+// Typescript fucked up flat https://stackoverflow.com/a/61420611/743464
+export var combineWs = function (a, b) { return [a, b].flat(Infinity); };
+// Move whitespace from one node to another. Since whitespace is trailing, if a
+// node is injected after a previous node, move the whitespace from the earlier
+// node to the later one. This keeps comments in the same place.
+export var transferWhitespace = function (to, from) {
+    return 'semi' in to && 'semi' in from
+        ? [
+            __assign(__assign({}, to), { semi: __assign(__assign({}, to.semi), { whitespace: from.semi.whitespace }) }),
+            __assign(__assign({}, from), { semi: __assign(__assign({}, from.semi), { whitespace: '\n' }) }),
+        ]
+        : 'whitespace' in to && 'semi' in from
+            ? [
+                __assign(__assign({}, to), { whitespace: combineWs(to.whitespace, from.semi.whitespace) }),
+                __assign(__assign({}, from), { semi: __assign(__assign({}, from.semi), { whitespace: '\n' }) }),
+            ]
+            : [to, from];
+};
+export var getLiteralIndent = function (node) {
+    return [(node === null || node === void 0 ? void 0 : node.whitespace) || '']
+        .flat(Infinity)
+        .join('')
+        .split(/\r|\n/)
+        .sort()
+        .at(-1);
+};
+export var tryAddTrailingWhitespace = function (node, ws) {
+    return 'semi' in node
+        ? __assign(__assign({}, node), { semi: __assign(__assign({}, node.semi), { whitespace: combineWs(node.semi.whitespace, ws) }) }) : node;
+};
+export var guessFnIndent = function (fnBody) {
+    return getLiteralIndent(fnBody.body.lb) ||
+        fnBody.body.statements.reduce(function (ws, n) {
+            return ws || getLiteralIndent(n.semi);
+        }, '');
+};
+export var addFnStmtWithIndent = function (fnBody, newNode) {
+    var statements = fnBody.body.statements;
+    var indent = guessFnIndent(fnBody);
+    return __spreadArray(__spreadArray([], __read(statements), false), [
+        // This simple hack is way easier than trying to modify the function body
+        // opening brace and/or the previous statement
+        { type: 'literal', literal: '', whitespace: indent },
+        tryAddTrailingWhitespace(typeof newNode === 'string' ? makeFnStatement(newNode)[0] : newNode, "\n"),
+    ], false);
+};

package/util/whitespace.test.d.ts

@@ -0,0 +1 @@
+export {};

package/util/whitespace.test.js

@@ -0,0 +1,13 @@
+import { expect, it } from 'vitest';
+import { parser } from '@shaderfrog/glsl-parser';
+import { generate } from '@shaderfrog/glsl-parser';
+import { addFnStmtWithIndent } from './whitespace';
+import { findMainOrThrow } from './ast';
+it("addFnStmtWithIndent", function () {
+    var source = "void main() {\n  vec2 y;\n}\n";
+    var ast = parser.parse(source, { quiet: true });
+    var m = findMainOrThrow(ast);
+    m.body.statements = addFnStmtWithIndent(m, "return x");
+    // Should line up the whitespace properly!
+    expect(generate(m)).toBe("void main() {\n  vec2 y;\n  return x;\n}\n");
+});