FindInMap Bug Fix (#60)

* FindInMap bug fix, removed .todo from a now passing test

* Removed unecessary debug logs

* FindInMap bug fix, removed .todo from a now passing test

* Removed unecessary debug logs

* Move getSecondLevelKeysDynamic into getSecondLevelKeys

* fix lint errors

Author: gemammercado

src/autocomplete/IntrinsicFunctionArgumentCompletionProvider.ts

@@ -71,6 +71,7 @@ export class IntrinsicFunctionArgumentCompletionProvider implements CompletionPr
                 provider: 'IntrinsicFunctionArgument',
                 context: context.record(),
                 intrinsicFunction: intrinsicFunction.type,
+                args: intrinsicFunction.args,
             },
             'Processing intrinsic function argument completion request',
         );
@@ -530,38 +531,98 @@ export class IntrinsicFunctionArgumentCompletionProvider implements CompletionPr
         context: Context,
     ): CompletionItem[] | undefined {
         // Validate arguments structure for second-level keys
-        if (!Array.isArray(args) || args.length < 2 || typeof args[0] !== 'string' || typeof args[1] !== 'string') {
+        if (!this.isValidSecondLevelKeyArgs(args)) {
             log.debug('Invalid arguments for second-level key completions');
             return undefined;
         }
 
         try {
             const mappingName = args[0];
-            const topLevelKey = args[1];
+            const topLevelKey = args[1] as string | { Ref: unknown } | { '!Ref': unknown };
 
             const mappingEntity = this.getMappingEntity(mappingsEntities, mappingName);
             if (!mappingEntity) {
                 log.debug(`Mapping entity not found: ${mappingName}`);
                 return undefined;
             }
 
-            const secondLevelKeys = mappingEntity.getSecondLevelKeys(topLevelKey);
+            const secondLevelKeys = this.getSecondLevelKeysForTopLevelKey(mappingEntity, topLevelKey);
             if (secondLevelKeys.length === 0) {
-                log.debug(`No second-level keys found for mapping: ${mappingName}, top-level key: ${topLevelKey}`);
+                log.debug(`No second-level keys found for mapping: ${mappingName}`);
                 return undefined;
             }
 
             const items = secondLevelKeys.map((key) =>
                 createCompletionItem(key, CompletionItemKind.EnumMember, { context }),
             );
 
-            return context.text.length > 0 ? this.fuzzySearch(items, context.text) : items;
+            return this.filterSecondLevelKeyItems(items, context, topLevelKey);
         } catch (error) {
-            log.error({ error }, 'Error creating second-level key completions');
+            log.debug({ error }, 'Error creating second-level key completions');
             return undefined;
         }
     }
 
+    private isValidSecondLevelKeyArgs(args: unknown): args is [string, string | object] {
+        if (!Array.isArray(args) || args.length < 2 || typeof args[0] !== 'string') {
+            return false;
+        }
+
+        // Second argument valid if it is a string i.e. 'us-east-1' or object '{Ref: AWS::Region}' or '{!Ref: AWS::Region}'
+        return typeof args[1] === 'string' || this.isRefObject(args[1]);
+    }
+
+    private getSecondLevelKeysForTopLevelKey(
+        mappingEntity: Mapping,
+        topLevelKey: string | { Ref: unknown } | { '!Ref': unknown },
+    ): string[] {
+        if (typeof topLevelKey === 'string') {
+            return mappingEntity.getSecondLevelKeys(topLevelKey);
+        } else {
+            // For dynamic references, get all possible keys
+            return mappingEntity.getSecondLevelKeys();
+        }
+    }
+
+    private filterSecondLevelKeyItems(
+        items: CompletionItem[],
+        context: Context,
+        topLevelKey: string | { Ref: unknown } | { '!Ref': unknown },
+    ): CompletionItem[] {
+        // Check if context.text contains the full FindInMap syntax (empty third argument case)
+        if (context.text.startsWith('[') && context.text.endsWith(']')) {
+            return items;
+        }
+
+        // If no text typed, return all items
+        if (context.text.length === 0) {
+            return items;
+        }
+
+        return this.applySecondLevelKeyFiltering(items, context, topLevelKey);
+    }
+
+    private applySecondLevelKeyFiltering(
+        items: CompletionItem[],
+        context: Context,
+        topLevelKey: string | { Ref: unknown } | { '!Ref': unknown },
+    ): CompletionItem[] {
+        // For dynamic keys, try prefix matching first, then fall back to fuzzy search
+        if (typeof topLevelKey === 'object') {
+            const prefixMatches = items.filter((item) =>
+                item.label.toLowerCase().startsWith(context.text.toLowerCase()),
+            );
+
+            if (prefixMatches.length === 0) {
+                return this.fuzzySearch(items, context.text);
+            }
+
+            return prefixMatches;
+        }
+
+        return this.fuzzySearch(items, context.text);
+    }
+
     private getMappingEntity(mappingsEntities: Map<string, Context>, mappingName: string): Mapping | undefined {
         try {
             const mappingContext = mappingsEntities.get(mappingName);
@@ -741,4 +802,8 @@ export class IntrinsicFunctionArgumentCompletionProvider implements CompletionPr
 
         return undefined;
     }
+
+    private isRefObject(value: unknown): value is { Ref: unknown } | { '!Ref': unknown } {
+        return typeof value === 'object' && value !== null && ('Ref' in value || '!Ref' in value);
+    }
 }

src/context/semantic/Entity.ts

@@ -140,9 +140,20 @@ export class Mapping extends Entity {
         return Object.keys(this.value);
     }
 
-    public getSecondLevelKeys(topLevelKey: string): string[] {
-        if (this.value[topLevelKey]) {
-            return Object.keys(this.value[topLevelKey]);
+    public getSecondLevelKeys(topLevelKey?: string): string[] {
+        if (topLevelKey === undefined) {
+            const allKeys = new Set<string>();
+            const topLevelKeys = this.getTopLevelKeys();
+
+            for (const tlKey of topLevelKeys) {
+                const keys = this.getSecondLevelKeys(tlKey);
+                for (const key of keys) allKeys.add(key);
+            }
+            return [...allKeys];
+        } else {
+            if (this.value[topLevelKey]) {
+                return Object.keys(this.value[topLevelKey]);
+            }
         }
         return [];
     }

tst/e2e/autocomplete/Autocomplete.test.ts

@@ -407,10 +407,6 @@ Rules:
                             position: { line: 114, character: 54 },
                             expectation: CompletionExpectationBuilder.create()
                                 .expectContainsItems(['InstanceType'])
-                                .todo(
-                                    `second level of Mapping not being suggested when using !Ref for first level key
-                                works using 'us-east-1'`,
-                                )
                                 .build(),
                         },
                     },
@@ -1438,13 +1434,7 @@ O`,
                         description: 'suggest Mapping second level key in deeply nested intrinsic function',
                         verification: {
                             position: { line: 471, character: 61 },
-                            expectation: CompletionExpectationBuilder.create()
-                                .expectItems(['AMI'])
-                                .todo(
-                                    `fix bug in FindInMap completion where using intrinsic in second arg breaks
-                             suggestion for third arg`,
-                                )
-                                .build(),
+                            expectation: CompletionExpectationBuilder.create().expectItems(['AMI']).build(),
                         },
                     },
                     {