[mlir-vscode] Added LIT test discovery functionality

mlir/utils/vscode/package.json

@@ -2,7 +2,7 @@
   "name": "vscode-mlir",
   "displayName": "MLIR",
   "description": "MLIR Language Extension",
-  "version": "0.0.12",
+  "version": "0.0.13",
   "publisher": "llvm-vs-code-extensions",
   "homepage": "https://mlir.llvm.org/",
   "icon": "icon.png",
@@ -25,7 +25,8 @@
     "onCustomEditor:mlir.bytecode",
     "onLanguage:mlir",
     "onLanguage:pdll",
-    "onLanguage:tablegen"
+    "onLanguage:tablegen",
+    "onTestingStart"
   ],
   "main": "./out/extension",
   "scripts": {
@@ -106,6 +107,14 @@
         "configuration": "./tablegen-language-configuration.json"
       }
     ],
+    "testing": {
+      "controllers": [
+        {
+          "id": "litTestController",
+          "label": "LIT Tests"
+        }
+      ]
+    },
     "grammars": [
       {
         "language": "mlir",
@@ -150,6 +159,16 @@
       "type": "object",
       "title": "MLIR",
       "properties": {
+        "lit.lit_path": {
+          "type": "string",
+          "default": "lit.py",
+          "description": "Path to the lit.py script."
+        },
+        "lit.test_root_folder": {
+          "type": "string",
+          "default": ".",
+          "description": "Path to the folder containing lit tests."
+        },
         "mlir.server_path": {
           "scope": "resource",
           "type": "string",
@@ -208,6 +227,10 @@
       }
     },
     "commands": [
+      {
+        "command": "lit.reconfigure",
+        "title": "Reconfigure LIT Test Settings"
+      },
       {
         "command": "mlir.restart",
         "title": "mlir: Restart language server"

mlir/utils/vscode/src/LIT/lit.ts

@@ -0,0 +1,178 @@
+import * as vscode from 'vscode';
+import * as fs from 'fs';
+import * as path from 'path';
+
+export class LitTestProvider implements vscode.Disposable {
+    private controller: vscode.TestController;
+    private testItemRoot: vscode.TestItem;
+
+    constructor(context: vscode.ExtensionContext) {
+        // Create the TestController and root test item
+        this.controller = vscode.tests.createTestController('litTestController', 'LIT Tests');
+        this.testItemRoot = this.controller.createTestItem('litTestsRoot', 'LIT Tests');
+        this.controller.items.add(this.testItemRoot);
+        context.subscriptions.push(this.controller);
+
+        // Discover tests initially on extension activation
+        this.discoverLitTests();
+
+        // Set up file open listener for MLIR files
+        vscode.workspace.onDidOpenTextDocument(document => {
+            if (document.uri.fsPath.endsWith('.mlir')) {
+                console.log(`MLIR file opened: ${document.uri.fsPath}`);
+                this.discoverLitTests();
+            }
+        });
+
+        // Set up run profile for running tests
+        this.controller.createRunProfile('Run Tests', vscode.TestRunProfileKind.Run, async (request, token) => {
+            const run = this.controller.createTestRun(request);
+            for (const test of request.include ?? []) {
+                await this.runTest(run, test, token);
+            }
+            run.end();
+        });
+    }
+
+    // Function to prompt the user to re-enter the LIT tool path and test folder path
+    public async reconfigureLitSettings() {
+        const config = vscode.workspace.getConfiguration('lit');
+
+        // Prompt for the lit tool path and update the configuration
+        const litToolPath = await this.promptForPath('Please re-enter the path to the lit tool');
+        if (litToolPath) {
+            await config.update('lit_path', litToolPath, vscode.ConfigurationTarget.Workspace);
+        }
+
+        // Prompt for the test folder path and update the configuration
+        const testFolderPath = await this.promptForPath('Please re-enter the path to the folder containing LIT tests');
+        if (testFolderPath) {
+            await config.update('test_root_folder', testFolderPath, vscode.ConfigurationTarget.Workspace);
+        }
+
+        // Rediscover tests after reconfiguration
+        this.discoverLitTests();
+    }
+
+    // Function to discover LIT tests and display them in the test explorer
+    private async discoverLitTests() {
+        const config = vscode.workspace.getConfiguration('lit');
+
+        // Get the lit tool path and test folder from the config
+        let litToolPath = config.get<string>('lit_path');
+        let testFolderPath = config.get<string>('test_root_folder');
+
+        // If the lit tool path or test folder is not set, prompt the user to enter them
+        if (!litToolPath) {
+            litToolPath = await this.promptForPath('Please enter the path to the lit tool');
+            if (litToolPath) {
+                await config.update('lit_path', litToolPath, vscode.ConfigurationTarget.Workspace);
+            }
+        }
+
+        if (!testFolderPath) {
+            testFolderPath = await this.promptForPath('Please enter the path to the folder containing LIT tests');
+            if (testFolderPath) {
+                await config.update('test_root_folder', testFolderPath, vscode.ConfigurationTarget.Workspace);
+            }
+        }
+
+        // Ensure both values are now set before proceeding
+        if (!litToolPath || !testFolderPath) {
+            vscode.window.showErrorMessage('LIT tool path or test folder path not set. Test discovery cannot proceed.');
+            return;
+        }
+
+        // Ensure the test folder path is absolute (relative to workspace)
+        const absoluteTestFolderPath = path.isAbsolute(testFolderPath)
+            ? testFolderPath
+            : path.join(vscode.workspace.workspaceFolders?.[0].uri.fsPath || '', testFolderPath);
+
+        if (!fs.existsSync(absoluteTestFolderPath)) {
+            vscode.window.showErrorMessage(`Test folder not found: ${absoluteTestFolderPath}`);
+            return;
+        }
+
+        // Clear previous test items before discovering new ones
+        this.testItemRoot.children.replace([]); // Use replace([]) to clear the test items
+
+        // Recursively scan the folder for LIT tests
+        this.scanDirectory(this.testItemRoot, absoluteTestFolderPath);
+    }
+
+    // Function to scan a directory for LIT tests
+    private scanDirectory(parent: vscode.TestItem, directory: string) {
+        const items = fs.readdirSync(directory, { withFileTypes: true });
+
+        for (const item of items) {
+            const itemPath = path.join(directory, item.name);
+
+            if (item.isDirectory()) {
+                // Create a new TestItem for the directory
+                const dirTestItem = this.controller.createTestItem(itemPath, item.name);
+                parent.children.add(dirTestItem);
+
+                // Recursively scan this subdirectory
+                this.scanDirectory(dirTestItem, itemPath);
+            } else if (item.isFile() && this.isLitTestFile(item.name)) {
+                // It's a file and we assume it's a LIT test file
+                const testItem = this.controller.createTestItem(itemPath, item.name, vscode.Uri.file(itemPath));
+                parent.children.add(testItem);
+            }
+        }
+    }
+
+    // A simple helper function to check if a file is a LIT test (now checks for .mlir files)
+    private isLitTestFile(filename: string): boolean {
+        return filename.endsWith('.mlir'); // Now only checks for .mlir files
+    }
+
+    // Function to run a LIT test
+    private async runTest(run: vscode.TestRun, test: vscode.TestItem, token: vscode.CancellationToken) {
+        run.started(test);
+
+        const config = vscode.workspace.getConfiguration('lit');
+        const litToolPath = config.get<string>('lit_path') || 'lit';  // Default to 'lit'
+
+        try {
+            const result = await this.runLitTest(litToolPath, test.uri!.fsPath);
+            if (result.passed) {
+                run.passed(test);
+            } else {
+                run.failed(test, new vscode.TestMessage(result.errorMessage));
+            }
+        } catch (error) {
+            run.errored(test, new vscode.TestMessage(error.message));
+        }
+
+        run.end();
+    }
+
+    // Function to execute the LIT test using the lit tool
+    private async runLitTest(litToolPath: string, testPath: string): Promise<{ passed: boolean, errorMessage?: string }> {
+        const { exec } = require('child_process');
+
+        return new Promise((resolve, reject) => {
+            exec(`${litToolPath} -v ${testPath}`, (error: any, stdout: string, stderr: string) => {
+                if (error) {
+                    resolve({ passed: false, errorMessage: `stdout: ${stdout}\stderr: ${stderr}` });
+                } else {
+                    resolve({ passed: true });
+                }
+            });
+        });
+    }
+
+    // Helper function to prompt the user for a path
+    private async promptForPath(promptMessage: string): Promise<string | undefined> {
+        return vscode.window.showInputBox({
+            prompt: promptMessage,
+            placeHolder: 'Enter a valid path'
+        });
+    }
+
+    // Implementing the dispose method to clean up resources
+    public dispose() {
+        this.controller.dispose();  // Dispose of the TestController
+    }
+}

mlir/utils/vscode/src/extension.ts

@@ -3,7 +3,7 @@ import * as vscode from 'vscode';
 import {registerMLIRExtensions} from './MLIR/mlir';
 import {MLIRContext} from './mlirContext';
 import {registerPDLLExtensions} from './PDLL/pdll';
-
+import { LitTestProvider } from './LIT/lit';
 /**
  *  This method is called when the extension is activated. The extension is
  *  activated the very first time a command is executed.
@@ -15,13 +15,19 @@ export function activate(context: vscode.ExtensionContext) {
   const mlirContext = new MLIRContext();
   context.subscriptions.push(mlirContext);
 
+  const litTests = new LitTestProvider(context);  // Instantiate the LitTestProvider
+  context.subscriptions.push(litTests);  // Push it to context.subscriptions to handle cleanup
+
   // Initialize the commands of the extension.
   context.subscriptions.push(
       vscode.commands.registerCommand('mlir.restart', async () => {
         // Dispose and reactivate the context.
         mlirContext.dispose();
         await mlirContext.activate(outputChannel);
       }));
+  context.subscriptions.push(vscode.commands.registerCommand('lit.reconfigure', () => {
+    litTests.reconfigureLitSettings();
+  }));
   registerMLIRExtensions(context, mlirContext);
   registerPDLLExtensions(context, mlirContext);