From 6e4e68cb3012b848d2a90bd2722ffa1ea0e735d8 Mon Sep 17 00:00:00 2001
From: Andrew Scott <atscott@google.com>
Date: Thu, 10 Dec 2020 11:13:11 -0800
Subject: [PATCH] refactor(compiler-cli): Add flag to TCB for non-diagnostic
 requests (#40071)

The TCB utility functions used to find nodes in the TCB are currently
configured to ignore results when an ignore marker is found. However,
these ignore markers are only meant to affect diagnostics requests. The
Language Service may have a need to find nodes with diagnostic ignore
markers. The most common example of this would be finding references for
generic directives. The reference appears to the generic directive's
class appears on the type ctor in the TCB, which is ignored for
diagnostic purposes.
These functions should only skip results when the request is in the
context of a larger request for _diagnostics_. In all other cases, we
should get matches, even if a diagnostic ignore marker is encountered.

PR Close #40071
---
 .../src/ngtsc/typecheck/src/checker.ts        |  7 ++---
 .../src/ngtsc/typecheck/src/diagnostics.ts    |  3 ++-
 .../src/ngtsc/typecheck/src/tcb_util.ts       | 26 +++++++++++--------
 .../ivy/test/references_spec.ts               | 17 ++++++++++++
 4 files changed, 38 insertions(+), 15 deletions(-)

diff --git a/packages/compiler-cli/src/ngtsc/typecheck/src/checker.ts b/packages/compiler-cli/src/ngtsc/typecheck/src/checker.ts
index 0badb22070..0a032a995c 100644
--- a/packages/compiler-cli/src/ngtsc/typecheck/src/checker.ts
+++ b/packages/compiler-cli/src/ngtsc/typecheck/src/checker.ts
@@ -121,12 +121,12 @@ export class TemplateTypeCheckerImpl implements TemplateTypeChecker {
       throw new Error(`Error: no shim file in program: ${shimPath}`);
     }
 
-    let tcb: ts.Node|null = findTypeCheckBlock(shimSf, id);
+    let tcb: ts.Node|null = findTypeCheckBlock(shimSf, id, /*isDiagnosticsRequest*/ false);
 
     if (tcb === null) {
       // Try for an inline block.
       const inlineSf = getSourceFileOrError(program, sfPath);
-      tcb = findTypeCheckBlock(inlineSf, id);
+      tcb = findTypeCheckBlock(inlineSf, id, /*isDiagnosticsRequest*/ false);
     }
 
     let data: TemplateData|null = null;
@@ -194,7 +194,8 @@ export class TemplateTypeCheckerImpl implements TemplateTypeChecker {
     if (shimSf === undefined) {
       return null;
     }
-    return getTemplateMapping(shimSf, positionInShimFile, fileRecord.sourceManager);
+    return getTemplateMapping(
+        shimSf, positionInShimFile, fileRecord.sourceManager, /*isDiagnosticsRequest*/ false);
   }
 
   generateAllTypeCheckBlocks() {
diff --git a/packages/compiler-cli/src/ngtsc/typecheck/src/diagnostics.ts b/packages/compiler-cli/src/ngtsc/typecheck/src/diagnostics.ts
index 02a22ef292..cd828872b6 100644
--- a/packages/compiler-cli/src/ngtsc/typecheck/src/diagnostics.ts
+++ b/packages/compiler-cli/src/ngtsc/typecheck/src/diagnostics.ts
@@ -91,7 +91,8 @@ export function translateDiagnostic(
   if (diagnostic.file === undefined || diagnostic.start === undefined) {
     return null;
   }
-  const fullMapping = getTemplateMapping(diagnostic.file, diagnostic.start, resolver);
+  const fullMapping = getTemplateMapping(
+      diagnostic.file, diagnostic.start, resolver, /*isDiagnosticsRequest*/ true);
   if (fullMapping === null) {
     return null;
   }
diff --git a/packages/compiler-cli/src/ngtsc/typecheck/src/tcb_util.ts b/packages/compiler-cli/src/ngtsc/typecheck/src/tcb_util.ts
index 7bdc189dcb..ac4f3cfb5e 100644
--- a/packages/compiler-cli/src/ngtsc/typecheck/src/tcb_util.ts
+++ b/packages/compiler-cli/src/ngtsc/typecheck/src/tcb_util.ts
@@ -54,10 +54,10 @@ export function requiresInlineTypeCheckBlock(node: ClassDeclaration<ts.ClassDecl
 
 /** Maps a shim position back to a template location. */
 export function getTemplateMapping(
-    shimSf: ts.SourceFile, position: number, resolver: TemplateSourceResolver): FullTemplateMapping|
-    null {
+    shimSf: ts.SourceFile, position: number, resolver: TemplateSourceResolver,
+    isDiagnosticRequest: boolean): FullTemplateMapping|null {
   const node = getTokenAtPosition(shimSf, position);
-  const sourceLocation = findSourceLocation(node, shimSf);
+  const sourceLocation = findSourceLocation(node, shimSf, isDiagnosticRequest);
   if (sourceLocation === null) {
     return null;
   }
@@ -72,9 +72,10 @@ export function getTemplateMapping(
   return {sourceLocation, templateSourceMapping: mapping, span};
 }
 
-export function findTypeCheckBlock(file: ts.SourceFile, id: TemplateId): ts.Node|null {
+export function findTypeCheckBlock(
+    file: ts.SourceFile, id: TemplateId, isDiagnosticRequest: boolean): ts.Node|null {
   for (const stmt of file.statements) {
-    if (ts.isFunctionDeclaration(stmt) && getTemplateId(stmt, file) === id) {
+    if (ts.isFunctionDeclaration(stmt) && getTemplateId(stmt, file, isDiagnosticRequest) === id) {
       return stmt;
     }
   }
@@ -84,12 +85,14 @@ export function findTypeCheckBlock(file: ts.SourceFile, id: TemplateId): ts.Node
 /**
  * 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.
+ * marker comment is found up the tree (and this is part of a diagnostic request), this function
+ * returns null.
  */
-export function findSourceLocation(node: ts.Node, sourceFile: ts.SourceFile): SourceLocation|null {
+export function findSourceLocation(
+    node: ts.Node, sourceFile: ts.SourceFile, isDiagnosticsRequest: boolean): SourceLocation|null {
   // Search for comments until the TCB's function declaration is encountered.
   while (node !== undefined && !ts.isFunctionDeclaration(node)) {
-    if (hasIgnoreForDiagnosticsMarker(node, sourceFile)) {
+    if (hasIgnoreForDiagnosticsMarker(node, sourceFile) && isDiagnosticsRequest) {
       // There's an ignore marker on this node, so the diagnostic should not be reported.
       return null;
     }
@@ -98,7 +101,7 @@ export function findSourceLocation(node: ts.Node, sourceFile: ts.SourceFile): So
     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);
+      const id = getTemplateId(node, sourceFile, isDiagnosticsRequest);
       if (id === null) {
         return null;
       }
@@ -111,10 +114,11 @@ export function findSourceLocation(node: ts.Node, sourceFile: ts.SourceFile): So
   return null;
 }
 
-function getTemplateId(node: ts.Node, sourceFile: ts.SourceFile): TemplateId|null {
+function getTemplateId(
+    node: ts.Node, sourceFile: ts.SourceFile, isDiagnosticRequest: boolean): TemplateId|null {
   // Walk up to the function declaration of the TCB, the file information is attached there.
   while (!ts.isFunctionDeclaration(node)) {
-    if (hasIgnoreForDiagnosticsMarker(node, sourceFile)) {
+    if (hasIgnoreForDiagnosticsMarker(node, sourceFile) && isDiagnosticRequest) {
       // There's an ignore marker on this node, so the diagnostic should not be reported.
       return null;
     }
diff --git a/packages/language-service/ivy/test/references_spec.ts b/packages/language-service/ivy/test/references_spec.ts
index ecc56b1e4c..eff89b6d50 100644
--- a/packages/language-service/ivy/test/references_spec.ts
+++ b/packages/language-service/ivy/test/references_spec.ts
@@ -713,6 +713,23 @@ describe('find references', () => {
       assertTextSpans(refs, ['<div dir>', 'Dir', 'Dir2']);
       assertFileNames(refs, ['app.ts', 'dir.ts', 'dir2.ts']);
     });
+
+    it('should be able to request references for generic directives', () => {
+      const {text, cursor} = extractCursorInfo(`
+        import {Component, NgModule} from '@angular/core';
+
+        @Component({template: '<div *ngF¦or="let item of items"></div>'})
+        export class AppCmp {
+          items = [];
+        }
+      `);
+      const appFile = {name: _('/app.ts'), contents: text};
+      env = createModuleWithDeclarations([appFile]);
+      const refs = getReferencesAtPosition(_('/app.ts'), cursor)!;
+      expect(refs.length).toBe(6);
+      assertTextSpans(refs, ['<div *ngFor="let item of items"></div>', 'NgForOf']);
+      assertFileNames(refs, ['index.d.ts', 'app.ts']);
+    });
   });
 
   describe('components', () => {