cesium/Source/Scene/ModelExperimental/CustomShaderPipelineStage.js

CesiumJS is a JavaScript library for creating 3D globes and 2D maps in a web browser without a plugin.

import combine from "../../Core/combine.js";
import defined from "../../Core/defined.js";
import oneTimeWarning from "../../Core/oneTimeWarning.js";
import ShaderDestination from "../../Renderer/ShaderDestination.js";
import Pass from "../../Renderer/Pass.js";
import CustomShaderStageVS from "../../Shaders/ModelExperimental/CustomShaderStageVS.js";
import CustomShaderStageFS from "../../Shaders/ModelExperimental/CustomShaderStageFS.js";
import AlphaMode from "../AlphaMode.js";
import CustomShaderMode from "./CustomShaderMode.js";
import FeatureIdPipelineStage from "./FeatureIdPipelineStage.js";
import MetadataPipelineStage from "./MetadataPipelineStage.js";
import ModelExperimentalUtility from "./ModelExperimentalUtility.js";

/**
 * The custom shader pipeline stage takes GLSL callbacks from the
 * {@link CustomShader} and inserts them into the overall shader code for the
 * {@link ModelExperimental}. The input to the callback is a struct with many
 * properties that depend on the attributes of the primitive. This shader code
 * is automatically generated by this stage.
 *
 * @namespace CustomShaderPipelineStage
 *
 * @private
 */
const CustomShaderPipelineStage = {};
CustomShaderPipelineStage.name = "CustomShaderPipelineStage"; // Helps with debugging

CustomShaderPipelineStage.STRUCT_ID_ATTRIBUTES_VS = "AttributesVS";
CustomShaderPipelineStage.STRUCT_ID_ATTRIBUTES_FS = "AttributesFS";
CustomShaderPipelineStage.STRUCT_NAME_ATTRIBUTES = "Attributes";
CustomShaderPipelineStage.STRUCT_ID_VERTEX_INPUT = "VertexInput";
CustomShaderPipelineStage.STRUCT_NAME_VERTEX_INPUT = "VertexInput";
CustomShaderPipelineStage.STRUCT_ID_FRAGMENT_INPUT = "FragmentInput";
CustomShaderPipelineStage.STRUCT_NAME_FRAGMENT_INPUT = "FragmentInput";
CustomShaderPipelineStage.FUNCTION_ID_INITIALIZE_INPUT_STRUCT_VS =
  "initializeInputStructVS";
CustomShaderPipelineStage.FUNCTION_SIGNATURE_INITIALIZE_INPUT_STRUCT_VS =
  "void initializeInputStruct(out VertexInput vsInput, ProcessedAttributes attributes)";
CustomShaderPipelineStage.FUNCTION_ID_INITIALIZE_INPUT_STRUCT_FS =
  "initializeInputStructFS";
CustomShaderPipelineStage.FUNCTION_SIGNATURE_INITIALIZE_INPUT_STRUCT_FS =
  "void initializeInputStruct(out FragmentInput fsInput, ProcessedAttributes attributes)";

/**
 * Process a primitive. This modifies the following parts of the render
 * resources:
 * <ul>
 *   <li>Modifies the shader to include the custom shader code to the vertex and fragment shaders</li>
 *   <li>Modifies the shader to include automatically-generated structs that serve as input to the custom shader callbacks </li>
 *   <li>Modifies the shader to include any additional user-defined uniforms</li>
 *   <li>Modifies the shader to include any additional user-defined varyings</li>
 *   <li>Adds any user-defined uniforms to the uniform map</li>
 *   <li>If the user specified a lighting model, the settings are overridden in the render resources</li>
 * </ul>
 * <p>
 * This pipeline stage is designed to fail gracefully where possible. If the
 * primitive does not have the right attributes to satisfy the shader code,
 * defaults will be inferred (when reasonable to do so). If not, the custom
 * shader will be disabled.
 * <p>
 *
 * @param {PrimitiveRenderResources} renderResources The render resources for the primitive
 * @param {ModelComponents.Primitive} primitive The primitive to be rendered
 * @param {FrameState} frameState The frame state.
 * @private
 */
