strands_codegen.js

import { sortCFG } from './ir_cfg';
import { structType } from './ir_types';

export function generateShaderCode(strandsContext) {
  const {
    cfg,
    backend,
    vertexDeclarations,
    fragmentDeclarations,
    computeDeclarations
  } = strandsContext;

  const hooksObj = {
    uniforms: {},
    storageUniforms: {},
    varyingVariables: [],
    shaderNameMap: strandsContext.shaderNameMap,
    shaderNameState: strandsContext.shaderNameState
  };

  for (const { name, typeInfo, defaultValue } of strandsContext.uniforms) {
    if (typeInfo.baseType === 'storage') {
      if (defaultValue !== null && defaultValue !== undefined) {
        hooksObj.storageUniforms[name] = defaultValue;
      }
    } else {
      const key = backend.generateHookUniformKey(name, typeInfo);
      if (key !== null) {
        hooksObj.uniforms[key] = defaultValue;
      }
    }
  }

  // Add texture bindings to declarations for WebGPU backend
  if (backend.addTextureBindingsToDeclarations) {
    backend.addTextureBindingsToDeclarations(strandsContext);
  }

  // Add storage buffer bindings to declarations for WebGPU backend
  if (backend.addStorageBufferBindingsToDeclarations) {
    backend.addStorageBufferBindingsToDeclarations(strandsContext);
  }

  for (const {
    hookType,
    rootNodeID,
    entryBlockID,
    shaderContext
  } of strandsContext.hooks) {
    const generationContext = {
      indent: 1,
      codeLines: [],
      write(line) {
        this.codeLines.push('  '.repeat(this.indent) + line);
      },
      tempNames: {},
      declarations: [],
      nextTempID: 0,
      visitedNodes: new Set(),
      shaderContext, // 'vertex' or 'fragment'
      strandsContext // For shared variable tracking
    };

    const blocks = sortCFG(cfg.outgoingEdges, entryBlockID);
    for (const blockID of blocks) {
      backend.generateBlock(blockID, strandsContext, generationContext);
    }

    const firstLine = backend.hookEntry(hookType);
    let returnType;
    if (hookType.returnType.properties) {
      returnType = structType(hookType.returnType);
    } else if (
      !hookType.returnType.dataType ||
      hookType.returnType.typeName?.trim() === 'void'
    ) {
      returnType = null;
    } else {
      returnType = hookType.returnType.dataType;
    }

    if (rootNodeID !== undefined) {
      backend.generateReturnStatement(
        strandsContext,
        generationContext,
        rootNodeID,
        returnType
      );
    }
    hooksObj[`${hookType.returnType.typeName} ${hookType.name}`] = [
      firstLine,
      ...generationContext.codeLines,
      '}'
    ].join('\n');
  }

  // Finalize shared variable declarations based on usage
  if (strandsContext.sharedVariables) {
    for (const [varName, varInfo] of strandsContext.sharedVariables) {
      if (varInfo.usedInVertex && varInfo.usedInFragment) {
        // Used in both shaders - this is a true varying variable
        hooksObj.varyingVariables.push(
          backend.generateVaryingVariable(varName, varInfo.typeInfo)
        );
      } else if (varInfo.usedInVertex) {
        // Only used in vertex shader - declare as local variable
        vertexDeclarations.add(
          backend.generateLocalDeclaration(varName, varInfo.typeInfo)
        );
      } else if (varInfo.usedInFragment) {
        // Only used in fragment shader - declare as local variable
        fragmentDeclarations.add(
          backend.generateLocalDeclaration(varName, varInfo.typeInfo)
        );
      }
      // If not used anywhere, don't declare it
    }
  }

  // Register instanceID varying if used in a fragment hook
  if (strandsContext._instanceIDUsedInFragment) {
    hooksObj.instanceIDVarying = backend.generateInstanceIDVarying();
  }

  hooksObj.vertexDeclarations = [...vertexDeclarations].join('\n');
  hooksObj.fragmentDeclarations = [...fragmentDeclarations].join('\n');
  hooksObj.computeDeclarations = [...computeDeclarations].join('\n');

  return hooksObj;
}