First commit

extract_nodes.js

@@ -0,0 +1,238 @@
+import { parse } 				from "@babel/parser";
+import * as babelTraverse 		from "@babel/traverse";
+import generator from "@babel/generator";
+const generate = generator.default;
+
+
+import Shader 					from "../framework/WebGpu.js";
+
+import { readFile } 			from "fs/promises";
+import { writeFile } 			from "fs/promises";
+
+const traverse = babelTraverse.default.default;
+
+async function loadFile( pathName ) {
+
+	const content = await readFile( pathName, "utf-8" );
+
+	return content;
+
+}
+
+async function saveToJsonFile( pathName, object ) {
+
+	const jsonContent = JSON.stringify( object, null, 4 );
+
+	await writeFile( pathName, jsonContent, "utf-8" );
+
+}
+
+function parseBindings( wgsl ) {
+
+	var bindings = [];
+
+	const regex = /@group\((\d+)\)\s+@binding\((\d+)\)\s+var<(\w+)(?:,\s*(\w+))?>\s+(\w+)\s*:\s*([^;]+);/g;
+
+	let match;
+
+	while ( ( match = regex.exec( wgsl ) ) !== null ) {
+
+		const group 		= Number( match[1] );
+		const binding 		= Number( match[2] );
+		const storageClass 	= match[3];
+		const access 		= match[4];
+		const varName 		= match[5];
+
+		let type;
+
+		if ( storageClass === "storage" ) {
+
+			type = access === "read" ? "read-only-storage" : "storage";
+
+		} else if ( storageClass === "uniform" ) {
+
+			type = "uniform";
+
+		} else {
+
+			type = "storage";
+
+		}
+
+		const varType = match[6];
+
+		bindings.push({ group, binding, type, varName, varType });
+
+	}
+
+	bindings.sort( ( a, b ) => a.binding - b.binding );
+
+	return bindings;
+
+}
+
+function extractShaderEdges(code) {
+
+	const shaderInstances = new Map();
+	const edges = [];
+	const ifStatements = [];
+	const definitions = [];
+
+	const ast = parse(code, {
+		sourceType: "module",
+		plugins: ["classProperties"]
+	});
+
+	traverse(ast, {
+
+		IfStatement(path) {
+
+			const { test, consequent, alternate, start, end } = path.node;
+
+			ifStatements.push({
+				test: code.slice(test.start, test.end),
+				consequent: code.slice(consequent.start, consequent.end),
+				alternate: alternate ? code.slice(alternate.start, alternate.end) : null,
+				start,
+				end
+			});
+		},
+
+		AssignmentExpression(path) {
+
+			const node = path.node;
+
+			if (node.left.type === "MemberExpression") {
+
+				const shaderName = node.left.property.name;
+
+				if (node.right.type === "NewExpression" && node.right.callee.name === "Shader") {
+
+					const args = node.right.arguments;
+
+					if (args.length >= 2 && args[1].type === "StringLiteral") {
+
+						const shaderPath = args[1].value;
+
+						// Store more precise location info for matching
+						shaderInstances.set(shaderName, {
+							name: shaderName,
+							path: shaderPath,
+							start: path.parentPath.node.start,
+							end: path.parentPath.node.end
+						});
+					}
+				}
+			}
+		},
+
+		CallExpression(path) {
+
+			const node = path.node;
+
+			if (node.callee.type === "MemberExpression") {
+
+				const object = node.callee.object;
+				const method = node.callee.property.name;
+
+				if (object.type === "MemberExpression" && object.object.type === "ThisExpression") {
+
+					const toShaderName = object.property.name;
+
+					if (method === "setBuffer" && node.arguments.length === 2) {
+
+						const bufferName = node.arguments[0].type === "StringLiteral"
+							? node.arguments[0].value
+							: null;
+
+						const arg = node.arguments[1];
+
+						if (
+							arg.type === "CallExpression" &&
+							arg.callee.type === "MemberExpression" &&
+							arg.callee.property.name === "getBuffer" &&
+							arg.callee.object.type === "MemberExpression" &&
+							arg.callee.object.object.type === "ThisExpression"
+						) {
+
+							const fromShaderName = arg.callee.object.property.name;
+
+							const fromBufferName = arg.arguments.length === 1 && arg.arguments[0].type === "StringLiteral"
+								? arg.arguments[0].value
+								: null;
+
+							edges.push({
+								from: fromShaderName,
+								to: toShaderName,
+								buffer: bufferName,
+								fromBuffer: fromBufferName
+							});
+						}
+					}
+				}
+			}
+		},
+
+		VariableDeclarator(path) {
+			const name = path.node.id.name;
+			const init = path.node.init ? generate(path.node.init).code : null;
+			definitions.push({
+				name,
+				init,
+				start: path.node.start,
+				end: path.node.end
+			});
+		},
+
+		ClassProperty(path) {
+			const name = path.node.key.name;
+			const init = path.node.value ? generate(path.node.value).code : null;
+			definitions.push({
+				name,
+				init,
+				start: path.node.start,
+				end: path.node.end
+			});
+		}
+
+	});
+
+	// Attach ifStatements to shaders they are declared within
+	for (const shader of shaderInstances.values()) {
+
+		for (const ifStmt of ifStatements) {
+
+			if (shader.start >= ifStmt.start && shader.end <= ifStmt.end) {
+
+				if (!shader.ifStatements) shader.ifStatements = [];
+
+				shader.ifStatements.push(ifStmt);
+			}
+		}
+	}
+
+	return {
+		shaders: Array.from(shaderInstances.values()),
+		edges: edges,
+		definitions: definitions
+	};
+}
+
+const code 		= await loadFile("../Graphics/demo.js");
+
+const result 	= extractShaderEdges(code);
+
+for (var i = 0; i < result.shaders.length; i++) {
+
+	var shaderInfo 	= result.shaders[i];
+
+	var source 		= await loadFile(shaderInfo.path);
+
+	var bindings	= parseBindings(source);
+
+	shaderInfo.bindings = bindings;
+}
+
+await saveToJsonFile("nodes.json", result);
+
+console.log(result);