From b6855fc54458c48cdd9347fea8648706cb93327a Mon Sep 17 00:00:00 2001
From: Mark Sujew <mark.sujew@typefox.io>
Date: Wed, 3 Sep 2025 14:17:09 +0200
Subject: [PATCH] Fix infix operator CST issues

---
 .../langium/src/parser/cst-node-builder.ts    |  6 +-
 packages/langium/src/parser/langium-parser.ts |  3 +
 .../langium/test/lsp/goto-definition.test.ts  | 63 ++++++++++++++++++-
 3 files changed, 68 insertions(+), 4 deletions(-)

diff --git a/packages/langium/src/parser/cst-node-builder.ts b/packages/langium/src/parser/cst-node-builder.ts
index e4bee1378..45d3393d1 100644
--- a/packages/langium/src/parser/cst-node-builder.ts
+++ b/packages/langium/src/parser/cst-node-builder.ts
@@ -87,11 +87,11 @@ export class CstNodeBuilder {
         }
     }
 
-    construct(item: { $type: string | symbol | undefined, $cstNode: CstNode }): void {
+    construct(item: { $type: string | symbol | undefined, $cstNode: CstNode, $infix?: boolean }): void {
         const current: CstNode = this.current;
-        // The specified item could be a datatype ($type is symbol) or a fragment ($type is undefined)
+        // The specified item could be a datatype ($type is symbol), fragment ($type is undefined) or infix rule ($infix is true)
         // Only if the $type is a string, we actually assign the element
-        if (typeof item.$type === 'string') {
+        if (typeof item.$type === 'string' && !item.$infix) {
             this.current.astNode = <AstNode>item;
         }
         item.$cstNode = current;
diff --git a/packages/langium/src/parser/langium-parser.ts b/packages/langium/src/parser/langium-parser.ts
index 379140863..cbe47588e 100644
--- a/packages/langium/src/parser/langium-parser.ts
+++ b/packages/langium/src/parser/langium-parser.ts
@@ -41,6 +41,7 @@ interface DataTypeNode {
 }
 
 interface InfixElement {
+    $infix: true;
     $type: string;
     $cstNode: CompositeCstNode;
     parts: AstNode[];
@@ -512,12 +513,14 @@ export class LangiumParser extends AbstractLangiumParser {
 
         // Create sub-expressions
         const leftInfix: InfixElement = {
+            $infix: true,
             $type: obj.$type,
             $cstNode: obj.$cstNode,
             parts: leftParts,
             operators: leftOperators
         };
         const rightInfix: InfixElement = {
+            $infix: true,
             $type: obj.$type,
             $cstNode: obj.$cstNode,
             parts: rightParts,
diff --git a/packages/langium/test/lsp/goto-definition.test.ts b/packages/langium/test/lsp/goto-definition.test.ts
index 3de3d9a5c..00d3cd0b8 100644
--- a/packages/langium/test/lsp/goto-definition.test.ts
+++ b/packages/langium/test/lsp/goto-definition.test.ts
@@ -6,7 +6,7 @@
 
 import { describe, test } from 'vitest';
 import { EmptyFileSystem } from 'langium';
-import { createLangiumGrammarServices } from 'langium/grammar';
+import { createLangiumGrammarServices, createServicesForGrammar } from 'langium/grammar';
 import { expectGoToDefinition } from 'langium/test';
 
 /**
@@ -114,3 +114,64 @@ describe('Definition Provider', () => {
         });
     });
 });
+
+describe('Definition Provider with Infix Operators', async () => {
+
+    const infixGrammar = `
+    grammar Test
+    entry Model: elements+=Element*;
+    Element: Statement | Item;
+    Item: 'item' name=ID ';';
+
+    Statement: value=InfixExpr ';';
+
+    infix InfixExpr on Primary:
+        '*' | '/' 
+        > '+' | '-';
+
+    Primary: '(' InfixExpr ')' | {infer ItemRef} ref=[Item];
+
+    terminal ID: /\\w+/;
+    hidden terminal WS: /\\s+/;
+    hidden terminal COMMENT: /\\/\\/.*/;
+    `;
+
+    const infixServices = await createServicesForGrammar({ grammar: infixGrammar });
+    const gotoDefinitionInfix = expectGoToDefinition(infixServices);
+
+    test('Simple infix operator expression should find Item from reference', async () => {
+        await gotoDefinitionInfix({
+            text: `
+            item <|a|>;
+            <|>a;
+            `,
+            index: 0,
+            rangeIndex: 0
+        });
+    });
+
+    test('Complex infix operator expression should find Item from reference', async () => {
+        const text = `
+            item <|a|>;
+            item <|b|>;
+            item <|c|>;
+            <|>a + <|>b * <|>c;
+        `;
+        await gotoDefinitionInfix({
+            text,
+            index: 0,
+            rangeIndex: 0
+        });
+        await gotoDefinitionInfix({
+            text,
+            index: 1,
+            rangeIndex: 1
+        });
+        await gotoDefinitionInfix({
+            text,
+            index: 2,
+            rangeIndex: 2
+        });
+    });
+
+});