CustomShaderPipelineStage.process = function (
  renderResources,
  primitive,
  frameState
) {
  const shaderBuilder = renderResources.shaderBuilder;
  const customShader = renderResources.model.customShader;

  // Check the lighting model and translucent options first, as sometimes
  // these are used even if there is no vertex or fragment shader text.

  // if present, the lighting model overrides the material's lighting model.
  if (defined(customShader.lightingModel)) {
    renderResources.lightingOptions.lightingModel = customShader.lightingModel;
  }

  const alphaOptions = renderResources.alphaOptions;
  if (customShader.isTranslucent) {
    alphaOptions.pass = Pass.TRANSLUCENT;
    alphaOptions.alphaMode = AlphaMode.BLEND;
  } else {
    // Use the default pass (either OPAQUE or 3D_TILES), regardless of whether
    // the material pipeline stage used translucent. The default is configured
    // in AlphaPipelineStage
    alphaOptions.pass = undefined;
    alphaOptions.alphaMode = AlphaMode.OPAQUE;
  }

  // Generate lines of code for the shader, but don't add them to the shader
  // yet.
  const generatedCode = generateShaderLines(customShader, primitive);

  // In some corner cases, the primitive may not be compatible with the
  // shader. In this case, skip the custom shader.
  if (!generatedCode.customShaderEnabled) {
    return;
  }
  addLinesToShader(shaderBuilder, customShader, generatedCode);

  // the input to the fragment shader may include a low-precision ECEF position
  if (generatedCode.shouldComputePositionWC) {
    shaderBuilder.addDefine(
      "COMPUTE_POSITION_WC",
      undefined,
      ShaderDestination.BOTH
    );
  }

  if (defined(customShader.vertexShaderText)) {
    shaderBuilder.addDefine(
      "HAS_CUSTOM_VERTEX_SHADER",
      undefined,
      ShaderDestination.VERTEX
    );
  }

  if (defined(customShader.fragmentShaderText)) {
    shaderBuilder.addDefine(
      "HAS_CUSTOM_FRAGMENT_SHADER",
      undefined,
      ShaderDestination.FRAGMENT
    );

    // add defines like CUSTOM_SHADER_MODIFY_MATERIAL
    const shaderModeDefine = CustomShaderMode.getDefineName(customShader.mode);
    shaderBuilder.addDefine(
      shaderModeDefine,
      undefined,
      ShaderDestination.FRAGMENT
    );
  }

  const uniforms = customShader.uniforms;
  for (const uniformName in uniforms) {
    if (uniforms.hasOwnProperty(uniformName)) {
      const uniform = uniforms[uniformName];
      shaderBuilder.addUniform(uniform.type, uniformName);
    }
  }

  const varyings = customShader.varyings;
  for (const varyingName in varyings) {
    if (varyings.hasOwnProperty(varyingName)) {
      const varyingType = varyings[varyingName];
      shaderBuilder.addVarying(varyingType, varyingName);
    }
  }

  renderResources.uniformMap = combine(
    renderResources.uniformMap,
    customShader.uniformMap
  );
};

function getAttributesByName(attributes) {
  const names = {};
  for (let i = 0; i < attributes.length; i++) {
    const attribute = attributes[i];
    const attributeInfo = ModelExperimentalUtility.getAttributeInfo(attribute);

    names[attributeInfo.variableName] = attributeInfo;
  }
  return names;
}

// GLSL types of standard attribute types when uniquely defined
const attributeTypeLUT = {
  position: "vec3",
  normal: "vec3",
  tangent: "vec3",
  bitangent: "vec3",
  texCoord: "vec2",
  color: "vec4",
  joints: "ivec4",
  weights: "vec4",
};

// Corresponding attribute values
const attributeDefaultValueLUT = {
  position: "vec3(0.0)",
  normal: "vec3(0.0, 0.0, 1.0)",
  tangent: "vec3(1.0, 0.0, 0.0)",
  bitangent: "vec3(0.0, 1.0, 0.0)",
  texCoord: "vec2(0.0)",
  color: "vec4(1.0)",
  joints: "ivec4(0)",
  weights: "vec4(0.0)",
};

function inferAttributeDefaults(attributeName) {
  // remove trailing set indices. E.g. "texCoord_0" -> "texCoord"
  let trimmed = attributeName.replace(/_[0-9]+$/, "");
  // also remove the MC/EC since they will have the same default value
  trimmed = trimmed.replace(/(MC|EC)$/, "");

  const glslType = attributeTypeLUT[trimmed];
  const value = attributeDefaultValueLUT[trimmed];

  // - _CUSTOM_ATTRIBUTE has an unknown type.
  if (!defined(glslType)) {
    return undefined;
  }

  return {
    attributeField: [glslType, attributeName],
    value: value,
  };
}

