fix(compiler): handle tracking expressions requiring temporary variables (#58520)

Currently when we generate the tracking expression for a `@for` block, we process its expression in the context of the creation block. This is incorrect, because the expression may require ops of its own for cases like nullish coalescing or safe reads. The result is that while we do generate the correct variable, they're added to the creation block rather than the tracking function which causes an error at runtime.

These changes address the issue by keeping track of a separate set of ops for the `track` expression that are prepended to the generated function, similarly to how we handle event listeners.

Fixes #56256.

PR Close #58520

packages/compiler-cli/test/compliance/test_cases/r3_view_compiler_control_flow/GOLDEN_PARTIAL.js

@@ -2666,3 +2666,41 @@ it('case 2', () => {
  ****************************************************************************************************/
 export {};
 
+/****************************************************************************************************
+ * PARTIAL FILE: for_track_by_temporary_variables.js
+ ****************************************************************************************************/
+import { Component } from '@angular/core';
+import * as i0 from "@angular/core";
+export class MyApp {
+    constructor() {
+        this.items = [];
+    }
+}
+MyApp.ɵfac = i0.ɵɵngDeclareFactory({ minVersion: "12.0.0", version: "0.0.0-PLACEHOLDER", ngImport: i0, type: MyApp, deps: [], target: i0.ɵɵFactoryTarget.Component });
+MyApp.ɵcmp = i0.ɵɵngDeclareComponent({ minVersion: "17.0.0", version: "0.0.0-PLACEHOLDER", type: MyApp, isStandalone: true, selector: "ng-component", ngImport: i0, template: `
+    @for (item of items; track item?.name?.[0]?.toUpperCase() ?? foo) {}
+    @for (item of items; track item.name ?? $index ?? foo) {}
+  `, isInline: true });
+i0.ɵɵngDeclareClassMetadata({ minVersion: "12.0.0", version: "0.0.0-PLACEHOLDER", ngImport: i0, type: MyApp, decorators: [{
+            type: Component,
+            args: [{
+                    template: `
+    @for (item of items; track item?.name?.[0]?.toUpperCase() ?? foo) {}
+    @for (item of items; track item.name ?? $index ?? foo) {}
+  `,
+                }]
+        }] });
+
+/****************************************************************************************************
+ * PARTIAL FILE: for_track_by_temporary_variables.d.ts
+ ****************************************************************************************************/
+import * as i0 from "@angular/core";
+export declare class MyApp {
+    foo: any;
+    items: {
+        name?: string;
+    }[];
+    static ɵfac: i0.ɵɵFactoryDeclaration<MyApp, never>;
+    static ɵcmp: i0.ɵɵComponentDeclaration<MyApp, "ng-component", never, {}, {}, never, never, true, never>;
+}
+

packages/compiler-cli/test/compliance/test_cases/r3_view_compiler_control_flow/TEST_CASES.json

@@ -690,6 +690,21 @@
           ]
         }
       ]
+    },
+    {
+      "description": "should support expressions requiring temporary variables inside `track`",
+      "inputFiles": ["for_track_by_temporary_variables.ts"],
+      "expectations": [
+        {
+          "failureMessage": "Incorrect generated output.",
+          "files": [
+            {
+              "expected": "for_track_by_temporary_variables_template.js",
+              "generated": "for_track_by_temporary_variables.js"
+            }
+          ]
+        }
+      ]
     }
   ]
 }

packages/compiler-cli/test/compliance/test_cases/r3_view_compiler_control_flow/for_track_by_temporary_variables.ts

@@ -0,0 +1,12 @@
+import {Component} from '@angular/core';
+
+@Component({
+  template: `
+    @for (item of items; track item?.name?.[0]?.toUpperCase() ?? foo) {}
+    @for (item of items; track item.name ?? $index ?? foo) {}
+  `,
+})
+export class MyApp {
+  foo: any;
+  items: {name?: string}[] = [];
+}

packages/compiler-cli/test/compliance/test_cases/r3_view_compiler_control_flow/for_track_by_temporary_variables_template.js

