dope2/lib/util/trie.js

// A purely functional Radix Trie
// https://en.wikipedia.org/wiki/Radix_tree


export const emptyTrie = {
  children: [],
};

// invariant: the keys of a trie node should always be sorted
export const isProperlySorted = trie => {
  for (let i=0; i<trie.children.length-1; i++) {
    if (trie.children[i] >= trie.children[i+1]) {
      return false; // not properly sorted!
    }
  }
  for (const [_, node] of trie.children) {
    if (!isProperlySorted(node)) {
      return false;
    }
  }
  return true;
}

// find maximal common prefix, and whether string A is smaller than B
const commonPrefix = (strA, strB) => {
  let i=0
  for (; i<strA.length && i<strB.length; i++) {
    const charA = strA[i], charB = strB[i];
    if (charA < charB) {
      return [strA.slice(0, i), true];
    }
    if (charA > charB) {
      return [strA.slice(0, i), false];
    }
  }
  return [strA.slice(0, i), false];
}

// a funny kind of binary search, that assumes 'ls' is a sorted list of strings, and none of the strings have a common prefix
const binarySearch = (ls, key) => {
  return __binarySearch(ls, key, 0, ls.length);
}
const __binarySearch = (ls, key, min, max) => {
  if (min === max) {
    return [max, ""]; // otherwise we go out of bounds
  }
  const middle = Math.floor((min+max)/2);
  const [prefix, smaller] = commonPrefix(key, ls[middle][0]);
  if (prefix.length > 0) {
    return [middle, prefix];
  }
  if (smaller) {
    // key was smaller than middle
    return __binarySearch(ls, key, min, middle);
  }
  return __binarySearch(ls, key, middle+1, max);
}

const check = trie => {
  // uncomment if you think shit is broken:
  // if (!isProperlySorted(trie)) {
  //   throw new Error('not properly sorted!')
  // }
  return trie;
}

// insert (key,value) into trie.
export const insert = trie => key => value => {
  if (key.length === 0) {
    // set value of current node
    return {
      value,
      children: trie.children,
    };
  }

  const [insertPos, prefix] = binarySearch(trie.children, key);
  if (insertPos === trie.children.length) {
    // insert node at end
    return check({
      value: trie.value,
      children: [
        ...trie.children,
        [key, {value, children:[]}],
      ],
    });
  }
  if (prefix.length === 0) {
    // nothing in common...
    // insert new node into children
    return check({
      value: trie.value,
      children: trie.children.toSpliced(
        insertPos, // insert position
        0, // delete nothing
        [key, {value, children:[]}],
      ),
    });
  }

  const [haveKey, haveChildNode] = trie.children[insertPos];
  if (prefix.length === haveKey.length) {
    // recurse
    return check({
      value: trie.value,
      children: trie.children.with(
        insertPos, // position to update
        [haveKey, insert(haveChildNode)(key.slice(prefix.length))(value)],
      )
    })
  }

  // otherwise, split entry:
  const havePostFix = haveKey.slice(prefix.length);
  const postFix = key.slice(prefix.length);
  return check({
    value: trie.value,
    children: trie.children.with(
      insertPos, // position to update
      [prefix, {
        children: (havePostFix < postFix)
          ? [
              [havePostFix, haveChildNode],
              [postFix, {value, children: []}],
            ]
          : [
              [postFix, {value, children: []}],
              [havePostFix, haveChildNode],
            ],
      }],
    ),
  });
};

// given a prefix, return a string X such that prefix+X is a possibly larger prefix for the same entries as the original prefix.
export const growPrefix = trie => key => {
  const [pos, prefix] = binarySearch(trie.children, key);
  if (prefix.length === 0) {
    return "";
  }
  if (pos === trie.children.length) {
    return "";
  }
  const [haveKey, haveChildNode] = trie.children[pos];
  if (key.length < haveKey.length) {
    if (haveKey.startsWith(key)) {
      return haveKey.slice(key.length);
    }
  }
  if (key.length > haveKey.length) {
    if (key.startsWith(haveKey)) {
      return growPrefix(haveChildNode)(key.slice(haveKey.length));
    }
  }
  return "";
}

// get array of (key, value) entries whose keys are prefixed by 'key'.
export const suggest = trie => key => maxSuggestions => {
  return __suggest(trie, "", key, maxSuggestions);
}

const __suggest = (trie, path, remaining, maxSuggestions) => {
  if (maxSuggestions === 0) {
    return [];
  }

  if (remaining === '') {
    const results = [];
    if (trie.value !== undefined) {
      results.push([path, trie.value]);
      maxSuggestions--;
    }
    for (const [haveKey, haveChildNode] of trie.children) {
      const moreSuggestions = __suggest(haveChildNode, path+haveKey, remaining, maxSuggestions);
      results.push(...moreSuggestions);
      maxSuggestions -= moreSuggestions.length;
      if (maxSuggestions === 0) {
        break;
      }
    }
    return results;
  }

  const [pos] = binarySearch(trie.children, remaining);
  if (pos < trie.children.length) {
    const [haveKey, haveChildNode] = trie.children[pos];
    if (!remaining.startsWith(haveKey) && !haveKey.startsWith(remaining)) {
      return [];
    }
    return __suggest(haveChildNode, path+haveKey, remaining.slice(haveKey.length), maxSuggestions);
  }

  return [];
}