NazaraEngine/src/Nazara/Graphics/Resources/Shaders/PhysicallyBasedMaterial.nzsl

[nzsl_version("1.0")]
module PhysicallyBasedMaterial;

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

import SkinLinearPosition, SkinLinearPositionNormal from Engine.SkinningLinear;

// Pass-specific options
option DepthPass: 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;
option BillboardCenterLocation: i32 = -1;
option BillboardColorLocation: i32 = -1;
option BillboardSizeRotLocation: i32 = -1;

// Vertex declaration related options
option VertexColorLoc: i32 = -1;
option VertexNormalLoc: i32 = -1;
option VertexPositionLoc: i32;
option VertexTangentLoc: i32 = -1;
option VertexUvLoc: i32 = -1;

option VertexJointIndicesLoc: i32 = -1;
option VertexJointWeightsLoc: i32 = -1;

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

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

	[tag("BaseColor")]
	BaseColor: vec4[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]
}

struct VertToFrag
{
	[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 DistributionGGX, GeometrySmith, FresnelSchlick from Math.CookTorrancePBR;
import Pi from Math.Constants;

struct FragOut
{
	[location(0)] RenderTarget0: vec4[f32]
}

[entry(frag), cond(!DepthPass || AlphaTest)]
fn main(input: VertToFrag) -> FragOut
{
	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;
	}

	const if (HasNormal && !DepthPass)
	{
		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: f32;
		let roughness: f32;

		const if (HasMetallicTexture)
			metallic = MaterialMetallicMap.Sample(input.uv).x;
		else
			metallic = 0.0;

		const if (HasRoughnessTexture)
			roughness = MaterialRoughnessMap.Sample(input.uv).x;
		else
			roughness = 0.8;

		let F0 = vec3[f32](0.04, 0.04, 0.04);
		F0 = albedo * metallic + F0 * (1.0 - metallic);

		let albedoFactor = albedo / Pi;

		/*for i in u32(0) -> lightData.lightCount
		{
			let light = lightData.lights[i];

			let attenuation = 1.0;

			// TODO: Add switch instruction
			let lightToPosNorm: vec3[f32];
			if (light.type == DirectionalLight)
				lightToPosNorm = -light.parameter1.xyz;
			else
			{
				// PointLight | SpotLight
				let lightPos = light.parameter1.xyz;
				let lightInvRadius = light.parameter1.w;

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

				attenuation = max(1.0 - dist * lightInvRadius, 0.0);
				lightToPosNorm = lightToPos / max(dist, 0.0001);

				if (light.type == SpotLight)
				{
					let lightDir = light.parameter2.xyz;
					let lightInnerAngle = light.parameter3.x;
					let lightOuterAngle = light.parameter3.y;

					let curAngle = dot(lightDir, lightToPosNorm);
					let innerMinusOuterAngle = lightInnerAngle - lightOuterAngle;
	
					attenuation *= max((curAngle - lightOuterAngle) / innerMinusOuterAngle, 0.0);			
				}
			}

			let radiance = light.color.rgb * attenuation;

			let halfDir = normalize(lightToPosNorm + eyeVec);

			// Cook-Torrance BRDF
			let NDF = DistributionGGX(normal, halfDir, roughness);
			let G = GeometrySmith(normal, eyeVec, lightToPosNorm, roughness);
			let F = FresnelSchlick(max(dot(halfDir, eyeVec), 0.0), F0);

			let kS = F;
			let diffuse = vec3[f32](1.0, 1.0, 1.0) - kS;
			diffuse *= 1.0 - metallic;

			let numerator = NDF * G * F;
			let denominator = 4.0 * max(dot(normal, eyeVec), 0.0) * max(dot(normal, lightToPosNorm), 0.0);
			let specular = numerator / max(denominator, 0.0001);

			let NdotL = max(dot(normal, lightToPosNorm), 0.0);
			lightRadiance += (diffuse * albedoFactor + specular) * radiance * NdotL;
		}*/

		let ambient = (0.03).rrr * albedo;

		let color = ambient + lightRadiance * color.rgb;
		color = color / (color + vec3[f32](1.0, 1.0, 1.0));
		color = pow(color, (1.0 / 2.2).xxx);

		let output: FragOut;
		output.RenderTarget0 = vec4[f32](color, 1.0);
		return output;
	}
	else
	{
		let output: FragOut;
		output.RenderTarget0 = color;
		return output;
	}
}

// Dummy fragment shader (TODO: Add a way to delete stage?)
[entry(frag), cond(DepthPass && !AlphaTest)]
fn main() {}

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

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

	[cond(HasUV), 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), location(BillboardCenterLocation)]
	billboardCenter: vec3[f32],

	[cond(Billboard), location(BillboardSizeRotLocation)]
	billboardSizeRot: vec4[f32], //< width,height,sin,cos

	[cond(Billboard), location(BillboardColorLocation)]
	billboardColor: vec4[f32]
}

[entry(vert), cond(Billboard)]
fn billboardMain(input: VertIn) -> VertToFrag
{
	let size = input.billboardSizeRot.xy;
	let sinCos = input.billboardSizeRot.zw;

	let rotatedPosition = vec2[f32](
		input.pos.x * sinCos.y - input.pos.y * sinCos.x,
		input.pos.y * sinCos.y + input.pos.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 vertexPos = input.billboardCenter;
	vertexPos += cameraRight * rotatedPosition.x;
	vertexPos += cameraUp * rotatedPosition.y;

	let output: VertToFrag;
	output.position = viewerData.viewProjMatrix * instanceData.worldMatrix * vec4[f32](vertexPos, 1.0);
	
	const if (HasColor)
		output.color = input.billboardColor;
	
	const if (HasUV)
		output.uv = input.pos.xy + vec2[f32](0.5, 0.5);

	return output;
}

[entry(vert), cond(!Billboard)]
fn main(input: VertIn) -> VertToFrag
{
	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;
	}

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

	let output: VertToFrag;
	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 (HasUV)
		output.uv = input.uv;

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

	return output;
}