@@ -0,0 +1,35 @@
+function _forTrack0($index, $item) {
+  let tmp_0_0;
+  return (tmp_0_0 =
+    $item == null
+      ? null
+      : $item.name == null
+        ? null
+        : $item.name[0] == null
+          ? null
+          : $item.name[0].toUpperCase()) !== null && tmp_0_0 !== undefined
+    ? tmp_0_0
+    : this.foo;
+}
+
+function _forTrack1($index, $item) {
+  let tmp_0_0;
+  return (tmp_0_0 = (tmp_0_0 = $item.name) !== null && tmp_0_0 !== undefined ? tmp_0_0 : $index) !==
+    null && tmp_0_0 !== undefined
+    ? tmp_0_0
+    : this.foo;
+}
+
+…
+
+function MyApp_Template(rf, ctx) {
+  if (rf & 1) {
+    $r3$.ɵɵrepeaterCreate(0, MyApp_For_1_Template, 0, 0, null, null, _forTrack0, true);
+    $r3$.ɵɵrepeaterCreate(2, MyApp_For_3_Template, 0, 0, null, null, _forTrack1, true);
+  }
+  if (rf & 2) {
+    $r3$.ɵɵrepeater(ctx.items);
+    $r3$.ɵɵadvance(2);
+    $r3$.ɵɵrepeater(ctx.items);
+  }
+}

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

@@ -50,7 +50,8 @@ export type Expression =
   | ConstCollectedExpr
   | TwoWayBindingSetExpr
   | ContextLetReferenceExpr
-  | StoreLetExpr;
+  | StoreLetExpr
+  | TrackContextExpr;
 
 /**
  * Transformer type which converts expressions into general `o.Expression`s (which may be an
@@ -1153,7 +1154,13 @@ export function transformExpressionsInOp(
         op.trustedValueFn && transformExpressionsInExpression(op.trustedValueFn, transform, flags);
       break;
     case OpKind.RepeaterCreate:
-      op.track = transformExpressionsInExpression(op.track, transform, flags);
+      if (op.trackByOps === null) {
+        op.track = transformExpressionsInExpression(op.track, transform, flags);
+      } else {
+        for (const innerOp of op.trackByOps) {
+          transformExpressionsInOp(innerOp, transform, flags | VisitorContextFlag.InChildOperation);
+        }
+      }
       if (op.trackByFn !== null) {
         op.trackByFn = transformExpressionsInExpression(op.trackByFn, transform, flags);
       }

packages/compiler/src/template/pipeline/ir/src/ops/create.ts

@@ -324,6 +324,12 @@ export interface RepeaterCreateOp extends ElementOpBase, ConsumesVarsTrait {
    */
   track: o.Expression;
 
