compressedPly.ts

// SuperSplat / PlayCanvas compressed PLY: a chunk element of quantization
// ranges per 256 splats (positions, log scales, colors), then four packed
// uint32 words per splat. Not standard 3DGS PLY; PlyFile.read reads both.
// Writers sort splats in Morton order of their centers so that each chunk's
// ranges are small.
import type { ArrayData } from "./ArrayData";
import { SplatMath } from "./math";
import { COMPRESSED_CHUNK_PROPERTIES, COMPRESSED_CHUNK_SIZE, COMPRESSED_VERTEX_PROPERTIES, PlyFile, encodeCompressedSh } from "./ply";
import { counter, type Progress } from "./progress";
import { Sh } from "./sh";
import { SortKeys } from "./sort";
import { Splats } from "./splats";

// Per splat: center, ln scales, quaternion w x y z, rgb, opacity
const VALUES = 14;
// splat-transform's clamp of ln scales
const LN_SCALE_LIMIT = 20;

// The splat bundle of `data` (any layout), in Morton order unless `sort` is
// false, and with `sh` (default true) its SH bundle as an sh element of
// uchar f_rest_* in splat-transform's mapping (clamped to [-4, 4] by the
// format; degree 3 at most). A LoD tree is not written.
function writeCompressedPly(data: ArrayData, { sort = true, sh = true, progress }: {
  sort?: boolean;
  sh?: boolean;
  progress?: Progress;
} = {}): Uint8Array<ArrayBuffer> {
  const { count } = data;
  const decode = Splats.decoder(Splats.get(data).bundle);
  const found = sh ? Sh.find(data.bundles) : undefined;
  const decodeSh = found && Sh.decoder(data, found);
  // The format stops at degree 3: higher degrees write their first 3 bands
  const perChannel = found ? Sh.COEFFICIENTS[Math.min(found.degree, 3)] : 0;
  const shValues = perChannel * 3;
  const tick = counter(progress, 2 * count, "encoding");

  const values = new Float32Array(count * VALUES);
  const s = Splats.makeExtSplat();
  const q = { x: 0, y: 0, z: 0, w: 1 };
  for (let i = 0; i < count; i += 1) {
    tick(i);
    decode(i, s);
    SplatMath.normalizeQuat(s.quaternion, q);
    const v = i * VALUES;
    values[v] = s.center.x;
    values[v + 1] = s.center.y;
    values[v + 2] = s.center.z;
    values[v + 3] = lnScale(s.scales.x);
    values[v + 4] = lnScale(s.scales.y);
    values[v + 5] = lnScale(s.scales.z);
    values[v + 6] = q.w;
    values[v + 7] = q.x;
    values[v + 8] = q.y;
    values[v + 9] = q.z;
    values[v + 10] = s.rgb.r;
    values[v + 11] = s.rgb.g;
    values[v + 12] = s.rgb.b;
    values[v + 13] = s.opacity;
  }
  const order = sort ? mortonOrder(values, count) : undefined;

  const chunks = Math.ceil(count / COMPRESSED_CHUNK_SIZE);
  const header = new TextEncoder().encode([
    "ply", "format binary_little_endian 1.0", "comment Generated by Radix",
    `element chunk ${chunks}`, ...COMPRESSED_CHUNK_PROPERTIES.map((name) => `property float ${name}`),
    `element vertex ${count}`, ...COMPRESSED_VERTEX_PROPERTIES.map((name) => `property uint ${name}`),
    ...(shValues > 0 ? [`element sh ${count}`, ...Array.from({ length: shValues }, (_, k) => `property uchar f_rest_${k}`)] : []),
    "end_header", "",
  ].join("\n"));
  const stride = COMPRESSED_CHUNK_PROPERTIES.length;
  const vertexStart = header.byteLength + chunks * stride * 4;
  const shStart = vertexStart + count * COMPRESSED_VERTEX_PROPERTIES.length * 4;
  const bytes = new Uint8Array(shStart + count * shValues);
  bytes.set(header);
  const view = new DataView(bytes.buffer);

  // Mins then maxes of the center, ln scales and rgb, as the chunk properties
  const range = new Float32Array(stride);
  const extend = (at: number, value: number) => {
    range[at] = Math.min(range[at], value);
    range[at + 3] = Math.max(range[at + 3], value);
  };
  for (let c = 0; c < chunks; c += 1) {
    const start = c * COMPRESSED_CHUNK_SIZE;
    const end = Math.min(count, start + COMPRESSED_CHUNK_SIZE);
    for (let g = 0; g < stride; g += 6) {
      range.fill(Infinity, g, g + 3);
      range.fill(-Infinity, g + 3, g + 6);
    }
    for (let j = start; j < end; j += 1) {
      const v = (order ? order[j] : j) * VALUES;
      for (let k = 0; k < 3; k += 1) {
        extend(k, values[v + k]);
        extend(6 + k, values[v + 3 + k]);
        extend(12 + k, values[v + 10 + k]);
      }
    }
    for (let k = 0; k < stride; k += 1) {
      view.setFloat32(header.byteLength + (c * stride + k) * 4, range[k], true);
    }

    for (let j = start; j < end; j += 1) {
      tick(count + j);
      const v = (order ? order[j] : j) * VALUES;
      const at = vertexStart + j * 16;
      view.setUint32(at, pack111011(values, v, range, 0), true);
      view.setUint32(at + 4, packRotation(values, v + 6), true);
      view.setUint32(at + 8, pack111011(values, v + 3, range, 6), true);
      view.setUint32(at + 12, packColor(values, v + 10, range), true);
    }
  }

  // f_rest_{c * N + k} (channel-major) of each splat in the written order
  if (decodeSh) {
    const coeffs = new Float32Array(Sh.VALUES);
    for (let j = 0; j < count; j += 1) {
      decodeSh(order ? order[j] : j, coeffs);
      const at = shStart + j * shValues;
      for (let c = 0; c < 3; c += 1) {
        for (let k = 0; k < perChannel; k += 1) {
          bytes[at + c * perChannel + k] = encodeCompressedSh(coeffs[k * 3 + c]);
        }
      }
    }
  }
  return bytes;
}

