This typings

package.json

@@ -8,7 +8,7 @@
   "license": "MIT",
   "scripts": {
     "build": "rimraf dist && tsc && copyfiles -u 1 src/type-collector-snippet.ts dist",
-    "format": "prettier --write src/**/*.ts packages/**/src/**.ts",
+    "format": "prettier --write src/*.ts src/**/*.ts packages/**/src/**.ts",
     "precommit": "lint-staged",
     "prepublish": "npm run build",
     "lint": "tslint -p .",

src/apply-types.ts

@@ -2,36 +2,17 @@ import * as fs from 'fs';
 import * as path from 'path';
 import * as ts from 'typescript';
 
-import { IExtraOptions } from './instrument';
+import { getProgram, ICompilerOptions } from './compiler-helper';
 import { applyReplacements, Replacement } from './replacement';
-import { ISourceLocation } from './type-collector-snippet';
+import { ICollectedTypeInfo, ISourceLocation } from './type-collector-snippet';
 
-export type ICollectedTypeInfo = Array<
-    [string, number, Array<[string | undefined, ISourceLocation | undefined]>, IExtraOptions]
->;
-
-export interface IApplyTypesOptions {
+export interface IApplyTypesOptions extends ICompilerOptions {
     /**
      * A prefix that will be added in front of each type applied. You can use a javascript comment
      * to mark the automatically added types. The prefix will be added after the colon character,
      * just before the actual type.
      */
     prefix?: string;
-
-    /**
-     * If given, all the file paths in the collected type info will be resolved relative to this directory.
-     */
-    rootDir?: string;
-
-    /**
-     * Path to your project's tsconfig file
-     */
-    tsConfig?: string;
-
-    // You probably never need to touch these two - they are used by the integration tests to setup
-    // a virtual file system for TS:
-    tsConfigHost?: ts.ParseConfigHost;
-    tsCompilerHost?: ts.CompilerHost;
 }
 
 function findType(program?: ts.Program, typeName?: string, sourcePos?: ISourceLocation) {
@@ -76,34 +57,26 @@ export function applyTypesToFile(
             continue;
         }
 
+        let thisPrefix = '';
         let suffix = '';
         if (opts && opts.parens) {
             replacements.push(Replacement.insert(opts.parens[0], '('));
             suffix = ')';
         }
-        replacements.push(Replacement.insert(pos, ': ' + prefix + sortedTypes.join('|') + suffix));
+        if (opts && opts.comma) {
+            suffix = ', ';
+        }
+        if (opts && opts.thisType) {
+            thisPrefix = 'this';
+        }
+        replacements.push(Replacement.insert(pos, thisPrefix + ': ' + prefix + sortedTypes.join('|') + suffix));
     }
     return applyReplacements(source, replacements);
 }
 
 export function applyTypes(typeInfo: ICollectedTypeInfo, options: IApplyTypesOptions = {}) {
     const files: { [key: string]: typeof typeInfo } = {};
-    let program: ts.Program | undefined;
-    if (options.tsConfig) {
-        const configHost = options.tsConfigHost || ts.sys;
-        const { config, error } = ts.readConfigFile(options.tsConfig, configHost.readFile);
-        if (error) {
-            throw new Error(`Error while reading ${options.tsConfig}: ${error.messageText}`);
-        }
-
-        const parsed = ts.parseJsonConfigFileContent(config, configHost, options.rootDir || '');
-        if (parsed.errors.length) {
-            const errors = parsed.errors.map((e) => e.messageText).join(', ');
-            throw new Error(`Error while parsing ${options.tsConfig}: ${errors}`);
-        }
-
-        program = ts.createProgram(parsed.fileNames, parsed.options, options.tsCompilerHost);
-    }
+    const program: ts.Program | undefined = getProgram(options);
     for (const entry of typeInfo) {
         const file = entry[0];
         if (!files[file]) {

src/apply-types.spec.ts → src/compiler-helper.spec.ts

@@ -1,9 +1,9 @@
 import * as ts from 'typescript';
-import { applyTypes } from './apply-types';
+import { getProgram } from './compiler-helper';
 
-describe('applyTypes', () => {
+describe('getProgram', () => {
     it('should throw an error if given non-existing tsconfig.json file', () => {
-        expect(() => applyTypes([], { tsConfig: 'not-found-file.json' })).toThrowError(
+        expect(() => getProgram({ tsConfig: 'not-found-file.json' })).toThrowError(
             `Error while reading not-found-file.json: The specified path does not exist: 'not-found-file.json'.`,
         );
     });
@@ -13,7 +13,7 @@ describe('applyTypes', () => {
             ...ts.sys,
             readFile: jest.fn(() => '<invalid json>'),
         };
-        expect(() => applyTypes([], { tsConfig: 'tsconfig.bad.json', tsConfigHost })).toThrowError(
+        expect(() => getProgram({ tsConfig: 'tsconfig.bad.json', tsConfigHost })).toThrowError(
             `Error while reading tsconfig.bad.json: '{' expected.`,
         );
         expect(tsConfigHost.readFile).toHaveBeenCalledWith('tsconfig.bad.json');
@@ -24,7 +24,7 @@ describe('applyTypes', () => {
             ...ts.sys,
             readFile: jest.fn(() => '{ "include": 123 }'),
         };
-        expect(() => applyTypes([], { tsConfig: 'tsconfig.invalid.json', tsConfigHost })).toThrowError(
+        expect(() => getProgram({ tsConfig: 'tsconfig.invalid.json', tsConfigHost })).toThrowError(
             `Error while parsing tsconfig.invalid.json: Compiler option 'include' requires a value of type Array.`,
         );
         expect(tsConfigHost.readFile).toHaveBeenCalledWith('tsconfig.invalid.json');

src/compiler-helper.ts

@@ -0,0 +1,38 @@
+import * as ts from 'typescript';
+
+export interface ICompilerOptions {
+    /**
+     * If given, all the file paths in the collected type info will be resolved relative to this directory.
+     */
+    rootDir?: string;
+
+    /**
+     * Path to your project's tsconfig file
+     */
+    tsConfig?: string;
+
+    // You probably never need to touch these two - they are used by the integration tests to setup
+    // a virtual file system for TS:
+    tsConfigHost?: ts.ParseConfigHost;
+    tsCompilerHost?: ts.CompilerHost;
+}
+
+export function getProgram(options: ICompilerOptions) {
+    let program: ts.Program | undefined;
+    if (options.tsConfig) {
+        const configHost = options.tsConfigHost || ts.sys;
+        const { config, error } = ts.readConfigFile(options.tsConfig, configHost.readFile);
+        if (error) {
+            throw new Error(`Error while reading ${options.tsConfig}: ${error.messageText}`);
+        }
+
+        const parsed = ts.parseJsonConfigFileContent(config, configHost, options.rootDir || '');
+        if (parsed.errors.length) {
+            const errors = parsed.errors.map((e) => e.messageText).join(', ');
+            throw new Error(`Error while parsing ${options.tsConfig}: ${errors}`);
+        }
+
+        program = ts.createProgram(parsed.fileNames, parsed.options, options.tsCompilerHost);
+    }
+    return program;
+}

src/instrument.ts

@@ -1,14 +1,18 @@
 import * as ts from 'typescript';
 
+import { getProgram, ICompilerOptions } from './compiler-helper';
 import { applyReplacements, Replacement } from './replacement';
 
-export interface IInstrumentOptions {
+export interface IInstrumentOptions extends ICompilerOptions {
     instrumentCallExpressions: boolean;
+    instrumentImplicitThis: boolean;
 }
 
 export interface IExtraOptions {
     arrow?: boolean;
     parens?: [number, number];
+    thisType?: boolean;
+    comma?: boolean;
 }
 
 function hasParensAroundArguments(node: ts.FunctionLike) {
@@ -25,9 +29,64 @@ function hasParensAroundArguments(node: ts.FunctionLike) {
     }
 }
 
-function visit(node: ts.Node, replacements: Replacement[], fileName: string, options: IInstrumentOptions) {
+function visit(
+    node: ts.Node,
+    replacements: Replacement[],
+    fileName: string,
+    options: IInstrumentOptions,
+    program?: ts.Program,
+    semanticDiagnostics?: ReadonlyArray<ts.Diagnostic>,
+) {
     const isArrow = ts.isArrowFunction(node);
     if (ts.isFunctionDeclaration(node) || ts.isMethodDeclaration(node) || ts.isArrowFunction(node)) {
+        if (node.body) {
+            const needsThisInstrumentation =
+                options.instrumentImplicitThis &&
+                program &&
+                semanticDiagnostics &&
+                semanticDiagnostics.find((diagnostic) => {
+                    if (
+                        diagnostic.code === 2683 &&
+                        diagnostic.file &&
+                        diagnostic.file.fileName === node.getSourceFile().fileName &&
+                        diagnostic.start
+                    ) {
+                        if (node.body && ts.isBlock(node.body)) {
+                            const body = node.body as ts.FunctionBody;
+                            return (
+                                body.statements.find((statement) => {
+                                    return (
+                                        diagnostic.start !== undefined &&
+                                        statement.pos <= diagnostic.start &&
+                                        diagnostic.start <= statement.end
+                                    );
+                                }) !== undefined
+                            );
+                        } else {
+                            const body = node.body as ts.Expression;
+                            return body.pos <= diagnostic.start && diagnostic.start <= body.end;
+                        }
+                    }
+                    return false;
+                }) !== undefined;
+            if (needsThisInstrumentation) {
+                const opts: IExtraOptions = { thisType: true };
+                if (node.parameters.length > 0) {
+                    opts.comma = true;
+                }
+                const params = [
+                    JSON.stringify('this'),
+                    'this',
+                    node.parameters.pos,
+                    JSON.stringify(fileName),
+                    JSON.stringify(opts),
+                ];
+                const instrumentExpr = `$_$twiz(${params.join(',')})`;
+
+                replacements.push(Replacement.insert(node.body.getStart() + 1, `${instrumentExpr};`));
+            }
+        }
+
         const isShortArrow = ts.isArrowFunction(node) && !ts.isBlock(node.body);
         for (const param of node.parameters) {
             if (!param.type && !param.initializer && node.body) {
@@ -105,7 +164,7 @@ function visit(node: ts.Node, replacements: Replacement[], fileName: string, opt
         }
     }
 
-    node.forEachChild((child) => visit(child, replacements, fileName, options));
+    node.forEachChild((child) => visit(child, replacements, fileName, options, program, semanticDiagnostics));
 }
 
 const declaration = `
@@ -118,11 +177,18 @@ const declaration = `
 export function instrument(source: string, fileName: string, options?: IInstrumentOptions) {
     const instrumentOptions: IInstrumentOptions = {
         instrumentCallExpressions: false,
+        instrumentImplicitThis: false,
         ...options,
     };
-    const sourceFile = ts.createSourceFile(fileName, source, ts.ScriptTarget.Latest, true);
+    const program: ts.Program | undefined = getProgram(instrumentOptions);
+    const sourceFile = program
+        ? program.getSourceFile(fileName)
+        : ts.createSourceFile(fileName, source, ts.ScriptTarget.Latest, true);
     const replacements = [] as Replacement[];
-    visit(sourceFile, replacements, fileName, instrumentOptions);
+    if (sourceFile) {
+        const semanticDiagnostics = program ? program.getSemanticDiagnostics(sourceFile) : undefined;
+        visit(sourceFile, replacements, fileName, instrumentOptions, program, semanticDiagnostics);
+    }
     if (replacements.length) {
         replacements.push(Replacement.insert(0, declaration));
     }