aggremagic.ts

interface InitialGroup<K, T> {
  key: K;
  keyString: string;
  items: T[];
}

export interface HierarchicalGroup<K, T, A> extends InitialGroup<K, T> {
  id: string;
  level: number;
  parent?: HierarchicalGroup<K, T, A>;
  children: AggregatedGroup<K, T, A>[];
}

export interface AggregatedGroup<K, T, A> extends HierarchicalGroup<K, T, A> {
  aggregate: A;
}

export type Grouper<K, T> = [(t: T) => K, (k: K) => string];
export type Aggregator<K, T, A> = (group: HierarchicalGroup<K, T, A>) => A;

const STRING_KEY_SEPARATOR = '\x01';

export function defaultKeyStringifier<T extends Array<any>>(t: T): string {
  return t.join(STRING_KEY_SEPARATOR);
}

function buildSimpleGroups<K, T>(input: readonly T[], grouper: Grouper<K, T>): InitialGroup<K, T>[] {
  const initialGroupMap: Record<string, InitialGroup<K, T>> = {};
  const [getKey, getKeyString] = grouper;
  input.forEach((item) => {
    const key = getKey(item);
    const keyString = getKeyString(key);
    if (!initialGroupMap[keyString]) {
      initialGroupMap[keyString] = {
        key,
        keyString,
        items: [],
      };
    }
    initialGroupMap[keyString].items.push(item);
  });
  return Object.values(initialGroupMap);
}

function buildHierarchy<K, T, A>(
  initialGroups: InitialGroup<K, T>[],
  parent: HierarchicalGroup<K, T, A> | undefined,
  nextGroupers: Grouper<K, T>[],
  aggregator: Aggregator<K, T, A>,
  makeChildren: typeof internalAggregate,
) {
  return initialGroups.map((ig, index) => {
    const hg: HierarchicalGroup<K, T, A> = {
      ...ig,
      id: [parent ? parent.id : '', `${index}`].filter(Boolean).join(STRING_KEY_SEPARATOR),
      level: parent ? parent.level + 1 : 0,
      parent,
      children: [],
    };
    if (nextGroupers.length) {
      hg.children = makeChildren(ig.items, nextGroupers, aggregator, hg, makeChildren);
    }
    return hg;
  });
}

function aggregateHierarchicalGroups<K, T, A>(
  hierarchicalGroups: HierarchicalGroup<K, T, A>[],
  aggregator: Aggregator<K, T, A>,
) {
  return hierarchicalGroups.map((hg) => ({
    ...hg,
    aggregate: aggregator(hg),
  }));
}

function internalAggregate<K, T, A>(
  input: readonly T[],
  groupers: readonly Grouper<K, T>[],
  aggregator: Aggregator<K, T, A>,
  parent: HierarchicalGroup<K, T, A> | undefined,
  makeChildren: typeof internalAggregate,
): AggregatedGroup<K, T, A>[] {
  if (!groupers.length) {
    return [];
  }
  const initialGroups = buildSimpleGroups(input, groupers[0]);
  const hierarchicalGroups = buildHierarchy(initialGroups, parent, groupers.slice(1), aggregator, makeChildren);
  return aggregateHierarchicalGroups(hierarchicalGroups, aggregator);
}

export function aggregate<K, T, A>(
  input: readonly T[],
  groupers: readonly Grouper<K, T>[],
  aggregator: Aggregator<K, T, A>,
): AggregatedGroup<K, T, A>[] {
  return internalAggregate(input, groupers, aggregator, undefined, internalAggregate);
}