NZSL 1.1 PBR shader

[nzsl_version("1.1")]
module PhysicallyBasedMaterial;

import InstanceData from Engine.InstanceData;
import LightData from Engine.LightData;
import SkeletalData from Engine.SkeletalData;
import ViewerData from Engine.ViewerData;

import * from Engine.LightShadow;
import SkinLinearPosition, SkinLinearPositionNormal from Engine.SkinningLinear;

// Pass-specific options
option DepthPass: bool = false;
option DistanceDepth: bool = false;
option ShadowPass: bool = false;

// Basic material options
option HasBaseColorTexture: bool = false;
option HasAlphaTexture: bool = false;
option AlphaTest: bool = false;

// Physically-based material options
option HasEmissiveTexture: bool = false;
option HasHeightTexture: bool = false;
option HasMetallicTexture: bool = false;
option HasNormalTexture: bool = false;
option HasRoughnessTexture: bool = false;
option HasSpecularTexture: bool = false;

// Billboard related options
option Billboard: bool = false;

// Vertex declaration related options
const InvalidLoc = u32.Max;

option VertexColorLoc: u32 = InvalidLoc;
option VertexNormalLoc: u32 = InvalidLoc;
option VertexPositionLoc: u32 = InvalidLoc;
option VertexSizeRotLocation: u32 = InvalidLoc;
option VertexTangentLoc: u32 = InvalidLoc;
option VertexUvLoc: u32 = InvalidLoc;

option VertexJointIndicesLoc: u32 = InvalidLoc;
option VertexJointWeightsLoc: u32 = InvalidLoc;

option MaxLightCount: u32 = 3;

const HasNormal = (VertexNormalLoc != InvalidLoc);
const HasVertexColor = (VertexColorLoc != InvalidLoc);
const HasColor = (HasVertexColor || Billboard);
const HasTangent = (VertexTangentLoc != InvalidLoc);
const HasVertexUV = (VertexUvLoc != InvalidLoc);
const HasUV = (HasVertexUV);
const HasNormalMapping = HasNormalTexture && HasNormal && HasTangent && !DepthPass;
const HasSkinning = (VertexJointIndicesLoc != InvalidLoc && VertexJointWeightsLoc != InvalidLoc);

[layout(std140)]
struct MaterialSettings
{
    // Basic settings
    [tag("AlphaTestThreshold")]
    AlphaThreshold: f32,

    [tag("ShadowMapNormalOffset")]
    ShadowMapNormalOffset: f32,

    [tag("ShadowPosScale")]
    ShadowPosScale: f32,

    [tag("BaseColor")]
    BaseColor: vec4[f32],

    // PBR settings
    [tag("MetallicFactor")]
    MetallicFactor: f32,

    [tag("RoughnessFactor")]
    RoughnessFactor: f32
}

// TODO: Add enums
const DirectionalLight = 0;
const PointLight = 1;
const SpotLight = 2;

[tag("Material")]
[auto_binding]
external
{
    [tag("Settings")] settings: uniform[MaterialSettings],
    [tag("BaseColorMap")] MaterialBaseColorMap: sampler2D[f32],
    [tag("AlphaMap")] MaterialAlphaMap: sampler2D[f32],
    [tag("EmissiveMap")] MaterialEmissiveMap: sampler2D[f32],
    [tag("HeightMap")] MaterialHeightMap: sampler2D[f32],
    [tag("MetallicMap")] MaterialMetallicMap: sampler2D[f32],
    [tag("NormalMap")] MaterialNormalMap: sampler2D[f32],
    [tag("RoughnessMap")] MaterialRoughnessMap: sampler2D[f32],
    [tag("SpecularMap")] MaterialSpecularMap: sampler2D[f32],
}

[tag("Engine")]
[auto_binding]
external
{
    [tag("TextureOverlay")] TextureOverlay: sampler2D[f32],
    [tag("InstanceData")] instanceData: uniform[InstanceData],
    [tag("ViewerData")] viewerData: uniform[ViewerData],
    [tag("SkeletalData")] skeletalData: uniform[SkeletalData],
    [tag("LightData")] lightData: uniform[LightData],
    [tag("ShadowMapsDirectional")] shadowMapsDirectional: array[depth_sampler2D_array[f32], MaxLightCount],
    [tag("ShadowMapsPoint")] shadowMapsPoint: array[sampler_cube[f32], MaxLightCount],
    [tag("ShadowMapsSpot")] shadowMapsSpot: array[depth_sampler2D[f32], MaxLightCount],
}

[export]
struct VertOut
{
    [location(0)] worldPos: vec3[f32],
    [location(1), cond(HasUV)] uv: vec2[f32],
    [location(2), cond(HasColor)] color: vec4[f32],
    [location(3), cond(HasNormal)] normal: vec3[f32],
    [location(4), cond(HasNormalMapping)] tangent: vec3[f32],
    [builtin(position)] position: vec4[f32],
}

// Fragment stage
import ComputeLightRadiance, DistributionGGX, GeometrySmith, FresnelSchlick from Math.CookTorrancePBR;
import Pi from Math.Constants;

