feat(language-service): enable get references for directive and compo…

packages/language-service/ivy/references.ts

@@ -5,14 +5,14 @@
  * 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 {TmplAstVariable} from '@angular/compiler';
+import {TmplAstBoundAttribute, TmplAstTextAttribute, TmplAstVariable} from '@angular/compiler';
 import {NgCompiler} from '@angular/compiler-cli/src/ngtsc/core';
 import {absoluteFrom, absoluteFromSourceFile, AbsoluteFsPath} from '@angular/compiler-cli/src/ngtsc/file_system';
-import {SymbolKind, TemplateTypeChecker, TypeCheckingProgramStrategy} from '@angular/compiler-cli/src/ngtsc/typecheck/api';
+import {DirectiveSymbol, SymbolKind, TemplateTypeChecker, TypeCheckingProgramStrategy} from '@angular/compiler-cli/src/ngtsc/typecheck/api';
 import * as ts from 'typescript';
 
 import {getTargetAtPosition} from './template_target';
-import {getTemplateInfoAtPosition, isWithin, TemplateInfo, toTextSpan} from './utils';
+import {getDirectiveMatchesForAttribute, getDirectiveMatchesForElementTag, getTemplateInfoAtPosition, isWithin, TemplateInfo, toTextSpan} from './utils';
 
 export class ReferenceBuilder {
   private readonly ttc = this.compiler.getTemplateTypeChecker();
@@ -43,19 +43,27 @@ export class ReferenceBuilder {
       return undefined;
     }
     switch (symbol.kind) {
-      case SymbolKind.Element:
       case SymbolKind.Directive:
       case SymbolKind.Template:
-      case SymbolKind.DomBinding:
         // References to elements, templates, and directives will be through template references
         // (#ref). They shouldn't be used directly for a Language Service reference request.
-        //
-        // Dom bindings aren't currently type-checked (see `checkTypeOfDomBindings`) so they don't
-        // have a shim location and so we cannot find references for them.
-        //
-        // TODO(atscott): Consider finding references for elements that are components as well as
-        // when the position is on an element attribute that directly maps to a directive.
         return undefined;
+      case SymbolKind.Element: {
+        const matches = getDirectiveMatchesForElementTag(symbol.templateNode, symbol.directives);
+        return this.getReferencesForDirectives(matches);
+      }
+      case SymbolKind.DomBinding: {
+        // Dom bindings aren't currently type-checked (see `checkTypeOfDomBindings`) so they don't
+        // have a shim location. This means we can't match dom bindings to their lib.dom reference,
+        // but we can still see if they match to a directive.
+        if (!(positionDetails.node instanceof TmplAstTextAttribute) &&
+            !(positionDetails.node instanceof TmplAstBoundAttribute)) {
+          return undefined;
+        }
+        const directives = getDirectiveMatchesForAttribute(
+            positionDetails.node.name, symbol.host.templateNode, symbol.host.directives);
+        return this.getReferencesForDirectives(directives);
+      }
       case SymbolKind.Reference: {
         const {shimPath, positionInShimFile} = symbol.referenceVarLocation;
         return this.getReferencesAtTypescriptPosition(shimPath, positionInShimFile);
@@ -94,6 +102,27 @@ export class ReferenceBuilder {
     }
   }
 
+  private getReferencesForDirectives(directives: Set<DirectiveSymbol>):
+      ts.ReferenceEntry[]|undefined {
+    const allDirectiveRefs: ts.ReferenceEntry[] = [];
+    for (const dir of directives.values()) {
+      const dirClass = dir.tsSymbol.valueDeclaration;
+      if (dirClass === undefined || !ts.isClassDeclaration(dirClass) ||
+          dirClass.name === undefined) {
+        continue;
+      }
+
+      const dirFile = dirClass.getSourceFile().fileName;
+      const dirPosition = dirClass.name.getStart();
+      const directiveRefs = this.getReferencesAtTypescriptPosition(dirFile, dirPosition);
+      if (directiveRefs !== undefined) {
+        allDirectiveRefs.push(...directiveRefs);
+      }
+    }
+
+    return allDirectiveRefs.length > 0 ? allDirectiveRefs : undefined;
+  }
+
   private getReferencesAtTypescriptPosition(fileName: string, position: number):
       ts.ReferenceEntry[]|undefined {
     const refs = this.tsLS.getReferencesAtPosition(fileName, position);

packages/language-service/ivy/test/references_spec.ts

@@ -679,6 +679,100 @@ describe('find references', () => {
       assertTextSpans(refs, ['<div dir>', 'Dir']);
       assertFileNames(refs, ['app.ts', 'dir.ts']);
     });
+
+    it('gets references to all matching directives when cursor is on an attribute', () => {
+      const dirFile = `
+      import {Directive} from '@angular/core';
+
+      @Directive({selector: '[dir]'})
+      export class Dir {}`;
+      const dirFile2 = `
+      import {Directive} from '@angular/core';
+
+      @Directive({selector: '[dir]'})
+      export class Dir2 {}`;
+      const {text, cursor} = extractCursorInfo(`
+        import {Component, NgModule} from '@angular/core';
+        import {Dir} from './dir';
+        import {Dir2} from './dir2';
+
+        @Component({template: '<div di¦r></div>'})
+        export class AppCmp {
+        }
+
+        @NgModule({declarations: [AppCmp, Dir, Dir2]})
+        export class AppModule {}
+      `);
+      env = LanguageServiceTestEnvironment.setup([
+        {name: _('/app.ts'), contents: text, isRoot: true},
+        {name: _('/dir.ts'), contents: dirFile},
+        {name: _('/dir2.ts'), contents: dirFile2},
+      ]);
+      const refs = getReferencesAtPosition(_('/app.ts'), cursor)!;
+      expect(refs.length).toBe(8);
+      assertTextSpans(refs, ['<div dir>', 'Dir', 'Dir2']);
+      assertFileNames(refs, ['app.ts', 'dir.ts', 'dir2.ts']);
+    });
+  });
+
+  describe('components', () => {
+    it('works for component classes', () => {
+      const {text, cursor} = extractCursorInfo(`
+      import {Component} from '@angular/core';
+
+      @Component({selector: 'my-comp', template: ''})
+      export class MyCo¦mp {}`);
+      const appFile = `
+        import {Component, NgModule} from '@angular/core';
+        import {MyComp} from './comp';
+
+        @Component({template: '<my-comp></my-comp>'})
+        export class AppCmp {
+        }
+
+        @NgModule({declarations: [AppCmp, MyComp]})
+        export class AppModule {}
+      `;
+      env = LanguageServiceTestEnvironment.setup([
+        {name: _('/app.ts'), contents: appFile, isRoot: true},
+        {name: _('/comp.ts'), contents: text},
+      ]);
+      const refs = getReferencesAtPosition(_('/comp.ts'), cursor)!;
+      // 4 references are:  class declaration, template usage, app import and use in declarations
+      // list.
+      expect(refs.length).toBe(4);
+      assertTextSpans(refs, ['<my-comp>', 'MyComp']);
+      assertFileNames(refs, ['app.ts', 'comp.ts']);
+    });
+
+    it('gets works when cursor is on element tag', () => {
+      const compFile = `
+      import {Component} from '@angular/core';
+
+      @Component({selector: 'my-comp', template: ''})
+      export class MyComp {}`;
+      const {text, cursor} = extractCursorInfo(`
+        import {Component, NgModule} from '@angular/core';
+        import {MyComp} from './comp';
+
+        @Component({template: '<my-c¦omp></my-comp>'})
+        export class AppCmp {
+        }
+
+        @NgModule({declarations: [AppCmp, MyComp]})
+        export class AppModule {}
+      `);
+      env = LanguageServiceTestEnvironment.setup([
+        {name: _('/app.ts'), contents: text, isRoot: true},
+        {name: _('/comp.ts'), contents: compFile},
+      ]);
+      const refs = getReferencesAtPosition(_('/app.ts'), cursor)!;
+      // 4 references are:  class declaration, template usage, app import and use in declarations
+      // list.
+      expect(refs.length).toBe(4);
+      assertTextSpans(refs, ['<my-comp>', 'MyComp']);
+      assertFileNames(refs, ['app.ts', 'comp.ts']);
+    });
   });
 
   function getReferencesAtPosition(fileName: string, position: number) {