Tree Utilities

tree.ts

import { map, filter, reduce, assign } from 'lodash';

/**
 * Map over a tree from the bottom-most nodes up. Applies transformation to deepest children,
 * then to parents, and so on. You can even rename the property that contains the children.
 *
 * Usage:
 * tmap(
 *   x => ({ ...x, value: x.value + 1 || 0 }),
 *   'children',
 *   [{value: 0, children: []}, {children: [{value: 5, children: [{value: 100}]}]}]
 * )
 * // This will increment all values by 1 or initialize value to 0.
 */
export const tmap = <T extends { [key: string]: T[] | any }, R>(
  transformation: (t: T) => R,
  children: keyof T,
  nodes: T[],
): R[] => {
  return map(nodes, node => {
    if (node[children] !== undefined) {
      const newNode = assign({}, node, { [children]: tmap<T, R>(transformation, children, node[children]) });
      return transformation(newNode);
    }
    return transformation(node);
  });
};

/**
 * Filter over a tree from bottom-most nodes up. Applies predicate to deepest children,
 * then to parents, and so on. So rolling up empty children arrays is possible.
 */
export const tfilter = <T extends { [key: string]: T[] | any }>(
  predicate: (t: T) => boolean,
  children: keyof T,
  nodes: T[],
): T[] => {
  if (Array.isArray(nodes)) {
    return reduce(
      nodes,
      (acc, node) => {
        if (node[children] !== undefined) {
          const newNode = assign({}, node, { [children]: tfilter<T>(predicate, children, node[children]) });
          return predicate(newNode) ? [...acc, newNode] : acc;
        }
        return predicate(node) ? [...acc, node] : acc;
      },
      <T[]>[],
    );
  }
  return nodes;
};