Over-engineered TypeScript Types - but I learned some stuff!

Hey team 🌟! This is another long rambling one, to try to document how my brain thought about this problem! My writing soundtrack was "Trust Me" by Sincere Engineer on repeat 🎶:

Hope you like it!

Background

At work we get a few days each month to work on whatever we like, to encourage learning/exploration/curiosity. Last month I decided I would try to solve an interesting JavaScript problem, but ended up uncovering a pretty interesting TypeScript problem 🤓...

My goal was to make the following work:

I wanted to create a function that works like require, but transforms the imported module to run in a separate thread. The implementation of the JavaScript part isn't the point of this post, but here's a simplified version of how I got it to work:

Comlink does most of the heavy lifting (thanks 🥰!) and the real inspiration for this post comes from a little note at the end of the Comlink docs:

While this type (Comlink.Remote<T>) has been battle-tested over some time now, it is implemented on a best-effort basis. There are some nuances that are incredibly hard if not impossible to encode correctly in TypeScript’s type system.

So what are those nuances and can TypeScript handle them? How should the types for workerRequire work?

Incredibly hard nuances

My original use case involves rendering a simple terminal-based UI using Ink. I want to run a task, and update the UI showing the task's progress. Unfortunately, running a CPU intensive task blocks UI from updating until the task is done, which makes the UI feel slow and clunky 😑. It would be much nicer to run the task in a Worker thread and update the UI in the main thread. This is like what you might do in browser-based app!

I want to go from something like this:

// BEFORE:
const { somethingCPUIntensive } = require('./cpu-intensive');
// The main thread is blocked
somethingCPUIntensive();

to something like this:

// AFTER:
const { somethingCPUIntensive } = workerRequire('./cpu-intensive');
// The main thread is no longer blocked
await somethingCPUIntensive();

I have a './cpu-intensive.js' file that exports a function:

// ./cpu-intensive.js
export function somethingCPUIntensive () {
    return fibonacci(50); // plain non-memoized fibonacci function which is pretty slow...
}

The types of this function make sense in isolation. When it is required with workerRequire the function is transformed from synchronous to asynchronous. That means that the types must also be:

// BEFORE:
const { somethingCPUIntensive } = require('./cpu-intensive');
// typeof somethingCPUIntensive is () => number

// AFTER:
const { somethingCPUIntensive } = workerRequire('./cpu-intensive');
// typeof somethingCPUIntensive is () => number

TypeScript can describe this transform succinctly:

type Func = (...args: Array<unknown>) => unknown;
type AsyncFunc<SyncFunc> = SyncFunc extends Func
    ? (...args: Parameters<SyncFunc>) => Promise<ReturnType<SyncFunc>>
    : never;

function fibonacci (n: number): number {
  if (n <= 1) return 1;
  return fibonacci(n - 1) + fibonacci(n - 2);
}

type SyncAdd = typeof fibonacci; // (n: number) => number
type AsyncAdd = AsyncFunc<typeof fibonacci>; // (n: number) => Promise<number>;

This is already quite deep in TypeScript stuff, so if you need a refresher check out Conditional Types, Parameters and ReturnType

Everything in this case works fine, because there is a clear distinction between which code runs in each thread. The fibonacci function runs synchronously in the worker thread, but the main thread receives the result asynchronously.

Unfortunately, the code doesn't actually solve the problem yet. The worker handles the intensive work, but it doesn't update the main thread as the task progresses. There needs to be some kind of callback:

// BEFORE:
const { somethingCPUIntensive } = require('./cpu-intensive');
// The main thread is blocked
somethingCPUIntensive((data) => updateProgress(data));

// AFTER:
const { somethingCPUIntensive } = workerRequire('./cpu-intensive');
await somethingCPUIntensive((data) => updateProgress(data));

In JavaScript land, this Just Works™️. Comlink handles passing arguments from the worker thread to the main thread, and then passing the result back to the worker thread. Seems kind of ✨ magic ✨ right? And it is! But from a types point of view this introduces even more weirdness 🙃...

Imagine that this is the original synchronous implementation of the somethingCPUIntensive function:

// ./cpu-intensive.js
export function somethingCPUIntensive(updateProgress) {
    const progress = doSomeWork();
    updateProgress(progress);
    const moreProgress = doSomeMoreWork();
    updateProgress(moreProgress);
    // ...
}

With a standard require this works as written, everything is synchronous and blocking:

const { somethingCPUIntensive } = require('./cpu-intensive');
somethingCPUIntensive((data) => updateProgress(data));

But with workerRequire everything is different:

