Deferred Shading v4

#pragma once

#include "Shader.h"

class DeferredRenderer
{
private:
  struct QuadVertex
  {
    float x, y, z;
    float tx, ty;

    QuadVertex( float ax, float ay, float az, float atx, float aty )
    {
      x = ax;
      y = ay;
      z = az;

      tx = atx;
      ty = aty;
    }
  };

  IDirect3DTexture9 * depthMap;
  IDirect3DTexture9 * normalMap;
  IDirect3DTexture9 * diffuseMap;

  IDirect3DSurface9 * positionSurface;
  IDirect3DSurface9 * normalSurface;
  IDirect3DSurface9 * diffuseSurface;
  IDirect3DSurface9 * backSurface;

  VertexShader * vertexShaderPassOne;
  PixelShader * pixelShaderPassOne;

  VertexShader * vertexShaderPassTwo;
  PixelShader * pixelShaderPassTwo;

  IDirect3DVertexBuffer9 * screenQuad;
  IDirect3DVertexDeclaration9 * vertexDeclaration;

  D3DXHANDLE v1World;
  D3DXHANDLE v1Proj;
  D3DXHANDLE v1View;

  D3DXHANDLE v2World;
  D3DXHANDLE v2Proj;
  D3DXHANDLE v2View;

  D3DXHANDLE p2LightPos;
  D3DXHANDLE p2LightRange;
  D3DXHANDLE p2CameraPos;
  D3DXHANDLE p2InvViewProj;



  DWORD icosphereFVF;
  DWORD icosphereStride;
  DWORD icosphereVertexCount;
  DWORD icosphereFaceCount;
  IDirect3DVertexBuffer9 * icosphereVB;
  IDirect3DIndexBuffer9 * icosphereIB;
public:

  void CreateScreenQuad()
  {
    g_device->CreateVertexBuffer( 6 * sizeof( QuadVertex ), D3DUSAGE_WRITEONLY, D3DFVF_XYZ | D3DFVF_TEX1, D3DPOOL_DEFAULT, &screenQuad, 0 );

    QuadVertex vertices[ ] = { QuadVertex( -0.5, -0.5, 0, 0, 0 ), QuadVertex( g_width - 0.5, 0 - 0.5, 0, 1, 0 ), QuadVertex ( 0 - 0.5, g_height - 0.5, 0, 0, 1 ),
      QuadVertex( g_width - 0.5, 0 - 0.5, 0, 1, 0 ), QuadVertex( g_width - 0.5, g_height - 0.5, 0, 1, 1 ), QuadVertex ( 0 - 0.5, g_height - 0.5, 0, 0, 1 ) };

    void * data = 0;

    screenQuad->Lock( 0, 0, &data, 0 );
    memcpy( data, vertices, sizeof( QuadVertex ) * 6 );
    screenQuad->Unlock( );

    D3DVERTEXELEMENT9 quadVertexDeclation[ ] =
    {
      { 0,  0, D3DDECLTYPE_FLOAT3, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_POSITION, 0 },
      { 0, 12, D3DDECLTYPE_FLOAT2, D3DDECLMETHOD_DEFAULT, D3DDECLUSAGE_TEXCOORD, 0 },
      D3DDECL_END()
    };

    g_device->CreateVertexDeclaration( quadVertexDeclation, &vertexDeclaration );
  }

