Type system of the React compiler

19 March 2024
Forget,
React,
Compiler

If you're wondering what the React compiler is, I recommend reading our recent update post for some background. This post is for anyone curious about the compiler theory behind it. Don't feel pressured to understand everything in this post in order to use the compiler.

Memo the props #

In React, a component wrapped in React.memo only re-renders if its props change.

const Greeting = memo(function Greeting({ user }) {
  return (
    <h1>
      Hello, {user.firstName} {user.lastName}!
    </h1>
  );
});

Greeting re-renders whenever its prop, user changes. React uses shallow comparison to determine if the props have changed.

In JavaScript, objects must maintain their identity to make the shallow comparison work, which is why memoization can be crucial. On the other hand, primitive values don't have any identity associated with them and can simply be compared directly.

Object.is({}, {}); // false
Object.is(3, 3); // true

Consider a simple component that calculates a total based on props:

function Price({ items, state }) {
  const subTotal = calculateSubTotal(items);
  const tax = calculateTax(subTotal, state);
  const total = subTotal + tax;
  return <Text text={total} />;
}

One naive way to prevent Text from re-rendering unnecessarily would be memoize everything like this:

function Price({ items, state }) {
  const subTotal = useMemo(() => calculateSubTotal(items), [items]);
  const tax = useMemo(() => calculateTax(subTotal, state), [subTotal, state]);
  const total = useMemo(() => subTotal + tax, [subTotal, tax]);
  return <Text text={total} />;
}

However, we don't really need to memoize primitives for shallow comparison. The memoization here is wasteful from both a memory and a bundle size perspective.

Can we teach the React compiler that these are primitive values? The React compiler could perform whole program analysis by compiling all the files, including the files containing calculateSubTotal and calculateTax to understand that they return a number. But, single file analysis has a bunch of really nice advantages like better performance, incremental rollout and lower compiler complexity.

Can the compiler infer that these are numbers from the way they are used?

Type inference #

One of the most classical type systems is the Hindley Milner type system that's typically used in functional languages. The type inference in the React compiler is inspired by the Algorithm W of this type system, but it's much simpler because JavaScript is too dynamic for sound typing. I'll briefly describe the various steps as implemented by the React compiler.

Initializing the type variables #

During the initial lowering from JavaScript source to the compiler's intermediate representation, every identifier gets an associated Type variable to store the type. A Type variable is just like any other variable but instead of storing values, we store types in them.

type Type = { kind: "type"; id: number } | { kind: "Primitive" };
// { kind: "Type", id: number } represents a type variable
// { kind: "Primitive" } represents a primitive type

let total; // identifier: { name: 'total', type: { kind: "Type", id: 0 } }

{ kind: "Type", id: 0 } is the Type variable associated with the identifier total.

Generating the type equations #

Rather than specifying the type system in scary and not so accessible formal notation, I'll try to explain one of the typing rules using our earlier example.

const total = subTotal + tax;

The above statement can be typed by saying that the operands of a BinaryExpression with an Arithmetic operator are primitives and the resulting value is also a primitive. In our example, subTotal and tax are the operands, the operator is + and the return value is total.

const total = subTotal + tax;
// subTotal -> primitive
// tax      -> primitive
// total    -> primitive

It's possible to use non-primitives as operands of a BinaryExpression in JavaScript but this has been a safe assumption in practice.

The first step of the type inference pass is to generate type equations based on the typing rules defined by the compiler. A type equation is simply a statement of equality between two types, like a mathematical equation. A simple type equation can be of the form "left = right", where left and right are types.

In code, it can be as simple as an object with two fields representing the left hand side and right hand side of the equation, like this:

type TypeEquation = {
  left: Type;
  right: Type;
};

And concretely, the typing rule we defined above can be generated like this:

function* generateTypeEquationsForBinaryExpression(instruction) {
  const { operands, lvalue } = instruction;

  yield { left: operands[0].type, right: { kind: "Primitive" } };
  // subTotal -> primitive

  yield { left: operands[1].type, right: { kind: "Primitive" } };
  // tax -> primitive

  yield { left: { lvalue.type }, right: { kind: "Primitive" } };
  // total -> primitive
}

Similarly type equations can be generated for other constructs in JavaScript like function application (ie, function call), and if-statements.

Solving the equations #

The process of solving these type equations is called unification. The unification process attempts to find a substitution of type variables that makes all the type equations true.

Solving the type equations for our example are pretty straightforward. The type variables for subTotal, tax and total can directly be substituted by the primitive type.

But consider the earlier statement that defines and initializes subTotal:

const subTotal = calculateSubTotal(items);

At the point of defining subTotal, we don't know its type. Only after we look at the usage of subTotal have we inferred that this is a primitive.

But in this type inference, the types flow back to the definition as well. We go back to the definition and notice that the return type of the function calculateSubTotal must be the same type as subTotal. And by solving this, we've now inferred that the return type of calculateSubTotal must be a primitive.

This is an example of how incredibly powerful type inference can be! We've inferred the return type of a function that exists in a separate compilation unit without having looked at it's implementation. Type systems often use this inference to boostrap and quickly start expanding inference across an untyped codebase.

But it does come with a significant downside -- if the inference is incorrect, it results in a surprising “action-at-a-distance” behaviors. This is why Flow moved to local type inference, trading more explicit type annotations for better errors.

We could've made the compiler use the type information from Typescript or Flow, but we wanted to make sure it worked well for untyped JavaScript too. We do plan to add support for these type systems in the future for more optimal memoization.

(Don't) Memo the props #

Now, going back to our original example of the Price component, the compiler can infer that all the values are primitives. There's actually no need to memoize total, subTotal or tax in this Price component, saving bundle size and memory!

function Price(t0) {
  const $ = useMemoCache(2);
  const { items, state } = t0;
  const subTotal = calculateSubTotal(items);
  const tax = calculateTax(subTotal, state);
  const total = subTotal + tax;
  let t1;

  if ($[0] !== total) {
    t1 = <Text text={total} />;
    $[0] = total;
    $[1] = t1;
  } else {
    t1 = $[1];
  }

  return t1;
}

Typing React #

Once we had the type system, it quickly became clear that we could use it as a platform to do various other analyses too.

Adding validations for certain rules of react became as easy as adding a few typing rules to our type system. For example, rather than building separate validations for each React API, just adding the typing rules for useState hook gives us this validation:

  1 | const [state, setState] = useState({ foo: { bar: {} } });
  2 | const foo = state.foo;
> 3 | foo.bar = 1;
    | ^^^ InvalidReact: Mutating a value returned from 'useState()',
          which should not be mutated. Use the setter function to
          update instead.

Note how even interior mutability is caught here -- not just simply a modification to state, but a modification to state.foo through an aliased variable (foo).