open-standards/joule-code/joule-code-rs/crates/joule-moq-web/web/room/compose.js

// Composition in the page (wai/extensions/live-compose.md): the room's participants laid
// out on a canvas by wai-live-web's integer layout, from the room's WCD1 descriptor and
// its participant set at the instant, and their audio mixed by its integer mixer into an
// AudioWorklet. The layout is recomputed on each lifecycle event (§5.8); between events
// the participant set does not change, and nor does the layout.

const MIX_FRAMES = 960; // 20 ms at 48 kHz

export class Compositor {
  constructor(wai, canvas, descriptor, { log = () => {}, bindingOf = null } = {}) {
    this.wai = wai;
    this.canvas = canvas;
    this.descriptorBytes = descriptor;
    this.desc = JSON.parse(wai.composeDescriptor(descriptor));
    if (this.desc.outcome !== "read") throw new Error(`the composition descriptor: ${this.desc.reason}`);
    const p = this.desc.params;
    canvas.width = p.width;
    canvas.height = p.height;
    this.ctx = canvas.getContext("2d", { willReadFrequently: true });
    const bg = this.desc.background;
    this.background = `#${bg}`;
    this.frame = null; // the composition at the last re-layout, as wai-live-web gives it
    this.layouts = []; // each re-layout: t, set digest, tiles
    this.video = new Map(); // principal -> latest VideoFrame
    this.drawn = new Map(); // principal -> the rectangles last drawn
    this.draws = 0;
    this.log = log;
    // The binding of each frame drawn (WAI browser-sink §10), when the page keeps them.
    this.bindingOf = bindingOf;
    this.audio = new Map(); // principal -> { chunks: Int16Array[], at, channels }
    this.mixed = { blocks: 0, voiced_blocks: 0, peak: 0 };
    this.worklet = null;
    this.workletStats = { played: 0, underruns: 0, peak: 0 };
  }

  relayout(room, t) {
    const out = JSON.parse(room.frame(this.descriptorBytes, BigInt(t)));
    if (out.outcome !== "compose") {
      this.log(`composition refused: ${out.reason} at ${JSON.stringify(out.at)}`);
      return null;
    }
    const f = out.frames[0];
    const prev = this.frame;
    this.frame = f;
    if (!prev || prev.set_digest !== f.set_digest) {
      this.layouts.push({ t: f.t, set_digest: f.set_digest, tiles: f.tiles.map((x) => ({ principal: x.principal, tile: x.tile })) });
      this.log(`re-layout at t=${f.t}: ${f.tiles.length} tile(s), set ${f.set_digest.slice(0, 12)}`);
    }
    return f;
  }

  setVideo(principal, vf) {
    const old = this.video.get(principal);
    if (old) old.close();
    this.video.set(principal, vf);
  }

  draw() {
    const c = this.ctx;
    c.fillStyle = this.background;
    c.fillRect(0, 0, this.canvas.width, this.canvas.height);
    if (!this.frame) return;
    // Tiles in drawing order: the grid's, then the picture-in-picture window (§5.7).
    for (const t of this.frame.tiles) {
      const vf = this.video.get(t.principal);
      if (!t.dst || !vf) continue;
      const [sx, sy, sw, sh] = t.src;
      const [dx, dy, dw, dh] = t.dst;
      c.drawImage(vf, sx, sy, sw, sh, dx, dy, dw, dh);
      const b = this.bindingOf ? this.bindingOf(vf) : null;
      this.drawn.set(t.principal, { tile: t.tile, src: t.src, dst: t.dst, frame: [vf.displayWidth, vf.displayHeight], binding: b ? b.binding : null, group: b ? b.group : null });
    }
    this.draws++;
  }

  // A pixel at the centre of each tile: what a test reads to see video was drawn there.
  samples() {
    if (!this.frame) return [];
    return this.frame.tiles.map((t) => {
      const [x, y, w, h] = t.tile;
      const d = this.ctx.getImageData(x + (w >> 1), y + (h >> 1), 1, 1).data;
      return { principal: t.principal, rgb: [d[0], d[1], d[2]] };
    });
  }

  // ------------------------------------------------------------------------------ audio

  async startAudio(url) {
    this.ac = new AudioContext({ sampleRate: this.desc.params.sample_rate, latencyHint: "interactive" });
    await this.ac.audioWorklet.addModule(url);
    this.worklet = new AudioWorkletNode(this.ac, "room-mix", { outputChannelCount: [2] });
    this.worklet.port.onmessage = (e) => { this.workletStats = e.data; };
    this.worklet.connect(this.ac.destination);
    try { await this.ac.resume(); } catch { /* resumes when allowed */ }
    this.mixTimer = setInterval(() => this.mixOnce(), 20);
  }

  addAudio(principal, audioData) {
    let v = this.audio.get(principal);
    if (!v) { v = { chunks: [], channels: audioData.numberOfChannels }; this.audio.set(principal, v); }
    const n = audioData.numberOfFrames;
    const ch = audioData.numberOfChannels;
    const f = new Float32Array(n * ch);
    audioData.copyTo(f, { planeIndex: 0, format: "f32" });
    const s = new Int16Array(n * ch);
    for (let i = 0; i < s.length; i++) s[i] = Math.max(-32768, Math.min(32767, Math.round(f[i] * 32767)));
    v.channels = ch;
    v.chunks.push(s);
    // At most a second queued per voice.
    let total = v.chunks.reduce((k, c) => k + c.length, 0);
    while (total > 48000 * ch && v.chunks.length > 1) total -= v.chunks.shift().length;
    audioData.close();
  }

  take(v, frames) {
    const want = frames * v.channels;
    const out = new Int16Array(want);
    let at = 0;
    while (at < want && v.chunks.length) {
      const c = v.chunks[0];
      const n = Math.min(c.length, want - at);
      out.set(c.subarray(0, n), at);
      at += n;
      if (n === c.length) v.chunks.shift(); else v.chunks[0] = c.subarray(n);
    }
    return { pcm: out, had: at };
  }

  // One block of the mix (§6): the voices of the layout, each with its gain and pan.
  mixOnce() {
    if (!this.frame || !this.worklet) return;
    const voices = this.frame.voices;
    const chans = [], gains = [], pans = [], parts = [];
    let had = 0;
    for (const v of voices) {
      const a = this.audio.get(v.principal);
      if (!a) continue;
      const t = this.take(a, MIX_FRAMES);
      had += t.had;
      chans.push(a.channels);
      gains.push(v.gain);
      pans.push(v.pan);
      parts.push(t.pcm);
    }
    if (!parts.length) return;
    const pcm = new Int16Array(parts.reduce((k, p) => k + p.length, 0));
    let at = 0;
    for (const p of parts) { pcm.set(p, at); at += p.length; }
    const out = this.wai.mixBlock(this.desc.params.channels, MIX_FRAMES, Uint8Array.from(chans), Uint32Array.from(gains), Int32Array.from(pans), pcm);
    let peak = 0;
    for (let i = 0; i < out.length; i++) { const a = Math.abs(out[i]); if (a > peak) peak = a; }
    this.mixed.blocks++;
    if (had) this.mixed.voiced_blocks++;
    if (peak > this.mixed.peak) this.mixed.peak = peak;
    this.worklet.port.postMessage({ pcm: out }, [out.buffer]);
  }
}