  void CreateRenderTargets()
  {
    D3DCAPS9 caps;  g_device->GetDeviceCaps( &caps );

    if( caps.NumSimultaneousRTs < 4 )
      MessageBoxA( 0, "Your video card does not support MRT.", 0, MB_OK | MB_ICONERROR );

    if( FAILED( D3DXCreateTexture( g_device, g_width, g_height, 0, D3DUSAGE_RENDERTARGET, D3DFMT_R32F, D3DPOOL_DEFAULT, &depthMap )))
      MessageBoxA( 0, "Failed to create 'Depth Map' texture.", 0, MB_OK | MB_ICONERROR );
    if( FAILED( D3DXCreateTexture( g_device, g_width, g_height, 0, D3DUSAGE_RENDERTARGET, D3DFMT_A8B8G8R8, D3DPOOL_DEFAULT, &normalMap )))
      MessageBoxA( 0, "Failed to create 'Normal Map' texture.", 0, MB_OK | MB_ICONERROR );
    if( FAILED( D3DXCreateTexture( g_device, g_width, g_height, 0, D3DUSAGE_RENDERTARGET, D3DFMT_A8B8G8R8, D3DPOOL_DEFAULT, &diffuseMap )))
      MessageBoxA( 0, "Failed to create 'Diffuse Map' texture.", 0, MB_OK | MB_ICONERROR );

    depthMap->GetSurfaceLevel( 0, &positionSurface );
    normalMap->GetSurfaceLevel( 0, &normalSurface );
    diffuseMap->GetSurfaceLevel( 0, &diffuseSurface );

    g_device->GetRenderTarget( 0, &backSurface );

    ID3DXMesh * icosphere = 0;
    D3DXCreateSphere( g_device, 1.0, 6, 6, &icosphere, 0 );

    void * dataSrc = 0;
    void * dataDest = 0;
    icosphereFVF = icosphere->GetFVF();
    icosphereStride = icosphere->GetNumBytesPerVertex();
    icosphereFaceCount = icosphere->GetNumFaces();
    icosphereVertexCount = icosphere->GetNumVertices();

    // copy vertex buffer
    int size = icosphere->GetNumBytesPerVertex() * icosphere->GetNumVertices();
    g_device->CreateVertexBuffer( size, D3DUSAGE_WRITEONLY, icosphereFVF, D3DPOOL_DEFAULT, &icosphereVB, 0 );
    icosphereVB->Lock( 0, 0, &dataDest, 0 );
    icosphere->LockVertexBuffer( 0, &dataSrc );
    memcpy( dataDest, dataSrc, size );
    icosphere->UnlockVertexBuffer();
    icosphereVB->Unlock();

    // copy index buffer
    size = icosphere->GetNumFaces() * sizeof( short ) * 3;
    g_device->CreateIndexBuffer( size, D3DUSAGE_WRITEONLY, D3DFMT_INDEX16, D3DPOOL_DEFAULT, &icosphereIB, 0 );
    icosphereIB->Lock( 0, 0, &dataDest, 0 );
    icosphere->LockIndexBuffer( 0, &dataSrc );
    memcpy( dataDest, dataSrc, size );
    icosphere->UnlockIndexBuffer();
    icosphereIB->Unlock();

    icosphere->Release();
  }

  void InitPassOneShaders()
  {
    // PASS 1 - Filling the G-Buffer
    string vertexSourcePassOne =
      "float4x4 g_world;\n"
      "float4x4 g_view;\n"
      "float4x4 g_projection;\n"

      "struct VS_INPUT {\n"
      "  float4 position : POSITION;\n"
      "  float4 normal   : NORMAL;\n"
      "  float2 texCoord : TEXCOORD0;\n"
      "};\n"

      "struct VS_OUTPUT {\n"
      "  float4 position : POSITION;\n"
      "  float4 screenPosition : TEXCOORD0;\n"
      "  float3 normal : TEXCOORD1;\n"
      "  float2 diffuseMapCoord : TEXCOORD2;\n"
      "};\n"

      "VS_OUTPUT main(VS_INPUT input) {\n"
      "  VS_OUTPUT output;\n"
      "  float4x4 worldViewProj = mul(mul(g_world, g_view), g_projection);\n"
      "  output.position = mul(input.position, worldViewProj);\n"
      "  output.normal = normalize(mul(input.normal, (float3x3)g_world));\n"
      "  output.screenPosition = output.position;\n"
      "  output.diffuseMapCoord = input.texCoord;\n"
      "  return output;\n"
      "};\n";

    vertexShaderPassOne = new VertexShader( vertexSourcePassOne );

    v1World = vertexShaderPassOne->GetConstantTable()->GetConstantByName( 0, "g_world" );
    v1Proj = vertexShaderPassOne->GetConstantTable()->GetConstantByName( 0, "g_projection" );
    v1View = vertexShaderPassOne->GetConstantTable()->GetConstantByName( 0, "g_view" );

    string pixelSourcePassOne =

      "texture diffuseMap;\n"

      "sampler diffuseSampler : register(s0) = sampler_state\n"
      "{\n"
      "   texture = <diffuseMap>;\n"
      "};\n"

      "struct PS_INPUT {\n"
      "  float4 screenPos : TEXCOORD0;\n"
      "  float3 normal : TEXCOORD1;\n"
      "  float2 diffuseMapCoord : TEXCOORD2;\n"
      "};\n"

      "struct MRT_OUTPUT {\n"
      "  float4 depth : COLOR0;\n"
      "  float4 normal : COLOR1;\n"
      "  float4 diffuseMap : COLOR2;\n"
      "};\n"

      "MRT_OUTPUT main( PS_INPUT input )\n"
      "{\n"
      "   MRT_OUTPUT output;\n"
      "   float3 normal = normalize( input.normal );\n"
      "   float depth = input.screenPos.z / input.screenPos.w;\n"
      "   output.depth = float4( depth, depth, depth, 1.0 );\n"
      "   output.normal = float4( normal.x, normal.y, normal.z, 1.0 );\n"
      "   output.diffuseMap = tex2D( diffuseSampler, input.diffuseMapCoord );\n"
      "   return output;"
      "};\n"
      ;

    pixelShaderPassOne = new PixelShader( pixelSourcePassOne );
  }

