Adding Lighting to a ShaderMaterial in three.js

I'd like to create a ShaderMaterial that has lighting like a MeshLambertMaterial. I've created a vertex and fragment shader, and included the uniforms from THREE.ShaderLib[ 'lambert' ].uniforms in the ShaderMaterial.

Correct me if I'm wrong, but I believe the next step to add lighting would be to merge the fragment and vertex shader code used in the MeshLambertMaterial with my custom shader code.

Merging the vertex shader would be straightforward. However, in the fragment shader, how do I provide a base color (which will be generated by my shader code) for the code in meshlambert_frag.glsl to apply the lighting calculations? The fragment shader would probably look something like this:

// my custom shader code before the main function
// meshlambert_frag.glsl shader code before the main function
void main {
  vec4 myBaseColor ... // set by my custom fragment shader code
  // meshlambert_frag.glsl code in the main function would use myBaseColor as a base 
  // color for the lighting calculations before setting gl_FragColor
}

Additionally, is copying such a large chunk of code (from the meshlambert shaders) bad practice in this case? What would be a better solution?

Solution

The way to pass a color to the fragment shader code used in MeshLambertMaterial was to modify the diffuseColor variable. So you can set the color in the fragment shader using diffuseColor.rgb *= aVec3. Here's what my shader code looks like:

// vertex
#define LAMBERT

varying vec3 vLightFront;

#ifdef DOUBLE_SIDED

    varying vec3 vLightBack;

#endif

#include <common>
#include <uv_pars_vertex>
#include <uv2_pars_vertex>
#include <envmap_pars_vertex>
#include <bsdfs>
#include <lights_pars>
#include <color_pars_vertex>
#include <morphtarget_pars_vertex>
#include <skinning_pars_vertex>
#include <shadowmap_pars_vertex>
#include <logdepthbuf_pars_vertex>
#include <clipping_planes_pars_vertex>

varying vec3 vertexColor;
varying vec2 vUv;
varying vec4 worldPosition;

void main() {
    vertexColor = vec3(255, 100, 0);
    vUv = uv;
    worldPosition = modelMatrix * vec4( position, 1.0 );

    #include <uv_vertex>
    #include <uv2_vertex>
    #include <color_vertex>

    #include <beginnormal_vertex>
    #include <morphnormal_vertex>
    #include <skinbase_vertex>
    #include <skinnormal_vertex>
    #include <defaultnormal_vertex>

    #include <begin_vertex>
    #include <morphtarget_vertex>
    #include <skinning_vertex>
    #include <project_vertex>
    #include <logdepthbuf_vertex>
    #include <clipping_planes_vertex>

    #include <worldpos_vertex>
    #include <envmap_vertex>
    #include <lights_lambert_vertex>
    #include <shadowmap_vertex>

}
// fragment
uniform float time;
varying vec2 vUv;

vec4 permute( vec4 x ) {

    return mod( ( ( x * 34.0 ) + 1.0 ) * x, 289.0 );

}

vec4 taylorInvSqrt( vec4 r ) {

    return 1.79284291400159 - 0.85373472095314 * r;

}

