@polygonjs/polygonjs/src/engine/nodes/gl/Subnet.ts

node-based WebGL 3D engine https://polygonjs.com

/**
 * a subnet can contain many nodes and is very useful to organise your shaders
 *
 *
 *
 */
import {GlType} from './../../poly/registers/nodes/types/Gl';
import {Constructor, valueof} from '../../../types/GlobalTypes';
import {TypedGlNode, BaseGlNodeType} from './_Base';
import {GlConnectionPointType, GL_CONNECTION_POINT_TYPES} from '../utils/io/connections/Gl';
import {NodeParamsConfig, ParamConfig} from '../utils/params/ParamsConfig';
import {ShadersCollectionController} from './code/utils/ShadersCollectionController';
import {NetworkChildNodeType, NetworkNodeType, NodeContext} from '../../poly/NodeContext';
import {GlNodeChildrenMap} from '../../poly/registers/nodes/Gl';
import {SubnetOutputGlNode} from './SubnetOutput';
import {ThreeToGl} from '../../../core/ThreeToGl';
import {SubnetInputGlNode} from './SubnetInput';
import {NodeCreateOptions} from '../utils/hierarchy/ChildrenController';
import {rangeStartEnd} from '../../../core/ArrayUtils';
import {IntegerParam} from '../../params/Integer';
import {StringParam} from '../../params/String';
import {TypedNodeTraverser} from '../utils/shaders/NodeTraverser';
import {CodeBuilder} from './code/utils/CodeBuilder';
import {LineType} from './code/utils/LineType';
import {BaseGLDefinition, FunctionGLDefinition} from './utils/GLDefinition';
import {CodeFormatter} from './code/utils/CodeFormatter';
import {ShaderName} from '../utils/shaders/ShaderName';
import {AddBodyLinesOptions} from './code/utils/LinesController';

export const ADD_BODY_LINES_OPTIONS: AddBodyLinesOptions = {
	makeUniq: false,
};

function visibleIfInputsCountAtLeast(index: number) {
	return {
		visibleIf: rangeStartEnd(index + 1, 10).map((i) => ({inputsCount: i})),
	};
}

function inputTypeParam(index: number) {
	return ParamConfig.INTEGER(GL_CONNECTION_POINT_TYPES.indexOf(GlConnectionPointType.FLOAT), {
		menu: {
			entries: GL_CONNECTION_POINT_TYPES.map((name, i) => {
				return {name: name, value: i};
			}),
		},
		separatorBefore: true,
		...visibleIfInputsCountAtLeast(index),
	});
}

function inputNameParam(index: number) {
	return ParamConfig.STRING(`input${index}`, {
		...visibleIfInputsCountAtLeast(index),
	});
}

export function TypedSubnetGlParamsConfigMixin<TBase extends Constructor>(Base: TBase) {
	return class Mixin extends Base {
		inputs = ParamConfig.FOLDER();
		inputsCount = ParamConfig.INTEGER(1, {
			range: [0, 10],
			rangeLocked: [true, true],
		});
		inputType0 = inputTypeParam(0);
		inputName0 = inputNameParam(0);
		inputType1 = inputTypeParam(1);
		inputName1 = inputNameParam(1);
		inputType2 = inputTypeParam(2);
		inputName2 = inputNameParam(2);
		inputType3 = inputTypeParam(3);
		inputName3 = inputNameParam(3);
		inputType4 = inputTypeParam(4);
		inputName4 = inputNameParam(4);
		inputType5 = inputTypeParam(5);
		inputName5 = inputNameParam(5);
		inputType6 = inputTypeParam(6);
		inputName6 = inputNameParam(6);
		inputType7 = inputTypeParam(7);
		inputName7 = inputNameParam(7);
		inputType8 = inputTypeParam(8);
		inputName8 = inputNameParam(8);
		inputType9 = inputTypeParam(9);
		inputName9 = inputNameParam(9);
		spare = ParamConfig.FOLDER();
	};
}
class TypedSubnetGlParamsConfig extends TypedSubnetGlParamsConfigMixin(NodeParamsConfig) {}
export class AbstractTypedSubnetGlNode<K extends NodeParamsConfig> extends TypedGlNode<K> {
	protected override _childrenControllerContext = NodeContext.GL;

