three-shader.js

/* @flow */
import * as THREE from 'three';
import AudioLoader from './audio-loader';
import MidiLoader from './midi-loader';
import VideoLoader from './video-loader';
import CameraLoader from './camera-loader';
import GamepadLoader from './gamepad-loader';
import KeyLoader from './key-loader';
import isVideo from 'is-video';
import type { Rc } from './config';
import { DEFAULT_VERTEX_SHADER, DEFAULT_FRAGMENT_SHADER } from './constants';

const DUMMY_TEXTURE = new THREE.Texture();

type RenderPassTarget = {
  name: string;
  targets: THREE.RenderTarget[];
}
type RenderPass = {
  scene: THREE.Scene;
  camera: THREE.Camera;
  target: ?RenderPassTarget;
}
export type Pass = {
  TARGET?: string;
  vs?: string;
  fs?: string;
}
type Uniforms = {
  [key: string]: {
    type: string;
    value: any;
  }
}

export default class ThreeShader {
  _pixelRatio: number;
  _frameskip: number;
  _start: number;
  _isPlaying: boolean;
  _frame: number;

  _passes: RenderPass[];

  _plane: THREE.Mesh;
  _renderer: THREE.Renderer;
  _targets: THREE.RenderTarget[];
  _textureLoader: THREE.TextureLoader;

  _audioLoader: AudioLoader;
  _cameraLoader: CameraLoader;
  _gamepadLoader: GamepadLoader;
  _keyLoader: KeyLoader;
  _midiLoader: MidiLoader;
  _videoLoader: VideoLoader;
  _uniforms: Uniforms;

  _vertexMode: string;

  constructor(rc: Rc) {
    this._pixelRatio = rc.pixelRatio;
    this._frameskip = rc.frameskip;
    this._vertexMode = rc.vertexMode;

    this._passes = [];

    // Create a target for backbuffer
    this._targets = [
      new THREE.WebGLRenderTarget(
        0, 0,
        { minFilter: THREE.LinearFilter, magFilter: THREE.NearestFilter, format: THREE.RGBAFormat }
      ),
      new THREE.WebGLRenderTarget(
        0, 0,
        { minFilter: THREE.LinearFilter, magFilter: THREE.NearestFilter, format: THREE.RGBAFormat }
      ),
    ];

    this._audioLoader = new AudioLoader();
    this._cameraLoader = new CameraLoader();
    this._gamepadLoader = new GamepadLoader();
    this._keyLoader = new KeyLoader();
    this._midiLoader = new MidiLoader();
    this._videoLoader = new VideoLoader();

    // Prepare uniforms
    this._start = Date.now();
    this._uniforms = {
      backbuffer: { type: 't', value: new THREE.Texture() },
      mouse: { type: 'v2', value: new THREE.Vector2() },
      resolution: { type: 'v2', value: new THREE.Vector2() },
      time: { type: 'f', value: 0.0 },
      vertexCount: { type: 'f', value: rc.vertexCount },
      PASSINDEX: { type: 'i', value: 0 },
    };

    this._textureLoader = new THREE.TextureLoader();
  }

  setPixelRatio(pixelRatio: number): void {
    this._pixelRatio = pixelRatio;
    this._renderer.setPixelRatio(1 / pixelRatio);
    this._resize();
  }

  setFrameskip(frameskip: number): void {
    this._frameskip = frameskip;
  }

  setVertexCount(count: number): void {
    this._uniforms.vertexCount.value = count;
  }

  setVertexMode(mode: string): void {
    this._vertexMode = mode;
  }

  setCanvas(canvas: HTMLCanvasElement): void {
    if (!canvas) {
      return;
    }

    this._renderer = new THREE.WebGLRenderer({ canvas });
    this._renderer.setPixelRatio(1 / this._pixelRatio);
    this._resize();
    window.addEventListener('resize', this._resize);
    window.addEventListener('mousemove', this.mousemove);

    this._frame = 0;
    this.animate();
  }

