hamt.js

//  A persistent hash array mapped trie (HAMT).
//
//  An immutable map from string keys to values. Every change returns a new map and leaves the old one
//  as it was, sharing every node the change didn't touch. It grows one small node at a time, so it
//  never has to be rehashed.


//  Hashing and indexing

const BITS = 5;               //  Bits of the hash used at each level
const WIDTH = 1 << BITS;      //  32 slots in a node
const MASK = WIDTH - 1;       //  0b11111

//  FNV-1a, a simple 32-bit string hash
function hash(key) {
  let h = 0x811c9dc5;
  for (let i = 0; i < key.length; i++) {
    h ^= key.charCodeAt(i);
    h = Math.imul(h, 0x01000193);
  }
  return h >>> 0;
}

//  The number of bits set in a 32-bit value
function popcount(x) {
  x = x - ((x >>> 1) & 0x55555555);
  x = (x & 0x33333333) + ((x >>> 2) & 0x33333333);
  x = (x + (x >>> 4)) & 0x0f0f0f0f;
  return Math.imul(x, 0x01010101) >>> 24;
}

//  Which of the 32 slots a hash falls in at this level
function fragment(h, shift) {
  return (h >>> shift) & MASK;
}

//  Where a slot's entry is in the packed array: the number of used slots below it
function index(bitmap, bit) {
  return popcount(bitmap & (bit - 1));
}


//  The nodes

//  A leaf holds one key, its value and the key's full hash
function Leaf(h, key, value) {
  return { type: 'leaf', hash: h, key, value };
}

//  A branch holds a bitmap of its used slots and an array of just those entries
function Branch(bitmap, entries) {
  return { type: 'branch', bitmap, entries };
}

//  Keys whose whole 32-bit hash is the same share a collision node
function Collision(h, leaves) {
  return { type: 'collision', hash: h, leaves };
}

const EMPTY = Branch(0, []);

//  Copies of an array with one entry replaced, inserted or removed. The original is never changed.
function arraySet(array, i, value) {
  const copy = array.slice();
  copy[i] = value;
  return copy;
}

function arrayInsert(array, i, value) {
  return [...array.slice(0, i), value, ...array.slice(i)];
}

function arrayRemove(array, i) {
  return [...array.slice(0, i), ...array.slice(i + 1)];
}


//  Looking up a key

function get(map, key) {
  const h = hash(key);
  let node = map;
  let shift = 0;

  while (true) {
    if (node.type === 'leaf')
      return node.key === key ? node.value : undefined;

    if (node.type === 'collision') {
      const leaf = node.leaves.find(l => l.key === key);
      return leaf ? leaf.value : undefined;
    }

    //  A branch: if the slot is empty the key isn't here, otherwise go down into it
    const bit = 1 << fragment(h, shift);
    if ((node.bitmap & bit) === 0)
      return undefined;

    node = node.entries[index(node.bitmap, bit)];
    shift += BITS;
  }
}


//  Adding a key

//  Returns a new map with key set to value
function set(map, key, value) {
  return insert(map, Leaf(hash(key), key, value), 0);
}

//  Inserts an entry (a leaf, or a collision node being pushed down) into a node at this level
function insert(node, entry, shift) {
  if (node.type === 'branch') {
    const bit = 1 << fragment(entry.hash, shift);
    const i = index(node.bitmap, bit);

    //  An empty slot: set its bit and put the entry in the packed array
    if ((node.bitmap & bit) === 0)
      return Branch(node.bitmap | bit, arrayInsert(node.entries, i, entry));

    //  A used slot: insert into whatever is there, and copy this branch to point at the result
    return Branch(node.bitmap, arraySet(node.entries, i, insert(node.entries[i], entry, shift + BITS)));
  }

  //  A leaf or collision node with a different hash: push it down into a new branch, then insert
  //  the entry beside it. If their next five bits are the same too, this happens again a level down.
  if (node.hash !== entry.hash)
    return insert(insert(EMPTY, node, shift), entry, shift);

  //  The same key: replace the value
  if (node.type === 'leaf' && node.key === entry.key)
    return entry;

  //  A different key with exactly the same hash: more bits won't separate them, so keep both
  if (node.type === 'leaf')
    return Collision(node.hash, [node, entry]);

  const i = node.leaves.findIndex(l => l.key === entry.key);
  return Collision(node.hash, i >= 0 ? arraySet(node.leaves, i, entry) : [...node.leaves, entry]);
}


//  Removing a key

//  Returns a new map without key, or the same map if key wasn't in it
function remove(map, key) {
  return removeFrom(map, key, hash(key), 0) || EMPTY;
}

//  Returns the node without the key: the same node if the key isn't there, or null if nothing is left
function removeFrom(node, key, h, shift) {
  if (node.type === 'leaf')
    return node.key === key ? null : node;

  if (node.type === 'collision') {
    const i = node.leaves.findIndex(l => l.key === key);
    if (i < 0)
      return node;

    //  Down to one leaf: it doesn't need a collision node any more
    if (node.leaves.length === 2)
      return node.leaves[1 - i];

    return Collision(node.hash, arrayRemove(node.leaves, i));
  }

  const bit = 1 << fragment(h, shift);
  if ((node.bitmap & bit) === 0)
    return node;

  const i = index(node.bitmap, bit);
  const child = node.entries[i];
  const newChild = removeFrom(child, key, h, shift + BITS);

  //  Not found further down: nothing changes, and nothing is copied
  if (newChild === child)
    return node;

  //  The slot is now empty: clear its bit and take its entry out
  if (newChild === null) {
    const bitmap = node.bitmap & ~bit;
    if (bitmap === 0)
      return null;

    const entries = arrayRemove(node.entries, i);

    //  A branch left with one leaf is replaced by the leaf, which moves up a level (but not at the root)
    if (shift > 0 && entries.length === 1 && entries[0].type !== 'branch')
      return entries[0];

    return Branch(bitmap, entries);
  }

  //  The child collapsed to a single leaf, and it's the only thing here: move it up again
  if (shift > 0 && node.entries.length === 1 && newChild.type !== 'branch')
    return newChild;

  return Branch(node.bitmap, arraySet(node.entries, i, newChild));
}


//  Going through every entry

//  Every [key, value] pair, depth first, in hash order
function* entries(node) {
  if (node.type === 'leaf') {
    yield [node.key, node.value];
  } else if (node.type === 'collision') {
    for (const leaf of node.leaves)
      yield [leaf.key, leaf.value];
  } else {
    for (const entry of node.entries)
      yield* entries(entry);
  }
}


export { EMPTY, get, set, remove, entries, hash, popcount };