Add inferring of node in MapFunction from tree

Related-to GH-3.
Closes GH-4.

Reviewed-by: Titus Wormer <[email protected]>

Author: JuanM04

.gitignore

@@ -1,5 +1,6 @@
 .DS_Store
-*.d.ts
+index.d.ts
+test.d.ts
 *.log
 coverage/
 node_modules/

complex-types.d.ts

@@ -0,0 +1,23 @@
+import type {Node, Parent} from 'unist'
+
+/**
+ * Internal utility to collect all descendants of in `Tree`.
+ * @see https://github.com/syntax-tree/unist-util-visit-parents/blob/18d36ad/complex-types.d.ts#L43
+ */
+export type InclusiveDescendant<
+  Tree extends Node = never,
+  Found = void
+> = Tree extends Parent
+  ?
+      | Tree
+      | InclusiveDescendant<
+          Exclude<Tree['children'][number], Found | Tree>,
+          Found | Tree
+        >
+  : Tree
+
+export type MapFunction<Tree extends Node = Node> = (
+  node: InclusiveDescendant<Tree>,
+  index: number | null,
+  parent: InclusiveDescendant<Tree> | null
+) => InclusiveDescendant<Tree>

index.js

@@ -1,46 +1,29 @@
 /**
- * @typedef {import('unist').Parent} Parent
- * @typedef {import('unist').Position} Position
  * @typedef {import('unist').Node} Node
- * @typedef {Record<string, unknown> & {type: string, position?: Position|undefined}} NodeLike
- */
-
-/**
- * Function called with a node to produce a new node.
- *
- * @callback MapFunction
- * @param {NodeLike|Node} node Current node being processed
- * @param {number} [index] Index of `node`, or `null`
- * @param {Parent} [parent] Parent of `node`, or `null`
- * @returns {NodeLike|Node} Node to be used in the new tree. Its children are not used: if the original node has children, those are mapped.
  */
 
 /**
  * Unist utility to create a new tree by mapping all nodes with the given function.
  *
- * @param {NodeLike|Node} tree Tree to map
- * @param {MapFunction} iteratee Function that returns a new node
- * @returns {NodeLike|Node} New mapped tree.
+ * @template {Node} Tree
+ * @param {Tree} tree Tree to map
+ * @param {import('./complex-types').MapFunction<Tree>} iteratee Function that returns a new node
+ * @returns {Tree} New mapped tree.
  */
 export function map(tree, iteratee) {
+  // @ts-expect-error Looks like a children.
   return preorder(tree, null, null)
 
-  /**
-   * @param {NodeLike|Node} node
-   * @param {number} [index]
-   * @param {Parent} [parent]
-   * @returns {Node}
-   */
+  /** @type {import('./complex-types').MapFunction<Tree>} */
   function preorder(node, index, parent) {
     var newNode = Object.assign({}, iteratee(node, index, parent))
 
     if ('children' in node) {
       // @ts-expect-error Looks like a parent.
       newNode.children = node.children.map(function (
-        /** @type {Node} */ child,
+        /** @type {import('./complex-types').InclusiveDescendant<Tree>} */ child,
         /** @type {number} */ index
       ) {
-        // @ts-expect-error Looks like a parent.
         return preorder(child, index, node)
       })
     }

package.json

@@ -52,7 +52,7 @@
   },
   "scripts": {
     "prepack": "npm run build && npm run format",
-    "build": "rimraf \"*.d.ts\" && tsc && type-coverage",
+    "build": "rimraf \"{index,test}.d.ts\" && tsc && type-coverage",
     "format": "remark . -qfo && prettier . -w --loglevel warn && xo --fix",
     "test-api": "node test.js",
     "test-coverage": "c8 --check-coverage --branches 100 --functions 100 --lines 100 --statements 100 --reporter lcov node test.js",

test.js

@@ -1,5 +1,8 @@
 /**
- * @typedef {import('unist').Node} Node
+ * @typedef {{type: 'leaf', value: string}} Leaf
+ * @typedef {{type: 'node', children: Array<Node | Leaf>}} Node
+ * @typedef {{type: 'root', children: Array<Node | Leaf>}} Root
+ * @typedef {Root | Node | Leaf} AnyNode
  */
 
 import test from 'tape'
@@ -15,13 +18,14 @@ test('unist-util-map', function (t) {
 
   t.deepEqual(
     map(
-      {
+      /** @type {Root} */
+      ({
         type: 'root',
         children: [
           {type: 'node', children: [{type: 'leaf', value: '1'}]},
           {type: 'leaf', value: '2'}
         ]
-      },
+      }),
       changeLeaf
     ),
     u('root', [u('node', [u('leaf', 'CHANGED')]), u('leaf', 'CHANGED')]),
@@ -45,8 +49,8 @@ test('unist-util-map', function (t) {
   t.end()
 
   /**
-   * @param {Node} node
-   * @returns {Node}
+   * @param {AnyNode} node
+   * @returns {AnyNode}
    */
   function changeLeaf(node) {
     return node.type === 'leaf'
@@ -55,8 +59,8 @@ test('unist-util-map', function (t) {
   }
 
   /**
-   * @param {Node} node
-   * @returns {Node?}
+   * @param {AnyNode} node
+   * @returns {AnyNode?}
    */
   function nullLeaf(node) {
     return node.type === 'leaf' ? null : node