From bbba752e98738c259e862ce0088dc8dda36cbec0 Mon Sep 17 00:00:00 2001
From: Sathya Gunasekaran <gsathya.ceg@gmail.com>
Date: Wed, 11 Jan 2023 11:30:58 +0000
Subject: [PATCH] [typer] Add type inference for Phis

---
 .../forget/src/TypeInference/InferTypes.ts    | 41 ++++++++++++++--
 .../hir/infer-phi-primitive.expect.md         | 49 +++++++++++++++++++
 .../fixtures/hir/infer-phi-primitive.js       | 11 +++++
 3 files changed, 97 insertions(+), 4 deletions(-)
 create mode 100644 compiler/forget/src/__tests__/fixtures/hir/infer-phi-primitive.expect.md
 create mode 100644 compiler/forget/src/__tests__/fixtures/hir/infer-phi-primitive.js

diff --git a/compiler/forget/src/TypeInference/InferTypes.ts b/compiler/forget/src/TypeInference/InferTypes.ts
index f73d788e96..b27ba32eb3 100644
--- a/compiler/forget/src/TypeInference/InferTypes.ts
+++ b/compiler/forget/src/TypeInference/InferTypes.ts
@@ -3,7 +3,6 @@ import invariant from "invariant";
 import {
   HIRFunction,
   Instruction,
-  Place,
   Type,
   typeEquals,
   TypeId,
@@ -46,6 +45,9 @@ export default function (func: HIRFunction) {
 
 function apply(func: HIRFunction, unifier: Unifier) {
   for (const [_, block] of func.body.blocks) {
+    for (const phi of block.phis) {
+      phi.type = unifier.get(phi.type);
+    }
     for (const instr of block.instructions) {
       for (const place of eachInstructionOperand(instr)) {
         place.identifier.type = unifier.get(place.identifier.type);
@@ -61,8 +63,20 @@ type TypeEquation = {
   right: Type;
 };
 
-function* generate(func: HIRFunction) {
+function* generate(
+  func: HIRFunction
+): Generator<TypeEquation, void, undefined> {
   for (const [_, block] of func.body.blocks) {
+    for (const phi of block.phis) {
+      yield {
+        left: phi.type,
+        right: {
+          kind: "Phi",
+          operands: [...phi.operands.values()].map((id) => id.type),
+        },
+      };
+    }
+
     for (const instr of block.instructions) {
       yield* generateTypeEquation(instr);
     }
@@ -148,6 +162,19 @@ class Unifier {
       return;
     }
 
+    if (type.kind === "Phi") {
+      const operands = new Set(type.operands.map((i) => this.get(i).kind));
+
+      invariant(operands.size > 0, "there should be at least one operand");
+      const kind = operands.values().next().value;
+
+      // there's only one unique type and it's not a type var
+      if (operands.size === 1 && kind !== "Type") {
+        this.unify(v, type.operands[0]);
+        return;
+      }
+    }
+
     if (this.occursCheck(v, type)) {
       throw new Error("cycle detected");
     }
@@ -162,6 +189,10 @@ class Unifier {
       return this.occursCheck(v, this.substitutions.get(type.id)!);
     }
 
+    if (type.kind === "Phi") {
+      return type.operands.some((o) => this.occursCheck(v, o));
+    }
+
     return false;
   }
 
@@ -169,11 +200,13 @@ class Unifier {
     if (type.kind === "Type") {
       if (this.substitutions.has(type.id)) {
         return this.get(this.substitutions.get(type.id)!);
-      } else {
-        return type;
       }
     }
 
+    if (type.kind === "Phi") {
+      return { kind: "Phi", operands: type.operands.map((o) => this.get(o)) };
+    }
+
     return type;
   }
 }
diff --git a/compiler/forget/src/__tests__/fixtures/hir/infer-phi-primitive.expect.md b/compiler/forget/src/__tests__/fixtures/hir/infer-phi-primitive.expect.md
new file mode 100644
index 0000000000..9065975ca8
--- /dev/null
+++ b/compiler/forget/src/__tests__/fixtures/hir/infer-phi-primitive.expect.md
@@ -0,0 +1,49 @@
+
+## Input
+
+```javascript
+function foo(a, b) {
+  let x;
+  if (a) {
+    x = 1;
+  } else {
+    x = 2;
+  }
+
+  let y = x;
+  return y;
+}
+
+```
+
+## Code
+
+```javascript
+function foo(a, b) {
+  const $ = React.useMemoCache();
+  const x = undefined;
+  const c_0 = $[0] !== a;
+  let x$0;
+  if (c_0) {
+    x$0 = undefined;
+
+    if (a) {
+      const x$1 = 1;
+      x$0 = x$1;
+    } else {
+      const x$2 = 2;
+      x$0 = x$2;
+    }
+
+    $[0] = a;
+    $[1] = x$0;
+  } else {
+    x$0 = $[1];
+  }
+
+  const y = x$0;
+  return y;
+}
+
+```
+      
\ No newline at end of file
diff --git a/compiler/forget/src/__tests__/fixtures/hir/infer-phi-primitive.js b/compiler/forget/src/__tests__/fixtures/hir/infer-phi-primitive.js
new file mode 100644
index 0000000000..faafa9ddd3
--- /dev/null
+++ b/compiler/forget/src/__tests__/fixtures/hir/infer-phi-primitive.js
@@ -0,0 +1,11 @@
+function foo(a, b) {
+  let x;
+  if (a) {
+    x = 1;
+  } else {
+    x = 2;
+  }
+
+  let y = x;
+  return y;
+}