const { somethingCPUIntensive } = workerRequire('./cpu-intensive');
await somethingCPUIntensive((data) => updateProgress(data));

The difference can be illustrated by adding types to somethingCPUIntensive:

// ./cpu-intensive.ts
export type Data = {};

export function somethingCPUIntensive(
    updateProgress: (data: Data) => void
) {
    const progress = doSomeWork();
    updateProgress(progress);
    const moreProgress = doSomeMoreWork();
    updateProgress(moreProgress);
    // ...
}

These types are entirely valid when used with require. But they are subtly invalid when used with workerRequire. Comlink will transform each call to updateProgress to being async!

Let me try to make it even more clear with another example. What if somethingCPUIntensive did something more complex, and the callback didn't return void?

// ./cpu-intensive.ts
export type Data = { ... };

export function somethingCPUIntensive(
    updateProgress: (data: Data) => Data
): Data {
    const progress = doSomeWork();
    const data = updateProgress(progress);
    const moreProgress = doSomeMoreWork(data);
    const moreData = updateProgress(moreProgress);
    // ...
}

Now the code is straight wrong! The code says that data and moreData to be of type Data, but when called via Comlink they would actually be Promise<Data> 🤯! The correct typing for this code would be this:

// ./cpu-intensive.ts
export type Data = { ... };

export function somethingCPUIntensive(
    updateProgress: (data: Data) => Promise<Data>
): Promise<Data> {
    const progress = doSomeWork();
    const data = await updateProgress(progress);
    const moreProgress = doSomeMoreWork(data);
    const moreData = await updateProgress(moreProgress);
    // ...
}

Where the callback function is originally declared, it appears to be synchronous. But where it is called it will be asynchronous! This is true if the function is passed from the main thread to the worker thread, or vice versa.

This is the aforementioned nuance and it is indeed quite mind-bending. All synchronous functions that are passed between threads will be transformed to asynchronous! This is true if the function is passed as a callback, or if it is a property on an object that is returned, or if it is an item in a Set Map or Array! The types on the receiving side must expect the function to be asynchronous. It's a pretty tricky mental model to work with.

The types for workerRequire function must encapsulate this confusion and ensure that required modules follow this model!

Typing workerRequire

The workerRequire API doesn't give much to work with:

workerRequire('./cpu-intensive');

The first problem is how to go from a string file path (or module name for a node_modules dependency) to the types of the resultant module:

Getting the Module type:

TypeScript has special handling for import statements, and also for require statements. But there is no way to tell the compiler that another function should inherit this behaviour.

The only way to go from a string path or module name to the type of that module is with the typeof operator combined with the import() syntax:

type Module = typeof import('./path/to/module');

A function that wraps require with correct types might look like this:

function myRequire<Result>(moduleId: string): Result {
    return require(moduleId);
}
const result = myRequire<typeof import('./path/to/module')>('./path/to/module');

Not exactly elegant, but it works 😅! The next step is to validate the Module type.

Validating the Module type:

I want TypeScript to validate that the module that lives in './cpu-intensive' handles asynchronicity correctly. It needs to search through every exported function in the module, and find all the parameters and return values that contain functions, and make sure that they return Promises... 🤔

To simplify things, I assume that the module only exports functions. So a valid module is an object that only has valid functions. And a valid function is a function where none of the arguments contain any synchronous functions, and the return value doesn't contain any synchronous functions. The actual exported function from the module doesn't have to be asynchronous!

