Implement 'organize imports' functionality

README.md

@@ -19,6 +19,7 @@ This extension provides syntax highlighting, type checking, and code formatting
 - Automatic imports
 - Snippets ([contributions welcome](https://github.com/dfinity/node-motoko/blob/main/contrib/snippets.json))
 - Go-to-definition
+- Organize imports
 - Documentation tooltips
 
 ## Installation
@@ -32,7 +33,7 @@ Get this extension through the [VS Marketplace](https://marketplace.visualstudio
 
 ## Extension Commands
 
-- `motoko.startService`: Starts (or restarts) the language service
+- `Motoko: Restart language server`: Starts (or restarts) the language server
 
 ## Extension Settings
 
@@ -41,6 +42,22 @@ Get this extension through the [VS Marketplace](https://marketplace.visualstudio
 - `motoko.formatter`: The formatter used by the extension
 - `motoko.legacyDfxSupport`: Uses legacy `dfx`-dependent features when a relevant `dfx.json` file is available
 
+## Advanced Configuration
+
+If you want VS Code to automatically format Motoko files on save, consider adding the following to your `settings.json` configuration:
+
+```json
+{
+  "[motoko]": {
+    "editor.defaultFormatter": "dfinity-foundation.vscode-motoko",
+    "editor.formatOnSave": true,
+    "editor.codeActionsOnSave": {
+        "source.organizeImports": true
+    }
+  }
+}
+```
+
 ## Recent Changes
 
 Projects using `dfx >= 0.11.1` use a new, experimental language server.

src/extension.ts

@@ -69,7 +69,7 @@ export function startServer(context: ExtensionContext) {
 
     // Cross-platform language server
     const module = context.asAbsolutePath(path.join('out', 'server.js'));
-    launchClient(context, {
+    restartLanguageServer(context, {
         run: { module, transport: TransportKind.ipc },
         debug: {
             module,
@@ -92,7 +92,10 @@ function launchDfxProject(context: ExtensionContext, dfxConfig: DfxConfig) {
             command: getDfxPath(),
             args: ['_language-service', canister],
         };
-        launchClient(context, { run: serverCommand, debug: serverCommand });
+        restartLanguageServer(context, {
+            run: serverCommand,
+            debug: serverCommand,
+        });
     };
 
     const canister = config.get<string>('canister');
@@ -114,7 +117,10 @@ function launchDfxProject(context: ExtensionContext, dfxConfig: DfxConfig) {
     }
 }
 
-function launchClient(context: ExtensionContext, serverOptions: ServerOptions) {
+function restartLanguageServer(
+    context: ExtensionContext,
+    serverOptions: ServerOptions,
+) {
     if (client) {
         console.log('Restarting Motoko language server');
         client.stop().catch((err) => console.error(err.stack || err));

src/server/imports.ts

@@ -2,8 +2,8 @@ import { pascalCase } from 'change-case';
 import { MultiMap } from 'mnemonist';
 import { AST, Node } from 'motoko/lib/ast';
 import { Context, getContext } from './context';
-import { Program, matchNode } from './syntax';
-import { getRelativeUri } from './utils';
+import { Import, Program, matchNode } from './syntax';
+import { formatMotoko, getRelativeUri } from './utils';
 
 interface ResolvedField {
     name: string;
@@ -198,3 +198,70 @@ function getImportInfo(
     }
     return [getImportName(uri), uri];
 }
+
+const importGroups: {
+    prefix: string;
+}[] = [
+    // IC imports
+    { prefix: 'ic:' },
+    // Canister alias imports
+    { prefix: 'canister:' },
+    // Package imports
+    { prefix: 'mo:' },
+    // Everything else
+    { prefix: '' },
+];
+
+export function organizeImports(imports: Import[]): string {
+    const groupParts: string[][] = importGroups.map(() => []);
+
+    // Combine imports with the same path
+    const combinedImports: Record<
+        string,
+        { names: string[]; fields: [string, string][] }
+    > = {};
+    imports.forEach((x) => {
+        const combined =
+            combinedImports[x.path] ||
+            (combinedImports[x.path] = { names: [], fields: [] });
+        if (x.name) {
+            combined.names.push(x.name);
+        }
+        combined.fields.push(...x.fields);
+    });
+
+    // Sort and print imports
+    Object.entries(combinedImports)
+        .sort(
+            // Sort by import path
+            (a, b) => a[0].localeCompare(b[0]),
+        )
+        .forEach(([path, { names, fields }]) => {
+            const parts =
+                groupParts[
+                    importGroups.findIndex((g) => path.startsWith(g.prefix))
+                ] || groupParts[groupParts.length - 1];
+            names.forEach((name) => {
+                parts.push(`import ${name} ${JSON.stringify(path)};`);
+            });
+            if (fields.length) {
+                parts.push(
+                    `import { ${fields
+                        .sort(
+                            // Sort by name, then alias
+                            (a, b) =>
+                                a[0].localeCompare(b[0]) ||
+                                (a[1] || a[0]).localeCompare(b[1] || b[0]),
+                        )
+                        .map(([name, alias]) =>
+                            !alias || name === alias
+                                ? name
+                                : `${name} = ${alias}`,
+                        )
+                        .join('; ')} } ${JSON.stringify(path)};`,
+                );
+            }
+        });
+
+    return formatMotoko(groupParts.map((p) => p.join('\n')).join('\n\n'));
+}

src/server/server.ts

@@ -19,6 +19,7 @@ import {
     MarkupKind,
     Position,
     ProposedFeatures,
+    Range,
     ReferenceParams,
     SignatureHelp,
     TextDocumentPositionParams,
@@ -46,13 +47,14 @@ import {
     rangeFromNode,
 } from './navigation';
 import { vesselSources } from './rust';
-import { Program, findNodes, asNode } from './syntax';
+import { Program, asNode, findNodes } from './syntax';
 import {
     formatMotoko,
     getFileText,
     resolveFilePath,
     resolveVirtualPath,
 } from './utils';
+import { organizeImports } from './imports';
 
 interface Settings {
     motoko: MotokoSettings;
@@ -387,8 +389,14 @@ connection.onInitialize((event): InitializeResult => {
             definitionProvider: true,
             // declarationProvider: true,
             // referencesProvider: true,
-            codeActionProvider: true,
+            codeActionProvider: {
+                codeActionKinds: [
+                    CodeActionKind.QuickFix,
+                    CodeActionKind.SourceOrganizeImports,
+                ],
+            },
             hoverProvider: true,
+            // executeCommandProvider: { commands: [] },
             // workspaceSymbolProvider: true,
             // diagnosticProvider: {
             //     documentSelector: ['motoko'],
@@ -795,11 +803,38 @@ function deleteVirtual(path: string) {
 }
 
 connection.onCodeAction((event) => {
+    const uri = event.textDocument.uri;
     const results: CodeAction[] = [];
 
-    // Automatic imports
+    // Organize imports
+    const status = getContext(uri).astResolver.request(uri);
+    const imports = status?.program?.imports;
+    if (imports?.length) {
+        const start = rangeFromNode(asNode(imports[0].ast))?.start;
+        const end = rangeFromNode(asNode(imports[imports.length - 1].ast))?.end;
+        if (!start || !end) {
+            console.warn('Unexpected import AST range format');
+            return;
+        }
+        const range = Range.create(
+            Position.create(start.line, 0),
+            Position.create(end.line + 1, 0),
+        );
+        const source = organizeImports(imports).trim() + '\n';
+        results.push({
+            title: 'Organize imports',
+            kind: CodeActionKind.SourceOrganizeImports,
+            isPreferred: true,
+            edit: {
+                changes: {
+                    [uri]: [TextEdit.replace(range, source)],
+                },
+            },
+        });
+    }
+
+    // Import quick-fix actions
     event.context?.diagnostics?.forEach((diagnostic) => {
-        const uri = event.textDocument.uri;
         const name = /unbound variable ([a-z0-9_]+)/i.exec(
             diagnostic.message,
         )?.[1];
@@ -808,9 +843,10 @@ connection.onCodeAction((event) => {
             context.importResolver.getImportPaths(name, uri).forEach((path) => {
                 // Add import suggestion
                 results.push({
+                    title: `Import "${path}"`,
                     kind: CodeActionKind.QuickFix,
                     isPreferred: true,
-                    title: `Import "${path}"`,
+                    diagnostics: [diagnostic],
                     edit: {
                         changes: {
                             [uri]: [