float snoise( vec3 v ) {

    const vec2 C = vec2( 1.0 / 6.0, 1.0 / 3.0 );
    const vec4 D = vec4( 0.0, 0.5, 1.0, 2.0 );

    // First corner

    vec3 i  = floor( v + dot( v, C.yyy ) );
    vec3 x0 = v - i + dot( i, C.xxx );

    // Other corners

    vec3 g = step( x0.yzx, x0.xyz );
    vec3 l = 1.0 - g;
    vec3 i1 = min( g.xyz, l.zxy );
    vec3 i2 = max( g.xyz, l.zxy );

    vec3 x1 = x0 - i1 + 1.0 * C.xxx;
    vec3 x2 = x0 - i2 + 2.0 * C.xxx;
    vec3 x3 = x0 - 1. + 3.0 * C.xxx;

    // Permutations

    i = mod( i, 289.0 );
    vec4 p = permute( permute( permute(
            i.z + vec4( 0.0, i1.z, i2.z, 1.0 ) )
        + i.y + vec4( 0.0, i1.y, i2.y, 1.0 ) )
        + i.x + vec4( 0.0, i1.x, i2.x, 1.0 ) );

    // Gradients
    // ( N*N points uniformly over a square, mapped onto an octahedron.)

    float n_ = 1.0 / 7.0; // N=7

    vec3 ns = n_ * D.wyz - D.xzx;

    vec4 j = p - 49.0 * floor( p * ns.z *ns.z );  //  mod(p,N*N)

    vec4 x_ = floor( j * ns.z );
    vec4 y_ = floor( j - 7.0 * x_ );    // mod(j,N)

    vec4 x = x_ *ns.x + ns.yyyy;
    vec4 y = y_ *ns.x + ns.yyyy;
    vec4 h = 1.0 - abs( x ) - abs( y );

    vec4 b0 = vec4( x.xy, y.xy );
    vec4 b1 = vec4( x.zw, y.zw );


    vec4 s0 = floor( b0 ) * 2.0 + 1.0;
    vec4 s1 = floor( b1 ) * 2.0 + 1.0;
    vec4 sh = -step( h, vec4( 0.0 ) );

    vec4 a0 = b0.xzyw + s0.xzyw * sh.xxyy;
    vec4 a1 = b1.xzyw + s1.xzyw * sh.zzww;

    vec3 p0 = vec3( a0.xy, h.x );
    vec3 p1 = vec3( a0.zw, h.y );
    vec3 p2 = vec3( a1.xy, h.z );
    vec3 p3 = vec3( a1.zw, h.w );

    // Normalise gradients

    vec4 norm = taylorInvSqrt( vec4( dot( p0, p0 ), dot( p1, p1 ), dot( p2, p2 ), dot( p3, p3 ) ) );
    p0 *= norm.x;
    p1 *= norm.y;
    p2 *= norm.z;
    p3 *= norm.w;

    // Mix final noise value

    vec4 m = max( 0.6 - vec4( dot( x0, x0 ), dot( x1, x1 ), dot( x2, x2 ), dot( x3, x3 ) ), 0.0 );
    m = m * m;
    return 42.0 * dot( m*m, vec4( dot( p0, x0 ), dot( p1, x1 ),
                                dot( p2, x2 ), dot( p3, x3 ) ) );

}

varying vec4 worldPosition;

uniform vec3 diffuse;
uniform vec3 emissive;
uniform float opacity;

varying vec3 vLightFront;

#ifdef DOUBLE_SIDED

    varying vec3 vLightBack;

#endif

#include <common>
#include <packing>
#include <color_pars_fragment>
#include <uv_pars_fragment>
#include <uv2_pars_fragment>
#include <map_pars_fragment>
#include <alphamap_pars_fragment>
#include <aomap_pars_fragment>
#include <lightmap_pars_fragment>
#include <emissivemap_pars_fragment>
#include <envmap_pars_fragment>
#include <bsdfs>
#include <lights_pars>
#include <fog_pars_fragment>
#include <shadowmap_pars_fragment>
#include <shadowmask_pars_fragment>
#include <specularmap_pars_fragment>
#include <logdepthbuf_pars_fragment>
#include <clipping_planes_pars_fragment>

float sigmoid(float x) {
    return x / (1.0 + abs(x));
}

varying vec3 vertexColor;

