Use variable type annotations to drive inference

Builds on the utilities added previously to infer types from type annotations on 
variable declarations. This is a limited form, where currently we only infer for 
local identifiers (not function parameters) and only infer a type for the 
variable initializer and not subsequent reassignments.

compiler/packages/babel-plugin-react-forget/src/HIR/BuildHIR.ts

@@ -1591,6 +1591,7 @@ function lowerExpression(
             kind: "StoreLocal",
             lvalue: { kind: InstructionKind.Const, place: { ...place } },
             value: last,
+            type: makeType(),
             loc: exprLoc,
           });
         }
@@ -1632,6 +1633,7 @@ function lowerExpression(
           kind: "StoreLocal",
           lvalue: { kind: InstructionKind.Const, place: { ...place } },
           value: consequent,
+          type: makeType(),
           loc: exprLoc,
         });
         return {
@@ -1650,6 +1652,7 @@ function lowerExpression(
           kind: "StoreLocal",
           lvalue: { kind: InstructionKind.Const, place: { ...place } },
           value: alternate,
+          type: makeType(),
           loc: exprLoc,
         });
         return {
@@ -1700,6 +1703,7 @@ function lowerExpression(
           kind: "StoreLocal",
           lvalue: { kind: InstructionKind.Const, place: { ...place } },
           value: { ...leftPlace },
+          type: makeType(),
           loc: leftPlace.loc,
         });
         return {
@@ -1716,6 +1720,7 @@ function lowerExpression(
           kind: "StoreLocal",
           lvalue: { kind: InstructionKind.Const, place: { ...place } },
           value: { ...right },
+          type: makeType(),
           loc: right.loc,
         });
         return {
@@ -1815,15 +1820,29 @@ function lowerExpression(
             right,
             loc: exprLoc,
           });
-          lowerValueToTemporary(builder, {
-            kind: getStoreKind(builder, leftExpr),
-            lvalue: {
-              place: { ...identifier },
-              kind: InstructionKind.Reassign,
-            },
-            value: { ...binaryPlace },
-            loc: exprLoc,
-          });
+          const kind = getStoreKind(builder, leftExpr);
+          if (kind === "StoreLocal") {
+            lowerValueToTemporary(builder, {
+              kind: "StoreLocal",
+              lvalue: {
+                place: { ...identifier },
+                kind: InstructionKind.Reassign,
+              },
+              value: { ...binaryPlace },
+              type: makeType(),
+              loc: exprLoc,
+            });
+          } else {
+            lowerValueToTemporary(builder, {
+              kind: "StoreContext",
+              lvalue: {
+                place: { ...identifier },
+                kind: InstructionKind.Reassign,
+              },
+              value: { ...binaryPlace },
+              loc: exprLoc,
+            });
+          }
           return { kind: "LoadLocal", place: identifier, loc: exprLoc };
         }
         case "MemberExpression": {
@@ -2272,6 +2291,7 @@ function lowerOptionalMemberExpression(
             kind: "StoreLocal",
             lvalue: { kind: InstructionKind.Const, place: { ...place } },
             value: { ...temp },
+            type: makeType(),
             loc,
           });
           return {
@@ -2326,6 +2346,7 @@ function lowerOptionalMemberExpression(
       kind: "StoreLocal",
       lvalue: { kind: InstructionKind.Const, place: { ...place } },
       value: { ...temp },
+      type: makeType(),
       loc,
     });
     return {
@@ -2382,6 +2403,7 @@ function lowerOptionalCallExpression(
             kind: "StoreLocal",
             lvalue: { kind: InstructionKind.Const, place: { ...place } },
             value: { ...temp },
+            type: makeType(),
             loc,
           });
           return {
@@ -2483,6 +2505,7 @@ function lowerOptionalCallExpression(
       kind: "StoreLocal",
       lvalue: { kind: InstructionKind.Const, place: { ...place } },
       value: { ...temp },
+      type: makeType(),
       loc,
     });
     return {
@@ -3215,10 +3238,22 @@ function lowerAssignment(
           loc,
         });
       } else {
+        const typeAnnotation = lvalue.get("typeAnnotation");
+        let type: Type;
+        if (typeAnnotation.isTSTypeAnnotation()) {
+          const typePath = typeAnnotation.get("typeAnnotation");
+          type = lowerType(builder, typePath);
+        } else if (typeAnnotation.isTypeAnnotation()) {
+          const typePath = typeAnnotation.get("typeAnnotation");
+          type = lowerType(builder, typePath);
+        } else {
+          type = makeType();
+        }
         temporary = lowerValueToTemporary(builder, {
           kind: "StoreLocal",
           lvalue: { place: { ...place }, kind },
           value,
+          type,
           loc,
         });
       }