struct FragOut
{
    [location(0)] RenderTarget0: vec4[f32],
    [builtin(frag_depth), cond(DistanceDepth)] fragdepth: f32,
}

[export]
fn ComputeColor(input: VertOut) -> vec4[f32]
{
    let color = settings.BaseColor;

    const if (HasUV)
        color.a *= TextureOverlay.Sample(input.uv).r;

    const if (HasColor)
        color *= input.color;

    const if (HasBaseColorTexture)
        color *= MaterialBaseColorMap.Sample(input.uv);

    const if (HasAlphaTexture)
        color.w *= MaterialAlphaMap.Sample(input.uv).x;

    const if (AlphaTest)
    {
        if (color.w < settings.AlphaThreshold)
            discard;
    }

    return color;
}

[export, cond(HasNormal)]
fn ComputeLighting(color: vec3[f32], input: VertOut) -> vec3[f32]
{
    let lightRadiance = vec3[f32](0.0, 0.0, 0.0);

    let eyeVec = normalize(viewerData.eyePosition - input.worldPos);

    let normal: vec3[f32];
    const if (HasNormalMapping)
    {
        let N = normalize(input.normal);
        let T = normalize(input.tangent);
        let B = cross(N, T);
        let tbnMatrix = mat3[f32](T, B, N);

        normal = normalize(tbnMatrix * (MaterialNormalMap.Sample(input.uv).xyz * 2.0 - vec3[f32](1.0, 1.0, 1.0)));
    }
    else
        normal = normalize(input.normal);

    let albedo = color.xyz;

    let metallic = settings.MetallicFactor;
    const if (HasMetallicTexture)
        metallic *= MaterialMetallicMap.Sample(input.uv).x;

    let roughness = settings.RoughnessFactor;
    const if (HasRoughnessTexture)
        roughness *= MaterialRoughnessMap.Sample(input.uv).x;

    let F0 = vec3[f32](0.04, 0.04, 0.04);
    F0 = lerp(F0, albedo, metallic.rrr);

    let albedoFactor = albedo / Pi;

    for lightIndex in u32(0) -> lightData.directionalLightCount
    {
        let light = lightData.directionalLights[lightIndex];

        let lambert = max(dot(normal, -light.direction), 0.0);

        let shadowFactor = ComputeDirectionalLightShadow(light, shadowMapsDirectional[lightIndex], input.worldPos, lambert, viewerData.viewMatrix);
        let radiance = ComputeLightRadiance(light.color.rgb, -light.direction, albedoFactor, eyeVec, F0, normal, metallic, roughness);

        lightRadiance += shadowFactor * radiance;
    }

    for lightIndex in u32(0) -> lightData.pointLightCount
    {
        let light = lightData.pointLights[lightIndex];

        let lightToPos = input.worldPos - light.position;
        let dist = length(lightToPos);

        let attenuation = max(1.0 - dist * light.invRadius, 0.0);
        let lightToPosNorm = lightToPos / max(dist, 0.0001);

        let shadowFactor = ComputePointLightShadow(light, shadowMapsPoint[lightIndex], dist, lightToPosNorm);
        let radiance = ComputeLightRadiance(light.color.rgb * attenuation, -lightToPosNorm, albedoFactor, eyeVec, F0, normal, metallic, roughness);

        lightRadiance += shadowFactor * radiance;
    }

    for lightIndex in u32(0) -> lightData.spotLightCount
    {
        let light = lightData.spotLights[lightIndex];

        let lightToPos = input.worldPos - light.position;
        let dist = length(lightToPos);
        let lightToPosNorm = lightToPos / max(dist, 0.0001);

        let curAngle = dot(light.direction, lightToPosNorm);
        let innerMinusOuterAngle = light.innerAngle - light.outerAngle;

        let attenuation = max(1.0 - dist * light.invRadius, 0.0);
        attenuation *= max((curAngle - light.outerAngle) / innerMinusOuterAngle, 0.0);

        let lambert = clamp(dot(normal, -lightToPosNorm), 0.0, 1.0);

        let shadowFactor = ComputeSpotLightShadow(light, shadowMapsSpot[lightIndex], input.worldPos, lambert);
        let radiance = ComputeLightRadiance(light.color.rgb * attenuation, -lightToPosNorm, albedoFactor, eyeVec, F0, normal, metallic, roughness);

        lightRadiance += shadowFactor * radiance;
    }

    let ambient = (0.03).rrr * albedo;

    let finalColor = ambient + lightRadiance;
    finalColor = finalColor / (finalColor + vec3[f32](1.0, 1.0, 1.0));

    return finalColor;
}

[export, entry(frag), cond(!DepthPass)]
fn FragMain(input: VertOut) -> FragOut
{
    let color = ComputeColor(input);

    const if (HasNormal)
        color.rgb = ComputeLighting(color.rgb, input);

    let output: FragOut;
    output.RenderTarget0 = color;
    return output;
}

// Shadow passes entries
[entry(frag), cond(DepthPass && DistanceDepth)]
[depth_write(replace)]
fn FragDepthDist(input: VertOut) -> FragOut
{
    let color = ComputeColor(input);

    let output: FragOut;
    output.RenderTarget0 = color;

    let dist = distance(viewerData.eyePosition, input.worldPos);
    output.fragdepth = dist / viewerData.farPlane;

    return output;
}

