Typed function to create output based on input with refinement and controlled return types

I have a bunch of different inputs:

interface StringInput {
  _tag: "_string";
}

interface NumberInput {
  _tag: "_number";
}

interface ListInput {
  _tag: "_list";
}

type Input = StringInput | NumberInput | ListInput;

Along with the corresponding outputs:

class StringOutput {
  _tag = "StringOutput" as const;
  constructor(_: StringInput) {}
}

class NumberOutput {
  _tag = "NumberOutput" as const;
  constructor(_: NumberInput) {}
}

class ListOutput {
  _tag = "ListOutput" as const;
  constructor(_: ListInput) {}
}

type Output = StringOutput | NumberOutput | ListOutput;

I'd like to have a function that creates the right output given an input. The function should be well-typed especially in its implementation.

 

I'm stuck with two attemps, each of which solves half of the problem.

Attempt 1: switch. It is a well-known improvement possibility for TypeScript. The input parameter i get refined, which is something I need in order to call constructor of the corresponding output object. On the other hand the generic I does not get refined, causing the problem with the return type.

namespace Attempt1 {

  type RemoveUn<S> = S extends `_${infer R}` ? R : S

  type InputOutputMap = {
    [K in Input["_tag"]]: Extract<Output,  { _tag: `${Capitalize<RemoveUn<K>>}Output` }>
  }

  function createValueOfObjKey<
    const I extends Input
  >(i: I): InputOutputMap[I["_tag"]] {
    switch(i._tag) {
      case "_string": {
        return new StringOutput(i); // error
      }
      case "_number": {
        return new NumberOutput(i); // error
      }
      case "_list": {
        return new ListOutput(i); // error
      }
    }
  }

}

Attempt 2: indexed access types. The compiler does quite well controlling them, so what I return must be in line with what I have stated in the signature. But I lost the refinement on i and I'm not able to restore it.

namespace Attempt2 {

  type RemoveUn<S> = S extends `_${infer R}` ? R : S

  type InputOutputMap = {
    [K in Input["_tag"]]: Extract<Output,  { _tag: `${Capitalize<RemoveUn<K>>}Output` }>
  }

  function createValueOfObjKey<
    const I extends Extract<Input, { _tag: K }>,
    const K extends Input["_tag"] = I["_tag"]
  >(i: I): InputOutputMap[K] {
    return {
      get _string() {
        return new StringOutput(i as StringInput);
      },
      get _number() {
        return new NumberOutput(i as NumberInput);
      },
      get _list() {
        return new ListOutput(i as ListInput);
      }
    }[i._tag];
  }

}

 

Any idea? TypeScript playground

The sort of relationship between i._tag and the input and output types isn't very easily supported by TypeScript. I consider this strongly related to what I call "correlated union types" as described in microsoft/TypeScript#30581. Conceptually you could just put the three constructors in an object called ctors and then have your function return new ctors[i._tag](i) and it should type check. In practice the compiler doesn't really understand it directly.

You can refactor to use generics and indexed access types as described in microsoft/TypeScript#47109, but it's not incredibly straightforward. The idea is to come up with a "basic" mapping type , like your InputOutputMap but for the input. I'll call it TagInputMap (and rename yours to TagOutputMap, since it's analogous):

type TagInputMap = { [T in Input as T["_tag"]]: T }

From this we can then start describing what we're doing. Let's define TagInput<K> as the intended input to createValueOfObjKey as follows:

type TagInput<K extends keyof TagInputMap> = { _tag: K } & TagInputMap[K]

This could almost just be TagInputMap[K], but then when you call createValueOfObjKey(i), the compiler would fail to infer K (see microsoft/TypeScript#51612). Hence the "redundant" intersection.

Now ctors can be explicitly written as a mapped type from TagInputMap to TagOutputMap:

const ctors: {
  [P in keyof TagInputMap]:
  new (_: TagInputMap[P]) => TagOutputMap[P]
} = {
  _list: ListOutput,
  _number: NumberOutput,
  _string: StringOutput
}

and finally the function implementation:

function createValueOfObjKey<K extends keyof TagInputMap>(i: TagInput<K>) {
  const tag: K = i._tag
  return new ctors[tag](i)
}

Again, this could almost be new ctors[i._tag](i), but the intersection confuses the compiler. We want i._tag to be seen only as K here, but the compiler sees it as K & InputMap[K]["_tag"] which causes things to break (it will fall back to a union like K & ("_string" | "_number" | "_list") and fail). The assignment to tag of type K is the most expedient and still type-safe way to fix it.

And let's try it out from the caller's side:

const ret = createValueOfObjKey({ _tag: "_list" })
// const ret: ListOutput

Looks good.

Playground link to code