refactor(compiler): add pure function support to the template pipeline

This commit adds support for generating pure functions in the output `ConstantPool` based on `ir.PureFunctionExpr`s. Note that nothing yet generates these pure function forms - in the future they will be used both in the implementation of the `pipeBindV` instruction as well as literal arrays and maps in expressions.
angular · alxhub · May 2, 2023 · Apr 25, 2023 · Apr 26, 2023 · May 1, 2023
commit 3f6e4ba43ea06d74610bb688a214f4dda61d7f30
diff --git a/packages/compiler/src/template/pipeline/ir/src/enums.ts b/packages/compiler/src/template/pipeline/ir/src/enums.ts
@@ -132,6 +132,16 @@ export enum ExpressionKind {
    * Runtime operation to reset the current view context after `RestoreView`.
    */
   ResetView,
+
+  /**
+   * Defines and calls a function with change-detected arguments.
+   */
+  PureFunctionExpr,
+
+  /**
+   * Indicates a positional parameter to a pure function definition.
+   */
+  PureFunctionParameterExpr,
 }
 
 /**

diff --git a/packages/compiler/src/template/pipeline/ir/src/expression.ts b/packages/compiler/src/template/pipeline/ir/src/expression.ts
@@ -10,7 +10,7 @@ import * as o from '../../../../output/output_ast';
 import type {ParseSourceSpan} from '../../../../parse_util';
 
 import {ExpressionKind, OpKind} from './enums';
-import {UsesSlotIndex, UsesSlotIndexTrait} from './traits';
+import {UsesSlotIndex, UsesSlotIndexTrait, UsesVarOffset, UsesVarOffsetTrait} from './traits';
 
 import type {XrefId} from './operations';
 import type {CreateOp} from './ops/create';
@@ -19,8 +19,9 @@ import type {UpdateOp} from './ops/update';
 /**
  * An `o.Expression` subtype representing a logical expression in the intermediate representation.
  */
-export type Expression = LexicalReadExpr|ReferenceExpr|ContextExpr|NextContextExpr|
-    GetCurrentViewExpr|RestoreViewExpr|ResetViewExpr|ReadVariableExpr;
+export type Expression =
+    LexicalReadExpr|ReferenceExpr|ContextExpr|NextContextExpr|GetCurrentViewExpr|RestoreViewExpr|
+    ResetViewExpr|ReadVariableExpr|PureFunctionExpr|PureFunctionParameterExpr;
 
 /**
  * Transformer type which converts IR expressions into general `o.Expression`s (which may be an
@@ -268,6 +269,96 @@ export class ReadVariableExpr extends ExpressionBase {
   override transformInternalExpressions(): void {}
 }
 
+export class PureFunctionExpr extends ExpressionBase implements ConsumesVarsTrait,
+                                                                UsesVarOffsetTrait {
+  override readonly kind = ExpressionKind.PureFunctionExpr;
+  readonly[ConsumesVarsTrait] = true;
+  readonly[UsesVarOffset] = true;
+
+  varOffset: number|null = null;
+
+  /**
+   * The expression which should be memoized as a pure computation.
+   *
+   * This expression contains internal `PureFunctionParameterExpr`s, which are placeholders for the
+   * positional argument expressions in `args.
+   */
+  body: o.Expression|null;
+
+  /**
+   * Positional arguments to the pure function which will memoize the `body` expression, which act
+   * as memoization keys.
+   */
+  args: o.Expression[];
+
+  /**
+   * Once extracted to the `ConstantPool`, a reference to the function which defines the computation
+   * of `body`.
+   */
+  fn: o.Expression|null = null;
+
+  constructor(expression: o.Expression, args: o.Expression[]) {
+    super();
+    this.body = expression;
+    this.args = args;
+  }
+
+  override visitExpression(visitor: o.ExpressionVisitor, context: any) {
+    this.body?.visitExpression(visitor, context);
+    for (const arg of this.args) {
+      arg.visitExpression(visitor, context);
+    }
+  }
+
+  override isEquivalent(other: o.Expression): boolean {
+    if (!(other instanceof PureFunctionExpr) || other.args.length !== this.args.length) {
+      return false;
+    }
+
+    return other.body !== null && this.body !== null && other.body.isEquivalent(this.body) &&
+        other.args.every((arg, idx) => arg.isEquivalent(this.args[idx]));
+  }
+
+  override isConstant(): boolean {
+    return false;
+  }
+
+  override transformInternalExpressions(transform: ExpressionTransform, flags: VisitorContextFlag):
+      void {
+    if (this.body !== null) {
+      // TODO: figure out if this is the right flag to pass here.
+      this.body = transformExpressionsInExpression(
+          this.body, transform, flags | VisitorContextFlag.InChildOperation);
+    } else if (this.fn !== null) {
+      this.fn = transformExpressionsInExpression(this.fn, transform, flags);
+    }
+
+    for (let i = 0; i < this.args.length; i++) {
+      this.args[i] = transformExpressionsInExpression(this.args[i], transform, flags);
+    }
+  }
+}
+
+export class PureFunctionParameterExpr extends ExpressionBase {
+  override readonly kind = ExpressionKind.PureFunctionParameterExpr;
+
+  constructor(public index: number) {
+    super();
+  }
+
+  override visitExpression(): void {}
+
+  override isEquivalent(other: o.Expression): boolean {
+    return other instanceof PureFunctionParameterExpr && other.index === this.index;
+  }
+
+  override isConstant(): boolean {
+    return true;
+  }
+
+  override transformInternalExpressions(): void {}
+}
+
 /**
  * Visits all `Expression`s in the AST of `op` with the `visitor` function.
  */