function generateVertexShaderLines(
  customShader,
  attributesByName,
  vertexLines
) {
  const categories = partitionAttributes(
    attributesByName,
    customShader.usedVariablesVertex.attributeSet,
    false
  );
  const addToShader = categories.addToShader;
  const needsDefault = categories.missingAttributes;

  let variableName;
  let vertexInitialization;
  const attributeFields = [];
  const initializationLines = [];
  for (variableName in addToShader) {
    if (addToShader.hasOwnProperty(variableName)) {
      const attributeInfo = addToShader[variableName];
      const attributeField = [attributeInfo.glslType, variableName];
      attributeFields.push(attributeField);

      // Initializing attribute structs are just a matter of copying the
      // attribute or varying: E.g.:
      // "    vsInput.attributes.position = a_position;"
      vertexInitialization = `vsInput.attributes.${variableName} = attributes.${variableName};`;
      initializationLines.push(vertexInitialization);
    }
  }

  for (let i = 0; i < needsDefault.length; i++) {
    variableName = needsDefault[i];
    const attributeDefaults = inferAttributeDefaults(variableName);
    if (!defined(attributeDefaults)) {
      CustomShaderPipelineStage._oneTimeWarning(
        "CustomShaderPipelineStage.incompatiblePrimitiveVS",
        `Primitive is missing attribute ${variableName}, disabling custom vertex shader`
      );
      // This primitive isn't compatible with the shader. Return early
      // to skip the vertex shader
      return;
    }

    attributeFields.push(attributeDefaults.attributeField);
    vertexInitialization = `vsInput.attributes.${variableName} = ${attributeDefaults.value};`;
    initializationLines.push(vertexInitialization);
  }

  vertexLines.enabled = true;
  vertexLines.attributeFields = attributeFields;
  vertexLines.initializationLines = initializationLines;
}

function generatePositionBuiltins(customShader) {
  const attributeFields = [];
  const initializationLines = [];
  const usedVariables = customShader.usedVariablesFragment.attributeSet;

  // Model space position is the same position as in the glTF accessor,
  // this is already added to the shader with other attributes.

  // World coordinates in ECEF coordinates. Note that this is
  // low precision (32-bit floats) on the GPU.
  if (usedVariables.hasOwnProperty("positionWC")) {
    attributeFields.push(["vec3", "positionWC"]);
    initializationLines.push(
      "fsInput.attributes.positionWC = attributes.positionWC;"
    );
  }

  // position in eye coordinates
  if (usedVariables.hasOwnProperty("positionEC")) {
    attributeFields.push(["vec3", "positionEC"]);
    initializationLines.push(
      "fsInput.attributes.positionEC = attributes.positionEC;"
    );
  }

  return {
    attributeFields: attributeFields,
    initializationLines: initializationLines,
  };
}

function generateFragmentShaderLines(
  customShader,
  attributesByName,
  fragmentLines
) {
  const categories = partitionAttributes(
    attributesByName,
    customShader.usedVariablesFragment.attributeSet,
    true
  );
  const addToShader = categories.addToShader;
  const needsDefault = categories.missingAttributes;

  let variableName;
  let fragmentInitialization;
  const attributeFields = [];
  const initializationLines = [];
  for (variableName in addToShader) {
    if (addToShader.hasOwnProperty(variableName)) {
      const attributeInfo = addToShader[variableName];

      const attributeField = [attributeInfo.glslType, variableName];
      attributeFields.push(attributeField);

      // Initializing attribute structs are just a matter of copying the
      // value from the processed attributes
      // "    fsInput.attributes.positionMC = attributes.positionMC;"
      fragmentInitialization = `fsInput.attributes.${variableName} = attributes.${variableName};`;
      initializationLines.push(fragmentInitialization);
    }
  }

  for (let i = 0; i < needsDefault.length; i++) {
    variableName = needsDefault[i];
    const attributeDefaults = inferAttributeDefaults(variableName);
    if (!defined(attributeDefaults)) {
      CustomShaderPipelineStage._oneTimeWarning(
        "CustomShaderPipelineStage.incompatiblePrimitiveFS",
        `Primitive is missing attribute ${variableName}, disabling custom fragment shader.`
      );

      // This primitive isn't compatible with the shader. Return early
      // so the fragment shader is skipped
      return;
    }

    attributeFields.push(attributeDefaults.attributeField);
    fragmentInitialization = `fsInput.attributes.${variableName} = ${attributeDefaults.value};`;
    initializationLines.push(fragmentInitialization);
  }

  // Built-ins for positions in various coordinate systems.
  const positionBuiltins = generatePositionBuiltins(customShader);

  fragmentLines.enabled = true;
  fragmentLines.attributeFields = attributeFields.concat(
    positionBuiltins.attributeFields
  );
  fragmentLines.initializationLines = positionBuiltins.initializationLines.concat(
    initializationLines
  );
}

