Typescript Trees

If you check out the yeet playground, one of the first things you might notice is this emoji-speckled sidebar.

What you're looking at is a hierarchical tree view (powered by the excellent react-arborist library). I'm guessing it's not such an unfamiliar sight. Data trees are an extremely natural and intuitive way to organize data -- whether you're a web dev or a third century Neoplatonist philosopher.

For us at yeet, the sidebar tree represents relationships between Linux devices, processes, and other system components that we monitor for analysis.

The TypeScript Challenge

When dealing with tree structures in TypeScript, the naive implementation looks something like this:

type SidebarNode = {
  type: "Host" | "Yeet" | "Process" | "Memory" | /* ...and more... */;
  data: any; // 🚩 Type safety lost here!
  parent?: SidebarNode;
  children: SidebarNode[];
}

This gets the job done, but there's one red flag: The any type. As soon as we access node.data, we lose all the safety and guarantees that TypeScript gives us.

For yeet this is particularly important. Different node types have different kinds of associated data. For example, "Memory" nodes track memory utilization while "BpfMap" nodes can access streams of events.

We can do better with discriminated unions. Assuming we have types defined for each variant of data:

type SidebarNode<Type extends string, Data> = {
  type: Type,
  data: Data,
  parent?: SidebarNode;
  children: SidebarNode[];
}

type HostNode = SidebarNode<"Host", HostData>;
type YeetNode = SidebarNode<"Yeet", YeetData>;
type ProcessNode = SidebarNode<"Process", ProcessData>;
/* ... other node types ... */

type SidebarNodeInstance = HostNode | YeetNode | ProcessNode | /* ... */;

Now we can use TypeScript's discriminated union support to automatically infer the correct type of data in context:

const n: SidebarNode = {...};
n.data; // type is "HostData" | "YeetData" | "ProcessData" | ...
if (n.type === "Host") n.data; // type is "HostData"

Better. But we're still missing something crucial: parent-child relationship rules.

The Problem with Tree Relationships

In our system, not every node type can be a child of every other node type. We have specific rules like:

• A Host node can have Yeet, DeviceTree, or ProcessTree nodes as children
• A Memory node can only have Collection nodes as children
• A Collection node cannot have any children

We want TypeScript to enforce these rules at compile time, but our current approach doesn't capture these constraints.

Encoding the Hierarchy in the Type System

The key insight is to encode the valid parent-child relationships directly into the type system:

type Schema =
  | ["Host"                    , "Yeet" | "DeviceTree" | "ProcessTree"]
  | ["Yeet"                    , "BpfProgram" | "BpfMap"              ]
  | ["DeviceTree"              , "NetworkTree" | "Memory" | "Cpu"     ]
  | ["NetworkTree"             , "NetworkInterface"                   ]
  | ["Cpu"                     , "CpuCore"                            ]
  | ["ProcessTree" | "Process" , "Process" | "Thread"                 ]
  | ["Memory" | "BpfMap"       , "Collection"                         ]

Each tuple in this union type represents a valid parent-child relationship. The first element is the parent type, and the second element is the child type (or union of types).

Extracting Parent and Child Types

Once we have this hierarchy defined, we can use TypeScript's type system to extract the valid parent and child types for any given node type:

type KidType<T> = (Schema & [T, string])[1];
type MomType<T> = (Schema & [string, T])[0];

Let's break down how these work:

• For KidType<T>:

  1. We intersect Schema with a tuple pattern [T, string]
  2. This finds all tuples where the first element matches type T
  3. Then we extract the second element [1], which gives us all valid child types

• For MomType<T>:

  1. We intersect with [string, T] to find tuples where the second element is T
  2. Then extract the first element [0] to get all valid parent types

The next step is to set up mappings between these type names and the actual node types.

type SidebarNodeType = SidebarNodeInstance["type"]; // "Host" | "Yeet" | ...

// To search SidebarNodeInstance by type name, preprocess it to a map.
// In our code base, this improved typechecking performance.
type SidebarNodeByType = {
  [K in SidebarNodeType]: Extract<SidebarNodeInstance, { type: K }>
};

Finally, let's amend the SidebarNode type from above with the new parent-child constraints.

type SidebarNode<Type extends string, Data> = {
  type: Type,
  data: Data,
  parent?: SidebarNodeByType[MomType<Type>];
  children: SidebarNodeByType[KidType<Type>][];
}

Now the Schema we defined is statically enforced:

hostNode.children.push(yeetNode); // ✅ Compiles
yeetNode.parent = hostNode; // ✅ Compiles

hostNode.children.push(cpuNode); // ❌ Error!
cpuNode.parent = hostNode; // ❌ Error!

Conclusion

TypeScript's type system is more powerful than many realize. By encoding domain rules like tree hierarchy constraints directly into the type system, we can replace verbose runtime checks and make certain bugs impossible to write.

This approach provides type safety, autocompletion, and self-documentation. It's valuable if you have many node types with a specific structure and type safety is a priority.

If you want to see these techniques and more in action, check out yeet, our high-performance Linux monitoring tool. Or check out our blog post about using yeet to audit github actions.