fix: resolve to the correct enum in field attribute when there's ambiguity (#513)

Author: ymc9

packages/schema/src/language-server/zmodel-linker.ts

@@ -6,6 +6,7 @@ import {
     DataModel,
     DataModelField,
     DataModelFieldType,
+    Enum,
     EnumField,
     Expression,
     FunctionDecl,
@@ -16,6 +17,7 @@ import {
     isDataModel,
     isDataModelField,
     isDataModelFieldType,
+    isEnum,
     isReferenceExpr,
     LiteralExpr,
     MemberAccessExpr,
@@ -164,6 +166,10 @@ export class ZModelLinker extends DefaultLinker {
                 this.resolveDataModel(node as DataModel, document, extraScopes);
                 break;
 
+            case DataModelField:
+                this.resolveDataModelField(node as DataModelField, document, extraScopes);
+                break;
+
             default:
                 this.resolveDefault(node, document, extraScopes);
                 break;
@@ -451,6 +457,49 @@ export class ZModelLinker extends DefaultLinker {
         return this.resolveDefault(node, document, extraScopes);
     }
 
+    private resolveDataModelField(
+        node: DataModelField,
+        document: LangiumDocument<AstNode>,
+        extraScopes: ScopeProvider[]
+    ) {
+        // Field declaration may contain enum references, and enum fields are pushed to the global
+        // scope, so if there're enums with fields with the same name, an arbitrary one will be
+        // used as resolution target. The correct behavior is to resolve to the enum that's used
+        // as the declaration type of the field:
+        //
+        // enum FirstEnum {
+        //     E1
+        //     E2
+        // }
+
+        // enum SecondEnum  {
+        //     E1
+        //     E3
+        //     E4
+        // }
+
+        // model M {
+        //     id Int @id
+        //     first  SecondEnum @default(E1) <- should resolve to SecondEnum
+        //     second FirstEnum @default(E1) <- should resolve to FirstEnum
+        // }
+        //
+
+        // make sure type is resolved first
+        this.resolve(node.type, document, extraScopes);
+
+        let scopes = extraScopes;
+
+        // if the field has enum declaration type, resolve the rest with that enum's fields on top of the scopes
+        if (node.type.reference?.ref && isEnum(node.type.reference.ref)) {
+            const contextEnum = node.type.reference.ref as Enum;
+            const enumScope: ScopeProvider = (name) => contextEnum.fields.find((f) => f.name === name);
+            scopes = [enumScope, ...scopes];
+        }
+
+        this.resolveDefault(node, document, scopes);
+    }
+
     private resolveDefault(node: AstNode, document: LangiumDocument<AstNode>, extraScopes: ScopeProvider[]) {
         for (const [property, value] of Object.entries(node)) {
             if (!property.startsWith('$')) {

packages/schema/tests/schema/parser.test.ts

@@ -40,7 +40,7 @@ describe('Parsing Tests', () => {
         expect((ds.fields[1].value as InvocationExpr).args[0].value.$type).toBe(LiteralExpr);
     });
 
-    it('enum', async () => {
+    it('enum simple', async () => {
         const content = `
             enum UserRole {
                 USER
@@ -64,6 +64,39 @@ describe('Parsing Tests', () => {
         expect((attrVal.value as ReferenceExpr).target.ref?.name).toBe('USER');
     });
 
+    it('enum dup name resolve', async () => {
+        const content = `
+            datasource db {
+                provider = "postgresql"
+                url      = env("DATABASE_URL")
+            }
+            
+            enum FirstEnum {
+                E1 // used in both ENUMs
+                E2
+            }
+            
+            enum SecondEnum  {
+                E1 // used in both ENUMs
+                E3
+                E4
+            }
+            
+            model M {
+                id Int @id
+                first  SecondEnum @default(E1)
+                second FirstEnum @default(E1)
+            }        
+            `;
+
+        const doc = await loadModel(content);
+        const firstEnum = doc.declarations.find((d) => d.name === 'FirstEnum');
+        const secondEnum = doc.declarations.find((d) => d.name === 'SecondEnum');
+        const m = doc.declarations.find((d) => d.name === 'M') as DataModel;
+        expect(m.fields[1].attributes[0].args[0].value.$resolvedType?.decl).toBe(secondEnum);
+        expect(m.fields[2].attributes[0].args[0].value.$resolvedType?.decl).toBe(firstEnum);
+    });
+
     it('model field types', async () => {
         const content = `
             model User {