  _createPlane(fs: ?string, vs: ?string) {
    let plane;
    if (vs) {
      // Create an object for vertexMode
      const geometry = new THREE.BufferGeometry();
      var vertices = new Float32Array(this._uniforms.vertexCount.value * 3);
      geometry.addAttribute('position', new THREE.BufferAttribute(vertices, 3));
      const vertexIds = new Float32Array(this._uniforms.vertexCount.value);
      vertexIds.forEach((v, i) => {
        vertexIds[i] = i;
      });
      geometry.addAttribute('vertexId', new THREE.BufferAttribute(vertexIds, 1));

      const material = new THREE.ShaderMaterial({
        uniforms: this._uniforms,
        vertexShader: vs,
        fragmentShader: fs || DEFAULT_FRAGMENT_SHADER,
        extensions: {
          derivatives: true,
          drawBuffers: false,
          fragDepth: false,
          shaderTextureLOD: false,
        },
      });

      if (this._vertexMode === 'POINTS') {
        plane = new THREE.Points(geometry, material);
      } else if (this._vertexMode === 'LINE_LOOP') {
        plane = new THREE.LineLoop(geometry, material);
      } else if (this._vertexMode === 'LINE_STRIP') {
        plane = new THREE.Line(geometry, material);
      } else if (this._vertexMode === 'LINES') {
        plane = new THREE.LineSegments(geometry, material);
      } else if (this._vertexMode === 'TRI_STRIP') {
        plane = new THREE.Mesh(geometry, material);
        plane.setDrawMode(THREE.TrianglesStripDrawMode);
      } else if (this._vertexMode === 'TRI_FAN') {
        plane = new THREE.Mesh(geometry, material);
        plane.setDrawMode(THREE.TrianglesFanDrawMode);
      } else {
        plane = new THREE.Mesh(geometry, material);
      }
    } else {
      // Create plane
      const geometry = new THREE.PlaneGeometry(2, 2);
      const material = new THREE.ShaderMaterial({
        uniforms: this._uniforms,
        vertexShader: DEFAULT_VERTEX_SHADER,
        fragmentShader: fs,
        extensions: {
          derivatives: true,
          drawBuffers: false,
          fragDepth: false,
          shaderTextureLOD: false,
        },
      });
      plane = new THREE.Mesh(geometry, material);
    }

    return plane;
  }

  _createRenderPass(pass: Pass): RenderPass {
    const scene = new THREE.Scene();
    const camera = new THREE.OrthographicCamera(-1, 1, 1, -1, 0.1, 10);
    camera.position.set(0, 0, 1);
    camera.lookAt(scene.position);

    const plane = this._createPlane(pass.fs, pass.vs);
    scene.add(plane);

    let target: RenderPassTarget;
    if (pass.TARGET) {
      const targetName = pass.TARGET;
      this._uniforms[targetName] = DUMMY_TEXTURE;
      target = {
        name: targetName,
        targets: [
          new THREE.WebGLRenderTarget(
            window.innerWidth / this._pixelRatio, window.innerHeight / this._pixelRatio,
            { minFilter: THREE.LinearFilter, magFilter: THREE.NearestFilter, format: THREE.RGBAFormat }
          ),
          new THREE.WebGLRenderTarget(
            window.innerWidth / this._pixelRatio, window.innerHeight / this._pixelRatio,
            { minFilter: THREE.LinearFilter, magFilter: THREE.NearestFilter, format: THREE.RGBAFormat }
          ),
        ],
      };
    }

    return { scene, camera, target };
  }

  loadShader(passes: Pass[]): void {
    // Dispose old targets
    this._passes.forEach(pass => {
      const target = pass.target;
      if (target) {
        target.targets[0].texture.dispose();
        target.targets[1].texture.dispose();
      }
    });

    // Create new Passes
    this._passes = passes.map(pass => this._createRenderPass(pass));
  }

  loadTexture(name: string, textureUrl: string): void {
    const texture = isVideo(textureUrl) ? this._videoLoader.load(name, textureUrl) : this._textureLoader.load(textureUrl);
    this._uniforms[name] = {
      type: 't',
      value: texture,
    };
  }

  unloadTexture(name: string, textureUrl: string, remove: boolean): void {
    const texture = this._uniforms[name];
    texture.value.dispose();

    if (remove && isVideo(textureUrl)) {
      this._videoLoader.unload(textureUrl);
    }
  }

  mousemove = (e: MouseEvent) => {
    this._uniforms.mouse.value.x = e.clientX / window.innerWidth;
    this._uniforms.mouse.value.y = 1 - e.clientY / window.innerHeight;
  }

