Fn::ForEach Autocomplete support (#194)

* autocomplete support for Fn::ForEach resource type
* autocomplete support for Fn::ForEach resource properties

---------

Co-authored-by: Deep Furiya <[email protected]>

Author: deepfuriya

src/autocomplete/EntityFieldCompletionProvider.ts

@@ -9,7 +9,17 @@ import { createCompletionItem } from './CompletionUtils';
 /* eslint-disable no-restricted-syntax -- Entire class depends on Entity */
 export class EntityFieldCompletionProvider<T extends Entity> implements CompletionProvider {
     public getCompletions(context: Context, _: CompletionParams): CompletionItem[] {
-        const entity = context.entity as T;
+        // Extract the actual entity (handle both regular and ForEach resources)
+        let entity;
+        if (context.getEntityType() === EntityType.ForEachResource) {
+            entity = context.getResourceEntity() as unknown as T;
+        } else {
+            entity = context.entity as T;
+        }
+
+        if (!entity) {
+            return [];
+        }
 
         const items = this.getFieldsAsCompletionItems(entity);
         if (context.text.length > 0) {

src/autocomplete/ResourceEntityCompletionProvider.ts

@@ -26,21 +26,24 @@ export class ResourceEntityCompletionProvider implements CompletionProvider {
     getCompletions(context: Context, params: CompletionParams): CompletionItem[] | undefined {
         const entityCompletions = this.entityFieldProvider.getCompletions(context, params);
 
+        // Extract the actual resource entity (handle both regular and ForEach resources)
+        const resource = context.getResourceEntity();
+        if (!resource) {
+            return entityCompletions;
+        }
+
         // Enhance the "Properties" completion with a snippet
         if (entityCompletions) {
             const propertiesIndex = entityCompletions.findIndex((item) => item.label === 'Properties');
 
-            if (propertiesIndex !== -1) {
-                const resource = context.entity as Resource;
-                if (resource.Type) {
-                    const schema = this.schemaRetriever.getDefault().schemas.get(resource.Type);
-                    if (schema) {
-                        entityCompletions[propertiesIndex] = this.createPropertiesSnippetCompletion(
-                            schema,
-                            context,
-                            params,
-                        );
-                    }
+            if (propertiesIndex !== -1 && resource.Type) {
+                const schema = this.schemaRetriever.getDefault().schemas.get(resource.Type);
+                if (schema) {
+                    entityCompletions[propertiesIndex] = this.createPropertiesSnippetCompletion(
+                        schema,
+                        context,
+                        params,
+                    );
                 }
             }
         }

src/autocomplete/ResourcePropertyCompletionProvider.ts

@@ -22,7 +22,7 @@ import {
 import { Context } from '../context/Context';
 import { ResourceAttribute, TopLevelSection, ResourceAttributesSet } from '../context/ContextType';
 import { Resource } from '../context/semantic/Entity';
-import { CfnValue } from '../context/semantic/SemanticTypes';
+import { CfnValue, EntityType } from '../context/semantic/SemanticTypes';
 import { NodeType } from '../context/syntaxtree/utils/NodeType';
 import { CommonNodeTypes } from '../context/syntaxtree/utils/TreeSitterTypes';
 import { propertyTypesToMarkdown } from '../hover/HoverFormatter';
@@ -66,7 +66,10 @@ export class ResourcePropertyCompletionProvider implements CompletionProvider {
      * Uses robust schema resolution approach from hover provider
      */
     private getPropertyCompletions(context: Context): PropertyCompletionsResult {
-        const resource = context.entity as Resource;
+        const resource = context.getResourceEntity();
+        if (!resource) {
+            return { completions: [], skipFuzzySearch: false };
+        }
         let completions: CompletionItem[] = [];
         let skipFuzzySearch = false;
 
@@ -138,7 +141,12 @@ export class ResourcePropertyCompletionProvider implements CompletionProvider {
     }
 
     private getSchemaPath(context: Context): string {
-        let segments = context.propertyPath.slice(3);
+        const propertiesIndex = context.propertyPath.indexOf('Properties');
+        if (propertiesIndex === -1) {
+            return '/properties';
+        }
+
+        let segments = context.propertyPath.slice(propertiesIndex + 1);
 
         // For key completions (except SYNTHETIC_KEY_OR_VALUE), remove last segment
         if (
@@ -271,11 +279,16 @@ export class ResourcePropertyCompletionProvider implements CompletionProvider {
         const lastSegment = propertyPath[propertyPath.length - 1];
         const isArrayItemContext = typeof lastSegment === 'number' || lastSegment === '';
 
-        if (propertyPath.length > 3 && isArrayItemContext) {
-            const entity = context.entity as Resource;
-            if (entity?.Properties) {
-                const pathSegments = propertyPath.slice(3); // Remove ['Resources', 'LogicalId', 'Properties']
-                let current: Record<string, CfnValue> | CfnValue | undefined = entity.Properties;
+        // Find the Properties index dynamically
+        const startIndex = context.getEntityType() === EntityType.ForEachResource ? 4 : 2;
+        const propertiesIndex = propertyPath.indexOf('Properties', startIndex);
+
+        if (propertiesIndex !== -1 && isArrayItemContext) {
+            const resource = context.getResourceEntity();
+
+            if (resource?.Properties) {
+                const pathSegments = propertyPath.slice(propertiesIndex + 1);
+                let current: Record<string, CfnValue> | CfnValue | undefined = resource.Properties;
 
                 for (let i = 0; i < pathSegments.length - 1; i++) {
                     if (current && typeof current === 'object' && pathSegments[i] in current) {

src/autocomplete/ResourceSectionCompletionProvider.ts

@@ -1,7 +1,7 @@
 import { CompletionItem, CompletionParams, CompletionTriggerKind } from 'vscode-languageserver';
 import { Context } from '../context/Context';
-import { ResourceAttributesSet, TopLevelSection } from '../context/ContextType';
-import { Resource } from '../context/semantic/Entity';
+import { ResourceAttributesSet } from '../context/ContextType';
+import { EntityType } from '../context/semantic/SemanticTypes';
 import { CfnExternal } from '../server/CfnExternal';
 import { CfnInfraCore } from '../server/CfnInfraCore';
 import { CfnLspProviders } from '../server/CfnLspProviders';
@@ -39,25 +39,23 @@ export class ResourceSectionCompletionProvider implements CompletionProvider {
             return this.resourceProviders
                 .get(ResourceCompletionType.Entity)
                 ?.getCompletions(context, params) as CompletionItem[];
-        } else if (context.entitySection === 'Type') {
+        } else if (context.entitySection === 'Type' || this.isAtResourceTypeField(context)) {
             return this.resourceProviders
                 .get(ResourceCompletionType.Type)
                 ?.getCompletions(context, params) as CompletionItem[];
         } else if (
             context.entitySection === 'Properties' ||
             ResourceAttributesSet.has(context.entitySection as string) ||
-            (context.matchPathWithLogicalId(TopLevelSection.Resources, 'Properties') && context.propertyPath.length > 3)
+            this.isInPropertiesSection(context)
         ) {
             const schemaPropertyCompletions = this.resourceProviders
                 .get(ResourceCompletionType.Property)
                 ?.getCompletions(context, params) as CompletionItem[];
 
-            if (
-                params.context?.triggerKind === CompletionTriggerKind.Invoked &&
-                context.matchPathWithLogicalId(TopLevelSection.Resources, 'Properties')
-            ) {
-                const resource = context.entity as Resource;
-                if (resource.Type) {
+            if (params.context?.triggerKind === CompletionTriggerKind.Invoked && this.isInPropertiesSection(context)) {
+                const resource = context.getResourceEntity();
+
+                if (resource?.Type) {
                     const stateCompletionPromise = this.resourceProviders
                         .get(ResourceCompletionType.State)
                         ?.getCompletions(context, params) as Promise<CompletionItem[]>;
@@ -77,6 +75,22 @@ export class ResourceSectionCompletionProvider implements CompletionProvider {
         }
         return [];
     }
+
+    private isInPropertiesSection(context: Context): boolean {
+        // Find 'Properties' starting after the resource structure
+        const startIndex = context.getEntityType() === EntityType.ForEachResource ? 4 : 2;
+        const propertiesIndex = context.propertyPath.indexOf('Properties', startIndex);
+        return propertiesIndex !== -1 && context.propertyPath.length >= propertiesIndex + 1;
+    }
+
+    private isAtResourceTypeField(context: Context): boolean {
+        const propertyPathLength = context.getEntityType() === EntityType.ForEachResource ? 5 : 3;
+
+        return (
+            context.propertyPath.length === propertyPathLength &&
+            context.propertyPath[context.propertyPath.length - 1] === 'Type'
+        );
+    }
 }
 
 export function createResourceCompletionProviders(

src/context/Context.ts

@@ -13,7 +13,7 @@ import {
     TopLevelSectionsWithLogicalIdsSet,
 } from './ContextType';
 import { IntrinsicContext } from './IntrinsicContext';
-import { Entity } from './semantic/Entity';
+import { Entity, ForEachResource, Resource } from './semantic/Entity';
 import { entityTypeFromSection, nodeToEntity } from './semantic/EntityBuilder';
 import { normalizeIntrinsicFunction } from './semantic/Intrinsics';
 import { EntityType } from './semantic/SemanticTypes';
@@ -76,6 +76,18 @@ export class Context {
         return entityTypeFromSection(this.section, this.logicalId);
     }
 
+    public getResourceEntity(): Resource | undefined {
+        const entityType = this.getEntityType();
+        if (entityType === EntityType.Resource) {
+            return this.entity as Resource;
+        }
+        if (entityType === EntityType.ForEachResource) {
+            const forEachResource = this.entity as ForEachResource;
+            return forEachResource.resource;
+        }
+        return undefined;
+    }
+
     public get intrinsicContext(): IntrinsicContext {
         this._intrinsicContext ??= this.telemetry.measure(
             'create.intrinsicContext',
@@ -260,19 +272,24 @@ export class Context {
             return false;
         }
 
-        // Case 1: If we are over 3 we know for sure we are beyond the entity level
-        if (this.propertyPath.length > 3) {
+        // Determine the entity key level based on entity type
+        // Regular: ['Resources', 'LogicalId', 'Key'] - level 3
+        // ForEachResource: ['Resources', 'Fn::ForEach::Name', 2, 'ResourceKey', 'Key'] - level 5
+        const entityKeyLevel = this.getEntityType() === EntityType.ForEachResource ? 5 : 3;
+
+        // Case 1: If we are beyond the entity key level
+        if (this.propertyPath.length > entityKeyLevel) {
             return false;
         }
 
         // Case 2: Two situations exist that we need to account for:
         // isKey and isValue can be True when at the first key inside a value
         // when we are at level 2 this means we are at Entity/LogicalId as the first key
-        // when we are at level 3 this means we are at Entity/LogicalId/Properties as the first key
+        // when we are at level 3 (or 5 for ForEach) this means we are at Entity/LogicalId/Properties as the first key
         if (this.isKey() && this.isValue()) {
             if (this.propertyPath.length === 2) {
                 return true;
-            } else if (this.propertyPath.length === 3) {
+            } else if (this.propertyPath.length === entityKeyLevel) {
                 return false;
             }
         }

tst/unit/autocomplete/EntityFieldCompletionProvider.test.ts

@@ -2,7 +2,7 @@ import { describe, expect, test } from 'vitest';
 import { CompletionParams } from 'vscode-languageserver';
 import { EntityFieldCompletionProvider } from '../../../src/autocomplete/EntityFieldCompletionProvider';
 import { Output, Parameter } from '../../../src/context/semantic/Entity';
-import { createOutputContext, createParameterContext } from '../../utils/MockContext';
+import { createForEachResourceContext, createOutputContext, createParameterContext } from '../../utils/MockContext';
 
 describe('EntityFieldCompletionProvider', () => {
     const parameterFieldCompletionProvider = new EntityFieldCompletionProvider<Parameter>();
@@ -153,4 +153,92 @@ describe('EntityFieldCompletionProvider', () => {
             expect(result?.at(0)?.label).equal('Export');
         });
     });
+
+    describe('Fn::ForEach Resource', () => {
+        const resourceFieldCompletionProvider = new EntityFieldCompletionProvider();
+
+        test('should suggest resource fields for ForEach resource', () => {
+            const mockContext = createForEachResourceContext('Fn::ForEach::Buckets', 'S3Bucket${BucketName}', {
+                text: '',
+                propertyPath: ['Resources', 'Fn::ForEach::Buckets', 2, 'S3Bucket${BucketName}', ''],
+                data: {
+                    Type: 'AWS::S3::Bucket',
+                },
+            });
+
+            const result = resourceFieldCompletionProvider.getCompletions(mockContext, mockParams);
+            expect(result).toBeDefined();
+            expect(result.length).toBeGreaterThan(0);
+
+            const resultLabels = result?.map((item) => item.label);
+            expect(resultLabels).toContain('Properties');
+            expect(resultLabels).toContain('DependsOn');
+            expect(resultLabels).toContain('Metadata');
+        });
+
+        test('should filter resource fields with partial text for ForEach resource', () => {
+            const mockContext = createForEachResourceContext('Fn::ForEach::Instances', 'Instance${Name}', {
+                text: 'Prop',
+                propertyPath: ['Resources', 'Fn::ForEach::Instances', 2, 'Instance${Name}', 'Prop'],
+                data: {
+                    Type: 'AWS::EC2::Instance',
+                },
+            });
+
+            const result = resourceFieldCompletionProvider.getCompletions(mockContext, mockParams);
+            expect(result).toBeDefined();
+            expect(result.length).toBeGreaterThan(0);
+
+            const propertiesItem = result?.find((item) => item.label === 'Properties');
+            expect(propertiesItem).toBeDefined();
+        });
+
+        test('should not suggest already defined fields for ForEach resource', () => {
+            const mockContext = createForEachResourceContext('Fn::ForEach::Buckets', 'S3Bucket${BucketName}', {
+                text: '',
+                propertyPath: ['Resources', 'Fn::ForEach::Buckets', 2, 'S3Bucket${BucketName}', ''],
+                data: {
+                    Type: 'AWS::S3::Bucket',
+                    Properties: {},
+                    DependsOn: 'SomeResource',
+                },
+            });
+
+            const result = resourceFieldCompletionProvider.getCompletions(mockContext, mockParams);
+            expect(result).toBeDefined();
+
+            const resultLabels = result?.map((item) => item.label);
+            expect(resultLabels).not.toContain('Type');
+            expect(resultLabels).not.toContain('Properties');
+            expect(resultLabels).not.toContain('DependsOn');
+        });
+
+        test('should suggest all fields when ForEach resource has no fields defined', () => {
+            const mockContext = createForEachResourceContext('Fn::ForEach::Buckets', 'S3Bucket${BucketName}', {
+                text: '',
+                propertyPath: ['Resources', 'Fn::ForEach::Buckets', 2, 'S3Bucket${BucketName}', ''],
+                data: {},
+            });
+
+            const result = resourceFieldCompletionProvider.getCompletions(mockContext, mockParams);
+            expect(result).toBeDefined();
+            expect(result.length).toBeGreaterThan(0);
+
+            const resultLabels = result?.map((item) => item.label);
+            expect(resultLabels).toContain('Type');
+            expect(resultLabels).toContain('Properties');
+        });
+
+        test('should return empty when ForEach resource has no resource property', () => {
+            const mockContext = createForEachResourceContext('Fn::ForEach::Buckets', 'S3Bucket${BucketName}', {
+                text: '',
+                propertyPath: ['Resources', 'Fn::ForEach::Buckets', 2, 'S3Bucket${BucketName}', ''],
+                data: undefined,
+            });
+
+            const result = resourceFieldCompletionProvider.getCompletions(mockContext, mockParams);
+            expect(result).toBeDefined();
+            expect(result.length).toBe(0);
+        });
+    });
 });