	override initializeNode() {
		this.childrenController?.setOutputNodeFindMethod(() => {
			return this.nodesByType(SubnetOutputGlNode.type())[0];
		});

		this.io.connection_points.set_input_name_function(this._expectedInputName.bind(this));

		this.io.connection_points.set_expected_input_types_function(this._expectedInputTypes.bind(this));
		this.io.connection_points.set_expected_output_types_function(this._expectedOutputTypes.bind(this));
		this.io.connection_points.set_output_name_function(this._expectedOutputName.bind(this));
	}
	protected _expectedInputTypes(): GlConnectionPointType[] {
		return [];
	}
	protected _expectedInputName(index: number) {
		return 'default';
	}

	protected _expectedOutputTypes() {
		return this._expectedInputTypes();
	}

	protected _expectedOutputName(index: number) {
		return this._expectedInputName(index);
	}
	//
	//
	// defines the outputs for the child subnet input
	//
	//
	childExpectedInputConnectionPointTypes() {
		return this._expectedInputTypes();
	}
	childExpectedOutputConnectionPointTypes() {
		return this._expectedOutputTypes();
	}
	childExpectedInputConnectionPointName(index: number) {
		return this._expectedInputName(index);
	}
	childExpectedOutputConnectionPointName(index: number) {
		return this._expectedOutputName(index);
	}

	//
	//
	// CHILDREN
	//
	//
	override createNode<S extends keyof GlNodeChildrenMap>(
		node_class: S,
		options?: NodeCreateOptions
	): GlNodeChildrenMap[S];
	override createNode<K extends valueof<GlNodeChildrenMap>>(
		node_class: Constructor<K>,
		options?: NodeCreateOptions
	): K;
	override createNode<K extends valueof<GlNodeChildrenMap>>(
		node_class: Constructor<K>,
		options?: NodeCreateOptions
	): K {
		return super.createNode(node_class, options) as K;
	}
	override children() {
		return super.children() as BaseGlNodeType[];
	}
	override nodesByType<K extends keyof GlNodeChildrenMap>(type: K): GlNodeChildrenMap[K][] {
		return super.nodesByType(type) as GlNodeChildrenMap[K][];
	}

	//
	//
	// set_lines
	//
	//
	protected _setLinesPreBlock(shadersCollectionController: ShadersCollectionController) {
		const bodyLines: string[] = [];
		const connection_points = this.io.inputs.namedInputConnectionPoints();
		if (!connection_points) {
			return;
		}
		for (let i = 0; i < connection_points.length; i++) {
			const connection_point = connection_points[i];
			const gl_type = connection_point.type();
			const out = this.glVarName(connection_point.name());
			const in_value = ThreeToGl.any(this.variableForInput(connection_point.name()));
			const body_line = `${gl_type} ${out} = ${in_value}`;
			bodyLines.push(body_line);
		}

		shadersCollectionController.addBodyLines(this, bodyLines);
	}
	protected setLinesBlockStart(shadersCollectionController: ShadersCollectionController) {
		shadersCollectionController.addBodyLines(this, [`if(true){`]);
	}
	protected setLinesBlockEnd(shadersCollectionController: ShadersCollectionController) {
		shadersCollectionController.addBodyLines(this, ['}']);
	}
	setSubnetInputLines(shadersCollectionController: ShadersCollectionController, childNode: SubnetInputGlNode) {
		const connections = this.io.connections.inputConnections();
		if (!connections) {
			return;
		}
		const bodyLines: string[] = [];
		for (const connection of connections) {
			if (connection) {
				const connection_point = connection.destConnectionPoint();
				if (connection_point) {
					const in_value = ThreeToGl.any(this.variableForInput(connection_point.name()));
					const gl_type = connection_point.type();
					const out = childNode.glVarName(connection_point.name());
					const body_line = `	${gl_type} ${out} = ${in_value}`;
					bodyLines.push(body_line);
				}
			}
		}
		shadersCollectionController.addBodyLines(childNode, bodyLines, undefined, ADD_BODY_LINES_OPTIONS);
	}
	private subnetOutputLines(childNode: SubnetOutputGlNode): string[] {
		const connections = childNode.io.connections.inputConnections();
		if (!connections) {
			return [];
		}
		const bodyLines: string[] = [];

		for (const connection of connections) {
			if (connection) {
				const connectionPoint = connection.destConnectionPoint();
				if (connectionPoint) {
					const in_value = ThreeToGl.any(childNode.variableForInput(connectionPoint.name()));
					const out = this.glVarName(connectionPoint.name());
					// const body_line = `${gl_type} ${out} = ${in_value}`;
					// do not use the type, to avoid re-defining a variable that should be defined in the parent node
					const bodyLine = `	${out} = ${in_value}`;
					bodyLines.push(bodyLine);
				}
			}
		}
		return bodyLines;
	}
	setSubnetOutputLines(shadersCollectionController: ShadersCollectionController, childNode: SubnetOutputGlNode) {
		const bodyLines: string[] = this.subnetOutputLines(childNode);
		shadersCollectionController.addBodyLines(childNode, bodyLines, undefined, ADD_BODY_LINES_OPTIONS);
	}

