@codeyam/codeyam-cli 0.1.0-staging.bbe4da9 → 0.1.0-staging.dd216e0

package/analyzer-template/packages/ai/src/lib/dataStructure/equivalencyManagers/frameworks/ReactFrameworkManager.ts

@@ -17,6 +17,12 @@ interface UseStateInfo {
 export class ReactFrameworkManager implements EquivalencyManager {
   private useStateInfos: UseStateInfo[] = [];
   private useCallbackEquivalents: Record<string, string> = {};
+  // Track variables used as ref props - they should be typed as 'null'
+  // Format: { scopeNodeName -> Set of variable paths }
+  private refVariables: Map<string, Set<string>> = new Map();
+  // Track ref variable source paths for updating external function call schemas in finalize
+  // Format: { functionCallName -> Set of source paths (e.g., 'useAutoAnimate().functionCallReturnValue[]') }
+  private refVariableSources: Map<string, Set<string>> = new Map();
 
   postProcess(scopeNode: ScopeNode, scopeDataStructure: ScopeDataStructure) {
     this.handleUseState(scopeNode, scopeDataStructure);
@@ -139,7 +145,27 @@ export class ReactFrameworkManager implements EquivalencyManager {
     'main',
   ]);
 
-  finalize(_scopeNode: ScopeNode, _scopeDataStructure: ScopeDataStructure) {}
+  finalize(scopeNode: ScopeNode, scopeDataStructure: ScopeDataStructure) {
+    // Set all tracked ref variables to 'null' type
+    // This ensures mock data generation uses null instead of placeholder strings
+    const refVarsForScope = this.refVariables.get(scopeNode.name);
+    if (refVarsForScope) {
+      for (const refVar of refVarsForScope) {
+        scopeNode.schema[refVar] = 'null';
+      }
+    }
+
+    // Update external function call schemas with ref source paths
+    // This runs AFTER schema building is complete, so the null values won't be overwritten
+    for (const efc of scopeDataStructure.externalFunctionCalls) {
+      const refSourcePaths = this.refVariableSources.get(efc.name);
+      if (refSourcePaths) {
+        for (const sourcePath of refSourcePaths) {
+          efc.schema[sourcePath] = 'null';
+        }
+      }
+    }
+  }
 
   private findUseStateInfo(useStateCall: string) {
     return this.useStateInfos.find(
@@ -251,9 +277,9 @@ export class ReactFrameworkManager implements EquivalencyManager {
               equivalentPathParts[0].replace(existingFunctionName, ''),
             ...equivalentPathParts.slice(1),
           ]);
-          for (const equivalentValue of scopeNode.equivalencies[
-            equivalentPath
-          ]) {
+          const equivalenciesToReroute =
+            scopeNode.equivalencies[equivalentPath] ?? [];
+          for (const equivalentValue of equivalenciesToReroute) {
             scopeDataStructure.removeEquivalency(
               equivalentPath,
               equivalentValue,
@@ -273,7 +299,9 @@ export class ReactFrameworkManager implements EquivalencyManager {
           equivalentPath = newEquivalentPath;
         }
 
-        for (const equivalentValue of scopeNode.equivalencies[equivalentPath]) {
+        const equivalenciesToModify =
+          scopeNode.equivalencies[equivalentPath] ?? [];
+        for (const equivalentValue of equivalenciesToModify) {
           if (equivalentValue.scopeNodeName === useCallbackScopeInfoName) {
             modifiedEquivalencies = true;
             equivalentValue.scopeNodeName =
@@ -562,6 +590,102 @@ export class ReactFrameworkManager implements EquivalencyManager {
     for (const [equivalentPath, equivalentValues] of Object.entries(
       scopeNode.equivalencies,
     )) {
+      // Check if the equivalency KEY (e.g., "div().signature[0].ref") ends with .ref
+      // If so, the VALUE (e.g., "parent") should be marked as null
+      const equivalentPathParts = scopeDataStructure.splitPath(equivalentPath);
+      const pathLength = equivalentPathParts.length;
+
+      if (pathLength >= 2) {
+        const secondToLast = equivalentPathParts[pathLength - 2];
+        const last = equivalentPathParts[pathLength - 1];
+
+        if (secondToLast === 'signature[0]' && last === 'ref') {
+          // This is a ref prop usage like div().signature[0].ref -> parent
+          for (const equivalentValue of equivalentValues) {
+            propagated = true;
+
+            // Track this variable as a ref - we'll set its type to 'null' in finalize()
+            // after the schema is built (setting directly here would be overwritten)
+            // IMPORTANT: The ref variable might be defined in a parent scope, not the
+            // current scope. We need to find where the variable is actually defined by
+            // walking up the scope tree.
+            const varName = equivalentValue.schemaPath;
+            let targetScope: ScopeNode | null = scopeNode;
+
+            // Walk up the scope tree to find where this variable is defined
+            while (targetScope) {
+              // Check if this scope has the variable in its schema or equivalencies
+              const hasVariable =
+                varName in (targetScope.schema || {}) ||
+                varName in (targetScope.equivalencies || {});
+
+              if (hasVariable) {
+                if (!this.refVariables.has(targetScope.name)) {
+                  this.refVariables.set(targetScope.name, new Set());
+                }
+                this.refVariables.get(targetScope.name)!.add(varName);
+
+                // IMPORTANT: Also update the source of the ref variable in the external
+                // function call's schema. This ensures dataForMocks produces [null]
+                // instead of ["string"] for hooks like useAutoAnimate.
+                const varEquivalencies = targetScope.equivalencies?.[varName];
+                if (varEquivalencies) {
+                  for (const varEquivalency of varEquivalencies) {
+                    const sourcePath = varEquivalency.schemaPath;
+                    // Check if this comes from an external function call
+                    // (contains functionCallReturnValue which indicates it's a hook/function return)
+                    if (sourcePath.includes('functionCallReturnValue')) {
+                      // Store the source path for later - we'll update the external function
+                      // call's schema in finalize() which runs AFTER schema building is complete.
+                      // Setting it here would be overwritten by later schema building.
+                      for (const efc of scopeDataStructure.externalFunctionCalls) {
+                        if (sourcePath.startsWith(efc.name)) {
+                          if (!this.refVariableSources.has(efc.name)) {
+                            this.refVariableSources.set(efc.name, new Set());
+                          }
+                          // Store both the specific indexed path and the generic [] path
+                          this.refVariableSources
+                            .get(efc.name)!
+                            .add(sourcePath);
+                          const genericPath = sourcePath.replace(
+                            /\[\d+\]/g,
+                            '[]',
+                          );
+                          if (genericPath !== sourcePath) {
+                            this.refVariableSources
+                              .get(efc.name)!
+                              .add(genericPath);
+                          }
+                        }
+                      }
+                    }
+                  }
+                }
+              }
+
+              // Move up to parent scope
+              if (targetScope.tree && targetScope.tree.length > 0) {
+                const parentScopeName =
+                  targetScope.tree[targetScope.tree.length - 1];
+                targetScope = scopeDataStructure.getScopeOrFunctionCallInfo(
+                  parentScopeName,
+                ) as ScopeNode | null;
+              } else {
+                targetScope = null;
+              }
+            }
+
+            scopeDataStructure.removeEquivalency(
+              equivalentPath,
+              equivalentValue,
+              scopeNode,
+            );
+          }
+          continue;
+        }
+      }
+
+      // Original logic for checking the value side
       for (const equivalentValue of equivalentValues) {
         const equivalentValueParts = scopeDataStructure.splitPath(
           equivalentValue.schemaPath,

package/analyzer-template/packages/ai/src/lib/dataStructure/helpers/cleanKnownObjectFunctions.ts

@@ -64,6 +64,14 @@ const protoMethodSets: Record<PrimitiveKind, Set<string>> = {
   ),
 };
 
+/** Built-in properties (not methods) that should be removed from primitives.
+ * These are properties that exist on the prototype and shouldn't appear
+ * as schema entries (they cause malformed mock data structures). */
+const protoPropertySets: Partial<Record<PrimitiveKind, Set<string>>> = {
+  array: new Set(['length']),
+  string: new Set(['length']),
+};
+
 /**
  * Extract primitive kind from a type string, handling union types.
  * E.g., "number | undefined" -> "number", "string | null" -> "string"
@@ -118,11 +126,21 @@ function scrub(
       const rootKind = rootTable.get(prefix);
       if (!rootKind) continue;
 
-      const rawMethod = keyParts[cut].split('(')[0];
+      const rawToken = keyParts[cut];
+      const rawMethod = rawToken.split('(')[0];
+
+      // Check if it's a prototype method (e.g., map, filter, toString)
       if (protoMethodSets[rootKind].has(rawMethod)) {
         onRemove(keyPath);
         break;
       }
+
+      // Check if it's a built-in property (e.g., length on arrays/strings)
+      const propertySet = protoPropertySets[rootKind];
+      if (propertySet && propertySet.has(rawToken)) {
+        onRemove(keyPath);
+        break;
+      }
     }
   }
 }

package/analyzer-template/packages/ai/src/lib/dataStructure/helpers/cleanNonObjectFunctions.ts

@@ -702,13 +702,16 @@ function handleFunctionMapping(
     }
 
     // Don't add [] element for ambiguous methods without other array evidence
+    // Also don't add [] if the path was already identified as a string type
     const methodName = getMethodName(lastPart);
     const isAmbiguousMethod = ambiguousArrayMethodsSet.has(methodName);
+    const isStringType = keyMapping[previousPath] === 'string';
     if (
       !isValidFunctionCall(previousPart) &&
       arrayMethodsSet.has(methodName) &&
       !keyMapping[previousPath + '[]'] &&
-      !isAmbiguousMethod
+      !isAmbiguousMethod &&
+      !isStringType
     ) {
       keyMapping[previousPath + '[]'] = 'unknown';
     }

package/analyzer-template/packages/ai/src/lib/dataStructure/helpers/convertDotNotation.ts

@@ -26,7 +26,6 @@ export default function convertDotNotation(
     const parts = splitOutsideParenthesesAndArrays(rawPath);
     if (!parts || parts.length === 0) continue;
 
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
     let cursor: any = result;
 
     for (let i = 0; i < parts.length; i++) {
@@ -181,8 +180,9 @@ export default function convertDotNotation(
           if (!isSkippableLeafType(typ)) {
             cursor[key] = [typ];
           } else {
-            // structural/placeholder subtype — ensure array exists but don't clobber
-            if (!Array.isArray(cursor[key])) cursor[key] = [];
+            // structural/placeholder subtype — ensure array exists but don't clobber existing values
+            // Only create an array if nothing exists yet; preserve concrete types like 'string'
+            if (cursor[key] === undefined) cursor[key] = [];
           }
           i++; // consume '[]'
         } else {

package/analyzer-template/packages/ai/src/lib/dataStructure/helpers/fillInSchemaGapsAndUnknowns.ts

@@ -377,6 +377,34 @@ export function fillInDirectSchemaGapsAndUnknowns({
       });
     }
 
+    // Remove .length paths when the parent is typed as 'array' or 'string' (or nullable variants).
+    // These are built-in properties that were used for type inference but shouldn't appear
+    // in the final schema (they create { length: N } objects instead of proper arrays).
+    // However, if the parent could be an 'object' or other type, keep .length as it may be
+    // a custom property.
+    for (const key of Object.keys(schema)) {
+      const parts = splitOutsideParenthesesAndArrays(key);
+      if (parts[parts.length - 1] === 'length') {
+        const parentPath = joinParenthesesAndArrays(parts.slice(0, -1));
+        const parentType = schema[parentPath];
+        if (parentType) {
+          // Split union types and filter out nullable parts
+          const typeParts = parentType.split(' | ').map((t) => t.trim());
+          const nonNullableTypes = typeParts.filter(
+            (t) => t !== 'undefined' && t !== 'null',
+          );
+          // Only remove if ALL non-nullable types are array or string
+          // This ensures we don't remove .length from potential custom objects
+          if (
+            nonNullableTypes.length > 0 &&
+            nonNullableTypes.every((t) => t === 'array' || t === 'string')
+          ) {
+            delete schema[key];
+          }
+        }
+      }
+    }
+
     return schema;
   } catch (error) {
     console.info(
@@ -454,6 +482,8 @@ function checkIfKnownType(
   let isString: boolean | undefined;
   let hasStrongArrayEvidence = false;
   let hasOnlyAmbiguousArrayMethods = true;
+  let hasLengthAccess = false;
+  let hasArrayStringAmbiguousMethod = false;
 
   for (const functionOrAttributeName of functionAndAttributesKeysMapping[key] ??
     []) {
@@ -467,11 +497,18 @@ function checkIfKnownType(
       knownMethodCalls(functionOrAttributeName).includes('string') ||
       functionOrAttributeName === 'length';
 
-    // Check if this is strong evidence of an array
-    if (
-      STRONG_ARRAY_METHODS.has(methodName) ||
-      functionOrAttributeName === 'length'
-    ) {
+    // Track .length access
+    if (functionOrAttributeName === 'length') {
+      hasLengthAccess = true;
+    }
+
+    // Track methods that exist on both arrays and strings (like includes, indexOf)
+    if (couldBeArray && couldBeString && functionOrAttributeName !== 'length') {
+      hasArrayStringAmbiguousMethod = true;
+    }
+
+    // Check if this is strong evidence of an array (unambiguous array methods)
+    if (STRONG_ARRAY_METHODS.has(methodName)) {
       hasStrongArrayEvidence = true;
     }
 
@@ -495,10 +532,24 @@ function checkIfKnownType(
     }
   }
 
+  // Strong array evidence (unambiguous methods like .map, .filter) -> definitely array
+  if (isArray && hasStrongArrayEvidence) {
+    return 'array';
+  }
+
+  // When we have BOTH .length AND an ambiguous method like .includes(),
+  // prefer array. This pattern (checking .length AND calling .includes)
+  // is very common with arrays (e.g., arr.length > 0 && arr.includes(x))
+  // but rare with strings (you'd typically just call str.includes(x))
+  if (isArray && hasLengthAccess && hasArrayStringAmbiguousMethod) {
+    return 'array';
+  }
+
   if (isString) return 'string';
+
   // Only infer array if we have strong evidence OR not only ambiguous methods
   // This prevents router.push() from being incorrectly inferred as an array
-  if (isArray && (hasStrongArrayEvidence || !hasOnlyAmbiguousArrayMethods)) {
+  if (isArray && !hasOnlyAmbiguousArrayMethods) {
     return 'array';
   }
 }

package/analyzer-template/packages/ai/src/lib/dataStructure/helpers/getFunctionCallRoot.ts

@@ -1,11 +1,41 @@
 import { splitOutsideParenthesesAndArrays } from '../../splitOutsideParentheses';
 
+/**
+ * Find the position of the first '(' that is NOT inside angle brackets.
+ * Returns -1 if not found.
+ */
+function findFirstParenOutsideTypeParams(signature: string): number {
+  let angleDepth = 0;
+
+  for (let i = 0; i < signature.length; i++) {
+    const char = signature[i];
+
+    if (char === '<') {
+      angleDepth++;
+    } else if (char === '>') {
+      angleDepth--;
+    } else if (char === '(' && angleDepth === 0) {
+      return i;
+    }
+  }
+
+  return -1;
+}
+
 export default function getFunctionCallRoot(functionCallSignature: string) {
   const functionCallSignatureParts = splitOutsideParenthesesAndArrays(
     functionCallSignature,
   );
 
-  return functionCallSignatureParts[0]
-    .replace(/(cyScope\d+)F/, '$1')
-    .split('(')[0];
+  const firstPart = functionCallSignatureParts[0];
+
+  // Find the first '(' that's not inside type parameters (angle brackets)
+  // This handles cases like "useFetcher<{ data: UserData | null }>()"
+  // which should return "useFetcher<{ data: UserData | null }>"
+  // (preserving type params for scope node matching)
+  const parenPos = findFirstParenOutsideTypeParams(firstPart);
+
+  const result = parenPos >= 0 ? firstPart.substring(0, parenPos) : firstPart;
+
+  return result.replace(/(cyScope\d+)F/, '$1');
 }

package/analyzer-template/packages/ai/src/lib/gatherRelevantDependentKeyAttributes.ts

@@ -1,4 +1,9 @@
-import { Entity, ReadonlyAnalysis, ReadonlyAnalysisMap } from '~codeyam/types';
+import {
+  DataStructure,
+  Entity,
+  ReadonlyAnalysis,
+  ReadonlyAnalysisMap,
+} from '~codeyam/types';
 import {
   joinParenthesesAndArrays,
   splitOutsideParenthesesAndArrays,
@@ -7,6 +12,7 @@ import { gatherAttributesMap } from './promptGenerators/gatherAttributesMap';
 
 export default function gatherRelevantDependentKeyAttributes(
   entity: Pick<Entity, 'name' | 'metadata'>,
+  mergedDataStructure: Omit<DataStructure, 'equivalentSignatureVariables'>,
   dependentAnalyses: ReadonlyAnalysisMap,
 ) {
   // console.info(
@@ -52,10 +58,16 @@ export default function gatherRelevantDependentKeyAttributes(
   //   ),
   // );
 
-  const { associationMap } = gatherAttributesMap(entity);
+  const isolatedDataStructure = entity.metadata?.isolatedDataStructure;
+  const equivalentSignatureVariables =
+    isolatedDataStructure?.equivalentSignatureVariables ?? {};
+  const { associationMap } = gatherAttributesMap(
+    isolatedDataStructure ?? mergedDataStructure,
+    mergedDataStructure,
+    equivalentSignatureVariables,
+  );
 
-  const dependencySchemas =
-    entity.metadata?.isolatedDataStructure?.dependencySchemas;
+  const dependencySchemas = mergedDataStructure?.dependencySchemas;
   const relevantDependentKeyAttributes: ReadonlyAnalysis['metadata']['keyAttributes'][0][] =
     [];
   Object.keys(dependentAnalyses).forEach((path) => {

package/analyzer-template/packages/ai/src/lib/generateChangesEntityKeyAttributes.ts

@@ -29,8 +29,17 @@ export default async function generateChangesEntityKeyAttributes({
     existingKeyAttributes: existingKeyAttributes?.map((a) => a.externalPath),
   });
 
+  const isolatedDataStructure = entity.metadata?.isolatedDataStructure;
+  const mergedDataStructure = analysis.metadata.mergedDataStructure;
+  const equivalentSignatureVariables =
+    isolatedDataStructure?.equivalentSignatureVariables ?? {};
   const { attributesMap, associationMap, dataStructurePathMap } =
-    gatherAttributesMap(entity, true);
+    gatherAttributesMap(
+      isolatedDataStructure ?? mergedDataStructure,
+      mergedDataStructure,
+      equivalentSignatureVariables,
+      true,
+    );
 
   const prompt = generateEntityKeyAttributesGenerator({
     code: entity.code,
@@ -102,6 +111,9 @@ export default async function generateChangesEntityKeyAttributes({
       description: addedKeyAttribute.description,
       validValueOptions: addedKeyAttribute.validValueOptions,
       errorValueOptions: addedKeyAttribute.errorValueOptions,
+      // Preserve new fields from LLM response
+      valueType: addedKeyAttribute.valueType,
+      dependencies: addedKeyAttribute.dependencies,
     };
     if (index !== -1) {
       keyAttributes.splice(index, 1, fullKeyAttribute);
@@ -167,6 +179,14 @@ Note: A conditional that shows a different UI state (e.g., loading screen, redir
 
 Only return attributes from the "Attributes and Value Types" list—these are the only values we can control. Use the exact path provided. Use \`functionCallReturnValue\` for function return values. Use \`[]\` for array items (e.g., \`items[]\`).
 
+## Dependencies
+
+If this key attribute is only used if one or more other attributes have specific values, include those dependencies in a \`dependencies\` array.
+
+## Value Type
+
+Should reiterate the type of the attribute (e.g., "string", "object", "array", "boolean", etc.) based on the "Attributes and Value Types" list.
+
 ## Response Format
 
 For \`validValueOptions\`, include boundary/extreme values that trigger different behavior—but NOT values that would cause errors.
@@ -177,6 +197,10 @@ For \`validValueOptions\`, include boundary/extreme values that trigger differen
     {
       "path": "user",
       "description": "Controls whether the main form or redirect screen displays",
+      "dependencies": [
+        { "projectId": "must be a valid project ID for 'user' to have effect" }
+      ],
+      "valueType": "object",
       "validValueOptions": ["truthy user object", "null/undefined"]
     }
   ],
@@ -214,6 +238,14 @@ NOT key attributes:
 
 Only return attributes from the "Attributes and Value Types" list—these are the only values we can control. Use the exact path provided. Use \`functionCallReturnValue\` for function return values. Use \`[]\` for array items (e.g., \`items[]\`).
 
+## Dependencies
+
+If this key attribute is only used if one or more other attributes have specific values, include those dependencies in a \`dependencies\` array.
+
+## Value Type
+
+Should reiterate the type of the attribute (e.g., "string", "object", "array", "boolean", etc.) based on the "Attributes and Value Types" list.
+
 ## Response Format
 
 For \`validValueOptions\`, include boundary/extreme values that trigger different behavior—but NOT values that would cause errors.
@@ -224,6 +256,10 @@ For \`validValueOptions\`, include boundary/extreme values that trigger differen
     {
       "path": "config.enabled",
       "description": "Controls whether the main processing logic runs",
+      "dependencies": [
+        { "analysisMode": "must be 'full' for 'enabled' to have effect" }
+      ],
+      "valueType": "boolean",
       "validValueOptions": ["true", "false"]
     }
   ],

package/analyzer-template/packages/ai/src/lib/generateChangesEntityScenarioData.ts

@@ -328,7 +328,7 @@ const generateChangesSystemMessage = (
 Preserve existing data and only modify what's necessary for the new structure.
 - If structure unchanged and data sufficient, return empty data: \`{ "data": {} }\`
 - If new fields added, fill them with realistic values
-- Fill ALL fields with realistic values (except error attributes)
+- Fill ALL fields with realistic values (except error attributes and key attributes set to null or undefined)
 - NEVER include "error" fields—omit entirely`
     : `## Non-Default Scenario (Incremental Update)
 Preserve existing data and only modify what changed.
@@ -365,8 +365,11 @@ Use simple elements only (\`<div>\`, \`<span>\`). No custom components.
 { "onClick": { "~~codeyam-code~~": "() => console.log('clicked')" } }
 \`\`\`
 
-## Function Keys
-Preserve function call syntax exactly as written in the structure.
+## Mock Data Keys
+Preserve keys exactly as written in the structure. There are two formats:
+- Standard function calls: \`"useUser()"\`, \`"from().select()"\`
+- Variable-qualified calls (for multiple calls to the same function): \`"entityDiffFetcher <- useFetcher"\`
+  This reads as "entityDiffFetcher receives from useFetcher". Each variable gets distinct mock data.
 
 ## Response Format
 \`\`\`json

package/analyzer-template/packages/ai/src/lib/generateChangesEntityScenarios.ts

@@ -46,6 +46,7 @@ export default async function generateChangesEntityScenarios({
 }: GenerateChangesEntityScenariosArgs) {
   const relevantDependentKeyAttributes = gatherRelevantDependentKeyAttributes(
     entity,
+    analysis.metadata.mergedDataStructure,
     dependentAnalyses,
   );
   if (
@@ -71,6 +72,7 @@ export default async function generateChangesEntityScenarios({
 
   const prompt = generateChangesEntityScenariosGenerator({
     entity,
+    mergedDataStructure: analysis.metadata.mergedDataStructure,
     keyAttributes,
     addedKeyAttributes,
     removedKeyAttributes,

package/analyzer-template/packages/ai/src/lib/generateEntityDataStructure.ts

@@ -3,6 +3,7 @@ import isolateScopes from './isolateScopes';
 import analyzeScope from './analyzeScope';
 import { FileAnalyzer } from '~codeyam/analyze';
 import { AI, SerializableDataStructure } from '~codeyam/ai';
+import { resetScopeDataStructureMetrics } from './dataStructure/ScopeDataStructure';
 
 // import { awsLog } from '~codeyam/utils';
 
@@ -17,15 +18,23 @@ export default async function generateEntityDataStructure({
   fileAnalyzer,
   model,
 }: GenerateEntityDataStructureArgs): Promise<SerializableDataStructure> {
+  // Reset metrics counters to prevent false "infinite loop" detection when
+  // processing multiple entities in the same run
+  resetScopeDataStructureMetrics();
+
   // awsLog('CodeYam: Generating entity data structure', {
   //   filePath: entity.filePath,
   //   entityName: entity.name,
   //   // code: entity.code.length,
   // });
 
-  const fileScope = isolateScopes(fileAnalyzer.sourceFile.getText(), {
+  // Use getFullText() to include leading trivia (comments/whitespace) so that
+  // positions from isolateScopes match AST node positions which include trivia
+  const fileText = fileAnalyzer.sourceFile.getFullText();
+  const fileScope = isolateScopes(fileText, {
     entityNames: [entity.name],
   });
+
   let entityScope =
     fileScope.name === entity.name
       ? fileScope
@@ -39,12 +48,10 @@ export default async function generateEntityDataStructure({
   }
 
   if (!entityScope) {
-    if (entity.code === fileAnalyzer.sourceFile.getText()) {
+    if (entity.code === fileText) {
       entityScope = fileScope;
     } else {
-      const entityStartIndex = fileAnalyzer.sourceFile
-        .getText()
-        .indexOf(entity.code);
+      const entityStartIndex = fileText.indexOf(entity.code);
       const entityIndex = entityStartIndex + entity.code.length / 2;
       entityScope = Object.values(fileScope.childScopes).find(
         (s) => entityIndex > s.start && entityIndex < s.end,