+  /**
+   * Some kinds of expressions (e.g. safe reads or nullish coalescing) require additional ops
+   * in order to work. This OpList keeps track of those ops, if they're necessary.
+   */
+  trackByOps: OpList<UpdateOp> | null;
+
   /**
    * `null` initially, then an `o.Expression`. Might be a track expression, or might be a reference
    * into the constant pool.
@@ -393,6 +399,7 @@ export function createRepeaterCreateOp(
     emptyView,
     track,
     trackByFn: null,
+    trackByOps: null,
     tag,
     emptyTag,
     emptyAttributes: null,

packages/compiler/src/template/pipeline/src/compilation.ts

@@ -190,6 +190,10 @@ export abstract class CompilationUnit {
         for (const listenerOp of op.handlerOps) {
           yield listenerOp;
         }
+      } else if (op.kind === ir.OpKind.RepeaterCreate && op.trackByOps !== null) {
+        for (const trackOp of op.trackByOps) {
+          yield trackOp;
+        }
       }
     }
     for (const op of this.update) {

packages/compiler/src/template/pipeline/src/emit.ts

@@ -74,7 +74,6 @@ import {saveAndRestoreView} from './phases/save_restore_view';
 import {allocateSlots} from './phases/slot_allocation';
 import {specializeStyleBindings} from './phases/style_binding_specialization';
 import {generateTemporaryVariables} from './phases/temporary_variables';
-import {generateTrackFns} from './phases/track_fn_generation';
 import {optimizeTrackFns} from './phases/track_fn_optimization';
 import {generateTrackVariables} from './phases/track_variables';
 import {countVariables} from './phases/var_counting';
@@ -148,7 +147,6 @@ const phases: Phase[] = [
   {kind: Kind.Tmpl, fn: resolveI18nElementPlaceholders},
   {kind: Kind.Tmpl, fn: resolveI18nExpressionPlaceholders},
   {kind: Kind.Tmpl, fn: extractI18nMessages},
-  {kind: Kind.Tmpl, fn: generateTrackFns},
   {kind: Kind.Tmpl, fn: collectI18nConsts},
   {kind: Kind.Tmpl, fn: collectConstExpressions},
   {kind: Kind.Both, fn: collectElementConsts},

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

@@ -57,6 +57,9 @@ function recursivelyProcessView(view: ViewCompilationUnit, parentScope: Scope |
         if (op.emptyView) {
           recursivelyProcessView(view.job.views.get(op.emptyView)!, scope);
         }
+        if (op.trackByOps !== null) {
+          op.trackByOps.prepend(generateVariablesInScopeForView(view, scope, false));
+        }
         break;
       case ir.OpKind.Listener:
       case ir.OpKind.TwoWayListener:

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

@@ -391,7 +391,7 @@ function reifyCreateOperations(unit: CompilationUnit, ops: ir.OpList<ir.CreateOp
             op.vars!,
             op.tag,
             op.attributes,
-            op.trackByFn!,
+            reifyTrackBy(unit, op),
             op.usesComponentInstance,
             emptyViewFnName,
             emptyDecls,
@@ -632,6 +632,8 @@ function reifyIrExpression(expr: o.Expression): o.Expression {
       return ng.readContextLet(expr.targetSlot.slot!);
     case ir.ExpressionKind.StoreLet:
       return ng.storeLet(expr.value, expr.sourceSpan);
+    case ir.ExpressionKind.TrackContext:
+      return o.variable('this');
     default:
       throw new Error(
         `AssertionError: Unsupported reification of ir.Expression kind: ${
@@ -675,3 +677,46 @@ function reifyListenerHandler(
 
   return o.fn(params, handlerStmts, undefined, undefined, name);
 }
+
+/** Reifies the tracking expression of a `RepeaterCreateOp`. */
+function reifyTrackBy(unit: CompilationUnit, op: ir.RepeaterCreateOp): o.Expression {
+  // If the tracking function was created already, there's nothing left to do.
+  if (op.trackByFn !== null) {
+    return op.trackByFn;
+  }
+
+  const params: o.FnParam[] = [new o.FnParam('$index'), new o.FnParam('$item')];
+  let fn: o.FunctionExpr | o.ArrowFunctionExpr;
+
+  if (op.trackByOps === null) {
+    // If there are no additional ops related to the tracking function, we just need
+    // to turn it into a function that returns the result of the expression.
+    fn = op.usesComponentInstance
+      ? o.fn(params, [new o.ReturnStatement(op.track)])
+      : o.arrowFn(params, op.track);
+  } else {
+    // Otherwise first we need to reify the track-related ops.
+    reifyUpdateOperations(unit, op.trackByOps);
+
+    const statements: o.Statement[] = [];
+    for (const trackOp of op.trackByOps) {
+      if (trackOp.kind !== ir.OpKind.Statement) {
+        throw new Error(
+          `AssertionError: expected reified statements, but found op ${ir.OpKind[trackOp.kind]}`,
+        );
+      }
+      statements.push(trackOp.statement);
+    }
+
+    // Afterwards we can create the function from those ops.
+    fn =
+      op.usesComponentInstance ||
+      statements.length !== 1 ||
+      !(statements[0] instanceof o.ReturnStatement)
+        ? o.fn(params, statements)
+        : o.arrowFn(params, statements[0].value);
+  }
+
+  op.trackByFn = unit.job.pool.getSharedFunctionReference(fn, '_forTrack');
+  return op.trackByFn;
+}

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

@@ -51,6 +51,11 @@ function processLexicalScope(
       case ir.OpKind.TwoWayListener:
         processLexicalScope(view, op.handlerOps);
         break;
+      case ir.OpKind.RepeaterCreate:
+        if (op.trackByOps !== null) {
+          processLexicalScope(view, op.trackByOps);
+        }
+        break;
     }
   }
 

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

@@ -80,6 +80,11 @@ function processLexicalScope(
         // lexical scope.
         processLexicalScope(unit, op.handlerOps, savedView);
         break;
+      case ir.OpKind.RepeaterCreate:
+        if (op.trackByOps !== null) {
+          processLexicalScope(unit, op.trackByOps, savedView);
+        }
+        break;
     }
   }
 

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

@@ -83,6 +83,8 @@ function generateTemporaries(
 
     if (op.kind === ir.OpKind.Listener || op.kind === ir.OpKind.TwoWayListener) {
       op.handlerOps.prepend(generateTemporaries(op.handlerOps) as ir.UpdateOp[]);
+    } else if (op.kind === ir.OpKind.RepeaterCreate && op.trackByOps !== null) {
+      op.trackByOps.prepend(generateTemporaries(op.trackByOps) as ir.UpdateOp[]);
     }
   }
 

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