	// set_lines_block_end(shadersCollectionController: ShadersCollectionController, childNode: SubnetOutputGlNode) {
	// 	shadersCollectionController.addBodyLines(childNode, ['}']);
	// }

	override setLines(shadersCollectionController: ShadersCollectionController) {
		this._setLinesPreBlock(shadersCollectionController);
		this.setLinesBlockStart(shadersCollectionController);
		this._setLinesBlockContent(shadersCollectionController);
		this.setLinesBlockEnd(shadersCollectionController);
	}
	protected linesBlockContent(shadersCollectionController: ShadersCollectionController) {
		const codeBuilder = this._runCodeBuilder(shadersCollectionController);
		if (!codeBuilder) {
			return;
		}
		const shadername = shadersCollectionController.currentShaderName();
		const bodyLines = codeBuilder.lines(shadername, LineType.BODY);
		return this._sanitizeBodyLines(bodyLines);
	}
	private _setLinesBlockContent(shadersCollectionController: ShadersCollectionController) {
		const bodyLines = this.linesBlockContent(shadersCollectionController);
		if (!bodyLines) {
			return;
		}
		shadersCollectionController.addBodyLines(this, bodyLines, undefined, ADD_BODY_LINES_OPTIONS);
	}
	protected _runCodeBuilder(shadersCollectionController: ShadersCollectionController) {
		// I potentially could look for attribute nodes to use as output,
		// but for now, I'll enforce a rule that attribute nodes must be at the top level
		const outputNodes: SubnetOutputGlNode[] = this.nodesByType(NetworkChildNodeType.OUTPUT);
		const matNode = this.materialNode();
		if (!matNode) {
			return;
		}
		if (outputNodes.length == 0) {
			matNode.states.error.set(`${this.path()}:one output node is required`);
		}
		if (outputNodes.length > 1) {
			matNode.states.error.set(`${this.path()}:only one output node allowed`);
		}
		const subnetOutput = outputNodes[0];
		const subnetOutputInputConnectionPoints = subnetOutput.io.inputs.namedInputConnectionPoints();

		const subnetOutputInputNames = subnetOutputInputConnectionPoints
			? subnetOutputInputConnectionPoints.map((cp) => cp.name())
			: [];

		const assembler = shadersCollectionController.assembler();

		const nodeTraverser = new TypedNodeTraverser<NodeContext.GL>(
			this,
			shadersCollectionController.shaderNames(),
			(rootNode, shaderName) => {
				return subnetOutputInputNames;
			}
		);
		const codeBuilder = new CodeBuilder(
			nodeTraverser,
			(shaderName, rootNodes) => {
				// return [subnetOutput];
				return assembler.rootNodesByShaderName(shaderName, rootNodes);
			},
			assembler
		);
		const paramNodes: BaseGlNodeType[] = [];
		codeBuilder.buildFromNodes(outputNodes, paramNodes);
		this._addCodeBuilderDefinition(codeBuilder, shadersCollectionController);
		return codeBuilder;
	}
	private _addCodeBuilderDefinition(
		codeBuilder: CodeBuilder,
		shadersCollectionController: ShadersCollectionController
	) {
		const internalShadersCollectionController = codeBuilder.shadersCollectionController();
		if (!internalShadersCollectionController) {
			return;
		}
		const currentShaderName = shadersCollectionController.currentShaderName();
		internalShadersCollectionController.setCurrentShaderName(currentShaderName);

		// 1- add all definitions for each shaderName
		const shaderNames = shadersCollectionController.shaderNames();
		for (const shaderName of shaderNames) {
			const definitions: BaseGLDefinition[] = [];
			internalShadersCollectionController.traverseDefinitions(shaderName, (definition) => {
				// only add function if it is for the current shader
				const isNotFunction = !(definition instanceof FunctionGLDefinition);
				const isCurrentShader = shaderName == currentShaderName;
				if (isNotFunction || isCurrentShader) {
					definitions.push(definition);
				}
			});
			shadersCollectionController.addDefinitions(this, definitions, shaderName);
		}
		// 2- add vertex body lines if current shader name is fragment
		if (currentShaderName != ShaderName.VERTEX) {
			const attribNodes = this.nodesByType(GlType.ATTRIBUTE);
			const bodyLines: string[] = [];
			for (const attribNode of attribNodes) {
				const linesForNode = internalShadersCollectionController.bodyLines(ShaderName.VERTEX, attribNode);
				if (linesForNode) {
					bodyLines.push(...linesForNode);
				}
			}
			shadersCollectionController.addBodyLines(this, bodyLines, ShaderName.VERTEX, ADD_BODY_LINES_OPTIONS);
		}
	}

