js map: add jsdoc for typechecking

1 files changed, 268 insertions, 20 deletions

diff --git a/src/persistent-hash-map.mjs b/src/persistent-hash-map.mjs
index e7c76fe..bf5ef53 100644
--- a/src/persistent-hash-map.mjs
+++ b/src/persistent-hash-map.mjs
@@ -3,8 +3,11 @@ import { isEqual } from "./gleam.mjs";
 const referenceMap = new WeakMap();
 const tempDataView = new DataView(new ArrayBuffer(8));
 let referenceUID = 0;
-
-/** hash the object by reference using a weak map and incrementing uid */
+/**
+ * hash the object by reference using a weak map and incrementing uid
+ * @param {any} o
+ * @returns {number}
+ */
 function hashByReference(o) {
   const known = referenceMap.get(o);
   if (known !== undefined) {
@@ -17,11 +20,20 @@ function hashByReference(o) {
   referenceMap.set(o, hash);
   return hash;
 }
-/** merge two hashes in an order sensitive way */
+/**
+ * merge two hashes in an order sensitive way
+ * @param {number} a
+ * @param {number} b
+ * @returns {number}
+ */
 function hashMerge(a, b) {
   return (a ^ (b + 0x9e3779b9 + (a << 6) + (a >> 2))) | 0;
 }
-/** standard string hash popularised by java */
+/**
+ * standard string hash popularised by java
+ * @param {string} s
+ * @returns {number}
+ */
 function hashString(s) {
   let hash = 0;
   const len = s.length;
@@ -30,14 +42,30 @@ function hashString(s) {
   }
   return hash;
 }
-/** hash a number by converting to two integers and do some jumbling */
+/**
+ * hash a number by converting to two integers and do some jumbling
+ * @param {number} n
+ * @returns {number}
+ */
 function hashNumber(n) {
   tempDataView.setFloat64(0, n);
   const i = tempDataView.getInt32(0);
   const j = tempDataView.getInt32(4);
   return Math.imul(0x45d9f3b, (i >> 16) ^ i) ^ j;
 }
-/** hash any js object */
+/**
+ * hash a BigInt by converting it to a string and hashing that
+ * @param {BigInt} n
+ * @returns {number}
+ */
+function hashBigInt(n) {
+  return hashString(n.toString());
+}
+/**
+ * hash any js object
+ * @param {any} o
+ * @returns {number}
+ */
 function hashObject(o) {
   const proto = Object.getPrototypeOf(o);
   if (proto !== null && typeof proto.hashCode === "function") {
@@ -77,7 +105,11 @@ function hashObject(o) {
   }
   return h;
 }
-/** hash any js value */
+/**
+ * hash any js value
+ * @param {any} u
+ * @returns {number}
+ */
 export function getHash(u) {
   if (u === null) return 0x42108422;
   if (u === undefined) return 0x42108423;
@@ -89,15 +121,40 @@ export function getHash(u) {
     case "string":
       return hashString(u);
     case "bigint":
-      return hashNumber(u);
+      return hashBigInt(u);
     case "object":
       return hashObject(u);
     case "symbol":
       return hashByReference(u);
     case "function":
       return hashByReference(u);
+    default:
+      return 0; // should be unreachable
   }
 }
+/**
+ * @template K,V
+ * @typedef {ArrayNode<K,V> | IndexNode<K,V> | CollisionNode<K,V>} Node
+ */
+/**
+ * @template K,V
+ * @typedef {{ type: typeof ENTRY, k: K, v: V }} Entry
+ */
+/**
+ * @template K,V
+ * @typedef {{ type: typeof ARRAY_NODE, size: number, array: (undefined | Entry<K,V> | Node<K,V>)[] }} ArrayNode
+ */
+/**
+ * @template K,V
+ * @typedef {{ type: typeof INDEX_NODE, bitmap: number, array: (Entry<K,V> | Node<K,V>)[] }} IndexNode
+ */
+/**
+ * @template K,V
+ * @typedef {{ type: typeof COLLISION_NODE, hash: number, array: Entry<K, V>[] }} CollisionNode
+ */
+/**
+ * @typedef {{ val: boolean }} Flag
+ */
 const SHIFT = 5; // number of bits you need to shift by to get the next bucket
 const BUCKET_SIZE = Math.pow(2, SHIFT);
 const MASK = BUCKET_SIZE - 1; // used to zero out all bits not in the bucket
@@ -107,20 +164,35 @@ const ENTRY = 0;
 const ARRAY_NODE = 1;
 const INDEX_NODE = 2;
 const COLLISION_NODE = 3;
+/** @type {IndexNode<any,any>} */
 const EMPTY = {
   type: INDEX_NODE,
   bitmap: 0,
   array: [],
 };
-/** Mask the hash to get only the bucket corresponding to shift */
+/**
+ * Mask the hash to get only the bucket corresponding to shift
+ * @param {number} hash
+ * @param {number} shift
+ * @returns {number}
+ */
 function mask(hash, shift) {
   return (hash >>> shift) & MASK;
 }
-/** Set only the Nth bit where N is the masked hash */
+/**
+ * Set only the Nth bit where N is the masked hash
+ * @param {number} hash
+ * @param {number} shift
+ * @returns {number}
+ */
 function bitpos(hash, shift) {
   return 1 << mask(hash, shift);
 }
-/** Count the number of 1 bits in a number */
+/**
+ * Count the number of 1 bits in a number
+ * @param {number} x
+ * @returns {number}
+ */
 function bitcount(x) {
   x -= (x >> 1) & 0x55555555;
   x = (x & 0x33333333) + ((x >> 2) & 0x33333333);
@@ -129,11 +201,23 @@ function bitcount(x) {
   x += x >> 16;
   return x & 0x7f;
 }
-/** Calculate the array index of an item in a bitmap index node */
+/**
+ * Calculate the array index of an item in a bitmap index node
+ * @param {number} bitmap
+ * @param {number} bit
+ * @returns {number}
+ */
 function index(bitmap, bit) {
   return bitcount(bitmap & (bit - 1));
 }
-/** Efficiently copy an array and set one value at an index */
+/**
+ * Efficiently copy an array and set one value at an index
+ * @template T
+ * @param {T[]} arr
+ * @param {number} at
+ * @param {T} val
+ * @returns {T[]}
+ */
 function cloneAndSet(arr, at, val) {
   const len = arr.length;
   const out = new Array(len);
@@ -143,7 +227,14 @@ function cloneAndSet(arr, at, val) {
   out[at] = val;
   return out;
 }
-/** Efficiently copy an array and insert one value at an index */
+/**
+ * Efficiently copy an array and insert one value at an index
+ * @template T
+ * @param {T[]} arr
+ * @param {number} at
+ * @param {T} val
+ * @returns {T[]}
+ */
 function spliceIn(arr, at, val) {
   const len = arr.length;
   const out = new Array(len + 1);
@@ -158,7 +249,13 @@ function spliceIn(arr, at, val) {
   }
   return out;
 }
-/** Efficiently copy an array and remove one value at an index */
+/**
+ * Efficiently copy an array and remove one value at an index
+ * @template T
+ * @param {T[]} arr
+ * @param {number} at
+ * @returns {T[]}
+ */
 function spliceOut(arr, at) {
   const len = arr.length;
   const out = new Array(len - 1);
@@ -173,7 +270,17 @@ function spliceOut(arr, at) {
   }
   return out;
 }
-/** Create a new node containing two entries */
+/**
+ * Create a new node containing two entries
+ * @template K,V
+ * @param {number} shift
+ * @param {K} key1
+ * @param {V} val1
+ * @param {number} key2hash
+ * @param {K} key2
+ * @param {V} val2
+ * @returns {Node<K,V>}
+ */
 function createNode(shift, key1, val1, key2hash, key2, val2) {
   const key1hash = getHash(key1);
   if (key1hash === key2hash) {
@@ -196,7 +303,22 @@ function createNode(shift, key1, val1, key2hash, key2, val2) {
     addedLeaf
   );
 }
-/** Associate a node with a new entry, creating a new node. */
+/**
+ * @template T,K,V
+ * @callback AssocFunction
+ * @param {T} root
+ * @param {number} shift
+ * @param {number} hash
+ * @param {K} key
+ * @param {V} val
+ * @param {Flag} addedLeaf
+ * @returns {Node<K,V>}
+ */
+/**
+ * Associate a node with a new entry, creating a new node
+ * @template T,K,V
+ * @type {AssocFunction<Node<K,V>,K,V>}
+ */
 function assoc(root, shift, hash, key, val, addedLeaf) {
   switch (root.type) {
     case ARRAY_NODE:
@@ -207,6 +329,10 @@ function assoc(root, shift, hash, key, val, addedLeaf) {
       return assocCollision(root, shift, hash, key, val, addedLeaf);
   }
 }
+/**
+ * @template T,K,V
+ * @type {AssocFunction<ArrayNode<K,V>,K,V>}
+ */
 function assocArray(root, shift, hash, key, val, addedLeaf) {
   const idx = mask(hash, shift);
   const node = root.array[idx];
@@ -260,6 +386,10 @@ function assocArray(root, shift, hash, key, val, addedLeaf) {
     array: cloneAndSet(root.array, idx, n),
   };
 }
+/**
+ * @template T,K,V
+ * @type {AssocFunction<IndexNode<K,V>,K,V>}
+ */
 function assocIndex(root, shift, hash, key, val, addedLeaf) {
   const bit = bitpos(hash, shift);
   const idx = index(root.bitmap, bit);
@@ -350,6 +480,10 @@ function assocIndex(root, shift, hash, key, val, addedLeaf) {
     }
   }
 }
+/**
+ * @template T,K,V
+ * @type {AssocFunction<CollisionNode<K,V>,K,V>}
+ */
 function assocCollision(root, shift, hash, key, val, addedLeaf) {
   // if there is a hash collision
   if (hash === root.hash) {
@@ -389,7 +523,13 @@ function assocCollision(root, shift, hash, key, val, addedLeaf) {
     addedLeaf
   );
 }
-/** Find the index of a key in the collision node's array */
+/**
+ * Find the index of a key in the collision node's array
+ * @template K,V
+ * @param {CollisionNode<K,V>} root
+ * @param {K} key
+ * @returns {number}
+ */
 function collisionIndexOf(root, key) {
   const size = root.array.length;
   for (let i = 0; i < size; i++) {
@@ -399,7 +539,20 @@ function collisionIndexOf(root, key) {
   }
   return -1;
 }
-/** Return the found entry or undefined if not present in the root */
+/**
+ * @template T,K,V
+ * @callback FindFunction
+ * @param {T} root
+ * @param {number} shift
+ * @param {number} hash
+ * @param {K} key
+ * @returns {undefined | Entry<K,V>}
+ */
+/**
+ * Return the found entry or undefined if not present in the root
+ * @template K,V
+ * @type {FindFunction<Node<K,V>,K,V>}
+ */
 function find(root, shift, hash, key) {
   switch (root.type) {
     case ARRAY_NODE:
@@ -410,6 +563,10 @@ function find(root, shift, hash, key) {
       return findCollision(root, key);
   }
 }
+/**
+ * @template K,V
+ * @type {FindFunction<ArrayNode<K,V>,K,V>}
+ */
 function findArray(root, shift, hash, key) {
   const idx = mask(hash, shift);
   const node = root.array[idx];
@@ -421,6 +578,10 @@ function findArray(root, shift, hash, key) {
   }
   return find(node, shift + SHIFT, hash, key);
 }
+/**
+ * @template K,V
+ * @type {FindFunction<IndexNode<K,V>,K,V>}
+ */
 function findIndex(root, shift, hash, key) {
   const bit = bitpos(hash, shift);
   if ((root.bitmap & bit) === 0) {
@@ -436,6 +597,12 @@ function findIndex(root, shift, hash, key) {
   }
   return undefined;
 }
+/**
+ * @template K,V
+ * @param {CollisionNode<K,V>} root
+ * @param {K} key
+ * @returns {undefined | Entry<K,V>}
+ */
 function findCollision(root, key) {
   const idx = collisionIndexOf(root, key);
   if (idx < 0) {
@@ -444,8 +611,19 @@ function findCollision(root, key) {
   return root.array[idx];
 }
 /**
+ * @template T,K,V
+ * @callback WithoutFunction
+ * @param {T} root
+ * @param {number} shift
+ * @param {number} hash
+ * @param {K} key
+ * @returns {undefined | Node<K,V>}
+ */
+/**
  * Remove an entry from the root, returning the updated root.
  * Returns undefined if the node should be removed from the parent.
+ * @template K,V
+ * @type {WithoutFunction<Node<K,V>,K,V>}
  * */
 function without(root, shift, hash, key) {
   switch (root.type) {
@@ -457,6 +635,10 @@ function without(root, shift, hash, key) {
       return withoutCollision(root, key);
   }
 }
+/**
+ * @template K,V
+ * @type {WithoutFunction<ArrayNode<K,V>,K,V>}
+ */
 function withoutArray(root, shift, hash, key) {
   const idx = mask(hash, shift);
   const node = root.array[idx];
@@ -522,6 +704,10 @@ function withoutArray(root, shift, hash, key) {
     array: cloneAndSet(root.array, idx, n),
   };
 }
+/**
+ * @template K,V
+ * @type {WithoutFunction<IndexNode<K,V>,K,V>}
+ */
 function withoutIndex(root, shift, hash, key) {
   const bit = bitpos(hash, shift);
   if ((root.bitmap & bit) === 0) {
@@ -568,6 +754,12 @@ function withoutIndex(root, shift, hash, key) {
   }
   return root;
 }
+/**
+ * @template K,V
+ * @param {CollisionNode<K,V>} root
+ * @param {K} key
+ * @returns {undefined | Node<K,V>}
+ */
 function withoutCollision(root, key) {
   const idx = collisionIndexOf(root, key);
   // if the key not found, no changes
@@ -586,6 +778,12 @@ function withoutCollision(root, key) {
     array: spliceOut(root.array, idx),
   };
 }
+/**
+ * @template K,V
+ * @param {undefined | Node<K,V>} root
+ * @param {(value:V,key:K)=>void} fn
+ * @returns {void}
+ */
 function forEach(root, fn) {
   if (root === undefined) {
     return;
@@ -604,10 +802,19 @@ function forEach(root, fn) {
     forEach(item, fn);
   }
 }
-/** Extra wrapper to keep track of map size and clean up the API */
+/**
+ * Extra wrapper to keep track of map size and clean up the API
+ * @template K,V
+ */
 export default class PMap {
+  /**
+   * @template V
+   * @param {Record<string,V>} o
+   * @returns {PMap<string,V>}
+   */
   static fromObject(o) {
     const keys = Object.keys(o);
+    /** @type PMap<string,V> */
     let m = PMap.new();
     for (let i = 0; i < keys.length; i++) {
       const k = keys[i];
@@ -615,7 +822,13 @@ export default class PMap {
     }
     return m;
   }
+  /**
+   * @template K,V
+   * @param {Map<K,V>} o
+   * @returns {PMap<K,V>}
+   */
   static fromMap(o) {
+    /** @type PMap<K,V> */
     let m = PMap.new();
     o.forEach((v, k) => {
       m = m.set(k, v);
@@ -625,10 +838,20 @@ export default class PMap {
   static new() {
     return new PMap(undefined, 0);
   }
+  /**
+   * @param {undefined | Node<K,V>} root
+   * @param {number} size
+   */
   constructor(root, size) {
     this.root = root;
     this.size = size;
   }
+  /**
+   * @template NotFound
+   * @param {K} key
+   * @param {NotFound} notFound
+   * @returns {NotFound | V}
+   */
   get(key, notFound) {
     if (this.root === undefined) {
       return notFound;
@@ -639,6 +862,11 @@ export default class PMap {
     }
     return found.v;
   }
+  /**
+   * @param {K} key
+   * @param {V} val
+   * @returns {PMap<K,V>}
+   */
   set(key, val) {
     const addedLeaf = { val: false };
     const root = this.root === undefined ? EMPTY : this.root;
@@ -648,6 +876,10 @@ export default class PMap {
     }
     return new PMap(newRoot, addedLeaf.val ? this.size + 1 : this.size);
   }
+  /**
+   * @param {K} key
+   * @returns {PMap<K,V>}
+   */
   delete(key) {
     if (this.root === undefined) {
       return this;
@@ -661,20 +893,32 @@ export default class PMap {
     }
     return new PMap(newRoot, this.size - 1);
   }
+  /**
+   * @param {K} key
+   * @returns {boolean}
+   */
   has(key) {
     if (this.root === undefined) {
       return false;
     }
     return find(this.root, 0, getHash(key), key) !== undefined;
   }
+  /**
+   * @returns {[K,V][]}
+   */
   entries() {
     if (this.root === undefined) {
       return [];
     }
+    /** @type [K,V][] */
     const result = [];
     this.forEach((v, k) => result.push([k, v]));
     return result;
   }
+  /**
+   *
+   * @param {(val:V,key:K)=>void} fn
+   */
   forEach(fn) {
     forEach(this.root, fn);
   }
@@ -685,6 +929,10 @@ export default class PMap {
     });
     return h;
   }
+  /**
+   * @param {unknown} o
+   * @returns {boolean}
+   */
   equals(o) {
     if (!(o instanceof PMap)) {
       return false;