  _resize = () => {
    const [width, height] = [window.innerWidth, window.innerHeight];
    this._renderer.setSize(width, height);

    const [bufferWidth, bufferHeight] = [width / this._pixelRatio, height / this._pixelRatio];
    this._passes.forEach(p => {
      if (p.target) {
        p.target.targets.forEach(t => t.setSize(bufferWidth, bufferHeight));
      }
    });
    this._targets.forEach(t => t.setSize(bufferWidth, bufferHeight));
    this._uniforms.resolution.value.x = bufferWidth;
    this._uniforms.resolution.value.y = bufferHeight;
  }

  animate = () => {
    this._frame++;
    if (!this._isPlaying) {
      return;
    }

    requestAnimationFrame(this.animate);
    if (this._frame % this._frameskip === 0) {
      this._render();
    }
  }

  play(): void {
    this._isPlaying = true;
    this.animate();
  }

  stop(): void {
    this._isPlaying = false;
    this._audioLoader.disable();
    this._cameraLoader.disable();
    this._keyLoader.disable();
    this._midiLoader.disable();
    this._gamepadLoader.disable();
  }

  _render(): void {
    this._uniforms.time.value = (Date.now() - this._start) / 1000;
    this._targets = [this._targets[1], this._targets[0]];
    this._uniforms.backbuffer.value = this._targets[0].texture;

    if (this._audioLoader.isEnabled) {
      this._audioLoader.update();
      this._uniforms.volume.value = this._audioLoader.getVolume();
    }

    if (this._gamepadLoader.isEnabled) {
      this._gamepadLoader.update();
    }

    this._passes.forEach((pass: RenderPass, i: number) => {
      this._uniforms.PASSINDEX.value = i;

      const target = pass.target;
      if (target) {
        this._renderer.render(pass.scene, pass.camera, target.targets[1], true);

        // Swap buffers after render so that we can use the buffer in latter passes
        target.targets = [target.targets[1], target.targets[0]];
        this._uniforms[target.name].value = target.targets[0].texture;
      } else {
        this._renderer.render(pass.scene, pass.camera, null);
      }
    });

    const lastPass = this._passes[this._passes.length - 1];

    // Render last pass to canvas even if target is specified
    if (lastPass.target) {
      this._renderer.render(lastPass.scene, lastPass.camera, null);
    }

    // Render result to backbuffer
    this._renderer.render(lastPass.scene, lastPass.camera, this._targets[1], true);
  }

  toggleAudio(flag: boolean): void {
    if (flag) {
      this._audioLoader.enable();
      this._uniforms = {
        ...this._uniforms,
        volume: { type: 'f', value: 0 },
        spectrum: { type: 't', value: this._audioLoader.spectrum },
        samples: { type: 't', value: this._audioLoader.samples },
      };
    } else if (this._uniforms.spectrum) {
      this._uniforms.spectrum.value.dispose();
      this._uniforms.samples.value.dispose();
      this._audioLoader.disable();
    }
  }

  toggleMidi(flag: boolean): void {
    if (flag) {
      this._midiLoader.enable();
      this._uniforms = {
        ...this._uniforms,
        midi: { type: 't', value: this._midiLoader.midiTexture },
        note: { type: 't', value: this._midiLoader.noteTexture },
      };
    } else if (this._uniforms.midi) {
      this._uniforms.midi.value.dispose();
      this._uniforms.note.value.dispose();
      this._midiLoader.disable();
    }
  }

  toggleCamera(flag: boolean): void {
    if (flag) {
      this._cameraLoader.enable();
      this._uniforms = {
        ...this._uniforms,
        camera: { type: 't', value: this._cameraLoader.texture },
      };
    } else {
      this._cameraLoader.disable();
    }
  }

  toggleKeyboard(flag: boolean): void {
    if (flag) {
      this._keyLoader.enable();
      this._uniforms = {
        ...this._uniforms,
        key: { type: 't', value: this._keyLoader.texture },
      };
    } else {
      this._keyLoader.disable();
    }
  }

  toggleGamepad(flag: boolean): void {
    if (flag) {
      this._gamepadLoader.enable();
      this._uniforms = {
        ...this._uniforms,
        gamepad: { type: 't', value: this._gamepadLoader.texture },
      };
    } else {
      this._gamepadLoader.disable();
    }
  }
}