fix: MoLang optimization quick-fix range resolution to prevent malformed rewrites

ide/base/server/src/lsp/code-action/minecraft/molang/optimization.test.ts

@@ -0,0 +1,61 @@
+import { CodeActionParams, Diagnostic, DiagnosticSeverity, Range } from 'vscode-languageserver';
+import { CodeActionBuilder } from '../../builder';
+import { onCodeAction } from './optimization';
+
+function positionAt(text: string, offset: number) {
+  const lines = text.slice(0, offset).split('\n');
+  return { line: lines.length - 1, character: lines[lines.length - 1].length };
+}
+
+function rangeFromOffsets(text: string, startOffset: number, endOffset: number): Range {
+  return {
+    start: positionAt(text, startOffset),
+    end: positionAt(text, endOffset),
+  };
+}
+
+describe('Molang optimization code action', () => {
+  it('uses explicit offsets from diagnostic data for replacement range', () => {
+    const text = 'v.test=5*3*math.pi;';
+    const uri = 'file:///test.json';
+
+    const correctStart = text.indexOf('5*3*math.pi');
+    const correctEnd = correctStart + '5*3*math.pi'.length;
+    const badStart = text.indexOf('*math');
+    const badEnd = badStart + '*math'.length;
+
+    const diag: Diagnostic = {
+      range: rangeFromOffsets(text, badStart, badEnd),
+      message: 'Can rewrite the operation to...',
+      severity: DiagnosticSeverity.Information,
+      code: 'molang.optimization.constant-folding',
+      data: {
+        replacement: '15*math.pi',
+        startOffset: correctStart,
+        endOffset: correctEnd,
+      },
+      source: 'mc',
+    };
+
+    const params: CodeActionParams = {
+      textDocument: { uri },
+      range: diag.range,
+      context: { diagnostics: [diag] },
+    };
+
+    const context = {
+      document: {
+        uri,
+        positionAt: (offset: number) => positionAt(text, offset),
+      },
+    } as any;
+    const builder = new CodeActionBuilder(params, context);
+
+    onCodeAction(builder, diag);
+
+    expect(builder.out).toHaveLength(1);
+    const action = builder.out[0] as any;
+    expect(action.edit.changes[uri][0].range).toEqual(rangeFromOffsets(text, correctStart, correctEnd));
+    expect(action.edit.changes[uri][0].newText).toBe('15*math.pi');
+  });
+});

ide/base/server/src/lsp/code-action/minecraft/molang/optimization.ts

