refactor(compiler): transform all expressions in the template pipeline (#50118)

Previously the helper operations for transforming expressions in the
template pipeline would only operate against `ir.Expression`s. This commit
changes them to process `o.Expression`s instead, paving the way to use them
for transformations of native expressions in addition to IR expressions.

PR Close #50118

Author: alxhub

packages/compiler/src/template/pipeline/ir/src/expression.ts

@@ -24,15 +24,15 @@ export type Expression =
     ResetViewExpr|ReadVariableExpr|PureFunctionExpr|PureFunctionParameterExpr;
 
 /**
- * Transformer type which converts IR expressions into general `o.Expression`s (which may be an
+ * Transformer type which converts expressions into general `o.Expression`s (which may be an
  * identity transformation).
  */
-export type ExpressionTransform = (expr: Expression, flags: VisitorContextFlag) => o.Expression;
+export type ExpressionTransform = (expr: o.Expression, flags: VisitorContextFlag) => o.Expression;
 
 /**
  * Check whether a given `o.Expression` is a logical IR expression type.
  */
-export function isIrExpression(expr: o.Expression): boolean {
+export function isIrExpression(expr: o.Expression): expr is Expression {
   return expr instanceof ExpressionBase;
 }
 
@@ -363,7 +363,7 @@ export class PureFunctionParameterExpr extends ExpressionBase {
  * Visits all `Expression`s in the AST of `op` with the `visitor` function.
  */
 export function visitExpressionsInOp(
-    op: CreateOp|UpdateOp, visitor: (expr: Expression, flags: VisitorContextFlag) => void): void {
+    op: CreateOp|UpdateOp, visitor: (expr: o.Expression, flags: VisitorContextFlag) => void): void {
   transformExpressionsInOp(op, (expr, flags) => {
     visitor(expr, flags);
     return expr;
@@ -429,7 +429,6 @@ export function transformExpressionsInExpression(
     expr: o.Expression, transform: ExpressionTransform, flags: VisitorContextFlag): o.Expression {
   if (expr instanceof ExpressionBase) {
     expr.transformInternalExpressions(transform, flags);
-    return transform(expr as Expression, flags);
   } else if (expr instanceof o.BinaryOperatorExpr) {
     expr.lhs = transformExpressionsInExpression(expr.lhs, transform, flags);
     expr.rhs = transformExpressionsInExpression(expr.rhs, transform, flags);
@@ -459,7 +458,7 @@ export function transformExpressionsInExpression(
   } else {
     throw new Error(`Unhandled expression kind: ${expr.constructor.name}`);
   }
-  return expr;
+  return transform(expr, flags);
 }
 
 /**

packages/compiler/src/template/pipeline/src/phases/next_context_merging.ts

@@ -49,6 +49,10 @@ function mergeNextContextsInOps(ops: ir.OpList<ir.UpdateOp>): void {
     for (let candidate = op.next!; candidate.kind !== ir.OpKind.ListEnd && tryToMerge;
          candidate = candidate.next!) {
       ir.visitExpressionsInOp(candidate, (expr, flags) => {
+        if (!ir.isIrExpression(expr)) {
+          return expr;
+        }
+
         if (!tryToMerge) {
           // Either we've already merged, or failed to merge.
           return;

packages/compiler/src/template/pipeline/src/phases/reify.ts

@@ -140,7 +140,11 @@ function reifyUpdateOperations(_view: ViewCompilation, ops: ir.OpList<ir.UpdateO
   }
 }
 
-function reifyIrExpression(expr: ir.Expression): o.Expression {
+function reifyIrExpression(expr: o.Expression): o.Expression {
+  if (!ir.isIrExpression(expr)) {
+    return expr;
+  }
+
   switch (expr.kind) {
     case ir.ExpressionKind.NextContext:
       return ng.nextContext(expr.steps);

packages/compiler/src/template/pipeline/src/phases/slot_allocation.ts

@@ -77,6 +77,10 @@ export function phaseSlotAllocation(cpl: ComponentCompilation): void {
 
       // Process all `ir.Expression`s within this view, and look for `usesSlotIndexExprTrait`.
       ir.visitExpressionsInOp(op, expr => {
+        if (!ir.isIrExpression(expr)) {
+          return;
+        }
+
         if (!ir.hasUsesSlotIndexTrait(expr) || expr.slot !== null) {
           return;
         }

packages/compiler/src/template/pipeline/src/phases/var_counting.ts

@@ -23,6 +23,10 @@ export function phaseVarCounting(cpl: ComponentCompilation): void {
       }
 
       ir.visitExpressionsInOp(op, expr => {
+        if (!ir.isIrExpression(expr)) {
+          return;
+        }
+
         // Some expressions require knowledge of the number of variable slots consumed.
         if (ir.hasUsesVarOffsetTrait(expr)) {
           expr.varOffset = varCount;

packages/compiler/src/template/pipeline/src/phases/variable_optimization.ts

@@ -281,6 +281,10 @@ function collectOpInfo(op: ir.CreateOp|ir.UpdateOp): OpInfo {
   let fences = Fence.None;
   const variablesUsed = new Set<ir.XrefId>();
   ir.visitExpressionsInOp(op, expr => {
+    if (!ir.isIrExpression(expr)) {
+      return;
+    }
+
     switch (expr.kind) {
       case ir.ExpressionKind.ReadVariable:
         variablesUsed.add(expr.xref);
@@ -300,6 +304,10 @@ function countVariableUsages(
     op: ir.CreateOp|ir.UpdateOp, varUsages: Map<ir.XrefId, number>,
     varRemoteUsage: Set<ir.XrefId>): void {
   ir.visitExpressionsInOp(op, (expr, flags) => {
+    if (!ir.isIrExpression(expr)) {
+      return;
+    }
+
     if (expr.kind !== ir.ExpressionKind.ReadVariable) {
       return;
     }
@@ -323,6 +331,10 @@ function countVariableUsages(
 function uncountVariableUsages(
     op: ir.CreateOp|ir.UpdateOp, varUsages: Map<ir.XrefId, number>): void {
   ir.visitExpressionsInOp(op, expr => {
+    if (!ir.isIrExpression(expr)) {
+      return;
+    }
+
     if (expr.kind !== ir.ExpressionKind.ReadVariable) {
       return;
     }
@@ -377,6 +389,10 @@ function tryInlineVariableInitializer(
   let inliningAllowed = true;
 
   ir.transformExpressionsInOp(target, (expr, flags) => {
+    if (!ir.isIrExpression(expr)) {
+      return expr;
+    }
+
     if (inlined || !inliningAllowed) {
       // Either the inlining has already succeeded, or we've passed a fence that disallows inlining
       // at this point, so don't try.