diff --git a/src/trackByFunctionRule.ts b/src/trackByFunctionRule.ts
index f00734b9f..8a0828457 100644
--- a/src/trackByFunctionRule.ts
+++ b/src/trackByFunctionRule.ts
@@ -1,65 +1,68 @@
-import * as Lint from 'tslint';
-import * as ts from 'typescript';
+import { BoundDirectivePropertyAst } from '@angular/compiler';
+import { IRuleMetadata, RuleFailure, Rules } from 'tslint/lib';
+import { SourceFile } from 'typescript/lib/typescript';
 import { NgWalker } from './angular/ngWalker';
-import * as ast from '@angular/compiler';
 import { BasicTemplateAstVisitor } from './angular/templates/basicTemplateAstVisitor';
 
-export class Rule extends Lint.Rules.AbstractRule {
-  public static metadata: Lint.IRuleMetadata = {
-    ruleName: 'trackBy-function',
-    type: 'functionality',
+export class Rule extends Rules.AbstractRule {
+  static readonly metadata: IRuleMetadata = {
     description: 'Ensures a trackBy function is used.',
-    rationale: "The use of 'trackBy' is considered a good practice.",
     options: null,
     optionsDescription: 'Not configurable.',
+    rationale: "The use of 'trackBy' is considered a good practice.",
+    ruleName: 'trackBy-function',
+    type: 'functionality',
     typescriptOnly: true
   };
 
-  public apply(sourceFile: ts.SourceFile): Lint.RuleFailure[] {
-    return this.applyWithWalker(
-      new NgWalker(sourceFile, this.getOptions(), {
-        templateVisitorCtrl: TrackByTemplateVisitor
-      })
-    );
+  static readonly FAILURE_STRING = 'Missing trackBy function in ngFor directive';
+
+  apply(sourceFile: SourceFile): RuleFailure[] {
+    return this.applyWithWalker(new NgWalker(sourceFile, this.getOptions(), { templateVisitorCtrl: TrackByTemplateVisitor }));
   }
 }
 
-const ngForExpressionRe = new RegExp(/\*ngFor\s*=\s*(?:'|")(.+)(?:'|")/);
-const trackByRe = new RegExp(/trackBy\s*:/);
+export const getFailureMessage = (): string => {
+  return Rule.FAILURE_STRING;
+};
+
+class TrackByFunctionTemplateVisitor extends BasicTemplateAstVisitor {
+  visitDirectiveProperty(prop: BoundDirectivePropertyAst, context: any): any {
+    this.validateDirective(prop, context);
+    super.visitDirectiveProperty(prop, context);
+  }
 
-class TrackByNgForTemplateVisitor extends BasicTemplateAstVisitor {
-  static Error = 'Missing trackBy function in ngFor directive';
+  private validateDirective(prop: BoundDirectivePropertyAst, context: any): any {
+    const { templateName } = prop;
 
-  visitDirectiveProperty(prop: ast.BoundDirectivePropertyAst, context: BasicTemplateAstVisitor): any {
-    if (prop.sourceSpan) {
-      const directive = (<any>prop.sourceSpan).toString();
+    if (templateName !== 'ngForOf') {
+      return;
+    }
 
-      if (directive.startsWith('*ngFor')) {
-        const directiveMatch = directive.match(ngForExpressionRe);
-        const expr = directiveMatch && directiveMatch[1];
+    const pattern = /trackBy\s*:|\[ngForTrackBy\]\s*=\s*['"].*['"]/;
 
-        if (expr && !trackByRe.test(expr)) {
-          const span = prop.sourceSpan;
-          context.addFailure(
-            context.createFailure(span.start.offset, span.end.offset - span.start.offset, TrackByNgForTemplateVisitor.Error)
-          );
-        }
-      }
+    if (pattern.test(context.codeWithMap.source)) {
+      return;
     }
-    super.visitDirectiveProperty(prop, context);
+
+    const {
+      sourceSpan: {
+        end: { offset: endOffset },
+        start: { offset: startOffset }
+      }
+    } = prop;
+    context.addFailureFromStartToEnd(startOffset, endOffset, getFailureMessage());
   }
 }
 
 class TrackByTemplateVisitor extends BasicTemplateAstVisitor {
-  private visitors: (BasicTemplateAstVisitor)[] = [
-    new TrackByNgForTemplateVisitor(this.getSourceFile(), this.getOptions(), this.context, this.templateStart)
-  ];
+  private readonly visitors: ReadonlySet<BasicTemplateAstVisitor> = new Set([
+    new TrackByFunctionTemplateVisitor(this.getSourceFile(), this.getOptions(), this.context, this.templateStart)
+  ]);
+
+  visitDirectiveProperty(prop: BoundDirectivePropertyAst, context: any): any {
+    this.visitors.forEach(visitor => visitor.visitDirectiveProperty(prop, this));
 
-  visitDirectiveProperty(prop: ast.BoundDirectivePropertyAst, context: any): any {
-    this.visitors
-      .map(v => v.visitDirectiveProperty(prop, this))
-      .filter(f => !!f)
-      .forEach(f => this.addFailure(f));
     super.visitDirectiveProperty(prop, context);
   }
 }
diff --git a/test/trackByFunctionRule.spec.ts b/test/trackByFunctionRule.spec.ts
new file mode 100644
index 000000000..19503e157
--- /dev/null
+++ b/test/trackByFunctionRule.spec.ts
@@ -0,0 +1,196 @@
+import { getFailureMessage, Rule } from '../src/trackByFunctionRule';
+import { assertAnnotated, assertMultipleAnnotated, assertSuccess } from './testHelper';
+
+const {
+  metadata: { ruleName }
+} = Rule;
+
+describe(ruleName, () => {
+  describe('failure', () => {
+    it('should fail when trackBy function is not present', () => {
+      const source = `
+        @Component({
+          template: \`
+            <ul>
+              <li *ngFor="let item of [1, 2, 3];">
+                  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+                {{ item }}
+              </li>
+            </ul>
+          \`
+        })
+        class Bar {}
+      `;
+      assertAnnotated({ message: getFailureMessage(), ruleName, source });
+    });
+
+    it('should fail when trackBy is missing colon', () => {
+      const source = `
+        @Component({
+          template: \`
+            <div *ngFor="let item of [1, 2, 3]; trackBy trackByFn">
+                 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+              {{ item }}
+            </div>
+          \`
+        })
+        class Bar {}
+      `;
+      assertAnnotated({ message: getFailureMessage(), ruleName, source });
+    });
+
+    it('should fail when [ngForTrackBy] is missing in ng-template', () => {
+      const source = `
+        @Component({
+          template: \`
+            <ng-template ngFor let-item [ngForOf]="[1, 2, 3]" let-i="index">
+                                        ~~~~~~~~~~~~~~~~~~~~~
+              {{ item }}
+            </ng-template>
+          \`
+        })
+        class Bar {}
+      `;
+      assertAnnotated({ message: getFailureMessage(), ruleName, source });
+    });
+
+    it('should fail when trackBy function is missing in multiple *ngFor', () => {
+      const source = `
+        @Component({
+          template: \`
+            <div *ngFor="let item of [1, 2, 3];">
+                 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+              {{ item }}
+            </div>
+
+            <ng-template ngFor let-item [ngForOf]="[1, 2, 3]" let-i="index">
+                                        ^^^^^^^^^^^^^^^^^^^^^
+              {{ item }}
+            </ng-template>
+          \`
+        })
+        class Bar {}
+      `;
+      assertMultipleAnnotated({
+        failures: [
+          {
+            char: '~',
+            msg: getFailureMessage()
+          },
+          {
+            char: '^',
+            msg: getFailureMessage()
+          }
+        ],
+        ruleName,
+        source
+      });
+    });
+  });
+
+  describe('success', () => {
+    it('should succeed when trackBy function is present', () => {
+      const source = `
+        @Component({
+          template: \`
+            <div *ngFor="let item of [1, 2, 3]; trackBy: trackByFn">
+              {{ item }}
+            </div>
+          \`
+        })
+        class Bar {}
+      `;
+      assertSuccess(ruleName, source);
+    });
+
+    it('should succeed when trackBy function and exported index value are present', () => {
+      const source = `
+        @Component({
+          template: \`
+            <div *ngFor="let item of [1, 2, 3]; let i = index; trackBy: trackByFn">
+              {{ item }}
+            </div>
+          \`
+        })
+        class Bar {}
+      `;
+      assertSuccess(ruleName, source);
+    });
+
+    it('should succeed when trackBy function is present and has trailing spaces', () => {
+      const source = `
+        @Component({
+          template: \`
+            <div *ngFor="let item of [1, 2, 3]; trackBy : trackByFn">
+              {{ item }}
+            </div>
+          \`
+        })
+        class Bar {}
+      `;
+      assertSuccess(ruleName, source);
+    });
+
+    it('should succeed when trackBy function is present and *ngFor uses single quotes', () => {
+      const source = `
+        @Component({
+          template: \`
+            <div *ngFor='let item of [1, 2, 3]; let i = index; trackBy: trackByFn'>
+              {{ item }}
+            </div>
+          \`
+        })
+        class Bar {}
+      `;
+      assertSuccess(ruleName, source);
+    });
+
+    it('should succeed when *ngFor is surrounded by a lot of whitespaces', () => {
+      const source = `
+        @Component({
+          template: \`
+            <div *ngFor  =  "let item of [1, 2, 3]; let i = index; trackBy : trackByFn">
+              {{ item }}
+            </div>
+          \`
+        })
+        class Bar {}
+      `;
+      assertSuccess(ruleName, source);
+    });
+
+    it('should succeed when [ngForTrackBy] is present in ng-template', () => {
+      const source = `
+        @Component({
+          template: \`
+            <ng-template ngFor let-item [ngForOf]="[1, 2, 3]" let-i="index"
+              [ngForTrackBy]="trackByFn">
+              {{ item }}
+            </ng-template>
+          \`
+        })
+        class Bar {}
+      `;
+      assertSuccess(ruleName, source);
+    });
+
+    it('should succeed when trackBy function is present in multiple *ngFor', () => {
+      const source = `
+        @Component({
+          template: \`
+            <div *ngFor="let item of ['a', 'b', 'c']; index as i; trackBy: trackByFn">
+              {{ item }}
+            </div>
+
+            <ng-template ngFor let-item [ngForOf]="[1, 2, 3]" let-i="index"
+              [ngForTrackBy]="trackByFn">
+              {{ item }}
+            </ng-template>
+          \`
+        })
+        class Bar {}
+      `;
+      assertSuccess(ruleName, source);
+    });
+  });
+});
diff --git a/test/trackByRule.spec.ts b/test/trackByRule.spec.ts
deleted file mode 100644
index 6f0f2fb3a..000000000
--- a/test/trackByRule.spec.ts
+++ /dev/null
@@ -1,131 +0,0 @@
-import { assertSuccess, assertAnnotated } from './testHelper';
-import { Replacement } from 'tslint';
-import { expect } from 'chai';
-
-describe('trackBy-function', () => {
-  describe('ngFor', () => {
-    it('should have a trackBy function', () => {
-      let source = `
-        @Component({
-          template: \`
-            <ul>
-              <li *ngFor="let person of persons; trackBy: trackByFn">
-                {{ person.name }}
-              </li>
-            </ul>
-          \`
-        })
-        class Bar {}
-      `;
-      assertSuccess('trackBy-function', source);
-    });
-
-    it('should have a trackBy function with exported index value', () => {
-      let source = `
-        @Component({
-          template: \`
-            <ul>
-              <li *ngFor="let person of persons; let i = index; trackBy: trackByFn">
-                {{ person.name }}
-              </li>
-            </ul>
-          \`
-        })
-        class Bar {}
-      `;
-      assertSuccess('trackBy-function', source);
-    });
-
-    it('should have a trackBy function', () => {
-      let source = `
-        @Component({
-          template: \`
-            <ul>
-              <li *ngFor="let person of persons; trackBy : trackByFn">
-                {{ person.name }}
-              </li>
-            </ul>
-          \`
-        })
-        class Bar {}
-      `;
-      assertSuccess('trackBy-function', source);
-    });
-  });
-
-  it('should have a trackBy function with different quotes', () => {
-    let source = `
-      @Component({
-        template: \`
-          <ul>
-            <li *ngFor='let person of persons; let i = index; trackBy: trackByFn'>
-              {{ person.name }}
-            </li>
-          </ul>
-        \`
-      })
-      class Bar {}
-    `;
-    assertSuccess('trackBy-function', source);
-  });
-
-  it('should have a trackBy function with different spacing', () => {
-    let source = `
-      @Component({
-        template: \`
-          <ul>
-            <li *ngFor  =  "let person of persons; let i = index; trackBy : trackByFn">
-              {{ person.name }}
-            </li>
-          </ul>
-        \`
-      })
-      class Bar {}
-    `;
-    assertSuccess('trackBy-function', source);
-  });
-
-  describe('failure', () => {
-    it('should fail when trackBy is misspelled', () => {
-      let source = `
-        @Component({
-          template: \`
-            <ul>
-              <li *ngFor="let person of persons; trackBy trackByFn">
-                  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-                {{ person.name }}
-              </li>
-            </ul>
-          \`
-        })
-        class Bar {}
-      `;
-      assertAnnotated({
-        ruleName: 'trackBy-function',
-        message: 'Missing trackBy function in ngFor directive',
-        source
-      });
-    });
-
-    it('should fail when the ngFor directive have not trackBy function', () => {
-      let source = `
-        @Component({
-          template: \`
-            <ul>
-              <li *ngFor="let person of persons;">
-                  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-                {{ person.name }}
-              </li>
-            </ul>
-          \`
-        })
-        class Bar {}
-      `;
-      assertAnnotated({
-        ruleName: 'trackBy-function',
-        message: 'Missing trackBy function in ngFor directive',
-        source
-      });
-    });
-  });
-});