diff --git a/packages/compiler/src/template/pipeline/src/emit.ts b/packages/compiler/src/template/pipeline/src/emit.ts
@@ -28,6 +28,7 @@ import {phaseChaining} from './phases/chaining';
 import {phaseMergeNextContext} from './phases/next_context_merging';
 import {phaseNgContainer} from './phases/ng_container';
 import {phaseSaveRestoreView} from './phases/save_restore_view';
+import {phasePureFunctionExtraction} from './phases/pure_function_extraction';
 
 /**
  * Run all transformation phases in the correct order against a `ComponentCompilation`. After this
@@ -48,6 +49,7 @@ export function transformTemplate(cpl: ComponentCompilation): void {
   phaseMergeNextContext(cpl);
   phaseNgContainer(cpl);
   phaseEmptyElements(cpl);
+  phasePureFunctionExtraction(cpl);
   phaseReify(cpl);
   phaseChaining(cpl);
 }

diff --git a/packages/compiler/src/template/pipeline/src/instruction.ts b/packages/compiler/src/template/pipeline/src/instruction.ts
@@ -154,6 +154,17 @@ export function textInterpolate(strings: string[], expressions: o.Expression[]):
   return callVariadicInstruction(TEXT_INTERPOLATE_CONFIG, [], interpolationArgs);
 }
 
+export function pureFunction(
+    varOffset: number, fn: o.Expression, args: o.Expression[]): o.Expression {
+  return callVariadicInstructionExpr(
+      PURE_FUNCTION_CONFIG,
+      [
+        o.literal(varOffset),
+        fn,
+      ],
+      args,
+  );
+}
 
 
 function call<OpT extends ir.CreateOp|ir.UpdateOp>(
@@ -195,6 +206,22 @@ const TEXT_INTERPOLATE_CONFIG: VariadicInstructionConfig = {
   },
 };
 
+const PURE_FUNCTION_CONFIG: VariadicInstructionConfig = {
+  constant: [
+    Identifiers.pureFunction0,
+    Identifiers.pureFunction1,
+    Identifiers.pureFunction2,
+    Identifiers.pureFunction3,
+    Identifiers.pureFunction4,
+    Identifiers.pureFunction5,
+    Identifiers.pureFunction6,
+    Identifiers.pureFunction7,
+    Identifiers.pureFunction8,
+  ],
+  variable: Identifiers.pureFunctionV,
+  mapping: n => n,
+};
+
 function callVariadicInstructionExpr(
     config: VariadicInstructionConfig, baseArgs: o.Expression[],
     interpolationArgs: o.Expression[]): o.Expression {

diff --git a/packages/compiler/src/template/pipeline/src/phases/pure_function_extraction.ts b/packages/compiler/src/template/pipeline/src/phases/pure_function_extraction.ts
@@ -0,0 +1,66 @@
+/**
+ * @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.io/license
+ */
+
+import {GenericKeyFn, SharedConstantDefinition} from '../../../../constant_pool';
+import * as o from '../../../../output/output_ast';
+import * as ir from '../../ir';
+
+import type {ComponentCompilation} from '../compilation';
+
+export function phasePureFunctionExtraction(cpl: ComponentCompilation): void {
+  for (const view of cpl.views.values()) {
+    for (const op of view.ops()) {
+      ir.visitExpressionsInOp(op, expr => {
+        if (!(expr instanceof ir.PureFunctionExpr) || expr.body === null) {
+          return;
+        }
+
+        const constantDef = new PureFunctionConstant(expr.args.length);
+        expr.fn = cpl.pool.getSharedConstant(constantDef, expr.body);
+        expr.body = null;
+      });
+    }
+  }
+}
+
+class PureFunctionConstant extends GenericKeyFn implements SharedConstantDefinition {
+  constructor(private numArgs: number) {
+    super();
+  }
+
+  override keyOf(expr: o.Expression): string {
+    if (expr instanceof ir.PureFunctionParameterExpr) {
+      return `param(${expr.index})`;
+    } else {
+      return super.keyOf(expr);
+    }
+  }
+
+  toSharedConstantDeclaration(declName: string, keyExpr: o.Expression): o.Statement {
+    const fnParams: o.FnParam[] = [];
+    for (let idx = 0; idx < this.numArgs; idx++) {
+      fnParams.push(new o.FnParam('_p' + idx));
+    }
+
+    // We will never visit `ir.PureFunctionParameterExpr`s that don't belong to us, because this
+    // transform runs inside another visitor which will visit nested pure functions before this one.
+    const returnExpr = ir.transformExpressionsInExpression(keyExpr, expr => {
+      if (!(expr instanceof ir.PureFunctionParameterExpr)) {
+        return expr;
+      }
+
+      return o.variable('_p' + expr.index);
+    }, ir.VisitorContextFlag.None);
+
+    return new o.DeclareFunctionStmt(
+        declName,
+        fnParams,
+        [new o.ReturnStatement(returnExpr)],
+    );
+  }
+}
diff --git a/packages/compiler/src/template/pipeline/src/phases/reify.ts b/packages/compiler/src/template/pipeline/src/phases/reify.ts
@@ -162,6 +162,13 @@ function reifyIrExpression(expr: ir.Expression): o.Expression {
         throw new Error(`Read of unnamed variable ${expr.xref}`);
       }
       return o.variable(expr.name);