@@ -1,62 +0,0 @@
-/**
- * @license
- * Copyright Google LLC All Rights Reserved.
- *
- * Use of this source code is governed by an MIT-style license that can be
- * found in the LICENSE file at https://angular.dev/license
- */
-
-import * as o from '../../../../output/output_ast';
-import * as ir from '../../ir';
-
-import type {CompilationJob} from '../compilation';
-
-/**
- * Generate track functions that need to be extracted to the constant pool. This entails wrapping
- * them in an arrow (or traditional) function, replacing context reads with `this.`, and storing
- * them in the constant pool.
- *
- * Note that, if a track function was previously optimized, it will not need to be extracted, and
- * this phase is a no-op.
- */
-export function generateTrackFns(job: CompilationJob): void {
-  for (const unit of job.units) {
-    for (const op of unit.create) {
-      if (op.kind !== ir.OpKind.RepeaterCreate) {
-        continue;
-      }
-      if (op.trackByFn !== null) {
-        // The final track function was already set, probably because it was optimized.
-        continue;
-      }
-
-      // Find all component context reads.
-      let usesComponentContext = false;
-      op.track = ir.transformExpressionsInExpression(
-        op.track,
-        (expr) => {
-          if (expr instanceof ir.PipeBindingExpr || expr instanceof ir.PipeBindingVariadicExpr) {
-            throw new Error(`Illegal State: Pipes are not allowed in this context`);
-          }
-          if (expr instanceof ir.TrackContextExpr) {
-            usesComponentContext = true;
-            return o.variable('this');
-          }
-          return expr;
-        },
-        ir.VisitorContextFlag.None,
-      );
-
-      let fn: o.FunctionExpr | o.ArrowFunctionExpr;
-
-      const fnParams = [new o.FnParam('$index'), new o.FnParam('$item')];
-      if (usesComponentContext) {
-        fn = new o.FunctionExpr(fnParams, [new o.ReturnStatement(op.track)]);
-      } else {
-        fn = o.arrowFn(fnParams, op.track);
-      }
-
-      op.trackByFn = job.pool.getSharedFunctionReference(fn, '_forTrack');
-    }
-  }
-}

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

@@ -58,7 +58,9 @@ export function optimizeTrackFns(job: CompilationJob): void {
         op.track = ir.transformExpressionsInExpression(
           op.track,
           (expr) => {
-            if (expr instanceof ir.ContextExpr) {
+            if (expr instanceof ir.PipeBindingExpr || expr instanceof ir.PipeBindingVariadicExpr) {
+              throw new Error(`Illegal State: Pipes are not allowed in this context`);
+            } else if (expr instanceof ir.ContextExpr) {
               op.usesComponentInstance = true;
               return new ir.TrackContextExpr(expr.view);
             }
@@ -66,6 +68,14 @@ export function optimizeTrackFns(job: CompilationJob): void {
           },
           ir.VisitorContextFlag.None,
         );
+
+        // Also create an OpList for the tracking expression since it may need
+        // additional ops when generating the final code (e.g. temporary variables).
+        const trackOpList = new ir.OpList<ir.UpdateOp>();
+        trackOpList.push(
+          ir.createStatementOp(new o.ReturnStatement(op.track, op.track.sourceSpan)),
+        );
+        op.trackByOps = trackOpList;
       }
     }
   }

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

@@ -36,6 +36,8 @@ export function optimizeVariables(job: CompilationJob): void {
     for (const op of unit.create) {
       if (op.kind === ir.OpKind.Listener || op.kind === ir.OpKind.TwoWayListener) {
         inlineAlwaysInlineVariables(op.handlerOps);
+      } else if (op.kind === ir.OpKind.RepeaterCreate && op.trackByOps !== null) {
+        inlineAlwaysInlineVariables(op.trackByOps);
       }
     }
 
@@ -45,6 +47,8 @@ export function optimizeVariables(job: CompilationJob): void {
     for (const op of unit.create) {
       if (op.kind === ir.OpKind.Listener || op.kind === ir.OpKind.TwoWayListener) {
         optimizeVariablesInOpList(op.handlerOps, job.compatibility);
+      } else if (op.kind === ir.OpKind.RepeaterCreate && op.trackByOps !== null) {
+        optimizeVariablesInOpList(op.trackByOps, job.compatibility);
       }
     }
   }