  DeferredRenderer()
  {
    CreateScreenQuad();
    CreateRenderTargets();
    InitPassOneShaders();
    InitPassTwoShaders();
  }

  void InitPassTwoShaders()
  {
    // PASS 2 - Final render
    string vertexSourcePassTwo =
      "float4x4 g_world;\n"
      "float4x4 g_view;\n"
      "float4x4 g_projection;\n"
      "struct VS_INPUT {\n"
      "  float4 position : POSITION;\n"
      "  float3 texcoord : TEXCOORD0;\n"
      "};\n"

      "struct VS_OUTPUT {\n"
      "  float4 position : POSITION;\n"
      "  float3 texcoord : TEXCOORD0;\n"
      "};\n"

      "VS_OUTPUT main(VS_INPUT input) {\n"
      "  VS_OUTPUT output;\n"
      "  float4x4 worldViewProj = mul(mul(g_world, g_view), g_projection);\n"
      "  output.position   = mul(input.position, worldViewProj);\n"
      "  output.texcoord   = input.texcoord;\n"
      "  return output;\n"
      "};\n";

    vertexShaderPassTwo = new VertexShader( vertexSourcePassTwo );

    v2World = vertexShaderPassTwo->GetConstantTable()->GetConstantByName( 0, "g_world" );
    v2Proj = vertexShaderPassTwo->GetConstantTable()->GetConstantByName( 0, "g_projection" );
    v2View = vertexShaderPassTwo->GetConstantTable()->GetConstantByName( 0, "g_view" );

    string pixelSourcePassTwo =

      "texture depthMap;\n"
      "texture normalMap;\n"
      "texture diffuseMap;\n"

      "sampler depthSampler : register(s0) = sampler_state\n"
      "{\n"
      "   texture = <depthMap>;\n"
      "};\n"

      "sampler normalSampler : register(s1) = sampler_state\n"
      "{\n"
      "   texture = <normalMap>;\n"
      "};\n"

      "sampler diffuseSampler : register(s2) = sampler_state\n"
      "{\n"
      "   texture = <diffuseMap>;\n"
      "};\n"

      // light props
      "float3 lightPos;\n"
      "float lightRange;\n"

      // camera props
      "float3 cameraPosition;\n"

      "float4x4 invViewProj;\n"
      "float4 main( float3 texcoord : TEXCOORD0 ) : COLOR0\n"
      "{\n"

      "   float4 diffuseTexel = tex2D( diffuseSampler, texcoord.xy );\n"
      "   float depth = tex2D( depthSampler, texcoord.xy ).r;\n"
      "   float3 n = tex2D( normalSampler, texcoord.xy ).xyz;\n"

      "   float4 screenPosition;\n"
      "   screenPosition.x = texcoord.x * 2.0f - 1.0f;\n"
      "   screenPosition.y = -(texcoord.y * 2.0f - 1.0f);\n"
      "   screenPosition.z = depth;\n"
      "   screenPosition.w = 1.0f;\n"

      "   float4 p = mul( screenPosition, invViewProj );\n"
      "   p /= p.w;\n"

      "   float3 lightDirection = lightPos - p;"
      "   float3 l = normalize( lightDirection );\n"

      // specular

      "   float3 v = normalize( cameraPosition - p );\n"
      "   float3 r = reflect( -v, n );\n"
      "   float spec = pow( saturate( dot( l, r ) ), 20.0 );\n"

      // diffuse
      "   float diff = saturate(dot( l, n ));\n"

      "   float falloff = pow( lightRange, 2 ) / pow( dot( lightDirection, lightDirection ), 2 );\n"

      "   float o = clamp( falloff, 0, 2 ) * (  diff + spec );\n"

      "   return float4( diffuseTexel.x * o, diffuseTexel.y * o, diffuseTexel.z * o, diffuseTexel.w );\n"
      "};\n";


    pixelShaderPassTwo = new PixelShader( pixelSourcePassTwo );

    p2LightPos = pixelShaderPassTwo->GetConstantTable()->GetConstantByName( 0, "lightPos" );
    p2LightRange = pixelShaderPassTwo->GetConstantTable()->GetConstantByName( 0, "lightRange" );
    p2CameraPos = pixelShaderPassTwo->GetConstantTable()->GetConstantByName( 0, "cameraPosition" );
    p2InvViewProj  = pixelShaderPassTwo->GetConstantTable()->GetConstantByName( 0, "invViewProj" );
  }

