Helper functions for deeply nested data structures. Can be used with GraphQL queries that have to return arbitrarily deep trees. "flatten" is called in the backend resolver to traverse the tree into an array, "unflatten" is called in the frontend resolver to bring it to original shape. The shape has a key "id" and an arbitrary key for nested chi…

deeply_nested_flatten_unflatten.ts

type IFlattenable<K extends string> = {
  [k in K]?: IFlattenable<K>[];
} & { id: number };

type IFlattened<T extends IFlattenable<K>, K extends string> = Omit<
  T,
  K | 'id'
> & {
  parentId: number;
  id: number;
};

const _flatten = <T extends IFlattenable<K>, K extends string>(
  element: T,
  childrenKey: K,
  parentId?: number,
): IFlattened<T, K>[] => {
  const { [childrenKey]: children, id, ...rest } = element;
  if (!children || children.length === 0) return [{ id, parentId, ...rest }];
  const flatChildren: IFlattened<T, K>[] = [];
  children.forEach((c) =>
    flatChildren.push(..._flatten<T, K>(c as T, childrenKey, id)),
  );
  return [{ id, parentId, ...rest }, ...flatChildren];
};

const flatten = <T extends IFlattenable<K>, K extends string>(
  elements: T[],
  childrenKey: K,
): IFlattened<T, K>[] => {
  const flatElements: IFlattened<T, K>[] = [];
  elements.forEach((c) => flatElements.push(..._flatten<T, K>(c, childrenKey)));
  return flatElements;
};

const _unflatten = <T extends IFlattenable<K>, K extends string>(
  parent: IFlattened<T, K>,
  flatElements: IFlattened<T, K>[],
  childrenKey: K,
): IFlattenable<K> => {
  const flatChildren = flatElements.filter((c) => c.parentId === parent.id);
  const { parentId, id, ...rest } = parent;
  if (flatChildren.length === 0) return { id, ...rest };
  const children = flatChildren.map((c) =>
    _unflatten<T, K>(c, flatElements, childrenKey),
  );
  return { ...rest, id, [childrenKey]: children };
};

const unflatten = <T extends IFlattenable<K>, K extends string>(
  flatElements: IFlattened<T, K>[],
  childrenKey: K,
): IFlattenable<K>[] => {
  const parents = flatElements.filter((c) => c.parentId === undefined);
  return parents.map((p) => _unflatten<T, K>(p, flatElements, childrenKey));
};

/** Trying it on a IComment deeply nested structure, similar to Reddit replies */

interface IComment {
  id: number;
  text: string;
  replies?: IComment[];
}

const comment1: IComment = {
  id: 1,
  text: 'first comment',
  replies: [
    {
      id: 2,
      text: 'first reply',
      replies: [
        {
          id: 3,
          text: 'reply back',
          replies: [{ id: 7, text: 'deeply nested reply' }],
        },
      ],
    },
    { id: 4, text: 'bybye' },
    { id: 5, text: 'hey' },
  ],
};

const comment2: IComment = {
  id: 9,
  text: 'dog',
  replies: [
    {
      id: 19,
      text: 'cat',
      replies: [{ id: 10, text: 'tiger', replies: [{ id: 7, text: 'lion' }] }],
    },
    { id: 17, text: 'eagle' },
    { id: 18, text: 'moose' },
  ],
};

const flattened = flatten([comment1, comment2], 'replies');
const unflattened = unflatten(flattened, 'replies');
console.log(unflattened);