// These attributes are derived from positionMC, and are handled separately
// from other attributes
const builtinAttributes = {
  positionWC: true,
  positionEC: true,
};

function partitionAttributes(
  primitiveAttributes,
  shaderAttributeSet,
  isFragmentShader
) {
  // shaderAttributes = set of all attributes used in the shader
  // primitiveAttributes = set of all the primitive's attributes
  // partition into three categories:
  // - addToShader = shaderAttributes intersect primitiveAttributes
  // - missingAttributes = shaderAttributes - primitiveAttributes - builtinAttributes
  // - unneededAttributes = (primitiveAttributes - shaderAttributes) U builtinAttributes
  //
  // addToShader are attributes that should be added to the shader.
  // missingAttributes are attributes for which we need to provide a default value
  // unneededAttributes are other attributes that can be skipped.

  let renamed;
  let attributeName;
  const addToShader = {};
  for (attributeName in primitiveAttributes) {
    if (primitiveAttributes.hasOwnProperty(attributeName)) {
      const attribute = primitiveAttributes[attributeName];

      // normals and tangents are in model coordinates in the attributes but
      // in eye coordinates in the fragment shader.
      renamed = attributeName;
      if (isFragmentShader && attributeName === "normalMC") {
        renamed = "normalEC";
      } else if (isFragmentShader && attributeName === "tangentMC") {
        renamed = "tangentEC";
      }

      if (shaderAttributeSet.hasOwnProperty(renamed)) {
        addToShader[renamed] = attribute;
      }
    }
  }

  const missingAttributes = [];
  for (attributeName in shaderAttributeSet) {
    if (shaderAttributeSet.hasOwnProperty(attributeName)) {
      if (builtinAttributes.hasOwnProperty(attributeName)) {
        // Builtins are handled separately from attributes, so skip them here
        continue;
      }

      // normals and tangents are in model coordinates in the attributes but
      // in eye coordinates in the fragment shader.
      renamed = attributeName;
      if (isFragmentShader && attributeName === "normalEC") {
        renamed = "normalMC";
      } else if (isFragmentShader && attributeName === "tangentEC") {
        renamed = "tangentMC";
      }

      if (!primitiveAttributes.hasOwnProperty(renamed)) {
        missingAttributes.push(attributeName);
      }
    }
  }

  return {
    addToShader: addToShader,
    missingAttributes: missingAttributes,
  };
}

function generateShaderLines(customShader, primitive) {
  // Assume shader code is disabled unless proven otherwise
  const vertexLines = {
    enabled: false,
  };
  const fragmentLines = {
    enabled: false,
  };

  // Attempt to generate vertex and fragment shader lines before adding any
  // code to the shader.
  const attributesByName = getAttributesByName(primitive.attributes);
  if (defined(customShader.vertexShaderText)) {
    generateVertexShaderLines(customShader, attributesByName, vertexLines);
  }

  if (defined(customShader.fragmentShaderText)) {
    generateFragmentShaderLines(customShader, attributesByName, fragmentLines);
  }

  // positionWC must be computed in the vertex shader
  // for use in the fragmentShader. However, this can be skipped if:
  // - positionWC isn't used in the fragment shader
  // - or the fragment shader is disabled
  const attributeSetFS = customShader.usedVariablesFragment.attributeSet;
  const shouldComputePositionWC =
    attributeSetFS.hasOwnProperty("positionWC") && fragmentLines.enabled;

  // Return any generated shader code along with some flags to indicate which
  // defines should be added.
  return {
    vertexLines: vertexLines,
    fragmentLines: fragmentLines,
    vertexLinesEnabled: vertexLines.enabled,
    fragmentLinesEnabled: fragmentLines.enabled,
    customShaderEnabled: vertexLines.enabled || fragmentLines.enabled,
    shouldComputePositionWC: shouldComputePositionWC,
  };
}