  ~DeferredRenderer()
  {
    if( depthMap )
      depthMap->Release();
    if( normalMap )
      normalMap->Release();
    if( diffuseMap )
      diffuseMap->Release();

    if( positionSurface )
      positionSurface->Release();
    if( normalSurface )
      normalSurface->Release();
    if( diffuseSurface )
      diffuseSurface->Release();
  }

  void Begin()
  {
    g_device->SetRenderTarget( 0, positionSurface );
    g_device->SetRenderTarget( 1, normalSurface );
    g_device->SetRenderTarget( 2, diffuseSurface );

    g_device->Clear( 0, 0, D3DCLEAR_TARGET | D3DCLEAR_ZBUFFER , D3DCOLOR_XRGB( 0, 0, 0 ), 1.0, 0 );
    g_device->BeginScene( );

    vertexShaderPassOne->Bind();
    pixelShaderPassOne->Bind();

    g_device->SetRenderState(D3DRS_ALPHABLENDENABLE, FALSE );
  }

  void RenderMesh( Mesh * mesh )
  {
    g_device->SetStreamSource( 0, mesh->vb, 0, sizeof( Vertex ));
    g_device->SetIndices( mesh->ib );

    D3DXMATRIX view; g_device->GetTransform( D3DTS_VIEW, &view );
    D3DXMATRIX proj; g_device->GetTransform( D3DTS_PROJECTION, &proj );
    D3DXMATRIX world; GetD3DMatrixFromBulletTransform( mesh->parent->globalTransform, world );

    vertexShaderPassOne->GetConstantTable()->SetMatrix( g_device, v1View, &view );
    vertexShaderPassOne->GetConstantTable()->SetMatrix( g_device, v1Proj, &proj );
    vertexShaderPassOne->GetConstantTable()->SetMatrix( g_device, v1World, &world );

    g_device->DrawIndexedPrimitive( D3DPT_TRIANGLELIST, 0, 0, mesh->vertexCount, 0, mesh->faceCount );
  }

