How to dynamically set the type of class property based on generic?

The type of property "props" in Class "MyNode" is determined by "type". Currently, this is how I implement it:

enum NodeType {
  Text,
  Rect,
  Triangle,
}

interface PropsTypeByNodeType {
  [NodeType.Text]: string;
  [NodeType.Rect]: Record<string, any>;
  [NodeType.Triangle]: string[];
}

class MyNode<T extends NodeType = NodeType> {
  type: T;
  props: PropsTypeByNodeType[T];
  alternate: MyNode<T>;
}

type TextNodeProps = MyNode<NodeType.Text>["props"]; // string
type RectNodeProps = MyNode<NodeType.Rect>["props"]; // Record<string, any>

But type narrowing not work

function compute(unknownNode: MyNode) {
  if (unknownNode.type === NodeType.Text) {
    unknownNode // expected type is narrowed down to MyNode<NodeType.Text>, actually get MyNode<NodeType>
  }
}

So I tried to solve it by using type distribution:

type NodeHelper<T extends NodeType = NodeType> = T extends NodeType ? MyNode<T> : never;

function compute(unknownNode: NodeHelper) {
  if (unknownNode.type === NodeType.Text) {
    unknownNode // finally get MyNode<NodeType.Text>!
  }
}

This seems to have been solved, but an error occurred when I wanted to write the following code:

// Return a node of the same type. One case is to return node.alternate
// but here it seems that the generic T is distributed as the default value NodeType
function clone<T extends NodeType>(node: NodeHelper<T>) {
  // expected type is MyNode<T>
  // actually get MyNode<NodeType.Text> | MyNode<NodeType.Rect> | MyNode<NodeType.Triangle>
  return node.alternate; 
}

Are there any ways to fix it?

You do indeed need to distribute MyNode<T> over unions in T if you want to be able to correlate the type, props, and alternate properties of MyNode with each other. Your NodeHelper<T> type is a distributive conditional type, which works to distribute and produces the union you want when T is a specific type like NodeType. This at least makes a discriminated union and you can do narrowing inside compute().

But TypeScript does not understand how NodeHelper<T> works for generic T. Evaluation of generic conditional types in TypeScript is either completely deferred (so TypeScript has no idea what might be assignable to them) or prematurely evaluated by widening the generic to its constraint. Inside your clone() it's doing the latter, widening to the NodeHelper<NodeType> union, which isn't what you want.

TypeScript doesn't directly support doing correlated/generic operations over discriminated unions. This is the subject of microsoft/TypeScript#30581. And again, you do need to distribute MyNode<T> over unions in T, but the way you did it isn't going to work for generics. The supported way of distributing a type over key-like unions (that is, unions of enums, or string literals) where TypeScript can keep it generic is via a distributive object type as coined in microsoft/TypeScript#. This is when you index into a mapped type, like so:

type NodeHelper<T extends NodeType = NodeType> = { [K in T]: MyNode<K> }[T]

You can verify that this produces the same output as your NodeHelper for any specific T type. But now inside clone(), TypeScript sees operations on NodeHelper<T> as being generic, and so node.alternate is of type MyNode<T> as expected.

Playground link to code