From 45b3bd48991772df7d014bcfde60e183074dc25b Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Tue, 3 Jan 2023 16:59:35 -0800
Subject: [PATCH] Foundation for IndexStore/IndexLoad (no lowering)

Adds new types for IndexLoad/IndexStore (renamed later in the stack to
ComputedLoad/ComputedStore) which will be used to represent computed property
access/update. The actual lowering to use these is later in the stack.
---
 compiler/forget/src/HIR/Codegen.ts            | 20 +++++++++++
 compiler/forget/src/HIR/HIR.ts                |  5 +++
 .../forget/src/HIR/InferReferenceEffects.ts   | 36 +++++++++++++++++++
 compiler/forget/src/HIR/PrintHIR.ts           | 12 +++++++
 compiler/forget/src/HIR/visitors.ts           | 22 ++++++++++++
 .../InferReactiveScopeVariables.ts            |  2 ++
 .../PropagateScopeDependencies.ts             |  2 ++
 7 files changed, 99 insertions(+)

diff --git a/compiler/forget/src/HIR/Codegen.ts b/compiler/forget/src/HIR/Codegen.ts
index 0e6b776df1..d843b649a1 100644
--- a/compiler/forget/src/HIR/Codegen.ts
+++ b/compiler/forget/src/HIR/Codegen.ts
@@ -473,6 +473,26 @@ export function codegenInstructionValue(
       );
       break;
     }