	// align with the right number of tabs
	protected _sanitizeBodyLines(lines: string[]): string[] {
		const level = CodeFormatter.nodeDistanceToMaterial(this);
		const prefix = `\t`.repeat(level);

		return lines.map((line) => {
			const trimmed = line.trim();
			if (trimmed.length == 0) {
				return '';
			} else {
				return `${prefix}${trimmed}`;
			}
		});
	}
}

export class TypedSubnetGlNode<K extends TypedSubnetGlParamsConfig> extends AbstractTypedSubnetGlNode<K> {
	override initializeNode() {
		super.initializeNode();

		this.io.connection_points.set_input_name_function(this._expectedInputName.bind(this));

		this.io.connection_points.set_expected_input_types_function(this._expectedInputTypes.bind(this));
		this.io.connection_points.set_expected_output_types_function(this._expectedOutputTypes.bind(this));
		this.io.connection_points.set_output_name_function(this._expectedOutputName.bind(this));
	}
	protected _inputTypeParams(): IntegerParam[] {
		return [
			this.p.inputType0,
			this.p.inputType1,
			this.p.inputType2,
			this.p.inputType3,
			this.p.inputType4,
			this.p.inputType5,
			this.p.inputType6,
			this.p.inputType7,
			this.p.inputType8,
			this.p.inputType9,
		];
	}
	protected _inputNameParams(): StringParam[] {
		return [
			this.p.inputName0,
			this.p.inputName1,
			this.p.inputName2,
			this.p.inputName3,
			this.p.inputName4,
			this.p.inputName5,
			this.p.inputName6,
			this.p.inputName7,
			this.p.inputName8,
			this.p.inputName9,
		];
	}

	setInputType(index: number, type: GlConnectionPointType) {
		const param = this._inputTypeParams()[index];
		if (!param) {
			return;
		}
		param.set(GL_CONNECTION_POINT_TYPES.indexOf(type));
	}
	setInputName(index: number, inputName: string) {
		const param = this._inputNameParams()[index];
		if (!param) {
			return;
		}
		param.set(inputName);
	}

	protected _expectedInputsCount(): number {
		return this.pv.inputsCount;
	}

	protected override _expectedInputTypes(): GlConnectionPointType[] {
		const count = this.pv.inputsCount;
		const params: IntegerParam[] = this._inputTypeParams();
		return rangeStartEnd(0, count).map((value, i) => GL_CONNECTION_POINT_TYPES[params[i].value]);
	}
	protected override _expectedInputName(index: number) {
		const params: StringParam[] = this._inputNameParams();
		const param = params[index];
		return param ? param.value : GlConnectionPointType.FLOAT;
	}

	protected override _expectedOutputTypes() {
		const count = this.pv.inputsCount;
		const params: IntegerParam[] = this._inputTypeParams();
		return rangeStartEnd(0, count).map((value, i) => GL_CONNECTION_POINT_TYPES[params[i].value]);
	}

	protected override _expectedOutputName(index: number) {
		// return this._expected_input_name(index);
		const params: StringParam[] = this._inputNameParams();
		return params[index].value;
	}
}

class SubnetGlParamsConfig extends TypedSubnetGlParamsConfigMixin(NodeParamsConfig) {}
const ParamsConfig = new SubnetGlParamsConfig();

export class SubnetGlNode extends TypedSubnetGlNode<SubnetGlParamsConfig> {
	override paramsConfig = ParamsConfig;
	static override type() {
		return NetworkNodeType.SUBNET;
	}
}