feat: support function expressions (#177)

## PR Checklist - [x] Addresses an existing open issue: fixes #176 - [x] That issue was marked as [`status: accepting prs`](https://github.com/JoshuaKGoldberg/ts-function-inliner/issues?q=is%3Aopen+is%3Aissue+label%3A%22status%3A+accepting+prs%22) - [x] Steps in [CONTRIBUTING.md](https://github.com/JoshuaKGoldberg/ts-function-inliner/blob/main/.github/CONTRIBUTING.md) were taken ## Overview Refactors the retrieval of function declarations to a separate `getFunctionDeclarationFromCall`.
JoshuaKGoldberg · Jun 3, 2024 · 2ccbf68 · 2ccbf68
1 parent b76855b
commit 2ccbf68
Show file tree

Hide file tree

Showing 7 changed files with 171 additions and 111 deletions.
diff --git a/src/getFunctionDeclarationFromCall.ts b/src/getFunctionDeclarationFromCall.ts
@@ -0,0 +1,26 @@
+import ts from "typescript";
+
+import { isSmallFunctionLike } from "./isSmallFunctionLike.js";
+
+export function getFunctionDeclarationFromCall(
+	node: ts.CallExpression,
+	typeChecker: ts.TypeChecker,
+) {
+	let declaration = typeChecker.getSymbolAtLocation(node.expression)
+		?.valueDeclaration;
+
+	if (!declaration) {
+		return undefined;
+	}
+
+	if (
+		ts.isPropertyAssignment(declaration) ||
+		ts.isShorthandPropertyAssignment(declaration)
+	) {
+		declaration =
+			typeChecker.getTypeAtLocation(declaration).getSymbol()
+				?.valueDeclaration ?? declaration;
+	}
+
+	return isSmallFunctionLike(declaration) ? declaration : undefined;
+}
diff --git a/src/isFunctionDeclarationWithBody.ts b/src/isFunctionDeclarationWithBody.ts
@@ -1,9 +1,12 @@
 import ts from "typescript";
 
-import { FunctionDeclarationWithBody } from "./types.js";
+import { FunctionLikeWithBody } from "./types.js";
 
-export const isFunctionDeclarationWithBody = (
+export const isFunctionWithBody = (
 	node: ts.Node,
-): node is FunctionDeclarationWithBody => {
-	return ts.isFunctionDeclaration(node) && !!node.body;
+): node is FunctionLikeWithBody => {
+	return (
+		(ts.isFunctionDeclaration(node) || ts.isFunctionExpression(node)) &&
+		!!node.body
+	);
 };
diff --git a/src/isSmallFunctionLikeDeclaration.ts → src/isSmallFunctionLike.ts b/src/isSmallFunctionLikeDeclaration.ts → src/isSmallFunctionLike.ts
@@ -1,16 +1,13 @@
 import ts from "typescript";
 
 import { CollectedValue, collectValue } from "./collectValue.js";
-import { isFunctionDeclarationWithBody } from "./isFunctionDeclarationWithBody.js";
-import { SmallFunctionLikeDeclaration } from "./types.js";
+import { isFunctionWithBody } from "./isFunctionDeclarationWithBody.js";
+import { SmallFunctionLikeWithBody } from "./types.js";
 
-export const isSmallFunctionLikeDeclaration = (
+export const isSmallFunctionLike = (
 	node: ts.Node,
-): node is SmallFunctionLikeDeclaration => {
-	if (
-		!isFunctionDeclarationWithBody(node) ||
-		node.body.statements.length !== 1
-	) {
+): node is SmallFunctionLikeWithBody => {
+	if (!isFunctionWithBody(node) || node.body.statements.length !== 1) {
 		return false;
 	}
 

diff --git a/src/transformToInline.ts b/src/transformToInline.ts
@@ -1,14 +1,14 @@
 import ts from "typescript";
 
-import { SmallFunctionLikeDeclaration } from "./types.js";
+import { SmallFunctionLikeWithBody } from "./types.js";
 
 /**
  * Inlines a small function declaration into a call to that function,
  * replacing parameter identifiers in the body with the call's arguments.
  */
 export const transformToInline = (
 	callExpression: ts.CallExpression,
-	declaration: SmallFunctionLikeDeclaration,
+	declaration: SmallFunctionLikeWithBody,
 	context: ts.TransformationContext,
 ) => {
 	const parameters = new Map(

diff --git a/src/transformerProgram.test.ts b/src/transformerProgram.test.ts
@@ -59,91 +59,141 @@ function expectResultToBe(actual: string, expected: string) {
 }
 
 describe("transformerProgram", () => {
-	test("BinaryExpression", () => {
-		const result = getResult(`
-			function addToLength(base: string) {
-				return base.length + 3;
-			}
-			
-			addToLength("abc");
-		`);
-
-		expectResultToBe(
-			result,
-			`
-				function addToLength(base) {
+	describe("function contents", () => {
+		test("BinaryExpression", () => {
+			const result = getResult(`
+				function addToStringLength(base: string) {
 					return base.length + 3;
 				}
 				
-				"abc".length + 3;
-			`,
-		);
-	});
-
-	test("PostfixUnaryExpression", () => {
-		const result = getResult(`
-			function incrementCount(count: number) {
-				return count++;
-			}
-			
-			const value = 123;
-			incrementCount(value);
-		`);
-
-		expectResultToBe(
-			result,
-			`
-				function incrementCount(count) {
+				addToStringLength("abc");
+			`);
+
+			expectResultToBe(
+				result,
+				`
+					function addToStringLength(base) {
+						return base.length + 3;
+					}
+					
+					"abc".length + 3;
+				`,
+			);
+		});
+
+		test("PostfixUnaryExpression", () => {
+			const result = getResult(`
+				function incrementCount(count: number) {
 					return count++;
 				}
 				
 				const value = 123;
-				value++;
-			`,
-		);
-	});
-
-	test("PrefixUnaryExpression", () => {
-		const result = getResult(`
-			function isNotEmpty(text: string) {
-				return !!text.length;
-			}
-			
-			isNotEmpty("Boo! 👻");
-		`);
-
-		expectResultToBe(
-			result,
-			`
-				function isNotEmpty(text) {
+				incrementCount(value);
+			`);
+
+			expectResultToBe(
+				result,
+				`
+					function incrementCount(count) {
+						return count++;
+					}
+					
+					const value = 123;
+					value++;
+				`,
+			);
+		});
+
+		test("PrefixUnaryExpression", () => {
+			const result = getResult(`
+				function isNotEmpty(text: string) {
 					return !!text.length;
 				}
 				
-				!!"Boo! 👻".length;
-			`,
-		);
+				isNotEmpty("Boo! 👻");
+			`);
+
+			expectResultToBe(
+				result,
+				`
+					function isNotEmpty(text) {
+						return !!text.length;
+					}
+					
+					!!"Boo! 👻".length;
+				`,
+			);
+		});
 	});
 
-	test("User-defined type guard", () => {
-		const result = getResult(`
-			function isDefined<T extends string>(input: T | undefined): input is T {
-				return input !== undefined;
-			}
-
-			isDefined("");
-			isDefined(undefined);
-		`);
-
-		expectResultToBe(
-			result,
-			`
-				function isDefined(input) {
+	test("function kind", () => {
+		test("FunctionExpression in object property", () => {
+			const result = getResult(`
+				const Utils = {
+					isNotEmpty: function (text: string) {
+						return !!text.length;
+					}
+				}
+				
+				Utils.isNotEmpty("Boo! 👻");
+			`);
+
+			expectResultToBe(
+				result,
+				`
+					const Utils = {
+						isNotEmpty: function (text) {
+							return !!text.length;
+						}
+					}
+					
+					!!"Boo! 👻".length;
+				`,
+			);
+		});
+
+		test("FunctionExpression in variable", () => {
+			const result = getResult(`
+				const isNotEmpty = function (text: string) {
+					return !!text.length;
+				}
+				
+				isNotEmpty("Boo! 👻");
+			`);
+
+			expectResultToBe(
+				result,
+				`
+					const isNotEmpty = function (text) {
+						return !!text.length;
+					}
+					
+					!!"Boo! 👻".length;
+				`,
+			);
+		});
+
+		test("FunctionDeclaration with a user-defined type guard", () => {
+			const result = getResult(`
+				function inputValueIsNotUndefined<T extends string>(input: T | undefined): input is T {
 					return input !== undefined;
 				}
 
-				"" !== undefined;
-				undefined !== undefined;
-			`,
-		);
+				inputValueIsNotUndefined("");
+				inputValueIsNotUndefined(undefined);
+			`);
+
+			expectResultToBe(
+				result,
+				`
+					function inputValueIsNotUndefined(input) {
+						return input !== undefined;
+					}
+
+					"" !== undefined;
+					undefined !== undefined;
+				`,
+			);
+		});
 	});
 });
diff --git a/src/transformerProgram.ts b/src/transformerProgram.ts
@@ -1,8 +1,7 @@
 import ts from "typescript";
 
-import { isSmallFunctionLikeDeclaration } from "./isSmallFunctionLikeDeclaration.js";
+import { getFunctionDeclarationFromCall } from "./getFunctionDeclarationFromCall.js";
 import { transformToInline } from "./transformToInline.js";
-import { SmallFunctionLikeDeclaration } from "./types.js";
 
 export const transformerProgram = (program: ts.Program) => {
 	const transformerFactory: ts.TransformerFactory<ts.SourceFile> = (
@@ -11,28 +10,13 @@ export const transformerProgram = (program: ts.Program) => {
 		return (sourceFile) => {
 			const typeChecker = program.getTypeChecker();
 
-			const getDeclarationToInline = (
-				node: ts.CallExpression,
-			): SmallFunctionLikeDeclaration | undefined => {
-				const declaration = typeChecker.getSymbolAtLocation(node.expression)
-					?.valueDeclaration;
-
-				return declaration && isSmallFunctionLikeDeclaration(declaration)
-					? declaration
-					: undefined;
-			};
-
 			const visitor = (node: ts.Node): ts.Node => {
-				if (ts.isCallExpression(node)) {
-					const functionDeclaration = getDeclarationToInline(node);
-					if (functionDeclaration) {
-						const result = transformToInline(
-							node,
-							functionDeclaration,
-							context,
-						);
-						return result;
-					}
+				const functionDeclaration =
+					ts.isCallExpression(node) &&
+					getFunctionDeclarationFromCall(node, typeChecker);
+
+				if (functionDeclaration) {
+					return transformToInline(node, functionDeclaration, context);
 				}
 
 				return ts.visitEachChild(node, visitor, context);

diff --git a/src/types.ts b/src/types.ts
@@ -3,12 +3,12 @@ import ts from "typescript";
 export type MakeRequired<T, K extends keyof T> = Omit<T, K> &
 	Required<Pick<T, K>>;
 
-export type FunctionDeclarationWithBody = MakeRequired<
-	ts.FunctionDeclaration,
+export type FunctionLikeWithBody = MakeRequired<
+	ts.FunctionDeclaration | ts.FunctionExpression,
 	"body"
 >;
 
-export type SmallFunctionLikeDeclaration = ts.FunctionDeclaration & {
+export type SmallFunctionLikeWithBody = FunctionLikeWithBody & {
 	body: {
 		statements: [MakeRequired<ts.ReturnStatement, "expression">];
 	};