void main() {
    #include <clipping_planes_fragment>


    vec4 diffuseColor = vec4(diffuse, opacity );
    ReflectedLight reflectedLight = ReflectedLight( vec3( 0.0 ), vec3( 0.0 ), vec3( 0.0 ), vec3( 0.0 ) );
    vec3 totalEmissiveRadiance = emissive;

    #include <logdepthbuf_fragment>
    #include <map_fragment>
    if(worldPosition.y > 300.0) { //snow
        float noise = snoise(vec3(worldPosition.x * 4.0, worldPosition.y * 4.0, worldPosition.z * 4.0))/20.0;
        diffuseColor.rgb *= vec3(1.0 - noise, 1.0 - noise, 1.0 - noise);
    }
    else if (worldPosition.y > 100.0) { // dirt
        float scale = 5.0;
        float effectscale = 0.2;
        float noise = (snoise(vec3(worldPosition.x * scale, worldPosition.y * scale, worldPosition.z * scale)) - 0.2) * effectscale;
        noise = sigmoid(noise);
        diffuseColor.rgb *= vec3(0.54 + noise, 0.27 + noise, 0.07 + noise);
    }
    else { // grass
        float scale = 4.0;
        float effectscale = 0.08;
        float noise = (snoise(vec3(worldPosition.x * scale, worldPosition.y * scale, worldPosition.z * scale)) - 0.2) * effectscale;
        diffuseColor.rgb *= vec3(0, 0.48 + noise, 0.05 + noise);

    }

    #include <color_fragment>
    #include <alphamap_fragment>
    #include <alphatest_fragment>
    #include <specularmap_fragment>
    #include <emissivemap_fragment>

    // accumulation
    reflectedLight.indirectDiffuse = getAmbientLightIrradiance( ambientLightColor );

    #include <lightmap_fragment>

    reflectedLight.indirectDiffuse *= BRDF_Diffuse_Lambert( diffuseColor.rgb );

    #ifdef DOUBLE_SIDED

        reflectedLight.directDiffuse = ( gl_FrontFacing ) ? vLightFront : vLightBack;

    #else

        reflectedLight.directDiffuse = vLightFront;

    #endif

    reflectedLight.directDiffuse *= BRDF_Diffuse_Lambert( diffuseColor.rgb ) * getShadowMask();

    // modulation
    #include <aomap_fragment>

    vec3 outgoingLight = reflectedLight.directDiffuse + reflectedLight.indirectDiffuse + totalEmissiveRadiance;

    #include <normal_flip>
    #include <envmap_fragment>

    gl_FragColor = vec4( outgoingLight, diffuseColor.a );

    #include <premultiplied_alpha_fragment>
    #include <tonemapping_fragment>
    #include <encodings_fragment>
    #include <fog_fragment>
}

The result doesn't look great but the concept of having a custom shader with lighting is there. All you have to do is copy over the code used in the MeshLambertShader and modify diffuseColor.rgb. To reduce the amount of code in your shader, you could add shader chunks (using the THREE.ShaderChunk array) of the MeshLambert shader code. Then you could add #includes in your shader that correspond to the names of the shader chunks you just pushed. For example, you could add the following to your javascript file:

THREE.ShaderChunk["meshlambert_premain_fragment"] = `
        uniform vec3 diffuse;
        uniform vec3 emissive;
        uniform float opacity;

        varying vec3 vLightFront;

        #ifdef DOUBLE_SIDED

            varying vec3 vLightBack;

        #endif

        #include <common>
        #include <packing>
        #include <color_pars_fragment>
        #include <uv_pars_fragment>
        #include <uv2_pars_fragment>
        #include <map_pars_fragment>
        #include <alphamap_pars_fragment>
        #include <aomap_pars_fragment>
        #include <lightmap_pars_fragment>
        #include <emissivemap_pars_fragment>
        #include <envmap_pars_fragment>
        #include <bsdfs>
        #include <lights_pars>
        #include <fog_pars_fragment>
        #include <shadowmap_pars_fragment>
        #include <shadowmask_pars_fragment>
        #include <specularmap_pars_fragment>
        #include <logdepthbuf_pars_fragment>
        #include <clipping_planes_pars_fragment>
`;

And later, in your shader that you want to use the MeshLambert lighting, you could write the following, greatly reducing the amount of code in your shader file:

#include <meshlambert_premain_fragment>