Type Only HeapSort Implementation

Nathan Northcutt — Thu, 25 Jul 2024 09:19:44 +0000

Today we're going to be learning how to create a HeapSort implementation in TypeScript entirely through the type system. Why? Because it's a lot of fun to try to do these things and I needed it for my type only Sudoku Solver!

High Level Process

For our algorithm to work we are going to need a few components that can help us out.

A type to extract the target number property from our object array.
A type that can use that property to "heapify" the array.
A type that can sift up/down based on comparing the values.
A way to swap elements in a fixed type array.

There are a few other things we'll need for allowing TypeScript to understand actual numerical values including incrementing, decrementing, addition and comparisons. Since they don't exist natively in the type system I had to build them. For now, we'll just take advantage of the fact that is a solved problem.

Extracting numeric properties

This is actually a pretty easy problem to solve and you're probably doing it already in your code day to day.

/**
 * Defines keys that are valid for sorting
 */
type SortKeys<T extends object> = {
  [Key in keyof T]: T[Key] extends number ? Key : never;
}[Extract<keyof T, string>];

/**
 * An item to use in the heap sort
 */
type HeapSortItem<N extends number = number, T extends object = object> = {
  value: N;
  item: T;
};

/**
 * Map the given items into a HeapSortItem array
 */
type MapItems<
  T extends object,
  Prop extends SortKeys<T>,
  Values
> = Values extends [infer Next extends T, ...infer Rest]
  ? Next[Prop] extends number
    ? Rest extends never[]
      ? [HeapSortItem<Next[Prop], Next>]
      : [HeapSortItem<Next[Prop], Next>, ...MapItems<T, Prop, Rest>]
    : never
  : never;

/**
 * Extract the original items
 */
type ExtractItems<
  T extends object,
  Items extends HeapSortItem[]
> = Items extends [infer Next extends HeapSortItem<number, T>, ...infer Rest]
  ? Rest extends never[]
    ? [Next["item"]]
    : Rest extends HeapSortItem[]
    ? [Next["item"], ...ExtractItems<T, Rest>]
    : never
  : never;

Using those types, we can translate between the original objects, our HeapSortItem type and finally extract the originals back out. Nothing too fancy here, but it's one of our building blocks and I'll take an easy win.

Heapify

For the "heapify" process, we need to inspect each of the elements in the array and check to verify that no child has a value greater than it's parent. Procedurally, we just need to do a for loop in reverse over the items and check them against their parents. While there isn't a built in looping construct for our types, we can create one with a recursive type and an incrementing counter as part of our type parameter.

/**
 * Heapify the array one element at a time
 */
type Heapify<
  Arr extends HeapSortItem[],
  Limit extends number,
  N extends number = 0
> = N extends Limit
  ? Arr // Done
  : SiftUp<Arr, N> extends infer _ extends HeapSortItem[]
  ? Heapify<_, Limit, Increment<N>>
  : never;

There are a couple of things that are worth calling out here. First, we're leveraging a type Increment that comes from our math package in the other link. It simply adds 1 to the number to generate a new type. When N becomes the same size as the Limit value, the process simply returns the Arr as the final type.

A key step in making this work without the compiler generating a HUGE number of types is where we take the result of the SiftUp type and force it to be evaluated extends infer _ extends HeapSortItem[]. If we simply did Heapify<SiftUp<Arr, N>, Limit, Increment<N>> the compiler can defer the type expansion until the end, dangling huge numbers of types which slows everything down.

Finding Parents and Children

For the sifting to work, we need to be able to define parent and child relationships.

/**
 * "Mulitiply" the number by 2
 */
type TwoN<Index extends number> = Add<Index, Index>;

/**
 * "Divide" the number by 2
 */
type DivBy2<
  Index extends number,
  C extends number = 0
> = TwoN<C> extends infer N extends number
  ? N extends Index
    ? C
    : GT<Index, N> extends true
    ? DivBy2<Index, Increment<C>>
    : Decrement<C>
  : never;

/**
 * Get the left child of the index
 */
type LeftChild<Index extends number> =
  TwoN<Index> extends infer N extends number
    ? Add<N, 1> extends infer TwoNPlus1 extends number
      ? TwoNPlus1
      : never
    : never;

/**
 * Get the right child of the index
 */
type RightChild<Index extends number> =
  TwoN<Index> extends infer N extends number
    ? Add<N, 2> extends infer TwoNPlus2 extends number
      ? TwoNPlus2
      : never
    : never;

/**
 * Get the parent of the given index
 */
type Parent<Index extends number> =
  Decrement<Index> extends infer N extends number
    ? DivBy2<N> extends infer _ extends number
      ? _
      : never
    : never;

Again, leveraging some of our math primitives for Add we can generate the 2n+1, 2n+2 and (n-1)/2 values that we need to explore up and down the heap when doing our sifting below.

