Collect into a Set and a Map

"use typeshade";

/* @example
{
  "title": "Host: collect into a Set and a Map",
  "blurb": "Four kernel functions that take JavaScript's own `Set` and `Map` (change 0080), and the host that calls them. `labels` adds each nonzero label once, `latest` keeps the last value set for each key, `lookup` reads the map where `has` guards `get`, and `present` adds in one loop and reads in the next. The host passes `new Set()` and `new Map()`, and when each promise resolves they hold the entries a sequential run leaves: each key once, in the order first added, the value set last. On WebGPU each add or set pushes a record the call applies in order, and a lookup is a binary search of the keys the call uploads; WebGL2 runs none.",
  "renderable": false,
  "reason": "no entry point"
}
*/

// Each nonzero label once, in the order first seen. The function returns how many there are.
export function labels(pixels: array<u32>, seen: Set<u32>): u32 {
  for (let i: u32 = 0; i < pixels.length; i++) {
    if (pixels[i] !== 0) {
      seen.add(pixels[i]);
    }
  }
  return seen.size;
}

// The last value set for each key, in the place of the key's first set.
export function latest(keys: array<u32>, values: array<f32>, last: Map<u32, f32>) {
  for (let i: u32 = 0; i < keys.length; i++) {
    last.set(keys[i], values[i]);
  }
}

// One call adds in its first loop and reads in its second: each probe is 1 when an earlier pixel
// of the image holds its label, 0 otherwise.
export function present(pixels: array<u32>, seen: Set<u32>, probes: array<u32>, out: array<u32>) {
  for (let i: u32 = 0; i < pixels.length; i++) {
    seen.add(pixels[i]);
  }
  for (let i: u32 = 0; i < probes.length; i++) {
    out[i] = seen.has(probes[i]) ? 1 : 0;
  }
}

// A later call reads the map: a key it does not hold reads as -1.
export function lookup(keys: array<u32>, last: Map<u32, f32>, out: array<f32>) {
  for (let i: u32 = 0; i < keys.length; i++) {
    out[i] = last.has(keys[i]) ? last.get(keys[i])! : -1.;
  }
}

// A closed set of string keys uses ids on the GPU and strings on the host (change 0081).
type Channel = "r" | "g" | "b";
export function channels(xs: array<f32>, seen: Set<Channel>, last: Map<Channel, f32>, out: array<f32>): u32 {
  for (let i: u32 = 0; i < xs.length; i++) {
    const c: Channel = i % 3 === 0 ? "g" : i % 3 === 1 ? "r" : "b";
    seen.add(c);
  }
  for (let i: u32 = 0; i < xs.length; i++) {
    const c: Channel = i % 3 === 0 ? "g" : i % 3 === 1 ? "r" : "b";
    last.set(c, xs[i]);
  }
  for (let i: u32 = 0; i < out.length; i++) {
    const c: Channel = i % 3 === 0 ? "g" : i % 3 === 1 ? "r" : "b";
    out[i] = seen.has(c) && last.has(c) ? last.get(c)! : -1.;
  }
  return seen.size;
}

// Keys used as Map values also cross the boundary as strings, including an initial lookup.
export function channelNames(xs: array<f32>, names: Map<Channel, Channel>, out: array<u32>) {
  for (let i: u32 = 0; i < xs.length; i++) {
    names.set("g", "r");
    names.set("b", "g");
  }
  for (let i: u32 = 0; i < out.length; i++) {
    out[i] = u32(names.has("r") && names.get("r")! === "b" ? 1 : 0) + u32(names.has("g") && names.get("g")! === "r" ? 2 : 0);
  }
}

// Collect into JavaScript's own Set and Map from kernel functions (surface §74, change 0080).
//
// The host passes a `Set` and a `Map` as it passes an array, by reference. When the promise
// resolves they hold the entries a sequential run of the loop leaves, in insertion order.
import {
  channelNames,
  channels,
  labels,
  latest,
  lookup,
  present,
} from './host-set-unique.shade.ts';

/** The unique nonzero labels of a small image, the last weight set for each label, and the
 *  weights of labels 0 to 9 read back from the map, -1 for a label it does not hold. */
export async function run(n = 64): Promise<{
  seen: Set<number>;
  count: number;
  last: Map<number, number>;
  found: Float32Array;
  present: Uint32Array;
  names: Map<string, string>;
  nameMarks: Uint32Array;
  channels: { seen: Set<string>; last: Map<string, number>; found: Float32Array; count: number };
}> {
  // Labels 0 to 6, where 0 is the background.
  const pixels = Uint32Array.from({ length: n }, (_, i) => (i * 5) % 7);
  const seen = new Set<number>();
  const count = await labels(pixels, seen); // seen holds 5, 3, 1, 6, 4, 2

  const weights = Float32Array.from({ length: n }, (_, i) => i / 4);
  const last = new Map<number, number>();
  await latest(pixels, weights, last); // the weight of each label's last pixel

  const found = new Float32Array(10);
  await lookup(
    Uint32Array.from({ length: 10 }, (_, i) => i),
    last,
    found,
  );

  // One call that adds, then reads what it added: labels 0 to 9 against the image's.
  const marks = new Uint32Array(10);
  await present(
    pixels,
    new Set<number>(),
    Uint32Array.from({ length: 10 }, (_, i) => i),
    marks,
  );
  const channelSeen = new Set<'r' | 'g' | 'b'>(['r']);
  const channelLast = new Map<'r' | 'g' | 'b', number>([
    ['b', -1],
    ['r', -2],
  ]);
  const channelFound = new Float32Array(6);
  const channelCount = await channels(weights, channelSeen, channelLast, channelFound);
  const names = new Map<'r' | 'g' | 'b', 'r' | 'g' | 'b'>([['r', 'b']]);
  const nameMarks = new Uint32Array(3);
  await channelNames(weights, names, nameMarks);
  return {
    names,
    nameMarks,
    seen,
    count,
    last,
    found,
    present: marks,
    channels: { seen: channelSeen, last: channelLast, found: channelFound, count: channelCount },
  };
}