@@ -3525,6 +3560,7 @@ function lowerAssignment(
           kind: "StoreLocal",
           lvalue: { kind: InstructionKind.Const, place: { ...temp } },
           value: { ...defaultValue },
+          type: makeType(),
           loc,
         });
         return {
@@ -3541,6 +3577,7 @@ function lowerAssignment(
           kind: "StoreLocal",
           lvalue: { kind: InstructionKind.Const, place: { ...temp } },
           value: { ...value },
+          type: makeType(),
           loc,
         });
         return {

compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts

@@ -676,6 +676,7 @@ export type InstructionValue =
       kind: "StoreLocal";
       lvalue: LValue;
       value: Place;
+      type: Type;
       loc: SourceLocation;
     }
   | {

compiler/packages/babel-plugin-react-forget/src/HIR/Types.ts

@@ -90,6 +90,50 @@ export function makeType(): TypeVar {
   };
 }
 
+/**
+ * Duplicates the given type, copying types that are exact while creating fresh
+ * type identifiers for any abstract types.
+ */
+export function duplicateType(type: Type): Type {
+  switch (type.kind) {
+    case "Function": {
+      return {
+        kind: "Function",
+        return: duplicateType(type.return),
+        shapeId: type.shapeId,
+      };
+    }
+    case "Object": {
+      return { kind: "Object", shapeId: type.shapeId };
+    }
+    case "ObjectMethod": {
+      return { kind: "ObjectMethod" };
+    }
+    case "Phi": {
+      return {
+        kind: "Phi",
+        operands: type.operands.map((operand) => duplicateType(operand)),
+      };
+    }
+    case "Poly": {
+      return { kind: "Poly" };
+    }
+    case "Primitive": {
+      return { kind: "Primitive" };
+    }
+    case "Property": {
+      return {
+        kind: "Property",
+        object: duplicateType(type.object),
+        propertyName: type.propertyName,
+      };
+    }
+    case "Type": {
+      return makeType();
+    }
+  }
+}
+
 export function typeEquals(tA: Type, tB: Type): boolean {
   if (tA.kind !== tB.kind) return false;
   return (

compiler/packages/babel-plugin-react-forget/src/Inference/InlineImmediatelyInvokedFunctionExpressions.ts

@@ -245,6 +245,7 @@ function rewriteBlock(
       kind: "StoreLocal",
       lvalue: { kind: InstructionKind.Reassign, place: { ...returnValue } },
       value: terminal.value,
+      type: makeType(),
       loc: terminal.loc,
     },
   });

compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/ExtractScopeDeclarationsFromDestructuring.ts

@@ -15,6 +15,7 @@ import {
   ReactiveFunction,
   ReactiveInstruction,
   ReactiveScopeBlock,
+  makeType,
 } from "../HIR";
 import { eachPatternOperand, mapPatternOperands } from "../HIR/visitors";
 import { ReactiveFunctionTransform, visitReactiveFunction } from "./visitors";
@@ -177,6 +178,7 @@ function transformDestructuring(
           place: original,
         },
         value: temporary,
+        type: makeType(),
         loc: destructure.loc,
       },
       loc: instr.loc,

compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PruneHoistedContexts.ts

@@ -11,6 +11,7 @@ import {
   ReactiveFunction,
   ReactiveInstruction,
   ReactiveStatement,
+  makeType,
 } from "../HIR";
 import {
   ReactiveFunctionTransform,
@@ -59,6 +60,7 @@ class Visitor extends ReactiveFunctionTransform<HoistedIdentifiers> {
                 ...instruction.value.lvalue,
                 kind: InstructionKind.Const,
               },
+              type: makeType(),
               kind: "StoreLocal",
             },
           },

compiler/packages/babel-plugin-react-forget/src/TypeInference/InferTypes.ts

@@ -137,11 +137,12 @@ function* generateInstructionTypes(
     }
 
     case "StoreLocal": {
-      yield equation(left, value.value.identifier.type);
       yield equation(
         value.lvalue.place.identifier.type,
         value.value.identifier.type
       );
+      yield equation(value.type, value.lvalue.place.identifier.type);
+      yield equation(left, value.type);
       break;
     }
 

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/type-annotations/todo_type-annotations-props.expect.md