+    case "IndexStore": {
+      value = t.assignmentExpression(
+        "=",
+        t.memberExpression(
+          codegenPlace(temp, instrValue.object),
+          codegenPlace(temp, instrValue.property),
+          true
+        ),
+        codegenPlace(temp, instrValue.value)
+      );
+      break;
+    }
+    case "IndexLoad": {
+      value = t.memberExpression(
+        codegenPlace(temp, instrValue.object),
+        codegenPlace(temp, instrValue.property),
+        true
+      );
+      break;
+    }
     case "Identifier": {
       value = codegenPlace(temp, instrValue);
       break;
diff --git a/compiler/forget/src/HIR/HIR.ts b/compiler/forget/src/HIR/HIR.ts
index a3eaf3916a..b671f056ec 100644
--- a/compiler/forget/src/HIR/HIR.ts
+++ b/compiler/forget/src/HIR/HIR.ts
@@ -298,6 +298,11 @@ export type InstructionData =
   // load `object.property`
   | { kind: "PropertyLoad"; object: Place; property: string }
 
+  // store `object[index] = value` - like PropertyStore but with a dynamic property
+  | { kind: "IndexStore"; object: Place; property: Place; value: Place }
+  // load `object[index]` - like PropertyLoad but with a dynamic property
+  | { kind: "IndexLoad"; object: Place; property: Place }
+
   /**
    * Catch-all for statements such as type imports, nested class declarations, etc
    * which are not directly represented, but included for completeness and to allow
diff --git a/compiler/forget/src/HIR/InferReferenceEffects.ts b/compiler/forget/src/HIR/InferReferenceEffects.ts
index a873245abb..a929eb1599 100644
--- a/compiler/forget/src/HIR/InferReferenceEffects.ts
+++ b/compiler/forget/src/HIR/InferReferenceEffects.ts
@@ -607,6 +607,42 @@ function inferBlock(env: Environment, block: BasicBlock) {
         }
         continue;
       }
+      case "IndexStore": {
+        const effect = isObjectType(instrValue.object.identifier)
+          ? Effect.Store
+          : Effect.Mutate;
+        env.reference(instrValue.value, Effect.Read);
+        env.reference(instrValue.property, Effect.Read);
+        env.reference(instrValue.object, effect);
+
+        const lvalue = instr.lvalue;
+        if (lvalue !== null) {
+          env.alias(lvalue.place, instrValue.value);
+          lvalue.place.effect = Effect.Store;
+        }
+        continue;
+      }
+      case "IndexLoad": {
+        if (!env.isDefined(instrValue.object)) {
+          // TODO @josephsavona: improve handling of globals
+          const value: InstructionValue = {
+            kind: "Primitive",
+            loc: instrValue.loc,
+            value: undefined,
+          };
+          env.initialize(value, ValueKind.Frozen);
+          env.define(instrValue.object, value);
+        }
+
+        env.reference(instrValue.object, Effect.Read);
+        env.reference(instrValue.property, Effect.Read);
+        const lvalue = instr.lvalue;
+        if (lvalue !== null) {
+          env.initialize(instrValue, env.kind(instrValue.object));
+          env.define(lvalue.place, instrValue);
+        }
+        continue;
+      }
       case "Identifier": {
         env.reference(instrValue, Effect.Read);
         const lvalue = instr.lvalue;
diff --git a/compiler/forget/src/HIR/PrintHIR.ts b/compiler/forget/src/HIR/PrintHIR.ts
index 80527b34f1..46e0b04870 100644
--- a/compiler/forget/src/HIR/PrintHIR.ts
+++ b/compiler/forget/src/HIR/PrintHIR.ts
@@ -277,6 +277,18 @@ export function printInstructionValue(instrValue: InstructionValue): string {
       } = ${printPlace(instrValue.value)}`;
       break;
     }
+    case "IndexLoad": {
+      value = `IndexLoad ${printPlace(instrValue.object)}[${printPlace(
+        instrValue.property
+      )}]`;
+      break;
+    }
+    case "IndexStore": {
+      value = `IndexStore ${printPlace(instrValue.object)}[${printPlace(
+        instrValue.property
+      )}] = ${printPlace(instrValue.value)}`;
+      break;
+    }
     default: {
       assertExhaustive(
         instrValue,
diff --git a/compiler/forget/src/HIR/visitors.ts b/compiler/forget/src/HIR/visitors.ts
index d2d87cc066..8bf81c821c 100644
--- a/compiler/forget/src/HIR/visitors.ts
+++ b/compiler/forget/src/HIR/visitors.ts
@@ -50,6 +50,17 @@ export function* eachInstructionValueOperand(
       yield instrValue.value;
       break;
     }
+    case "IndexLoad": {
+      yield instrValue.object;
+      yield instrValue.property;
+      break;
+    }
+    case "IndexStore": {
+      yield instrValue.object;
+      yield instrValue.property;
+      yield instrValue.value;
+      break;
+    }
     case "UnaryExpression": {
       yield instrValue.value;
       break;
@@ -110,6 +121,17 @@ export function mapInstructionOperands(
       instrValue.value = fn(instrValue.value);
       break;
     }
+    case "IndexLoad": {
+      instrValue.object = fn(instrValue.object);
+      instrValue.property = fn(instrValue.property);
+      break;
+    }
+    case "IndexStore": {
+      instrValue.object = fn(instrValue.object);
+      instrValue.property = fn(instrValue.property);
+      instrValue.value = fn(instrValue.value);
+      break;
+    }
     case "Identifier": {
       instr.value = fn(instrValue);
       break;
diff --git a/compiler/forget/src/ReactiveScopes/InferReactiveScopeVariables.ts b/compiler/forget/src/ReactiveScopes/InferReactiveScopeVariables.ts
index 546521500b..a8435b261e 100644
--- a/compiler/forget/src/ReactiveScopes/InferReactiveScopeVariables.ts
+++ b/compiler/forget/src/ReactiveScopes/InferReactiveScopeVariables.ts
@@ -161,11 +161,13 @@ function mayAllocate(value: InstructionValue): boolean {
     case "BinaryExpression":
     case "Identifier":
     case "PropertyLoad":
+    case "IndexLoad":
     case "JSXText":
     case "Primitive": {
       return false;
     }
     case "PropertyStore":
+    case "IndexStore":
     case "ArrayExpression":
     case "CallExpression":
     case "JsxExpression":
diff --git a/compiler/forget/src/ReactiveScopes/PropagateScopeDependencies.ts b/compiler/forget/src/ReactiveScopes/PropagateScopeDependencies.ts
index 6a497bf083..172ab29dba 100644
--- a/compiler/forget/src/ReactiveScopes/PropagateScopeDependencies.ts
+++ b/compiler/forget/src/ReactiveScopes/PropagateScopeDependencies.ts
@@ -302,6 +302,8 @@ function valueKind(value: InstructionValue): DeclKind {
     case "Primitive": {
       return DeclKind.Const;
     }
+    case "IndexLoad":
+    case "IndexStore":
     case "PropertyStore":
     case "PropertyLoad":
     case "Identifier":