@@ -4,6 +4,8 @@ import { CodeActionBuilder } from '../../builder';
 /** The shape of data attached to molang optimization diagnostics */
 interface OptimizationData {
   replacement: string;
+  startOffset?: number;
+  endOffset?: number;
 }
 
 function isOptimizationData(data: unknown): data is OptimizationData {
@@ -13,7 +15,14 @@ function isOptimizationData(data: unknown): data is OptimizationData {
 export function onCodeAction(builder: CodeActionBuilder, diag: Diagnostic): void {
   if (!isOptimizationData(diag.data)) return;
 
-  const replacement = diag.data.replacement;
+  const { replacement, startOffset, endOffset } = diag.data;
+  const hasOffsetRange = typeof startOffset === 'number' && typeof endOffset === 'number' && endOffset >= startOffset;
+  const editRange = hasOffsetRange
+    ? {
+        start: builder.context.document.positionAt(startOffset),
+        end: builder.context.document.positionAt(endOffset),
+      }
+    : diag.range;
 
   builder.push({
     title: `Rewrite to: ${replacement}`,
@@ -22,7 +31,7 @@ export function onCodeAction(builder: CodeActionBuilder, diag: Diagnostic): void
     isPreferred: true,
     edit: {
       changes: {
-        [builder.context.document.uri]: [TextEdit.replace(diag.range, replacement)],
+        [builder.context.document.uri]: [TextEdit.replace(editRange, replacement)],
       },
     },
   });

packages/bedrock-diagnoser/src/diagnostics/molang/optimizations/registry.ts

@@ -72,7 +72,10 @@ export class OptimizationRegistry {
       const nodeRange = OffsetWord.create('x'.repeat(length), startOffset);
 
       for (const optimizedNode of result) {
-        const data = optimizedNode.replacement !== undefined ? { replacement: optimizedNode.replacement } : undefined;
+        const data =
+          optimizedNode.replacement !== undefined
+            ? { replacement: optimizedNode.replacement, startOffset, endOffset }
+            : undefined;
         diagnoser.add(
           nodeRange,
           optimizedNode.message,

packages/bedrock-diagnoser/test/lib/diagnostics/molang/optimizations.test.ts

@@ -284,7 +284,7 @@ describe('Molang Optimization Diagnostics', () => {
 
       const diag = diagnoser.items.find((d) => d.code === 'molang.optimization.self-cancellation');
       expect(diag).toBeDefined();
-      expect(diag?.data).toEqual({ replacement: '0' });
+      expect(diag?.data).toEqual(expect.objectContaining({ replacement: '0' }));
     });
   });
 
@@ -320,7 +320,7 @@ describe('Molang Optimization Diagnostics', () => {
 
       const diag = diagnoser.items.find((d) => d.code === 'molang.optimization.self-division');
       expect(diag).toBeDefined();
-      expect(diag?.data).toEqual({ replacement: '1' });
+      expect(diag?.data).toEqual(expect.objectContaining({ replacement: '1' }));
     });
   });
 
@@ -462,7 +462,7 @@ describe('Molang Optimization Diagnostics', () => {
 
       const diag = diagnoser.items.find((d) => d.code === 'molang.optimization.identity-operation');
       expect(diag).toBeDefined();
-      expect(diag?.data).toEqual({ replacement: 'v.smooth_turn' });
+      expect(diag?.data).toEqual(expect.objectContaining({ replacement: 'v.smooth_turn' }));
     });
 
     it('should include replacement data for identity operation: addition with 0', () => {
@@ -471,7 +471,7 @@ describe('Molang Optimization Diagnostics', () => {
 
       const diag = diagnoser.items.find((d) => d.code === 'molang.optimization.identity-operation');
       expect(diag).toBeDefined();
-      expect(diag?.data).toEqual({ replacement: 'v.speed' });
+      expect(diag?.data).toEqual(expect.objectContaining({ replacement: 'v.speed' }));
     });
 
     it('should include replacement data for constant folding', () => {
@@ -490,7 +490,7 @@ describe('Molang Optimization Diagnostics', () => {
 
       const diag = diagnoser.items.find((d) => d.code === 'molang.optimization.constant-result');
       expect(diag).toBeDefined();
-      expect(diag?.data).toEqual({ replacement: '6' });
+      expect(diag?.data).toEqual(expect.objectContaining({ replacement: '6' }));
     });
 
     it('should not add parentheses to math constants like math.pi in replacement (regression)', () => {
@@ -499,10 +499,14 @@ describe('Molang Optimization Diagnostics', () => {
 
       const diag = diagnoser.items.find((d) => d.code === 'molang.optimization.constant-folding');
       expect(diag).toBeDefined();
-      const replacement = (diag?.data as { replacement: string })?.replacement;
+      const data = diag?.data as { replacement: string; startOffset?: number; endOffset?: number };
+      const replacement = data?.replacement;
       expect(replacement).toBeDefined();
       expect(replacement).toContain('math.pi');
       expect(replacement).not.toContain('math.pi()');
+      expect(typeof data?.startOffset).toBe('number');
+      expect(typeof data?.endOffset).toBe('number');
+      expect((data.endOffset ?? 0) - (data.startOffset ?? 0)).toBe('5*3*math.pi'.length);
     });
   });
 });