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

phenomnominal

Craig β˜ οΈπŸ’€πŸ‘»

Posted on January 26, 2021

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();
Enter fullscreen mode Exit fullscreen mode

to something like this:

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

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...
}
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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>;
Enter fullscreen mode Exit fullscreen mode

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));
Enter fullscreen mode Exit fullscreen mode

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);
    // ...
}
Enter fullscreen mode Exit fullscreen mode

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

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

But with workerRequire everything is different:

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

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);
    // ...
}
Enter fullscreen mode Exit fullscreen mode

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);
    // ...
}
Enter fullscreen mode Exit fullscreen mode

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);
    // ...
}
Enter fullscreen mode Exit fullscreen mode

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');
Enter fullscreen mode Exit fullscreen mode

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');
Enter fullscreen mode Exit fullscreen mode

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');
Enter fullscreen mode Exit fullscreen mode

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 = // ...
Enter fullscreen mode Exit fullscreen mode

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();
    }
}
Enter fullscreen mode Exit fullscreen mode

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);
    }
    // ...
}
Enter fullscreen mode Exit fullscreen mode

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);
    }
}
Enter fullscreen mode Exit fullscreen mode

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();
        }
    });
}

// ...
Enter fullscreen mode Exit fullscreen mode

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;
Enter fullscreen mode Exit fullscreen mode

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

const a = b > c ? d : e;
Enter fullscreen mode Exit fullscreen mode

Which is in turn equivalent to the more verbose JavaScript

let a;
if (b > c) {
    a = d;
} else {
    a = e;
}
Enter fullscreen mode Exit fullscreen mode

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;
Enter fullscreen mode Exit fullscreen mode

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
Enter fullscreen mode Exit fullscreen mode

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;
Enter fullscreen mode Exit fullscreen mode

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
};
Enter fullscreen mode Exit fullscreen mode

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('πŸ’₯');
}
Enter fullscreen mode Exit fullscreen mode

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;
Enter fullscreen mode Exit fullscreen mode

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;
};
Enter fullscreen mode Exit fullscreen mode

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')>;
Enter fullscreen mode Exit fullscreen mode

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'
};
Enter fullscreen mode Exit fullscreen mode

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'
Enter fullscreen mode Exit fullscreen mode

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'
// }
Enter fullscreen mode Exit fullscreen mode

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;
Enter fullscreen mode Exit fullscreen mode

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;
Enter fullscreen mode Exit fullscreen mode

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!

πŸ’– πŸ’ͺ πŸ™… 🚩
phenomnominal
Craig β˜ οΈπŸ’€πŸ‘»

Posted on January 26, 2021

Join Our Newsletter. No Spam, Only the good stuff.

Sign up to receive the latest update from our blog.

Related