// SuperSplat / PlayCanvas compressed PLY files of splats
export class CompressedPlyFile {
  // Either kind of PLY (the same as PlyFile.read)
  static readonly read = PlyFile.read;
  // The splat bundle of `data` (any layout)
  static readonly write = writeCompressedPly;
}

function lnScale(scale: number) {
  return Math.max(-LN_SCALE_LIMIT, Math.min(LN_SCALE_LIMIT, Math.log(scale)));
}

// `value` in [min, max] as a `bits`-bit code (0 for an empty range)
function unorm(value: number, min: number, max: number, bits: number) {
  const top = 2 ** bits - 1;
  const t = max > min ? (value - min) / (max - min) : 0;
  return Math.max(0, Math.min(top, Math.floor(t * top + 0.5)));
}

// Three values at `v` in the ranges at `at` (mins) and `at + 3` (maxes)
function pack111011(values: Float32Array, v: number, range: Float32Array, at: number) {
  return unorm(values[v], range[at], range[at + 3], 11) * 2 ** 21
    + unorm(values[v + 1], range[at + 1], range[at + 4], 10) * 2 ** 11
    + unorm(values[v + 2], range[at + 2], range[at + 5], 11);
}

// rgb in the chunk's color ranges, then linear opacity
function packColor(values: Float32Array, v: number, range: Float32Array) {
  return unorm(values[v], range[12], range[15], 8) * 2 ** 24
    + unorm(values[v + 1], range[13], range[16], 8) * 2 ** 16
    + unorm(values[v + 2], range[14], range[17], 8) * 2 ** 8
    + unorm(values[v + 3], 0, 1, 8);
}

// The index of the largest of w, x, y, z in 2 bits, then the other three
// (made positive by flipping the sign, each within +-1 / sqrt(2)) in 10
function packRotation(values: Float32Array, v: number) {
  let largest = 0;
  for (let k = 1; k < 4; k += 1) {
    if (Math.abs(values[v + k]) > Math.abs(values[v + largest])) {
      largest = k;
    }
  }
  const sign = values[v + largest] < 0 ? -1 : 1;
  let word = largest;
  for (let k = 0; k < 4; k += 1) {
    if (k !== largest) {
      word = word * 1024 + unorm(sign * values[v + k] * Math.SQRT1_2 + 0.5, 0, 1, 10);
    }
  }
  return word;
}

// Splat indices in Morton order of their centers on a 1024^3 grid over
// their bounds
function mortonOrder(values: Float32Array, count: number) {
  const lo = [Infinity, Infinity, Infinity];
  const hi = [-Infinity, -Infinity, -Infinity];
  for (let i = 0; i < count; i += 1) {
    for (let k = 0; k < 3; k += 1) {
      lo[k] = Math.min(lo[k], values[i * VALUES + k]);
      hi[k] = Math.max(hi[k], values[i * VALUES + k]);
    }
  }
  const cell = (value: number, k: number) => (hi[k] > lo[k] ? Math.min(1023, Math.floor((1024 * (value - lo[k])) / (hi[k] - lo[k]))) : 0);

  const keys = new Uint32Array(count);
  for (let i = 0; i < count; i += 1) {
    const v = i * VALUES;
    keys[i] = ((spread3(cell(values[v + 2], 2)) << 2) | (spread3(cell(values[v + 1], 1)) << 1) | spread3(cell(values[v], 0))) >>> 0;
  }
  return SortKeys.sort({ keys, count, bits: 32 }).ordering;
}

// The 10 bits of x at every third bit
function spread3(x: number) {
  let s = (x ^ (x << 16)) & 0xff0000ff;
  s = (s ^ (s << 8)) & 0x0300f00f;
  s = (s ^ (s << 4)) & 0x030c30c3;
  return (s ^ (s << 2)) & 0x09249249;
}