Fix organizing imports (#507)

* Fix organizing imports

The organize imports command broke documents where imports are present.
This was a regression caused by the logic to support auto imports.

Closes #502

Author: remcohaszing

.changeset/silly-loops-wait.md

@@ -0,0 +1,8 @@
+---
+'@mdx-js/language-service': patch
+'@mdx-js/language-server': patch
+'@mdx-js/typescript-plugin': patch
+'vscode-mdx': patch
+---
+
+Fix organizing imports

fixtures/node16/organize-imports.mdx

@@ -0,0 +1,16 @@
+---
+frontmatter: yaml
+---
+
+import { useState } from 'react'
+import { compile } from '@mdx-js/mdx'
+import { unified } from 'unified'
+import { createRoot } from 'react-dom/client'
+
+Paragraph text.
+
+export function whatever() {
+  compile;
+  useState;
+  createRoot;
+}

package.json

@@ -56,6 +56,7 @@
       "complexity": "off",
       "max-lines": "off",
       "max-params": "off",
+      "no-labels": "off",
       "@typescript-eslint/consistent-type-definitions": "off",
       "@typescript-eslint/naming-convention": "off",
       "unicorn/prevent-abbreviations": "off"

packages/language-server/test/code-action.test.js

@@ -0,0 +1,97 @@
+/**
+ * @import {LanguageServerHandle} from '@volar/test-utils'
+ */
+import assert from 'node:assert/strict'
+import {afterEach, beforeEach, test} from 'node:test'
+import {CodeAction, CodeActionTriggerKind} from '@volar/language-server'
+import {createServer, fixturePath, fixtureUri, tsdk} from './utils.js'
+
+/** @type {LanguageServerHandle} */
+let serverHandle
+
+beforeEach(async () => {
+  serverHandle = createServer()
+  await serverHandle.initialize(fixtureUri('node16'), {
+    typescript: {enabled: true, tsdk}
+  })
+})
+
+afterEach(() => {
+  serverHandle.connection.dispose()
+})
+
+test('organize imports', async () => {
+  const {uri} = await serverHandle.openTextDocument(
+    fixturePath('node16/organize-imports.mdx'),
+    'mdx'
+  )
+
+  const codeActions = await serverHandle.sendCodeActionsRequest(
+    uri,
+    {
+      start: {line: 6, character: 0},
+      end: {line: 6, character: 0}
+    },
+    {
+      diagnostics: [],
+      only: ['source.organizeImports'],
+      triggerKind: CodeActionTriggerKind.Invoked
+    }
+  )
+
+  assert.ok(codeActions)
+  const codeAction = codeActions
+    .filter((c) => CodeAction.is(c))
+    .find((c) => c.kind === 'source.organizeImports')
+  delete codeAction?.data
+
+  assert.deepEqual(codeAction, {
+    diagnostics: [],
+    edit: {
+      documentChanges: [
+        {
+          edits: [
+            {
+              newText:
+                "import { compile } from '@mdx-js/mdx';\n" +
+                "import { useState } from 'react';\n" +
+                "import { createRoot } from 'react-dom/client';\n" +
+                "import { unified } from 'unified';\n",
+              range: {
+                end: {character: 0, line: 5},
+                start: {character: 0, line: 4}
+              }
+            },
+            {
+              newText: '',
+              range: {
+                end: {character: 0, line: 6},
+                start: {character: 0, line: 5}
+              }
+            },
+            {
+              newText: '',
+              range: {
+                end: {character: 0, line: 7},
+                start: {character: 0, line: 6}
+              }
+            },
+            {
+              newText: '',
+              range: {
+                end: {character: 0, line: 8},
+                start: {character: 0, line: 7}
+              }
+            }
+          ],
+          textDocument: {
+            uri: fixtureUri('node16/organize-imports.mdx'),
+            version: null
+          }
+        }
+      ]
+    },
+    kind: 'source.organizeImports',
+    title: 'Organize Imports'
+  })
+})

packages/language-service/lib/virtual-code.js

@@ -24,12 +24,8 @@ import {isInjectableComponent, isInjectableEstree} from './jsx-utils.js'
  * @param {string} jsxImportSource
  *   The string to use for the JSX import source tag.
  */
-const jsPrefix = (
-  tsCheck,
-  jsxImportSource
-) => `${tsCheck ? '// @ts-check\n' : ''}/* @jsxRuntime automatic
+const jsPrefix = (tsCheck, jsxImportSource) => `/* @jsxRuntime automatic
 @jsxImportSource ${jsxImportSource} */
-import '${jsxImportSource}/jsx-runtime'
 `
 
 /**
@@ -216,86 +212,6 @@ function getPropsName(node) {
   return 'props'
 }
 
-/**
- * Process exports of an MDX ESM node.
- *
- * @param {string} mdx
- *   The full MDX code to process.
- * @param {MdxjsEsm} node
- *   The MDX ESM node to process.
- * @param {CodeMapping} mapping
- *   The Volar mapping to add offsets to.
- * @param {string} esm
- *   The virtual ESM code up to the point this function was called.
- * @returns {string}
- *   The updated virtual ESM code.
- */
-function processExports(mdx, node, mapping, esm) {
-  const start = node.position?.start?.offset
-  const end = node.position?.end?.offset
-
-  if (start === undefined || end === undefined) {
-    return esm
-  }
-
-  const body = node.data?.estree?.body
-
-  if (!body?.length) {
-    return addOffset(mapping, mdx, esm, start, end, true)
-  }
-
-  for (const child of body) {
-    if (child.type === 'ExportDefaultDeclaration') {
-      const propsName = getPropsName(child)
-      if (propsName) {
-        esm += layoutJsDoc(propsName)
-      }
-
-      esm = addOffset(
-        mapping,
-        mdx,
-        esm + '\nconst MDXLayout = ',
-        child.declaration.start,
-        child.end,
-        true
-      )
-      continue
-    }
-
-    if (child.type === 'ExportNamedDeclaration' && child.source) {
-      const {specifiers} = child
-      for (let index = 0; index < specifiers.length; index++) {
-        const specifier = specifiers[index]
-        if (
-          specifier.local.type === 'Identifier'
-            ? specifier.local.name === 'default'
-            : specifier.local.value === 'default'
-        ) {
-          esm = addOffset(mapping, mdx, esm, start, specifier.start)
-          const nextPosition =
-            index === specifiers.length - 1
-              ? specifier.end
-              : mdx.indexOf(',', specifier.end) + 1
-          return (
-            addOffset(mapping, mdx, esm, nextPosition, end, true) +
-            '\nimport {' +
-            (specifier.exported.type === 'Identifier'
-              ? specifier.exported.name
-              : JSON.stringify(specifier.exported.value)) +
-            ' as MDXLayout} from ' +
-            JSON.stringify(child.source.value) +
-            '\n'
-          )
-        }
-      }
-    }
-
-    esm = addOffset(mapping, mdx, esm, child.start, child.end, true)
-  }
-
-  return esm + '\n'
-}
-
 /**
  * Pad the generated offsets of a Volar code mapping.
  *
@@ -327,6 +243,7 @@ function getEmbeddedCodes(
   jsxImportSource
 ) {
   let hasAwait = false
+  let hasImports = false
   let esm = jsPrefix(checkMdx, jsxImportSource)
   let jsx = ''
   let jsxVariables = ''
@@ -345,9 +262,9 @@ function getEmbeddedCodes(
   const esmMapping = {
     // The empty mapping makes sure there’s always a valid mapping to insert
     // auto-imports.
-    sourceOffsets: [0],
-    generatedOffsets: [esm.length],
-    lengths: [0],
+    sourceOffsets: [],
+    generatedOffsets: [],
+    lengths: [],
     data: {
       completion: true,
       format: true,
@@ -493,6 +410,83 @@ function getEmbeddedCodes(
     updateMarkdownFromOffsets(startOffset, endOffset)
   }
 
+  /**
+   * Process exports of an MDX ESM node.
+   *
+   * @param {MdxjsEsm} node
+   *   The MDX ESM node to process.
+   * @returns {undefined}
+   */
+  function processExports(node) {
+    const start = node.position?.start?.offset
+    const end = node.position?.end?.offset
+
+    if (start === undefined || end === undefined) {
+      return
+    }
+
+    const body = node.data?.estree?.body
+
+    if (!body?.length) {
+      esm = addOffset(esmMapping, mdx, esm, start, end, true)
+      return
+    }
+
+    bodyLoop: for (const child of body) {
+      if (child.type === 'ExportDefaultDeclaration') {
+        const propsName = getPropsName(child)
+        if (propsName) {
+          esm += layoutJsDoc(propsName)
+        }
+
+        esm = addOffset(
+          esmMapping,
+          mdx,
+          esm + '\nconst MDXLayout = ',
+          child.declaration.start,
+          child.end,
+          true
+        )
+        continue
+      }
+
+      if (child.type === 'ExportNamedDeclaration' && child.source) {
+        const {specifiers} = child
+        for (let index = 0; index < specifiers.length; index++) {
+          const specifier = specifiers[index]
+          if (
+            specifier.local.type === 'Identifier'
+              ? specifier.local.name === 'default'
+              : specifier.local.value === 'default'
+          ) {
+            esm = addOffset(esmMapping, mdx, esm, start, specifier.start)
+            const nextPosition =
+              index === specifiers.length - 1
+                ? specifier.end
+                : mdx.indexOf(',', specifier.end) + 1
+            esm =
+              addOffset(esmMapping, mdx, esm, nextPosition, end, true) +
+              '\nimport {' +
+              (specifier.exported.type === 'Identifier'
+                ? specifier.exported.name
+                : JSON.stringify(specifier.exported.value)) +
+              ' as MDXLayout} from ' +
+              JSON.stringify(child.source.value)
+            continue bodyLoop
+          }
+        }
+      }
+
+      if (child.type === 'ImportDeclaration') {
+        hasImports = true
+      }
+
+      esm = addOffset(esmMapping, mdx, esm, child.start, child.end, true)
+    }
+
+    esm += '\n'
+  }
+
   /**
    * @param {Program} program
    * @param {number} lastIndex
@@ -663,7 +657,7 @@ function getEmbeddedCodes(
 
         case 'mdxjsEsm': {
           updateMarkdownFromNode(node)
-          esm = processExports(mdx, node, esmMapping, esm)
+          processExports(node)
           break
         }
 
@@ -808,6 +802,26 @@ function getEmbeddedCodes(
     esm += '\n' + plugin.finalize() + '\n'
   }
 
+  let prefix = ''
+  if (checkMdx) {
+    prefix += '// @ts-check\n'
+  }
+
+  if (!hasImports) {
+    prefix += `import '${jsxImportSource}/jsx-runtime'\n`
+  }
+
+  if (prefix) {
+    padOffsets(esmMapping, prefix.length)
+    esm = prefix + esm
+  }
+
+  if (!hasImports) {
+    esmMapping.sourceOffsets.unshift(0)
+    esmMapping.generatedOffsets.unshift(prefix.length)
+    esmMapping.lengths.unshift(0)
+  }
+
   updateMarkdownFromOffsets(mdx.length, mdx.length)
   esm += componentStart(hasAwait, programScope)