That type might begin like this (don't worry if this is scary 😱 or if you don't know all the syntax, I'm going to break it down):

export type WorkerModule<Module> = {
  [Key in keyof Module]: Module[Key] extends Func
    ? WorkerModuleFunction<Module[Key]>
    : never;
};

export type WorkerModuleFunction = // ...

I like to think about about complex types like functions, so what would that look like in JavaScript?

export function validateWorkerModule (module) {
    Object.keys(module).forEach((key) => {
        if (!isValidWorkerFunction(module[key]) {
            throw new Error();
        }
    });
}

function isValidWorkerFunction (func) {
    func.parameters.forEach(parameter => {
        if(!isValidWorkerValue(parameter) {
            throw new Error();
        }
    })
    if (!isValidWorkerValue(func.returnType)) {
        throw new Error();
    }
}

So that kind of makes sense for the first level, but how does isValidWorkerValue work? If I keep going in JavaScript:

function isValidWorkerValue (value) {
    if (isFunction(value)) {
        return isValidWorkerFunction(value);
    }
    // ...
}

This is the first bit of recursion. A function can have a function as an argument, so TypeScript will have to go back up to confirm that the argument is a valid function. Which in turn could have another argument that is a function. But arguments can be other things as well:

function isValidWorkerValue (value) {
    if (isFunction(value)) {
        return isValidWorkerFunction(value);
    }
    if (isObject(value)) {
        return isValidWorkerObject(value);
    }
    if (isArray(value)) {
        return isValidWorkerArray(value);
    }
    if (isSet(value)) {
        return isValidWorkerSet(value);
    }
    if (isMap(value)) {
        return isValidWorkerMap(value);
    }
    if (isPromise(value)) {
        return isValidWorkerPromise(value);
    }
}

The more specific validation functions introduce another level of recursion. There can be deep Objects, or Arrays of Arrays, or Maps of Sets of Promises that return Functions 🙃:

function isValidWorkerObject (obj) {
    Object.keys(obj).forEach(key => {
        const value = obj[key];
        if (!isValidWorkerValue(value) {
            throw new Error();
        }
    });
}

function isValidWorkerArray (arr) {
    arr.forEach(value => {
        if (!isValidWorkerValue(value) {
            throw new Error();
        }
    });
}

// ...

If you've got this far, have a think about how Set, Map, and Promise might work? Let me know in the comments!

There is a nested logic here with multiple layers of recursion and weirdness going on 🤯... How does it translate into types?

There is to be TypeScript type syntax that is equal to the above pseudo-JS:

Typed if statements

I've already mentioned Conditional Types once in this post, and I'm going to keep using them here! A conditional type is a type with the form:

type A = B extends C ? D : E;

This is equivalent (and syntactically similar to a JavaScript ternary statement:

const a = b > c ? d : e;

Which is in turn equivalent to the more verbose JavaScript

let a;
if (b > c) {
    a = d;
} else {
    a = e;
}

That means I can write the isValidWorkerValue function as the following type:

export type WorkerValue<Value> = Value extends Obj
  ? Value extends WorkerObject<Value>
    ? Value
    : WorkerObject<Value>
  : Value extends Array<unknown>
  ? Value extends WorkerArray<Value>
    ? Value
    : WorkerArray<Value>
  : Value extends Func
  ? Value extends WorkerFunction<Value>
    ? Value
    : WorkerFunction<Value>
  : Value extends Set<unknown>
  ? Value extends WorkerSet<Value>
    ? Value
    : WorkerSet<Value>
  : Value extends Map<unknown, unknown>
  ? Value extends WorkerMap<Value>
    ? Value
    : WorkerMap<Value>
  : Value extends Promise<unknown>
  ? Value extends WorkerPromise<Value>
    ? Value
    : WorkerPromise<Value>
  : Value;

Which is totally ridiculous, but valid 🤓! But how do the more specific validation types work? There are two parts to this. First the validation needs to check that the given type argument is approximately the right type. I'll use WorkerPromise as an example:

export type WorkerPromise<P extends Promise<unknown>> = 
  P extends Promise<unknown> // Check if it's a Promise
    ? // This is a Promise, keep validating
    : // Not a Promise, bail out

If it isn't a Promise , it should resolve to never so that the WorkerValue type will skip to the next branch of the if/else blocks.

If it is a Promise then it moves onto the second part. The type needs to extract the inner type, which is the type of the resolved value. This can be done with the infer keyword, which is a special power of Conditional Types. Then the inner type can be validated as a WorkerValue - even more recursion!

export type WorkerPromise<P extends Promise<unknown>> = 
  P extends Promise<infer Value>
    ? Value extends WorkerValue<Value>
      ? P
      : SomethingMagic
    : never;

I'll come back to the SomethingMagic in a bit! 😉

Typed loops

The types also need to be able to iterate over each key/value pair in an Object type, or each item in an Array type. TypeScript provides a way to do that as well by using keyof:

export type WorkerObject<O> = {
  [Key in keyof O]: O[Key] extends WorkerValue<O[Key]>
    ? O[Key]
    : SomethingMagic
};

export type WorkerArray<A extends Array<unknown>> = {
  [Index in keyof A]: A[Index] extends WorkerValue<A[Index]>
    ? A[Index]
    : SomethingMagic
};

The keyof syntax is like an iterator, except it iterates over the indices of a type. With an Object-like type, it will give the key, which can then be used to look up the type of the key (with the O[Key] syntax). For an Array type, it will give the numeric indices, which can then be used with the same syntax (A[Index]).

Again, there is more recursion introduced here!

WorkerSet and WorkerMap is implemented similarly to WorkerPromise. Have a crack yourself maybe?

Typed ✨Errors✨

The one thing that is missing from the original JavaScript pseudocode is Errors. If you're guessing this is what the SomethingMagic is about, then you're right. But I'll explain why first.

The standard way to show that a type is invalid in TypeScript is the never type.

One example where never could be used would be a function that always throws:

function explode (): never {
    throw new Error('💥');
}

This function never returns anything, so the return type is never. This is also often used with conditional types:

type isPromise<P> = P extends Promise<unknown> ? P : never;
type Yes = isPromise<Promise<boolean>>; // Promise<Boolean>
type No = isPromise<boolean>; // never;

So it kind of makes sense to use never as a "this isn't a valid module" type too, as I hinted in the first type I introduced:

export type WorkerModule<Module> = {
  [Key in keyof Module]: Module[Key] extends Func
    ? WorkerModuleFunction<Module[Key]>
    : never;
};

Hopefully you find this type a little less scary this time?

The WorkerModule type is the type that validates the Module, and is used like this:

type Valid = WorkerModule<typeof import('./path/to/module')>;

If typeof import('./path/to/module') follows the rules, then that is great, the WorkerModule type acts as an identity and resolves to the same input. If the module doesn't follow the rules, then it resolves to never...

That is going to make it basically impossible to figure out why the module isn't valid.

Fortunately, a type doesn't have to resolve to never. It can resolve to an entirely different type, for example a string:

export type WorkerModule<Module> = {
  [Key in keyof Module]: Module[Key] extends Func
    ? WorkerModuleFunction<Module[Key]>
    : 'Module should only export functions'
};

Now using the type, it will explain what the specific issue is:

type MyModule = { foo: 'bar' }
type Valid = WorkerModule<MyModule> // 'Module should only export functions'

But it can be even better:

export type WorkerModuleError<Type, Message, Details = null> = {
  message: Message;
  type: Type;
  details: Details;
};

export type WorkerModule<Module> = {
  [Key in keyof Module]: Module[Key] extends Func
    ? WorkerModuleFunction<Module[Key]>
    : WorkerModuleError<
        Module,
        'Module should only export functions',
        Module[Key]
      >;
};

type MyModule = { foo: 'bar' }
type Valid = WorkerModule<MyModule>
// { 
//   message: 'Module should only export functions'
//   type: { foo: 'bar' }
//   details: 'bar'
// }

And that can continue even further into the validation, like with the WorkerPromise above:

export type WorkerPromise<P extends Promise<unknown>> = 
   P extends Promise<infer Value>
    ? Value extends WorkerValue<Value>
      ? P
      : WorkerModuleError<
          P,
          'Promise value should not contain synchronous functions',
          WorkerValue<Value>
        >
    : never;

Or a more complex example like WorkerFunction:

export type WorkerFunction<F extends Func> = 
  F extends (...args: infer Args) => infer Return
  ? Args extends WorkerArray<Args>
    ? Return extends Promise<infer Value>
      ? Value extends WorkerValue<Value>
        ? F
        : WorkerModuleError<
            F,
            'Function return value should not contain synchronous functions',
            WorkerValue<Value>
          >
      : WorkerModuleError<F, 'Function should return a Promise'>
    : WorkerModuleError<
        F,
        'Function arguments should not contain synchronous functions',
        WorkerArray<Args>
      >
  : never;

This makes it much easier to fix issues in the problematic Module!

Putting it all together:

So, it is possible to use TypeScript to check that a module will behave correctly with Comlink - or my workerRequire function! It can even have nice error messages - I will be using that trick more in the future for sure.

There's a few more little edge-cases that I haven't detailed here, but you can see the full implementation here and check out some of the Type tests here!

The same Type tests are running here in a CodeSandbox if you want to try them out:

Or even better, try it out! See if you can get something running in a separate thread. And then delete everything and use Comlink instead cause this is a Very Bad Idea™️.

Wrapping up:

That's definitely some weird stuff, but I had fun playing with it. And I solved the problem I needed to solve, and I learned a bunch 🎉.

Have I got bugs in my implementation? Yes! Have I encapsulated all the complexity the Comlink authors imagine? Almost definitely not! Have I learned some things? Heck yeah! Have you? Hit me up in the comments or on the bird site.

If you got this far, and you meet me at a conference/meetup in the post-COVID future, tell me, I owe you a drink 🥤!

Thank you!

Comments

[stefanovualto]

Man,

You went very far. I am not sure that I got 10% of what you were trying to achieve.

But kudos on some typing aspects that I didn't know!

👏