function addVertexLinesToShader(shaderBuilder, vertexLines) {
  // Vertex Lines ---------------------------------------------------------

  let i;
  let structId = CustomShaderPipelineStage.STRUCT_ID_ATTRIBUTES_VS;
  shaderBuilder.addStruct(
    structId,
    CustomShaderPipelineStage.STRUCT_NAME_ATTRIBUTES,
    ShaderDestination.VERTEX
  );

  const attributeFields = vertexLines.attributeFields;
  for (i = 0; i < attributeFields.length; i++) {
    const field = attributeFields[i];
    const glslType = field[0];
    const variableName = field[1];
    shaderBuilder.addStructField(structId, glslType, variableName);
  }

  // This could be hard-coded, but the symmetry with other structs makes unit
  // tests more convenient
  structId = CustomShaderPipelineStage.STRUCT_ID_VERTEX_INPUT;
  shaderBuilder.addStruct(
    structId,
    CustomShaderPipelineStage.STRUCT_NAME_VERTEX_INPUT,
    ShaderDestination.VERTEX
  );
  shaderBuilder.addStructField(
    structId,
    CustomShaderPipelineStage.STRUCT_NAME_ATTRIBUTES,
    "attributes"
  );
  // Add FeatureIds struct from the Feature ID stage
  shaderBuilder.addStructField(
    structId,
    FeatureIdPipelineStage.STRUCT_NAME_FEATURE_IDS,
    "featureIds"
  );
  // Add Metadata struct from the metadata stage
  shaderBuilder.addStructField(
    structId,
    MetadataPipelineStage.STRUCT_NAME_METADATA,
    "metadata"
  );

  const functionId =
    CustomShaderPipelineStage.FUNCTION_ID_INITIALIZE_INPUT_STRUCT_VS;
  shaderBuilder.addFunction(
    functionId,
    CustomShaderPipelineStage.FUNCTION_SIGNATURE_INITIALIZE_INPUT_STRUCT_VS,
    ShaderDestination.VERTEX
  );

  const initializationLines = vertexLines.initializationLines;
  shaderBuilder.addFunctionLines(functionId, initializationLines);
}

function addFragmentLinesToShader(shaderBuilder, fragmentLines) {
  let i;
  let structId = CustomShaderPipelineStage.STRUCT_ID_ATTRIBUTES_FS;
  shaderBuilder.addStruct(
    structId,
    CustomShaderPipelineStage.STRUCT_NAME_ATTRIBUTES,
    ShaderDestination.FRAGMENT
  );

  let field;
  let glslType;
  let variableName;
  const attributeFields = fragmentLines.attributeFields;
  for (i = 0; i < attributeFields.length; i++) {
    field = attributeFields[i];
    glslType = field[0];
    variableName = field[1];
    shaderBuilder.addStructField(structId, glslType, variableName);
  }

  structId = CustomShaderPipelineStage.STRUCT_ID_FRAGMENT_INPUT;
  shaderBuilder.addStruct(
    structId,
    CustomShaderPipelineStage.STRUCT_NAME_FRAGMENT_INPUT,
    ShaderDestination.FRAGMENT
  );
  shaderBuilder.addStructField(
    structId,
    CustomShaderPipelineStage.STRUCT_NAME_ATTRIBUTES,
    "attributes"
  );
  // Add FeatureIds struct from the Feature ID stage
  shaderBuilder.addStructField(
    structId,
    FeatureIdPipelineStage.STRUCT_NAME_FEATURE_IDS,
    "featureIds"
  );
  // Add Metadata struct from the metadata stage
  shaderBuilder.addStructField(
    structId,
    MetadataPipelineStage.STRUCT_NAME_METADATA,
    "metadata"
  );

  const functionId =
    CustomShaderPipelineStage.FUNCTION_ID_INITIALIZE_INPUT_STRUCT_FS;
  shaderBuilder.addFunction(
    functionId,
    CustomShaderPipelineStage.FUNCTION_SIGNATURE_INITIALIZE_INPUT_STRUCT_FS,
    ShaderDestination.FRAGMENT
  );

  const initializationLines = fragmentLines.initializationLines;
  shaderBuilder.addFunctionLines(functionId, initializationLines);
}

function addLinesToShader(shaderBuilder, customShader, generatedCode) {
  const vertexLines = generatedCode.vertexLines;
  if (vertexLines.enabled) {
    addVertexLinesToShader(shaderBuilder, vertexLines);

    shaderBuilder.addVertexLines([
      "#line 0",
      customShader.vertexShaderText,
      CustomShaderStageVS,
    ]);
  }

  const fragmentLines = generatedCode.fragmentLines;
  if (fragmentLines.enabled) {
    addFragmentLinesToShader(shaderBuilder, fragmentLines);

    shaderBuilder.addFragmentLines([
      "#line 0",
      customShader.fragmentShaderText,
      CustomShaderStageFS,
    ]);
  }
}

// exposed for testing.
CustomShaderPipelineStage._oneTimeWarning = oneTimeWarning;

export default CustomShaderPipelineStage;