[entry(frag), cond(DepthPass && AlphaTest && !DistanceDepth)]
fn FragDepth(input: VertOut) -> FragOut
{
    let color = ComputeColor(input);

    let output: FragOut;
    output.RenderTarget0 = color;
    return output;
}

[entry(frag), cond(DepthPass && !AlphaTest && !DistanceDepth)]
fn FragDepthNoAlpha() {} //< dummy

// Vertex stage
[cond(!Billboard)]
struct VertIn
{
    [location(VertexPositionLoc)]
    pos: vec3[f32],

    [cond(HasVertexColor), location(VertexColorLoc)]
    color: vec4[f32],

    [cond(HasVertexUV), location(VertexUvLoc)]
    uv: vec2[f32],

    [cond(HasNormal), location(VertexNormalLoc)]
    normal: vec3[f32],

    [cond(HasTangent), location(VertexTangentLoc)]
    tangent: vec3[f32],

    [cond(HasSkinning), location(VertexJointIndicesLoc)]
    jointIndices: vec4[i32],

    [cond(HasSkinning), location(VertexJointWeightsLoc)]
    jointWeights: vec4[f32]
}

[cond(Billboard)]
struct BillboardVertIn
{
    [builtin(vertex_index)] vertIndex: i32,

    [location(VertexSizeRotLocation)]
    sizeRot: vec4[f32], //< width,height,sin,cos

    [cond(HasVertexUV), location(VertexUvLoc)]
    uv: vec2[f32],

    [cond(HasVertexColor), location(VertexColorLoc)]
    color: vec4[f32]
}

const billboardPos = array[vec2[f32]](
    vec2[f32](-0.5, -0.5),
    vec2[f32]( 0.5, -0.5),
    vec2[f32](-0.5, 0.5),
    vec2[f32]( 0.5, 0.5)
);

[entry(vert), cond(Billboard)]
fn VertBillboard(input: BillboardVertIn) -> VertOut
{
    let position = billboardPos[input.vertIndex % 4];

    let size = input.sizeRot.xy;
    let sinCos = input.sizeRot.zw;

    let rotatedPosition = vec2[f32](
        position.x * sinCos.y - position.y * sinCos.x,
        position.y * sinCos.y + position.x * sinCos.x
    );
    rotatedPosition *= size;

    let cameraRight = vec3[f32](viewerData.viewMatrix[0][0], viewerData.viewMatrix[1][0], viewerData.viewMatrix[2][0]);
    let cameraUp = vec3[f32](viewerData.viewMatrix[0][1], viewerData.viewMatrix[1][1], viewerData.viewMatrix[2][1]);

    let worldPosition = vec3[f32](instanceData.worldMatrix[3].xyz);
    worldPosition += cameraRight * rotatedPosition.x;
    worldPosition += cameraUp * rotatedPosition.y;

    let output: VertOut;
    output.worldPos = worldPosition;
    output.position = viewerData.viewProjMatrix * vec4[f32](worldPosition, 1.0);

    const if (HasVertexUV)
        output.uv = input.uv;

    const if (HasColor)
        output.color = input.color;

    return output;
}

[entry(vert), cond(!Billboard)]
fn VertMain(input: VertIn) -> VertOut
{
    let pos: vec3[f32];
    const if (HasNormal) let normal: vec3[f32];

    const if (HasSkinning)
    {
        let jointMatrices = array[mat4[f32]](
            skeletalData.jointMatrices[input.jointIndices[0]],
            skeletalData.jointMatrices[input.jointIndices[1]],
            skeletalData.jointMatrices[input.jointIndices[2]],
            skeletalData.jointMatrices[input.jointIndices[3]]
        );

        const if (HasNormal)
        {
            let skinningOutput = SkinLinearPositionNormal(jointMatrices, input.jointWeights, input.pos, input.normal);
            pos = skinningOutput.position;
            normal = skinningOutput.normal;
        }
        else
        {
            let skinningOutput = SkinLinearPosition(jointMatrices, input.jointWeights, input.pos);
            pos = skinningOutput.position;
        }
    }
    else
    {
        pos = input.pos;
        const if (HasNormal)
            normal = input.normal;
    }

    const if (ShadowPass)
    {
        pos *= settings.ShadowPosScale;
        const if (HasNormal)
            pos -= normal * settings.ShadowMapNormalOffset;
    }

    let worldPosition = instanceData.worldMatrix * vec4[f32](pos, 1.0);

    let output: VertOut;
    output.worldPos = worldPosition.xyz;
    output.position = viewerData.viewProjMatrix * worldPosition;

    let rotationMatrix = transpose(inverse(mat3[f32](instanceData.worldMatrix)));

    const if (HasColor)
        output.color = input.color;

    const if (HasNormal)
        output.normal = rotationMatrix * input.normal;

    const if (HasVertexUV)
        output.uv = input.uv;

    const if (HasNormalMapping)
        output.tangent = rotationMatrix * input.tangent;

    return output;
}