Strongly-typed array partitioning.md

/**
 * Split an array into buckets based on the predicates. The size of the returned
 * array of buckets is equal to the number of predicates + 1. The last bucket
 * contains all the items that did not match any predicate.
 * 
 * The first matching predicate wins.
 *
 * Overload 1: Strongly typed predicates for type-narrowing
 */
export function partitionBy<T, R extends T[]>(
	arr: T[],
	...predicates: { [K in keyof R]: (item: T) => item is R[K] }
): [...{ [K in keyof R]: R[K][] }, Exclude<T, R[number]>[]]
/**
 * Split an array into buckets based on the predicates. The size of the returned
 * array of buckets is equal to the number of predicates + 1. The last bucket
 * contains all the items that did not match any predicate.
 * 
 * The first matching predicate wins.
 *
 * Overload 2: When type-narrowing cannot be inferred, use the input type
 */
export function partitionBy<T>(
	arr: T[],
	...predicates: Array<(item: T) => boolean>
): T[][]
/** Implementation */
export function partitionBy<T>(
	arr: T[],
	...predicates: Array<(item: T) => boolean>
): T[][] {
	const output: T[][] = Array.from(
		{ length: predicates.length + 1 },
		() => []
	)

	for (const [index, item] of arr.entries()) {
		const index = predicates.findIndex((predicate) => predicate(item))
		output.at(index)!.push(item)
	}

	return output
}

Usage

interface SystemParam {
    role: "system"
}
interface UserParam {
    role: "user"
}
type ChatParam = SystemParam | UserParam

let messages: ChatParam[] = []

// Overload 1: `system` and `user` arrays are correctly typed
let [system, user] = partitionBy(
    messages,
    (item) => item.role === "system"
)

// Overload 2: When no type-narrowing is determined
let [odd, even] = partitionBy(
	[1, 2, 3, 4, 5, 6, 7, 8, 9, 0],
	(item) => item % 2 === 0
)