  void End()
  {
    g_device->EndScene();

    g_device->SetRenderTarget( 0, backSurface );
    g_device->SetRenderTarget( 1, 0 );
    g_device->SetRenderTarget( 2, 0 );

    g_device->Clear( 0, 0, D3DCLEAR_TARGET | D3DCLEAR_STENCIL, D3DCOLOR_XRGB( 0, 0, 0 ), 1.0, 0 );
    g_device->BeginScene( );

    IDirect3DStateBlock9 * state;
    g_device->CreateStateBlock( D3DSBT_ALL, &state );

    g_device->SetRenderState( D3DRS_LIGHTING, FALSE );


    D3DXMATRIX ortho; D3DXMatrixOrthoOffCenterLH ( &ortho, 0, g_width, g_height, 0, 0, 1024 );
    D3DXMATRIX identity; D3DXMatrixIdentity( &identity );

    vertexShaderPassTwo->GetConstantTable()->SetMatrix( g_device, v2World, &identity );
    vertexShaderPassTwo->GetConstantTable()->SetMatrix( g_device, v2View, &identity );
    vertexShaderPassTwo->GetConstantTable()->SetMatrix( g_device, v2Proj, &ortho );

    g_device->SetTexture( 0, depthMap );
    g_device->SetTexture( 1, normalMap );
    g_device->SetTexture( 2, diffuseMap );

    g_device->SetRenderState(D3DRS_ALPHATESTENABLE, FALSE );
    g_device->SetRenderState(D3DRS_ALPHABLENDENABLE, TRUE );
    g_device->SetRenderState(D3DRS_SRCBLEND, D3DBLEND_ONE );
    g_device->SetRenderState(D3DRS_DESTBLEND, D3DBLEND_ONE );

    float temp[ 3 ];

    btVector3 cameraPos = g_camera->base->globalTransform.getOrigin();

    temp[ 0 ] = cameraPos.x();
    temp[ 1 ] = cameraPos.y();
    temp[ 2 ] = cameraPos.z();

    pixelShaderPassTwo->GetConstantTable()->SetFloatArray( g_device, p2CameraPos, temp, 3 );

    D3DXMATRIX view; g_device->GetTransform( D3DTS_VIEW, &view );
    D3DXMATRIX proj; g_device->GetTransform( D3DTS_PROJECTION, &proj );

    D3DXMATRIX viewProj;
    D3DXMatrixMultiply( &viewProj, &view, &proj );

    D3DXMATRIX invViewProj;
    D3DXMatrixInverse( &invViewProj, 0, &viewProj );

    pixelShaderPassTwo->GetConstantTable()->SetMatrix( g_device, p2InvViewProj, &invViewProj );


    g_device->SetRenderState( D3DRS_ZWRITEENABLE, FALSE );

    g_device->SetRenderState(D3DRS_STENCILREF, 0x0 );
    g_device->SetRenderState(D3DRS_STENCILMASK, 0xffffffff);
    g_device->SetRenderState(D3DRS_STENCILWRITEMASK, 0xffffffff);

    g_device->SetRenderState( D3DRS_STENCILENABLE, TRUE );
    g_device->SetRenderState( D3DRS_TWOSIDEDSTENCILMODE, TRUE );

    for( int i = 0; i < g_lights.size(); i++ )
    {
      Light * light = g_lights.at( i );

      btVector3 lightPos = light->base->globalTransform.getOrigin();

      temp[ 0 ] = lightPos.x();
      temp[ 1 ] = lightPos.y();
      temp[ 2 ] = lightPos.z();

      float scl = sqrt( 5 * light->radius );
      D3DXMATRIX scale; D3DXMatrixScaling( &scale, scl,scl,scl );
      D3DXMATRIX trans; D3DXMatrixTranslation( &trans, lightPos.x(), lightPos.y(), lightPos.z() );
      D3DXMATRIX world; D3DXMatrixMultiply( &world, &scale, &trans );


      // use default shaders
      g_device->SetVertexShader( 0 );
      g_device->SetPixelShader( 0 );

      g_device->SetTransform( D3DTS_WORLD, &world );
      g_device->SetTransform( D3DTS_VIEW, &view );
      g_device->SetTransform( D3DTS_PROJECTION, &proj );

      // disable draw into color buffer
      g_device->SetRenderState( D3DRS_COLORWRITEENABLE, 0x00000000 );

      g_device->SetRenderState( D3DRS_STENCILFAIL, D3DSTENCILOP_KEEP );
      g_device->SetRenderState( D3DRS_STENCILPASS, D3DSTENCILOP_KEEP );

      g_device->SetStreamSource( 0, icosphereVB, 0, icosphereStride );
      g_device->SetIndices( icosphereIB );
      g_device->SetFVF( icosphereFVF );

      g_device->SetRenderState( D3DRS_STENCILFUNC, D3DCMP_ALWAYS );
      g_device->SetRenderState( D3DRS_STENCILZFAIL, D3DSTENCILOP_DECR );

      g_device->SetRenderState( D3DRS_CCW_STENCILFUNC, D3DCMP_ALWAYS );
      g_device->SetRenderState( D3DRS_CCW_STENCILZFAIL, D3DSTENCILOP_INCR );

      // draw a sphere bounds light into stencil buffer
      g_device->SetRenderState( D3DRS_CULLMODE, D3DCULL_NONE );
      g_device->DrawIndexedPrimitive( D3DPT_TRIANGLELIST, 0, 0, icosphereVertexCount, 0, icosphereFaceCount );

      vertexShaderPassTwo->Bind();
      pixelShaderPassTwo->Bind();

      pixelShaderPassTwo->GetConstantTable()->SetFloatArray( g_device, p2LightPos, temp, 3 );
      pixelShaderPassTwo->GetConstantTable()->SetFloat( g_device, p2LightRange, light->radius );

      // enable draw into color buffer
      g_device->SetRenderState( D3DRS_COLORWRITEENABLE, 0xFFFFFFFF );

      /////////////////////////////////////////////////////////////////////////////
      // draw a screen quad
      g_device->SetStreamSource( 0, screenQuad, 0, sizeof( QuadVertex ));
      g_device->SetVertexDeclaration( vertexDeclaration );

      // disable culling at all
      g_device->SetRenderState( D3DRS_CULLMODE, D3DCULL_NONE );

      g_device->SetRenderState(D3DRS_STENCILFUNC, D3DCMP_NOTEQUAL);
      g_device->SetRenderState(D3DRS_STENCILPASS, D3DSTENCILOP_ZERO ); // fill stencil with zeros

      g_device->DrawPrimitive( D3DPT_TRIANGLELIST, 0, 2 );
    }

    state->Apply();
    state->Release();

    g_device->SetTexture( 0, 0 );
    g_device->SetTexture( 1, 0 );
    g_device->SetTexture( 2, 0 );


  }
};