From 5e8f349437059ae4562cc0d3ae529b1fb3201471 Mon Sep 17 00:00:00 2001
From: Dawson Booth <dawson@dawsonbooth.com>
Date: Fri, 17 Jan 2025 15:14:18 -0600
Subject: [PATCH] fix(core): handle parameter counting in custom instance
 returning useCallback

---
 packages/core/src/generators/mutator.ts | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/packages/core/src/generators/mutator.ts b/packages/core/src/generators/mutator.ts
index 9f9b2a3cf..b3f8e1d50 100644
--- a/packages/core/src/generators/mutator.ts
+++ b/packages/core/src/generators/mutator.ts
@@ -259,11 +259,23 @@ const parseFunction = (
       (b: any) => b.type === 'ReturnStatement',
     );
 
+    // If the function directly returns an arrow function
     if (returnStatement?.argument?.params) {
       return {
         numberOfParams: node.params.length,
         returnNumberOfParams: returnStatement.argument.params.length,
       };
+      // If the function returns a CallExpression (e.g., return useCallback(...))
+    } else if (
+      returnStatement?.argument?.type === 'CallExpression' &&
+      returnStatement.argument.arguments?.[0]?.type ===
+        'ArrowFunctionExpression'
+    ) {
+      const arrowFn = returnStatement.argument.arguments[0];
+      return {
+        numberOfParams: node.params.length,
+        returnNumberOfParams: arrowFn.params.length,
+      };
     }
     return {
       numberOfParams: node.params.length,
@@ -292,6 +304,15 @@ const parseFunction = (
       numberOfParams: declaration.init.params.length,
       returnNumberOfParams: returnStatement.argument.params.length,
     };
+  } else if (
+    returnStatement?.argument?.type === 'CallExpression' &&
+    returnStatement.argument.arguments?.[0]?.type === 'ArrowFunctionExpression'
+  ) {
+    const arrowFn = returnStatement.argument.arguments[0];
+    return {
+      numberOfParams: declaration.init.params.length,
+      returnNumberOfParams: arrowFn.params.length,
+    };
   }
 
   return {