refactor(compiler-cli): Expose API for mappping from TCB to template location (#39715)

Consumers of the `TemplateTypeChecker` API could be interested in
mapping from a shim location back to the original source location in the
template. One concrete example of this use-case is for the "find
references" action in the Language Service. This will return locations
in the TypeScript shim file, and we will then need to be able to map the
result back to the template.

PR Close #39715

packages/compiler-cli/src/ngtsc/typecheck/api/api.ts

@@ -6,7 +6,7 @@
  * found in the LICENSE file at https://angular.io/license
  */
 
-import {BoundTarget, DirectiveMeta, SchemaMetadata} from '@angular/compiler';
+import {AbsoluteSourceSpan, BoundTarget, DirectiveMeta, ParseSourceSpan, SchemaMetadata} from '@angular/compiler';
 import * as ts from 'typescript';
 
 import {AbsoluteFsPath} from '../../file_system';
@@ -306,6 +306,24 @@ export interface ExternalTemplateSourceMapping {
   templateUrl: string;
 }
 
+/**
+ * A mapping of a TCB template id to a span in the corresponding template source.
+ */
+export interface SourceLocation {
+  id: TemplateId;
+  span: AbsoluteSourceSpan;
+}
+
+/**
+ * A representation of all a node's template mapping information we know. Useful for producing
+ * diagnostics based on a TCB node or generally mapping from a TCB node back to a template location.
+ */
+export interface FullTemplateMapping {
+  sourceLocation: SourceLocation;
+  templateSourceMapping: TemplateSourceMapping;
+  span: ParseSourceSpan;
+}
+
 /**
  * Abstracts the operation of determining which shim file will host a particular component's
  * template type-checking code.

packages/compiler-cli/src/ngtsc/typecheck/api/checker.ts

@@ -9,9 +9,10 @@
 import {AST, ParseError, TmplAstNode, TmplAstTemplate} from '@angular/compiler';
 import * as ts from 'typescript';
 
+import {FullTemplateMapping} from './api';
 import {GlobalCompletion} from './completion';
 import {DirectiveInScope, PipeInScope} from './scope';
-import {Symbol} from './symbols';
+import {ShimLocation, Symbol} from './symbols';
 
 /**
  * Interface to the Angular Template Type Checker to extract diagnostics and intelligence from the
@@ -66,6 +67,12 @@ export interface TemplateTypeChecker {
    */
   getDiagnosticsForFile(sf: ts.SourceFile, optimizeFor: OptimizeFor): ts.Diagnostic[];
 
+  /**
+   * Given a `shim` and position within the file, returns information for mapping back to a template
+   * location.
+   */
+  getTemplateMappingAtShimLocation(shimLocation: ShimLocation): FullTemplateMapping|null;
+
   /**
    * Get all `ts.Diagnostic`s currently available that pertain to the given component.
    *

packages/compiler-cli/src/ngtsc/typecheck/src/checker.ts

@@ -6,24 +6,24 @@
  * found in the LICENSE file at https://angular.io/license
  */
 
-import {AST, ParseError, parseTemplate, TmplAstNode, TmplAstReference, TmplAstTemplate, TmplAstVariable} from '@angular/compiler';
+import {AST, ParseError, parseTemplate, TmplAstNode, TmplAstTemplate,} from '@angular/compiler';
 import * as ts from 'typescript';
 
-import {absoluteFromSourceFile, AbsoluteFsPath, getSourceFileOrError} from '../../file_system';
+import {absoluteFrom, absoluteFromSourceFile, AbsoluteFsPath, getSourceFileOrError} from '../../file_system';
 import {ReferenceEmitter} from '../../imports';
 import {IncrementalBuild} from '../../incremental/api';
 import {isNamedClassDeclaration, ReflectionHost} from '../../reflection';
 import {ComponentScopeReader} from '../../scope';
 import {isShim} from '../../shims';
 import {getSourceFileOrNull} from '../../util/src/typescript';
-import {CompletionKind, DirectiveInScope, GlobalCompletion, OptimizeFor, PipeInScope, ProgramTypeCheckAdapter, Symbol, TemplateId, TemplateTypeChecker, TypeCheckingConfig, TypeCheckingProgramStrategy, UpdateMode} from '../api';
+import {DirectiveInScope, FullTemplateMapping, GlobalCompletion, OptimizeFor, PipeInScope, ProgramTypeCheckAdapter, ShimLocation, Symbol, TemplateId, TemplateTypeChecker, TypeCheckingConfig, TypeCheckingProgramStrategy, UpdateMode} from '../api';
 import {TemplateDiagnostic} from '../diagnostics';
 
-import {ExpressionIdentifier, findFirstMatchingNode} from './comments';
 import {CompletionEngine} from './completion';
 import {InliningMode, ShimTypeCheckingData, TemplateData, TypeCheckContextImpl, TypeCheckingHost} from './context';
-import {findTypeCheckBlock, shouldReportDiagnostic, TemplateSourceResolver, translateDiagnostic} from './diagnostics';
+import {shouldReportDiagnostic, translateDiagnostic} from './diagnostics';
 import {TemplateSourceManager} from './source';
+import {findTypeCheckBlock, getTemplateMapping, TemplateSourceResolver} from './tcb_util';
 import {SymbolBuilder} from './template_symbol_builder';
 
 /**
@@ -172,6 +172,31 @@ export class TemplateTypeCheckerImpl implements TemplateTypeChecker {
     return {nodes};
   }
 
+  private getFileAndShimRecordsForPath(shimPath: AbsoluteFsPath):
+      {fileRecord: FileTypeCheckingData, shimRecord: ShimTypeCheckingData}|null {
+    for (const fileRecord of this.state.values()) {
+      if (fileRecord.shimData.has(shimPath)) {
+        return {fileRecord, shimRecord: fileRecord.shimData.get(shimPath)!};
+      }
+    }
+    return null;
+  }
+
+  getTemplateMappingAtShimLocation({shimPath, positionInShimFile}: ShimLocation):
+      FullTemplateMapping|null {
+    const records = this.getFileAndShimRecordsForPath(absoluteFrom(shimPath));
+    if (records === null) {
+      return null;
+    }
+    const {fileRecord} = records;
+
+    const shimSf = this.typeCheckingStrategy.getProgram().getSourceFile(absoluteFrom(shimPath));
+    if (shimSf === undefined) {
+      return null;
+    }
+    return getTemplateMapping(shimSf, positionInShimFile, fileRecord.sourceManager);
+  }
+
   /**
    * Retrieve type-checking diagnostics from the given `ts.SourceFile` using the most recent
    * type-checking program.

packages/compiler-cli/src/ngtsc/typecheck/src/context.ts

@@ -20,7 +20,8 @@ import {DomSchemaChecker, RegistryDomSchemaChecker} from './dom';
 import {Environment} from './environment';
 import {OutOfBandDiagnosticRecorder, OutOfBandDiagnosticRecorderImpl} from './oob';
 import {TemplateSourceManager} from './source';
-import {generateTypeCheckBlock, requiresInlineTypeCheckBlock} from './type_check_block';
+import {requiresInlineTypeCheckBlock} from './tcb_util';
+import {generateTypeCheckBlock} from './type_check_block';
 import {TypeCheckFile} from './type_check_file';
 import {generateInlineTypeCtor, requiresInlineTypeCtor} from './type_constructor';
 

packages/compiler-cli/src/ngtsc/typecheck/src/diagnostics.ts

@@ -7,34 +7,10 @@
  */
 import {AbsoluteSourceSpan, ParseSourceSpan} from '@angular/compiler';
 import * as ts from 'typescript';
-
-import {getTokenAtPosition} from '../../util/src/typescript';
-import {TemplateId, TemplateSourceMapping} from '../api';
+import {TemplateId} from '../api';
 import {makeTemplateDiagnostic, TemplateDiagnostic} from '../diagnostics';
+import {getTemplateMapping, TemplateSourceResolver} from './tcb_util';
 
-import {hasIgnoreMarker, readSpanComment} from './comments';
-
-
-/**
- * Adapter interface which allows the template type-checking diagnostics code to interpret offsets
- * in a TCB and map them back to original locations in the template.
- */
-export interface TemplateSourceResolver {
-  getTemplateId(node: ts.ClassDeclaration): TemplateId;
-
-  /**
-   * For the given template id, retrieve the original source mapping which describes how the offsets
-   * in the template should be interpreted.
-   */
-  getSourceMapping(id: TemplateId): TemplateSourceMapping;
-
-  /**
-   * Convert an absolute source span associated with the given template id into a full
-   * `ParseSourceSpan`. The returned parse span has line and column numbers in addition to only
-   * absolute offsets and gives access to the original template source.
-   */
-  toParseSourceSpan(id: TemplateId, span: AbsoluteSourceSpan): ParseSourceSpan|null;
-}
 
 /**
  * Wraps the node in parenthesis such that inserted span comments become attached to the proper
@@ -115,91 +91,13 @@ export function translateDiagnostic(
   if (diagnostic.file === undefined || diagnostic.start === undefined) {
     return null;
   }
-
-  // Locate the node that the diagnostic is reported on and determine its location in the source.
-  const node = getTokenAtPosition(diagnostic.file, diagnostic.start);
-  const sourceLocation = findSourceLocation(node, diagnostic.file);
-  if (sourceLocation === null) {
+  const fullMapping = getTemplateMapping(diagnostic.file, diagnostic.start, resolver);
+  if (fullMapping === null) {
     return null;
   }
 
-  // Now use the external resolver to obtain the full `ParseSourceFile` of the template.
-  const span = resolver.toParseSourceSpan(sourceLocation.id, sourceLocation.span);
-  if (span === null) {
-    return null;
-  }
-
-  const mapping = resolver.getSourceMapping(sourceLocation.id);
+  const {sourceLocation, templateSourceMapping, span} = fullMapping;
   return makeTemplateDiagnostic(
-      sourceLocation.id, mapping, span, diagnostic.category, diagnostic.code,
+      sourceLocation.id, templateSourceMapping, span, diagnostic.category, diagnostic.code,
       diagnostic.messageText);
 }
-
-export function findTypeCheckBlock(file: ts.SourceFile, id: TemplateId): ts.Node|null {
-  for (const stmt of file.statements) {
-    if (ts.isFunctionDeclaration(stmt) && getTemplateId(stmt, file) === id) {
-      return stmt;
-    }
-  }
-  return null;
-}
-
-interface SourceLocation {
-  id: TemplateId;
-  span: AbsoluteSourceSpan;
-}
-
-/**
- * Traverses up the AST starting from the given node to extract the source location from comments
- * that have been emitted into the TCB. If the node does not exist within a TCB, or if an ignore
- * marker comment is found up the tree, this function returns null.
- */
-function findSourceLocation(node: ts.Node, sourceFile: ts.SourceFile): SourceLocation|null {
-  // Search for comments until the TCB's function declaration is encountered.
-  while (node !== undefined && !ts.isFunctionDeclaration(node)) {
-    if (hasIgnoreMarker(node, sourceFile)) {
-      // There's an ignore marker on this node, so the diagnostic should not be reported.
-      return null;
-    }
-
-    const span = readSpanComment(node, sourceFile);
-    if (span !== null) {
-      // Once the positional information has been extracted, search further up the TCB to extract
-      // the unique id that is attached with the TCB's function declaration.
-      const id = getTemplateId(node, sourceFile);
-      if (id === null) {
-        return null;
-      }
-      return {id, span};
-    }
-
-    node = node.parent;
-  }
-
-  return null;
-}
-
-function getTemplateId(node: ts.Node, sourceFile: ts.SourceFile): TemplateId|null {
-  // Walk up to the function declaration of the TCB, the file information is attached there.
-  while (!ts.isFunctionDeclaration(node)) {
-    if (hasIgnoreMarker(node, sourceFile)) {
-      // There's an ignore marker on this node, so the diagnostic should not be reported.
-      return null;
-    }
-    node = node.parent;
-
-    // Bail once we have reached the root.
-    if (node === undefined) {
-      return null;
-    }
-  }
-
-  const start = node.getFullStart();
-  return ts.forEachLeadingCommentRange(sourceFile.text, start, (pos, end, kind) => {
-    if (kind !== ts.SyntaxKind.MultiLineCommentTrivia) {
-      return null;
-    }
-    const commentText = sourceFile.text.substring(pos + 2, end - 2);
-    return commentText;
-  }) as TemplateId || null;
-}

packages/compiler-cli/src/ngtsc/typecheck/src/dom.ts

@@ -13,7 +13,7 @@ import {ErrorCode, ngErrorCode} from '../../diagnostics';
 import {TemplateId} from '../api';
 import {makeTemplateDiagnostic, TemplateDiagnostic} from '../diagnostics';
 
-import {TemplateSourceResolver} from './diagnostics';
+import {TemplateSourceResolver} from './tcb_util';
 
 const REGISTRY = new DomElementSchemaRegistry();
 const REMOVE_XHTML_REGEX = /^:xhtml:/;

packages/compiler-cli/src/ngtsc/typecheck/src/oob.ts

@@ -14,7 +14,7 @@ import {ClassDeclaration} from '../../reflection';
 import {TemplateId} from '../api';
 import {makeTemplateDiagnostic, TemplateDiagnostic} from '../diagnostics';
 
-import {TemplateSourceResolver} from './diagnostics';
+import {TemplateSourceResolver} from './tcb_util';
 
 
 

packages/compiler-cli/src/ngtsc/typecheck/src/source.ts

@@ -12,8 +12,8 @@ import * as ts from 'typescript';
 import {TemplateId, TemplateSourceMapping} from '../api';
 import {getTemplateId} from '../diagnostics';
 
-import {TemplateSourceResolver} from './diagnostics';
 import {computeLineStartsMap, getLineAndCharacterFromPosition} from './line_mappings';
+import {TemplateSourceResolver} from './tcb_util';
 
 /**
  * Represents the source of a template that was processed during type-checking. This information is

packages/compiler-cli/src/ngtsc/typecheck/src/tcb_util.ts

@@ -0,0 +1,135 @@
+/**
+ * @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 {AbsoluteSourceSpan, ParseSourceSpan} from '@angular/compiler';
+import {ClassDeclaration} from '@angular/compiler-cli/src/ngtsc/reflection';
+import * as ts from 'typescript';
+
+import {getTokenAtPosition} from '../../util/src/typescript';
+import {FullTemplateMapping, SourceLocation, TemplateId, TemplateSourceMapping} from '../api';
+
+import {hasIgnoreMarker, readSpanComment} from './comments';
+import {checkIfClassIsExported, checkIfGenericTypesAreUnbound} from './ts_util';
+
+/**
+ * Adapter interface which allows the template type-checking diagnostics code to interpret offsets
+ * in a TCB and map them back to original locations in the template.
+ */
+export interface TemplateSourceResolver {
+  getTemplateId(node: ts.ClassDeclaration): TemplateId;
+
+  /**
+   * For the given template id, retrieve the original source mapping which describes how the offsets
+   * in the template should be interpreted.
+   */
+  getSourceMapping(id: TemplateId): TemplateSourceMapping;
+
+  /**
+   * Convert an absolute source span associated with the given template id into a full
+   * `ParseSourceSpan`. The returned parse span has line and column numbers in addition to only
+   * absolute offsets and gives access to the original template source.
+   */
+  toParseSourceSpan(id: TemplateId, span: AbsoluteSourceSpan): ParseSourceSpan|null;
+}
+
+export function requiresInlineTypeCheckBlock(node: ClassDeclaration<ts.ClassDeclaration>): boolean {
+  // In order to qualify for a declared TCB (not inline) two conditions must be met:
+  // 1) the class must be exported
+  // 2) it must not have constrained generic types
+  if (!checkIfClassIsExported(node)) {
+    // Condition 1 is false, the class is not exported.
+    return true;
+  } else if (!checkIfGenericTypesAreUnbound(node)) {
+    // Condition 2 is false, the class has constrained generic types
+    return true;
+  } else {
+    return false;
+  }
+}
+
+/** Maps a shim position back to a template location. */
+export function getTemplateMapping(
+    shimSf: ts.SourceFile, position: number, resolver: TemplateSourceResolver): FullTemplateMapping|
+    null {
+  const node = getTokenAtPosition(shimSf, position);
+  const sourceLocation = findSourceLocation(node, shimSf);
+  if (sourceLocation === null) {
+    return null;
+  }
+
+  const mapping = resolver.getSourceMapping(sourceLocation.id);
+  const span = resolver.toParseSourceSpan(sourceLocation.id, sourceLocation.span);
+  if (span === null) {
+    return null;
+  }
+  return {sourceLocation, templateSourceMapping: mapping, span};
+}
+
+export function findTypeCheckBlock(file: ts.SourceFile, id: TemplateId): ts.Node|null {
+  for (const stmt of file.statements) {
+    if (ts.isFunctionDeclaration(stmt) && getTemplateId(stmt, file) === id) {
+      return stmt;
+    }
+  }
+  return null;
+}
+
+/**
+ * Traverses up the AST starting from the given node to extract the source location from comments
+ * that have been emitted into the TCB. If the node does not exist within a TCB, or if an ignore
+ * marker comment is found up the tree, this function returns null.
+ */
+export function findSourceLocation(node: ts.Node, sourceFile: ts.SourceFile): SourceLocation|null {
+  // Search for comments until the TCB's function declaration is encountered.
+  while (node !== undefined && !ts.isFunctionDeclaration(node)) {
+    if (hasIgnoreMarker(node, sourceFile)) {
+      // There's an ignore marker on this node, so the diagnostic should not be reported.
+      return null;
+    }
+
+    const span = readSpanComment(node, sourceFile);
+    if (span !== null) {
+      // Once the positional information has been extracted, search further up the TCB to extract
+      // the unique id that is attached with the TCB's function declaration.
+      const id = getTemplateId(node, sourceFile);
+      if (id === null) {
+        return null;
+      }
+      return {id, span};
+    }
+
+    node = node.parent;
+  }
+
+  return null;
+}
+
+function getTemplateId(node: ts.Node, sourceFile: ts.SourceFile): TemplateId|null {
+  // Walk up to the function declaration of the TCB, the file information is attached there.
+  while (!ts.isFunctionDeclaration(node)) {
+    if (hasIgnoreMarker(node, sourceFile)) {
+      // There's an ignore marker on this node, so the diagnostic should not be reported.
+      return null;
+    }
+    node = node.parent;
+
+    // Bail once we have reached the root.
+    if (node === undefined) {
+      return null;
+    }
+  }
+
+  const start = node.getFullStart();
+  return ts.forEachLeadingCommentRange(sourceFile.text, start, (pos, end, kind) => {
+    if (kind !== ts.SyntaxKind.MultiLineCommentTrivia) {
+      return null;
+    }
+    const commentText = sourceFile.text.substring(pos + 2, end - 2);
+    return commentText;
+  }) as TemplateId || null;
+}

packages/compiler-cli/src/ngtsc/typecheck/src/type_check_block.ts

@@ -21,7 +21,7 @@ import {Environment} from './environment';
 import {astToTypescript, NULL_AS_ANY} from './expression';
 import {OutOfBandDiagnosticRecorder} from './oob';
 import {ExpressionSemanticVisitor} from './template_semantics';
-import {checkIfClassIsExported, checkIfGenericTypesAreUnbound, tsCallMethod, tsCastToAny, tsCreateElement, tsCreateTypeQueryForCoercedInput, tsCreateVariable, tsDeclareVariable} from './ts_util';
+import {tsCallMethod, tsCastToAny, tsCreateElement, tsCreateTypeQueryForCoercedInput, tsCreateVariable, tsDeclareVariable} from './ts_util';
 
 /**
  * Given a `ts.ClassDeclaration` for a component, and metadata regarding that component, compose a
@@ -1867,18 +1867,3 @@ class TcbEventHandlerTranslator extends TcbExpressionTranslator {
     return super.resolve(ast);
   }
 }
-
-export function requiresInlineTypeCheckBlock(node: ClassDeclaration<ts.ClassDeclaration>): boolean {
-  // In order to qualify for a declared TCB (not inline) two conditions must be met:
-  // 1) the class must be exported
-  // 2) it must not have constrained generic types
-  if (!checkIfClassIsExported(node)) {
-    // Condition 1 is false, the class is not exported.
-    return true;
-  } else if (!checkIfGenericTypesAreUnbound(node)) {
-    // Condition 2 is false, the class has constrained generic types
-    return true;
-  } else {
-    return false;
-  }
-}

packages/compiler-cli/src/ngtsc/typecheck/test/type_checker__get_symbol_of_template_node_spec.ts

@@ -99,6 +99,11 @@ runInEachFileSystem(() => {
       expect(
           (symbol.bindings[0].tsSymbol!.declarations[0] as ts.PropertyDeclaration).name.getText())
           .toEqual('name');
+
+      // Ensure we can go back to the original location using the shim location
+      const mapping =
+          templateTypeChecker.getTemplateMappingAtShimLocation(symbol.bindings[0].shimLocation)!;
+      expect(mapping.span.toString()).toEqual('name');
     });
 
     describe('templates', () => {
@@ -155,6 +160,14 @@ runInEachFileSystem(() => {
           assertVariableSymbol(symbol);
           expect(program.getTypeChecker().typeToString(symbol.tsType!)).toEqual('any');
           expect(symbol.declaration.name).toEqual('contextFoo');
+
+          // Ensure we can map the shim locations back to the template
+          const initializerMapping =
+              templateTypeChecker.getTemplateMappingAtShimLocation(symbol.initializerLocation)!;
+          expect(initializerMapping.span.toString()).toEqual('bar');
+          const localVarMapping =
+              templateTypeChecker.getTemplateMappingAtShimLocation(symbol.localVarLocation)!;
+          expect(localVarMapping.span.toString()).toEqual('contextFoo');
         });
 
         it('should get a symbol for local ref which refers to a directive', () => {
@@ -170,6 +183,11 @@ runInEachFileSystem(() => {
           assertReferenceSymbol(symbol);
           expect(program.getTypeChecker().symbolToString(symbol.tsSymbol)).toEqual('TestDir');
           assertDirectiveReference(symbol);
+
+          // Ensure we can map the var shim location back to the template
+          const localVarMapping =
+              templateTypeChecker.getTemplateMappingAtShimLocation(symbol.referenceVarLocation);
+          expect(localVarMapping!.span.toString()).toEqual('ref1');
         });
 
         function assertDirectiveReference(symbol: ReferenceSymbol) {