+    case ir.ExpressionKind.PureFunctionExpr:
+      if (expr.fn === null) {
+        throw new Error(`AssertionError: expected PureFunctions to have been extracted`);
+      }
+      return ng.pureFunction(expr.varOffset!, expr.fn, expr.args);
+    case ir.ExpressionKind.PureFunctionParameterExpr:
+      throw new Error(`AssertionError: expected PureFunctionParameterExpr to have been extracted`);
     default:
       throw new Error(`AssertionError: Unsupported reification of ir.Expression kind: ${
           ir.ExpressionKind[(expr as ir.Expression).kind]}`);

diff --git a/packages/compiler/src/template/pipeline/src/phases/var_counting.ts b/packages/compiler/src/template/pipeline/src/phases/var_counting.ts
@@ -69,5 +69,11 @@ function varsUsedByOp(op: (ir.CreateOp|ir.UpdateOp)&ir.ConsumesVarsTrait): numbe
 }
 
 function varsUsedByIrExpression(expr: ir.Expression&ir.ConsumesVarsTrait): number {
-  return 0;
+  switch (expr.kind) {
+    case ir.ExpressionKind.PureFunctionExpr:
+      return 1 + expr.args.length;
+    default:
+      throw new Error(
+          `AssertionError: unhandled ConsumesVarsTrait expression ${expr.constructor.name}`);
+  }
 }