@@ -0,0 +1,51 @@
+
+## Input
+
+```javascript
+function useArray(items: Array<number>) {
+  // With type information we know that the callback cannot escape
+  // and does not need to be memoized, only the result needs to be
+  // memoized:
+  return items.filter((x) => x !== 0);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useArray,
+  params: [[1, 0, 2, 0, 3, 0, 42]],
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function useArray(items) {
+  const $ = useMemoCache(3);
+  let t1;
+  if ($[0] !== items) {
+    let t0;
+    if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+      t0 = (x) => x !== 0;
+      $[2] = t0;
+    } else {
+      t0 = $[2];
+    }
+    t1 = items.filter(t0);
+    $[0] = items;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useArray,
+  params: [[1, 0, 2, 0, 3, 0, 42]],
+};
+
+```
+      
+### Eval output
+(kind: ok) [1,2,3,42]

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/type-annotations/todo_type-annotations-props.ts

@@ -0,0 +1,11 @@
+function useArray(items: Array<number>) {
+  // With type information we know that the callback cannot escape
+  // and does not need to be memoized, only the result needs to be
+  // memoized:
+  return items.filter((x) => x !== 0);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useArray,
+  params: [[1, 0, 2, 0, 3, 0, 42]],
+};

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/type-annotations/type-annotation-var-array.expect.md

@@ -0,0 +1,70 @@
+
+## Input
+
+```javascript
+function Component(props: { id: number }) {
+  const x: number[] = makeArray(props.id);
+  const y = x.at(0);
+  return y;
+}
+
+function makeArray<T>(x: T): Array<T> {
+  return [x];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ id: 42 }],
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+function Component(props) {
+  const $ = useMemoCache(4);
+  let t0;
+  if ($[0] !== props.id) {
+    t0 = makeArray(props.id);
+    $[0] = props.id;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  const x = t0;
+  let t1;
+  if ($[2] !== x) {
+    t1 = x.at(0);
+    $[2] = x;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  const y = t1;
+  return y;
+}
+
+function makeArray(x) {
+  const $ = useMemoCache(2);
+  let t0;
+  if ($[0] !== x) {
+    t0 = [x];
+    $[0] = x;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ id: 42 }],
+};
+
+```
+      
+### Eval output
+(kind: ok) 42

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/type-annotations/type-annotation-var-array.ts

@@ -0,0 +1,14 @@
+function Component(props: { id: number }) {
+  const x: number[] = makeArray(props.id);
+  const y = x.at(0);
+  return y;
+}
+
+function makeArray<T>(x: T): Array<T> {
+  return [x];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ id: 42 }],
+};

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/type-annotations/type-annotation-var-array_.flow.expect.md

@@ -0,0 +1,75 @@
+
+## Input
+
+```javascript
+// @flow
+import { identity } from "shared-runtime";
+
+function Component(props: { id: number }) {
+  const x: Array<number> = makeArray(props.id);
+  const y = x.at(0);
+  return y;
+}
+
+function makeArray<T>(x: T): Array<T> {
+  return [x];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ id: 42 }],
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+import { identity } from "shared-runtime";
+
+function Component(props) {
+  const $ = useMemoCache(4);
+  let t0;
+  if ($[0] !== props.id) {
+    t0 = makeArray(props.id);
+    $[0] = props.id;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  const x = t0;
+  let t1;
+  if ($[2] !== x) {
+    t1 = x.at(0);
+    $[2] = x;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  const y = t1;
+  return y;
+}
+
+function makeArray(x) {
+  const $ = useMemoCache(2);
+  let t0;
+  if ($[0] !== x) {
+    t0 = [x];
+    $[0] = x;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ id: 42 }],
+};
+
+```
+      
+### Eval output
+(kind: ok) 42

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/type-annotations/type-annotation-var-array_.flow.js

@@ -0,0 +1,17 @@
+// @flow
+import { identity } from "shared-runtime";
+
+function Component(props: { id: number }) {
+  const x: Array<number> = makeArray(props.id);
+  const y = x.at(0);
+  return y;
+}
+
+function makeArray<T>(x: T): Array<T> {
+  return [x];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ id: 42 }],
+};