Sifting Up/Down

Now that we can move around in our heap from parent to child and child to parent, we can implement our sifting types.

/**
 * Ensure the child moves up the heap if it's larger than it's parent
 */
type SiftUp<Arr extends HeapSortItem[], Idx extends number> = Idx extends 0
  ? Arr
  : Parent<Idx> extends infer _ extends number
  ? GT<Arr[Idx]["value"], Arr[_]["value"]> extends true
    ? Swap<Arr, _, Idx> extends infer Sifted extends HeapSortItem[]
      ? SiftUp<Sifted, _>
      : never
    : Arr
  : never;

/**
 * Move values down until the heap is back in a valid state
 */
type SiftDown<
  Arr extends HeapSortItem[],
  Idx extends number,
  Limit extends number
> = LeftChild<Idx> extends infer LC extends number
  ? LTE<LC, Limit> extends true
    ? RightChild<Idx> extends infer RC extends number
      ? LTE<RC, Limit> extends true
        ? GT<Arr[LC]["value"], Arr[RC]["value"]> extends true // Have to check GT of left or right
          ? LT<Arr[Idx]["value"], Arr[LC]["value"]> extends true
            ? Swap<Arr, Idx, LC> extends infer _ extends HeapSortItem[]
              ? SiftDown<_, LC, Limit>
              : never
            : Arr
          : LT<Arr[Idx]["value"], Arr[RC]["value"]> extends true
          ? Swap<Arr, Idx, RC> extends infer _ extends HeapSortItem[]
            ? SiftDown<_, RC, Limit>
            : never
          : Arr
        : LT<Arr[Idx]["value"], Arr[LC]["value"]> extends true
        ? Swap<Arr, Idx, LC> extends infer _ extends HeapSortItem[]
          ? SiftDown<_, LC, Limit>
          : never
        : Arr // Already sorted
      : never
    : Arr // No further children
  : never; // left child is always a number

While the SiftDown may look more complex than the SiftUp type, they are essentially doing the same thing, just with more conditional edge cases. We simply check the current value against it's parent/child and if that relationship is in violation of our heap properties (largest item always highest in the heap), we swap them and continue processing from that new location. There are a few more cases in the SiftDown since we want to swap with the largest of the two children and a node may have 0, 1 or 2 children which require additional conditional checks, at the end it's always a Swap and then recursive call.

Swapping Objects

The last component we need is a way to swap two elements in an array. Since the type system is immutable, we actually have to generate a new type with those values changed. We can easily handle that with our recursive exploration of array types and two utility methods.

/**
 * Swap two elements in the array
 */
type Swap<
  Arr extends HeapSortItem[],
  Left extends number,
  Right extends number
> = Replace<Replace<Arr, Arr[Right], Left>, Arr[Left], Right>;

/**
 * Replace the value at the given location
 */
type Replace<
  Arr,
  Item extends HeapSortItem,
  Idx extends number,
  N extends number = 0
> = Arr extends [infer Next extends HeapSortItem, ...infer Rest]
  ? N extends Idx
    ? [Item, ...Rest]
    : Rest extends never[]
    ? [Next]
    : [Next, ...Replace<Rest, Item, Idx, Increment<N>>]
  : never;

Putting It All Together

Now that we have all the pieces in place, we can implement our HeapSort that we expose to the world.

/**
 * Perform the in place heap sort
 */
type HeapSort<
  T extends object,
  Items extends HeapSortItem[],
  N extends number = Decrement<Items["length"]>
> = N extends 0
  ? ExtractItems<T, Items>
  : Swap<Items, 0, N> extends infer Partial extends HeapSortItem[]
  ? SiftDown<
      Partial,
      0,
      Decrement<N>
    > extends infer Sifted extends HeapSortItem[]
    ? HeapSort<T, Sifted, Decrement<N>>
    : never
  : never;

/**
 * Perform a heap sort of the items using the property key defined
 */
export type Sort<
  T extends object,
  Values extends T[],
  Prop extends SortKeys<T>
> = MapItems<T, Prop, Values> extends infer Items extends HeapSortItem[]
  ? Heapify<
      Items,
      Items["length"]
    > extends infer HeapItems extends HeapSortItem[]
    ? HeapSort<T, HeapItems>
    : never
  : never;

The main Sort type is there to ensure that we call Heapify on the resulting mapped HeapSortItem[] we generate in MapItems while the HeapSort type is responsible for running the Swap + SiftDown iterations until we only have 1 item left in the heap. At that point we can return the extracted original items, giving us a fully working HeapSort!

I hope you enjoyed that and you can also check out some of the other fun I've been having with LeetCode and TypeScript or writing a SQL Parser from scratch (including raw string -> SQL types).

Forem: Nathan Northcutt