fix: MoLang optimization quick-fix range resolution to prevent malformed rewrites (#497)

MoLang optimization diagnostics were generating correct replacement text
(`math.pi` no longer `math.pi()`), but quick-fix application still used
an incorrect edit range in some cases. This caused partial-token
replacement and corrupted output such as `5*3(math.pi() * 15).pi`.

- **Root cause**
- Optimization code actions relied on `diag.range`, which could point to
an inner/misaligned span instead of the full expression intended for
rewrite.

- **Diagnostic payload now carries source offsets**
-
`packages/bedrock-diagnoser/src/diagnostics/molang/optimizations/registry.ts`
- Added `startOffset` and `endOffset` to optimization diagnostic `data`
when a `replacement` is present.
- Keeps existing replacement behavior while giving code actions an
unambiguous source span.

- **Code action now prefers offset-derived range**
- `ide/base/server/src/lsp/code-action/minecraft/molang/optimization.ts`
- Quick fix computes `editRange` from `startOffset/endOffset` when
available; falls back to `diag.range` for compatibility.
- Ensures rewrite applies to the full expression (`5*3*math.pi`) rather
than a partial segment.

- **Regression coverage**
  - Added server-side code-action regression test:
-
`ide/base/server/src/lsp/code-action/minecraft/molang/optimization.test.ts`
- Verifies malformed `diag.range` does not affect replacement when
offsets are present.
- Updated optimization tests to allow expanded diagnostic data shape and
assert offset fields for the `5*3*math.pi` regression path:
-
`packages/bedrock-diagnoser/test/lib/diagnostics/molang/optimizations.test.ts`

```ts
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;
```

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: DaanV2 <2393905+DaanV2@users.noreply.github.com>

Author: Copilot

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);
     });
   });
 });