Type checked async function graph execution in TypeScript

At a previous job of mine, a core part of our system involved tying together a bunch of HTTP calls and other processing in a particular way, where calls could depend on both up-front inputs and the results of previous calls. Some of the processing also yielded intermediate values that needed to be passed up and yielded out to the rest of the system before the final return value (i.e., there are async generators involved, not just plain functions). This was all written out in a hardcoded way with a bunch of bespoke code, with some vague thoughts of generalizing it at some point in the future. Though we never got around to working on that while I was there, I have since put together something in that vein that is a pretty neat little demonstration of type-level computation in TypeScript.1

(While working on this post, I came across the very recent https://blog.howardjohn.info/posts/bash-dag/, which describes how to do essentially the same thing in Bash! Though no types there, of course.)

The goal

Given some async generators and a description of how they fit together, we want to run them such that some of their inputs come from the outputs of others as described, with one of the generators coming at the very end and providing the output of the whole thing.

As a simple example, we might start with the following.

typescript
type Func<T> = AsyncGenerator<string, T>;

async function* base(): Func<{ greetingIndex: number; name: string }> {
  yield "returning base values";
  return { greetingIndex: 1, name: "world" };
}

async function* selectGreeting(args: { greetingIndex: number }): Func<string> {
  yield "selecting greeting";
  return ["Howdy", "Hello"][args.greetingIndex];
}

async function* modifyName(args: { name: string }): Func<string> {
  yield "modifying name";
  return `${name}!`;
}

async function* out(greeting: string, name: string): Func<string> {
  return `${greeting}, ${name}`;
}

const structure = {
  base: [base, []],
  selectGreeting: [selectGreeting, ["base"]],
  modifyName: [modifyName, ["base"]],
  out: [out, ["selectGreeting", "modifyName"]],
} as const;

We have some individual functions that perform parts of the whole computation, with the structure object at the bottom describing how values flow between the functions. Each function to be executed corresponds to one key-value pair in the structure; the key provides a name for the function and the value contains the function object itself and the list of dependencies specified by the corresponding keys in the record.2 The return values of the dependencies of a function are passed in as the arguments in order when that function is run.

The base function creates assembles a record containing some fixed values, which are passed to and modified by selectGreeting and modifyName. The resulting values are both passed to out, which computes one final value from them ("Hello, world!", naturally). Meanwhile, each function may also yield values before returning, which also need to be passed along.

Given all of that, we want to have some other code that can examine structure and determine at compile time whether it is consistent with the definitions of the individual functions, as well as run the functions to get the final value out.

The idea

TypeScript has very precise type inference for compile-time constants, making it possible to have quite a complex value translated directly into a type. When that’s combined with TypeScript’s powerful ways of manipulating types, it turns out to be possible to describe the functions and their dependencies and have the compiler check validity.

Breaking it down, we want to check the following properties for each function and its dependencies:

Ideally, we would also do a global check that the graph can actually be completed (i.e., it is acyclic and there is a path to get from inputs to the final function) and doesn’t contain any unnecessary functions, but I don’t know how to do that at the moment.

TypeScript type system background

I’m assuming that the reader is broadly familiar with JavaScript semantics and the basics of the TypeScript type system, but we’re going to need some slightly advanced stuff, which I’ll go over here.

Literal types and as const

Main article: https://www.typescriptlang.org/docs/handbook/2/everyday-types.html

When you see, for example, "foo" in TypeScript code, it is often a value: an actual thing that exists in JavaScript when the code is running, has methods, takes up memory, and so on. However, in the relevant contexts, it can instead be a type: a set of values (in this case, a set containing only the single value "foo") that the TypeScript compiler reasons about while compiling, but that doesn’t exist or have any effect at runtime.

By default, the inferred type for a literal is not the narrowest possible type. For example, writing let x = "foo"; is equivalent to let x: string = "foo";. The narrowest possible type would, of course, be the literal type "foo". TypeScript doesn’t infer that since variables may be reassigned later on3 and that would be too restrictive; the wider inferred type is a reasonable guess for what set of values might be acceptable.

If you do want the variable to take on the narrowest possible type, one way to tell that to the compiler is to explicitly write let x: "foo" = "foo";. To avoid the repetition, you can instead write let x = "foo" as const;. Appending as const also works for more complex objects; TypeScript will generate a type describing the specific value almost as closely as possible. For example, if you write let x = [() => 17, ["a", "b"]] as const;, then the type of x is [() => number, ["a", "b"]].4

Indexing into types

Main article: https://www.typescriptlang.org/docs/handbook/2/indexed-access-types.html

Just like object values, object types can be indexed into. For example, if we have type T = { a: number, b: string }, then T['a'] evaluates to number. (The dot notation T.a can’t be used on types.) Something that can be done with type indexing but not value indexing is indexing with a union type to get out a union of multiple items at a time: T['a' | 'b'] evaluates to string | number.

When T is an object type, keyof T is the type of all valid keys for T, so another equivalent way of writing the type above would be T[keyof T]. In general, T[keyof T] evaluates to the union of all the values (in the sense of key-value pair, rather than types versus values) in T.

Conditional types

Main article: https://www.typescriptlang.org/docs/handbook/2/conditional-types.html

A type of the form A extends B ? T : F evaluates to either T or F depending on whether the type A is assignable to the type B (i.e., whether A is a subtype of B, to use non-TypeScript-specific terminology5).

Unlike the similar-looking operator on values, the conditions can’t be conjoined with &&, so to check multiple conditions at once, you need to either nest those conditionals (A extends B ? C extends D ? true : false : false) or wrap the left sides and right sides into tuples and check the subtyping relationships all at once ([A, C] extends [B, D] ? true : false). The first is rather verbose, while the second moves the pairs of types being checked against each other far apart. I also didn’t think of the second one until after getting everything working, so I have the first one in the code as written.

Mapped record types

Main article: https://www.typescriptlang.org/docs/handbook/2/mapped-types.html

Something we’re going to need more than once is the ability to take an object type and transform its key-value pairs in some way. This example comes from the TypeScript documentation linked above:

typescript
type OptionsFlags<Type> = {
  [Property in keyof Type]: boolean;
};
type Record = { a: boolean; b: number; c: string };
type Options = OptionsFlags<Record>; // this is { a: boolean, b: boolean, c: boolean }

OptionsFlags<Type> is evaluated like a list comprehension: for each type within in the union keyof Type, the part after the colon is evaluated with Property bound to that type and a record type is constructed from everything that comes out. Of course, the part after the colon needn’t be a constant type:

typescript
type Arrayify<T> = {
  [I in keyof T]: T[I][];
};

type Record = { a: boolean; b: number; c: string };
type Arrayified = Arrayify<Record>; // this is { a: boolean[]; b: number[]; c: string[] }

Mapped tuple types

JavaScript doesn’t make a distinction between tuples and arrays/lists as some other languages do, but a fixed-size array of types is called a tuple type in TypeScript. Tuple types can also be mapped like record types, akin to how arrays in JavaScript are kind of like objects with keys length, 0, 1, 2, etc.

However, for reasons I don’t understand, there needs to be an intermediate helper type to get the right behavior. Applying the map directly to a concrete tuple type does something unhelpful; instead, we need to define a generic type and pass the input tuple type to that.

typescript
type Tuple = [boolean, number, string];
type Indirect = Arrayify<Tuple>; // this is [boolean[], number[], string[]] as desired
type Direct = { [I in keyof Tuple]: Tuple[I][] }; // this is... something else entirely

(Arrayify is defined as above.) Even though the definition of Direct is exactly what you would get by expanding out the definition of Indirect, it ends up evaluating to something quite different (the arrayification is applied to every single property of the underlying list type: length, toString(), push(), etc.). But when the two techniques are applied to a record type, they do the same thing. I don’t know what’s up with that.

typescript
type Record = { a: boolean; b: number; c: string };
type Indirect = Arrayify<Record>; // this is { a: boolean[]; b: number[]; c: string[] }
type Direct = { [I in keyof Record]: Record[I][] }; // this is the same as above

Type checking

The main check is implemented by the code below. For an object structure defined as in the code above, Check<typeof structure> evaluates to true or false depending on whether structure describes a valid function graph.

typescript
type StructureBase = Record<
  string,
  readonly [(...args: any[]) => AsyncGenerator<any, any>, readonly [...string[]]]
>;

type Writeable<T> = { -readonly [P in keyof T]: T[P] };

type AsyncGenReturnType<T extends (..._: never) => AsyncGenerator<unknown, unknown>> = T extends (
  ..._: never
) => AsyncGenerator<infer _, infer R> ? R
  : never;

type GetReturns<S extends StructureBase, K extends readonly [...string[]]> = {
  [ind in keyof K]: AsyncGenReturnType<S[K[ind]][0]>;
};

type Check<S extends StructureBase> = {
  [name in keyof S]:
    // Check that all names in the list are also in the argument object.
    S[name][1][number] extends keyof S
      // Check that parameter types are compatible with corresponding return types.
      ? Writeable<GetReturns<S, S[name][1]>> extends Parameters<S[name][0]> ? true : false
      : false;
}[keyof S] extends true ? true : false;

That may look a bit scary, but all it’s doing is combining the principles from before to check the conditions for each key-value pair in the structure, with some helpers to set things up.

StructureBase defines what a structure needs to look like to be potentially worth checking at all. It’s a record type that maps names to tuples containing functions and the list of their dependency names.

Writeable papers over a little hiccup in the types but doesn’t really change anything interesting. When an object is declared with as const, the inferred type has readonly attributes all over the place, indicating that the value can’t be modified. We want to compare that to something without readonly, and Writeable strips out the readonly to make that work.

AsyncGenReturnType uses the infer feature of conditional types, which pattern-matches on the arguments to a generic type. The syntax is a bit much, but just know that it extracts the return type of an async generator, unsurprisingly.

GetReturns is a mapped tuple type that applies AsyncGenReturnType to pull out the return types of all the functions in the given structure with the given names.

Finally, Check itself incorporates a mapped type that checks the necessary conditions for each function described in the structure. In the first extends check, S[name] evaluates to one of the value types in the structure, so S[name][1] evaluates to the type for the list of names, so S[name][1][number] evaluates to the union type of all the names. This needs to extend keyof S, i.e., every element of the list needs to also be a valid key in the structure. The second extends check uses the helpers to gather up all of the return types for the dependencies of one of the functions and then makes sure

Then the outermost [keyof S] extends true ? true : false combines the results of each item within the structure. The [keyof S] takes the union of all the values in the record; if every item evaluates to true, then that equals true, which extends true. But if there’s any appearance of false, then the union becomes true | false (or just false if nothing at all is valid), which does not extend true.

Using the check

With the ability to do the validity check out of the way, let’s see how we apply it:

typescript
async function* execute<
  S extends StructureBase & {
    out: readonly [(args: any) => AsyncGenerator<string, any>, readonly string[]];
  },
>(struct: Check<S> extends true ? S : never): AsyncGenerator<string, AsyncGenReturnType<S["out"][0]>> {
  // body to be examined later
}

The constraint on S consists of the general structure shape constraint, plus the requirement that there exists an "out" element in the record, which we take as the output of the overall execution.

That’s almost all there is to it, but there’s a trick in the parameter list that I came up with but am surprised kind of surprised by. Initially, I had a version that took two parameters, one of type S and one of type Check<S> extends true ? true : never, the idea being that type inference would determine S first based on the argument and then evaluate Check<S>; if the result was false, then the second parameter’s type would be never and so no value could safely be passed in. That did work, but the second parameter was annoying. Then I tried the current version on a whim, not really expecting much, but it worked! If the structure is valid, then Check<S> is true and type checking passes. If not, there is an error that the argument can’t be assigned to never.

This breaks my intuition for how the type checker seems like it ought to work: somehow it looks inside the type parameter to Check and matches that up with the concrete type of the argument, even though there’s no a priori reason that those should be the same. It seems like it’s doing some kind of reverse inference where it looks at the outputs of the conditional and determines that Check<S> must be true, then uses the fact that it now knows what S must be to turn around and evaluate Check<S>. Either that or there’s some heuristic where it just guesses that S should be the same as the argument type since there isn’t anything else around, and then it plows ahead and does the check from there.6

An alternative: records for arguments

Instead of each function taking values from its dependencies in order, it’s also possible for each function to take a single argument that is an object containing the return values of its dependencies keyed by name, and in fact that’s what I came up with first (before learning about mapped tuple types). That has the benefit that everything is keyed by name and so ordering doesn’t matter—indeed, I already almost screwed up the order of the two arguments to out just in the minimal example up top. But it adds a lot more noise to both the function definitions and the checking type, so I’m sticking to the order-based one here. If anyone were to ever use this for real, going with objects could be a reasonable tradeoff for robustness.

Finally running some code

…well, actually, this post is getting long enough that I’m going to leave the body of the execute function for a followup post. That’s worth going through for completeness, but the really interesting part, in my mind, is the type-level stuff.


Footnotes

1

I’m not entirely sure at this remove how well what I describe here would’ve been able to apply in that situation at work, and frankly this might be too clever for me to want to use it for real even if it could. But it’s cool! To me, at least.

2

The names in the structure don’t need to match the names of the functions themselves; in fact, the functions don’t need to be named at all, which will be handy later. But if the functions are named, having the names match the structure seems wise.

3

Using const x rather than let x achieves the same effect in some of the very simplest cases, but it does much less overall and is not enough for our purposes.

4

Actually, there are also some readonly modifiers in the full type, reflecting the fact that the object can’t be modified. We’re handwaving that away here (see Writeable later).

Also note that this is only almost as precisely as possible, since we could have () => 17 instead.

5

Technically, there are minor differences between subtype compatibility and assignment compatibility, but don’t worry about that.

6

I was initially leaning toward the second possibility, but now that I’ve written that out and thought about it, the first one seems a lot more plausible than before. I suppose I’ll have to go and actually look at the TypeScript compiler code at some point.

Which reminds me: at the time of writing this post, TypeScript 7.0 (the Go rewrite of the TypeScript compiler) is supposed to be just about to be declared prod-ready, so keep an eye on that if you’re using TypeScript!