fix(compiler): optimize track function that only passes $index (#55872)

Currently we optimize methods that pass both `$index` and the item into a method. We can take this a step further by also optimizing calls that only pass `$index` into the first parameter.

PR Close #55872

Author: crisbeto

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

@@ -1548,6 +1548,53 @@ export declare class MyApp {
     static ɵcmp: i0.ɵɵComponentDeclaration<MyApp, "ng-component", never, {}, {}, never, never, false, never>;
 }
 
+/****************************************************************************************************
+ * PARTIAL FILE: for_template_track_method_only_index.js
+ ****************************************************************************************************/
+import { Component } from '@angular/core';
+import * as i0 from "@angular/core";
+export class MyApp {
+    constructor() {
+        this.message = 'hello';
+        this.items = [{ name: 'one' }, { name: 'two' }, { name: 'three' }];
+    }
+    trackFn(index) {
+        return index;
+    }
+}
+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, selector: "ng-component", ngImport: i0, template: `
+    <div>
+      {{message}}
+      @for (item of items; track trackFn($index)) {}
+    </div>
+  `, isInline: true });
+i0.ɵɵngDeclareClassMetadata({ minVersion: "12.0.0", version: "0.0.0-PLACEHOLDER", ngImport: i0, type: MyApp, decorators: [{
+            type: Component,
+            args: [{
+                    template: `
+    <div>
+      {{message}}
+      @for (item of items; track trackFn($index)) {}
+    </div>
+  `,
+                }]
+        }] });
+
+/****************************************************************************************************
+ * PARTIAL FILE: for_template_track_method_only_index.d.ts
+ ****************************************************************************************************/
+import * as i0 from "@angular/core";
+export declare class MyApp {
+    message: string;
+    items: {
+        name: string;
+    }[];
+    trackFn(index: number): number;
+    static ɵfac: i0.ɵɵFactoryDeclaration<MyApp, never>;
+    static ɵcmp: i0.ɵɵComponentDeclaration<MyApp, "ng-component", never, {}, {}, never, never, false, never>;
+}
+
 /****************************************************************************************************
  * PARTIAL FILE: for_pure_track_reuse.js
  ****************************************************************************************************/

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

@@ -421,6 +421,21 @@
         }
       ]
     },
+    {
+      "description": "should optimize tracking function that calls a method on the component only with $index from the root template",
+      "inputFiles": ["for_template_track_method_only_index.ts"],
+      "expectations": [
+        {
+          "files": [
+            {
+              "expected": "for_template_track_method_only_index_template.js",
+              "generated": "for_template_track_method_only_index.js"
+            }
+          ],
+          "failureMessage": "Incorrect template"
+        }
+      ]
+    },
     {
       "description": "should reuse identical pure tracking functions",
       "inputFiles": ["for_pure_track_reuse.ts"],

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

@@ -0,0 +1,18 @@
+import {Component} from '@angular/core';
+
+@Component({
+  template: `
+    <div>
+      {{message}}
+      @for (item of items; track trackFn($index)) {}
+    </div>
+  `,
+})
+export class MyApp {
+  message = 'hello';
+  items = [{name: 'one'}, {name: 'two'}, {name: 'three'}];
+
+  trackFn(index: number) {
+    return index;
+  }
+}

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

@@ -0,0 +1 @@
+$r3$.ɵɵrepeaterCreate(2, MyApp_For_3_Template, 0, 0, null, null, ctx.trackFn, true);

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

@@ -79,7 +79,7 @@ function isTrackByFunctionCall(
     receiver: ir.ContextExpr;
   };
 } {
-  if (!(expr instanceof o.InvokeFunctionExpr) || expr.args.length !== 2) {
+  if (!(expr instanceof o.InvokeFunctionExpr) || expr.args.length === 0 || expr.args.length > 2) {
     return false;
   }
 
@@ -95,6 +95,8 @@ function isTrackByFunctionCall(
   const [arg0, arg1] = expr.args;
   if (!(arg0 instanceof o.ReadVarExpr) || arg0.name !== '$index') {
     return false;
+  } else if (expr.args.length === 1) {
+    return true;
   }
   if (!(arg1 instanceof o.ReadVarExpr) || arg1.name !== '$item') {
     return false;