three.jsのマテリアルを継承してシェーダーを追記する

three.jsでカスタムシェーダーを実装する際に「`MeshStandardMaterial`の実装を引き継いでシェーダーを追記したい」ケースがあり、マテリアルの継承をしつつシェーダーを追記する効率のいい方法を思いついたので備忘録として残しておきます。

[最終的なコードはこちら](https://codesandbox.io/s/extend-onbeforecompile-class-d4ie1?file=/resource/js/CheckerdNormalMeshStdMaterial.js)

安直な思考だと「シェーダーを書くなら`ShaderMaterial`か`RawShaderMaterial`を使う」となりがちですが、さすがにライブラリ側ですでに実装されている機能をわざわざ自前で再現しないといけないわけないよね、、と思い「three.js extend material」などで検索したら以下のページにたどりつきました。

[three.jsのデフォルトのマテリアルを拡張する | びわの家ブログ](https://biwanoie.tokyo/o2o1j/)

どうやら`onBeforeCompile`というのを使うとシェーダーを上書きできるっぽい。

# Material.onBeforeCompile

その名の通り`WebGLRenderer`側でマテリアルのシェーダープログラムがコンパイルされる直前に呼ばれる関数です。  
[Material#onBeforeCompile – three.js docs](https://threejs.org/docs/index.html#api/en/materials/Material.onBeforeCompile)

シェーダープログラムのコンパイルはレンダリング時に毎回実行されるわけではなく、`Material.needsUpdate`が`true`のときだけ呼ばれます。（不要なコンパイルを走らせないため）

`onBeforeCompile`の第一引数にはプロパティに`vertexShader` `fragmentShader`を持つ`Shader`オブジェクト[^1]が渡ってくるので、この関数内でシェーダーソースの追記（改変）が可能です。

シェーダーを追記するといっても、`vertexShader` `fragmaneShader`は`String`なので`String.replace`を使って任意の記述に上書きする形で追記します。追記するためにはどの行の位置に追記するかを把握しておかないと追加できませんね。

まず、変更前のシェーダーの中身を確認したいので`console.log`でブラウザのコンソールに出力してみます。

```javascript
const material = new MeshStandardMaterial({ roughness: 0.25, metalness: 0.85 });

// WebGLRednerer.render() の中でシェーダープログラムをコンパイルする直前に呼ばれる
material.onBeforeCompile = (shader) => {
  const { fragmentShader } = shader;
  console.log(fragmentShader);
};
```
**※ 長くなるのでレンダラーなどの記述は省いています**

`MeshStandardMaterial`がちゃんと機能しているか確認するために、環境マップとライトを使用しています。
![01.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/86966/b4ac14ed-f424-d6b8-e5b9-1749fe061221.png)

<details><summary>`MeshStandardMaterial`の`fragmentShader`の中身</summary><div>

```glsl
#define STANDARD

#ifdef PHYSICAL
	#define REFLECTIVITY
	#define CLEARCOAT
	#define TRANSMISSION
#endif

uniform vec3 diffuse;
uniform vec3 emissive;
uniform float roughness;
uniform float metalness;
uniform float opacity;

#ifdef TRANSMISSION
	uniform float transmission;
#endif

#ifdef REFLECTIVITY
	uniform float reflectivity;
#endif

#ifdef CLEARCOAT
	uniform float clearcoat;
	uniform float clearcoatRoughness;
#endif

#ifdef USE_SHEEN
	uniform vec3 sheen;
#endif

varying vec3 vViewPosition;

#ifndef FLAT_SHADED

	varying vec3 vNormal;

	#ifdef USE_TANGENT

		varying vec3 vTangent;
		varying vec3 vBitangent;

	#endif

#endif

#include <common>
#include <packing>
#include <dithering_pars_fragment>
#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 <transmissionmap_pars_fragment>
#include <bsdfs>
#include <cube_uv_reflection_fragment>
#include <envmap_common_pars_fragment>
#include <envmap_physical_pars_fragment>
#include <fog_pars_fragment>
#include <lights_pars_begin>
#include <lights_physical_pars_fragment>
#include <shadowmap_pars_fragment>
#include <bumpmap_pars_fragment>
#include <normalmap_pars_fragment>
#include <clearcoat_pars_fragment>
#include <roughnessmap_pars_fragment>
#include <metalnessmap_pars_fragment>
#include <logdepthbuf_pars_fragment>
#include <clipping_planes_pars_fragment>

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;

	#ifdef TRANSMISSION
		float totalTransmission = transmission;
	#endif

	#include <logdepthbuf_fragment>
	#include <map_fragment>
	#include <color_fragment>
	#include <alphamap_fragment>
	#include <alphatest_fragment>
	#include <roughnessmap_fragment>
	#include <metalnessmap_fragment>
	#include <normal_fragment_begin>
	#include <normal_fragment_maps>
	#include <clearcoat_normal_fragment_begin>
	#include <clearcoat_normal_fragment_maps>
	#include <emissivemap_fragment>
	#include <transmissionmap_fragment>

	// accumulation
	#include <lights_physical_fragment>
	#include <lights_fragment_begin>
	#include <lights_fragment_maps>
	#include <lights_fragment_end>

	// modulation
	#include <aomap_fragment>

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

	// this is a stub for the transmission model
	#ifdef TRANSMISSION
		diffuseColor.a *= mix( saturate( 1. - totalTransmission + linearToRelativeLuminance( reflectedLight.directSpecular + reflectedLight.indirectSpecular ) ), 1.0, metalness );
	#endif

	gl_FragColor = vec4( outgoingLight, diffuseColor.a );

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

}
```
</div></details>

three.jsで実装されている各マテリアルのシェーダーは、機能ごとに細かく分かれたシェーダーソースを`#include <shader_name>`というthree.js独自の記法でインクルードする仕様になっています。
`onBeforeCompile`の引数で渡ってくる時点では`vertexShader` `fragmentShader`ともにインクルード前のソースになっています。

[^1]: ドキュメントを読むと、`onBeforeCompile`の第一引数には Shader とだけ書いてありますが、`vertexShader` `fragmaneShader`以外にも各種テクスチャの有無の`boolean`など、レンダリングに必要なプロパティがいろいろ用意されています。[three.js/WebGLPrograms.js at 94f043c4e105eb73236529231388402da2b07cba · mrdoob/three.js](https://github.com/mrdoob/three.js/blob/94f043c4e105eb73236529231388402da2b07cba/src/renderers/webgl/WebGLPrograms.js#L155-L275)

# シェーダーを追記する

`fragmaneShader`の最終行に追記して上書きしてみます。  
変更前のシェーダーの中身を確認した結果、`MeshStandardMaterial`の最終行は`#include <dithering_fragment>`なので、これをキーワードとして`String.replace`で新しい記述と置換します。  
`gl_FragColor`の`rgb`値を強制的に法線の`xyz`にする処理を追記してみます。

```javascript
const material = new MeshStandardMaterial({ roughness: 0.25, metalness: 0.85 });

material.onBeforeCompile = (shader) => {
  // 置換してシェーダーに追記する
  const keyword = '#include <dithering_fragment>';
  shader.fragmentShader = shader.fragmentShader.replace(keyword, `${ keyword }
  gl_FragColor.rgb = vNormal;// 法線で上書き
`);
};
```

**※ `fragmentShader.raplece`の第二引数に`${ keyword }`を含めないと、キーワードとして設定した行に対応する処理がごっそり上書きされて機能しなくなるので注意しましょう！（該当箇所を意図的に上書きしたい場合を除く）**

**※ 文字通り色も上書きになるので追記部分のあとに適用させたい処理がある場合は残したい処理より上の行に挿入しましょう！**

![02.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/86966/a339dc49-ab37-4825-1212-9a06490ecc7d.png)


# uniform変数を追加する

uniform変数を追加して、`bool`で強制法線表示をオン/オフできるようにしてみます。
`onBeforeCompile`の引数の`Shader`オブジェクトに`uniforms`プロパティが用意されているので、そこに新しいプロパティを追加することでuniform変数を追加できます。

three.jsの順当な書き方であればマテリアルを`NormalMaterial`に切り替える必要がありますが[^2]、シェーダーの上書きだとマテリアルを切り替える必要なく法線表示切り替えが実現できます。

[^2]: 自分がこれ以外の方法を知らないだけかもしれないので、もし簡単な方法があれば教えてください。

```javascript
const material = new MeshStandardMaterial({ roughness: 0.25, metalness: 0.85 });

material.onBeforeCompile = (shader) => {
  Object.assign(shader.uniforms, { uShowNormal: { value: true } });// 法線表示切り替え用booleanを用意

  const keyword1 = 'void main() {';
  shader.fragmentShader = shader.fragmentShader.replace(keyword1, `
uniform bool uShowNormal;// uniform変数を追加
${ keyword1 }`);

  const keyword2 = '#include <dithering_fragment>';
  shader.fragmentShader = shader.fragmentShader.replace(keyword2, `${ keyword2 }
  if (uShowNormal) gl_FragColor.rgb = vNormal;// uShowNormal が true のときだけ法線で上書き
`);
};
```

**※ JS側でuniform変数を追加している部分は、`Object.assign(shader, { uniforms: { uShowNormal: { value: true } } });`や`shader.uniforms = { uShowNormal: { value: true } };`と書きたいところですが、`onBeforeCompile`が呼ばれている時点ですでに`shader.uniforms`の中身に`Material`で使用するuniform変数がセットされている[^3]ので、`shader.uniforms`ごと上書きすると`WebGLRenderer.render`でエラーになってしまいます。気をつけましょう！**

[^3]: [three.js/WebGLRenderer.js at master · mrdoob/three.js](https://github.com/mrdoob/three.js/blob/master/src/renderers/WebGLRenderer.js#L1377)


## Material.userDataの活用

`onBeforeCompile`の中で`uniforms`変数を定義してしまうと`onBeforeCompile`のスコープ外で変数が触れず不便なので、`userData`を使用してスコープ外でも`uniform`変数が触れるようにします。

```javascript
const material = new MeshStandardMaterial({ roughness: 0.25, metalness: 0.85 });
Object.assign(material.userData, {
  uniforms: { uShowNormal: { value: false } }// userDataにuniformsを追加
});

material.onBeforeCompile = (shader) => {
  Object.assign(shader.uniforms, material.userData.uniforms);// userDataのuniformsとリンクさせる

  const keyword1 = 'void main() {';
  shader.fragmentShader = shader.fragmentShader.replace(keyword1, `
uniform bool uShowNormal;
${ keyword1 }`);

  const keyword2 = '#include <dithering_fragment>';
  shader.fragmentShader = shader.fragmentShader.replace(keyword2, `${ keyword2 }
  if (uShowNormal) gl_FragColor.rgb = vNormal;
`);
};

material.userData.uniforms.uShowNormal = true;// onBeforeCompile 外から変更可能
```

uniform変数を`onBeforeCompile`外に持ってきたことで、`dat.GUI`で操作しやすくなります。

![03.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/86966/cfb42a34-698d-1b6d-0aa3-23c59f8effda.png)


# クラス化する

ここまでは`MeshStandardMaterial`のインスタンスの`onBeforeCompile`に関数を代入する形で書いてきましたが、簡素化するために`MeshStandardMaterial`を継承したクラスにしてみます。
今まで通り`this.onBeforeCompile = () => {};`と書いても良いですが、せっかく`class`構文を使うのでメンバ関数として書きます。（実行時に`material.onBeforeCompile()`となるので挙動はどちらも同じです）

```javascript
// MeshStandardMaterial を継承
class NormalMeshStandardMaterial extends MeshStandardMaterial {
  constructor(params) {
    super(params);

    Object.assign(this.userData, {
      uniforms: { uShowNormal: { value: false } }
    });
  }

  onBeforeCompile(shader) {
    Object.assign(shader.uniforms, material.userData.uniforms);

    const keyword1 = 'void main() {';
    shader.fragmentShader = shader.fragmentShader.replace(keyword1, `
uniform bool uShowNormal;
${ keyword1 }`);

    const keyword2 = '#include <dithering_fragment>';
    shader.fragmentShader = shader.fragmentShader.replace(keyword2, `${ keyword2 }
  if (uShowNormal) gl_FragColor.rgb = vNormal;
`);
  }
}
```

# 複数回の継承を想定してシェーダーを追記する

ここからさらにシェーダーに追記して別の機能を追加したい場合、上記の`NormalMeshStandardMaterial`を継承することになるのですが、今の書き方だと子クラスのシェーダーの挿入箇所の指定を親クラスのシェーダーの追記内容から抜粋する形になります。
もし子クラスを実装したあとに親クラスのシェーダー追記部分に変更があったら、子クラスの挿入箇所の指定も変える必要が出てきてめんどくさいですよね。

めんどくさいのはイヤなので、親クラスでシェーダーを追記する前後に特定のコメントを記述することで、継承時の柔軟性をもたせます。

```javascript
class NormalMeshStandardMaterial extends MeshStandardMaterial {
  ...

  onBeforeCompile(shader) {
    Object.assign(shader.uniforms, material.userData.uniforms);

    const keyword1 = 'void main() {';
    shader.fragmentShader = shader.fragmentShader.replace(keyword1, `
/* NormalMeshStandardMaterial fragment defines - begin */
uniform bool uShowNormal;
/* NormalMeshStandardMaterial fragment defines - end */
${ keyword1 }`);

    const keyword2 = '#include <dithering_fragment>';
    shader.fragmentShader = shader.fragmentShader.replace(keyword2, `${ keyword2 }
  /* NormalMeshStandardMaterial fragment main - begin */
  if (uShowNormal) gl_FragColor.rgb = vNormal;
  /* NormalMeshStandardMaterial fragment main - end */
`);
  }
}
```

これで親クラスのシェーダー追記部分の変数宣言部の前後、main追記部の前後どこでも自由に挿入でき、かつ親クラスのシェーダーの内容に変更があっても子クラスの挿入箇所の指定も変更する必要がありません。

`NormalMeshStandardMaterial`を継承して、チェッカー模様にする`CheckerdNormalMeshStandardMaterial`を作ったとすると以下のようになります。

シェーダーの追記を何度も行うので、先にシェーダー追記用の関数を用意しておきます。

```javascript
/**
 * シェーダーソースのキーワード直前に追記
 * @param {string} shaderSource 
 * @param {string} keyword 
 * @param {string} append 
 */
const insertShaderBefore = (shaderSource, keyword, append) => {
  return shaderSource.replace(keyword, `${ append }
${ keyword }`);
};

/**
 * シェーダーソースのキーワード直後に追記
 * @param {string} shaderSource 
 * @param {string} keyword 
 * @param {string} append 
 */
const insertShaderAfter = (shaderSource, keyword, append) => {
  return shaderSource.replace(keyword, `${ keyword }
${ append }`);
};
```

```javascript
// NormalMeshStandardMaterial を継承
class CheckerdNormalMeshStandardMaterial extends NormalMeshStandardMaterial {
  constructor(params) {
    super(params);

    Object.assign(this.userData.uniforms, {
      uShowChecker: { value: true },
      uTime: { value: 0 }
    });
  }

  onBeforeCompile(shader) {
    super.onBeforeCompile(shader);
    Object.assign(shader.uniforms, this.userData.uniforms);

    // 共通の変数宣言部
    const commonDefines = `
uniform bool uShowChecker;
uniform float uTime;
varying vec2 vUv;
`;

    // 頂点シェーダー : main の前に uniform 変数を追記
    const keywordVert1 = 'void main() {';
    shader.vertexShader = insertShaderBefore(shader.vertexShader, keywordVert1, `
/* CheckerdNormalMeshStandardMaterial vertex defines - begin */
${ commonDefines }
/* CheckerdNormalMeshStandardMaterial vertex defines - end */
`);

    // 頂点シェーダー : 最終行に追記
    const keywordVert2 = '#include <fog_vertex>';
    shader.vertexShader = insertShaderAfter(shader.vertexShader, keywordVert2, `
  /* CheckerdNormalMeshStandardMaterial vertex main - begin */
  vUv = uv;
  /* CheckerdNormalMeshStandardMaterial vertex main - end */
`);

    // フラグメントシェーダー : NormalMeshStandardMaterial defines の後に変数を追記
    const keywordFrag1 = '/* NormalMeshStandardMaterial fragment defines - end */';
    shader.fragmentShader = insertShaderAfter(shader.fragmentShader, keywordFrag1, `
/* CheckerdNormalMeshStandardMaterial fragment defines - begin */
${ commonDefines }
/* CheckerdNormalMeshStandardMaterial fragment defines - end */
`);

    // フラグメントシェーダー : NormalMeshStandardMaterial main の前に追記
    const keywordFrag2 = '/* NormalMeshStandardMaterial fragment main - begin */';
    shader.fragmentShader = insertShaderBefore(shader.fragmentShader, keywordFrag2, `
  /* CheckerdNormalMeshStandardMaterial fragment main - begin */
  if (uShowChecker) {
    float t = uTime * 6.0;
    float freq = PI * 20.0;
    float mask = sin(vUv.x * freq * 2.0 + t) * sin(vUv.y * freq + t);
    mask = sign(mask);
    gl_FragColor.rgb *= mask;
  }
  /* CheckerdNormalMeshStandardMaterial fragment main - end */
`);
  }
}
```

`MeshStandardMaterial`の実装を引き継ぎつつ、ひとつのマテリアルでエフェクト表示と法線表示のオン／オフができるようになりました。
子クラスでも挿入前後にコメントを記載することで何度でも継承して`onBeforeCompile`でシェーダーを効率よく上書きすることができます。

![04.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/86966/dcf6630f-71ec-7e41-4d92-7798f9ca7ebf.png)

# デモ

[最終的なコードはこちら](https://codesandbox.io/s/extend-onbeforecompile-class-d4ie1?file=/resource/js/CheckerdNormalMeshStdMaterial.js)


# おまけ：シェーダーソースの確認方法

`onBeforeCompile`でシェーダーを追記する以外のケースにも当てはまるのですが、three.jsでシェーダーを書いているとthree.js側のシェーダーの記述が見えないため、全容が把握し辛いですよね。  
なので、three.js側で実装されている部分も含めたほぼ全文のシェーダーソースの確認方法を紹介します

## インクルード前のシェーダーソース

### Material.onBeforeCompile の引数で確認

前述の通り、three.jsで実装されている各種マテリアルのシェーダーは、機能ごとに細かく分かれたシェーダーソースを`#include <shader_name>`というthree.js独自の記法でインクルードする仕様になっています。  
`onBeforeCompile`の引数で渡ってくる時点では`vertexShader` `fragmentShader`ともにインクルード前のソースになっています。  
`onBeforeCompile`で上書きする際は`console.log`で一度インクルード前のソースを確認して、挿入箇所を探るのが良いでしょう。

### ライブラリ側で確認

各マテリアルのインクルード前のシェーダーソースは[three/src/renderers/shaders/ShaderLib/](https://github.com/mrdoob/three.js/tree/master/src/renderers/shaders/ShaderLib)に格納されています。

インクルードされるシェーダーソースは[three/src/renderer/shaders/ShaderChunk](https://github.com/mrdoob/three.js/tree/master/src/renderers/shaders/ShaderChunk)にありますが、いちいち`node_modules`を潜るのも面倒ですよね。

## インクルード後のシェーダーソース

### わざとエラーを起こしてコンソールにエラー出力させる

いちばん手っ取り早く低カロリーでできる方法です。  
セミコロンを消したり、`float`の小数点を消したり、文の途中で改行を入れたり、お好みの方法でエラーを起こしましょう。

ブラウザのDevToolsのコンソールにエラーとともに実行時のシェーダーソースが出力されます。

![05.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/86966/95bbd92e-b8ce-84f5-ef08-b472ffcf3ef9.png)


コンソールからコピーして、テキストエディタで行番号を消せば、ある程度見やすくなります。

![06.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/86966/eda70539-449f-d77c-db34-f987113f7be7.png)

少々手荒な方法ですが、このやり方であればインクルード後のシェーダーソースのほぼ全文を確認することができます。