diff --git a/packages/schema/src/language-server/zmodel-workspace-manager.ts b/packages/schema/src/language-server/zmodel-workspace-manager.ts index b4146cfd0..79b5bfb5e 100644 --- a/packages/schema/src/language-server/zmodel-workspace-manager.ts +++ b/packages/schema/src/language-server/zmodel-workspace-manager.ts @@ -1,13 +1,17 @@ -import { DefaultWorkspaceManager, LangiumDocument } from 'langium'; +import { isPlugin, Model } from '@zenstackhq/language/ast'; +import { getLiteral } from '@zenstackhq/sdk'; +import { DefaultWorkspaceManager, interruptAndCheck, LangiumDocument } from 'langium'; import path from 'path'; -import { WorkspaceFolder } from 'vscode-languageserver'; -import { URI } from 'vscode-uri'; -import { STD_LIB_MODULE_NAME } from './constants'; +import { CancellationToken, WorkspaceFolder } from 'vscode-languageserver'; +import { URI, Utils } from 'vscode-uri'; +import { PLUGIN_MODULE_NAME, STD_LIB_MODULE_NAME } from './constants'; /** * Custom Langium WorkspaceManager implementation which automatically loads stdlib.zmodel */ export default class ZModelWorkspaceManager extends DefaultWorkspaceManager { + public pluginModels = new Set(); + protected async loadAdditionalDocuments( _folders: WorkspaceFolder[], _collector: (document: LangiumDocument) => void @@ -18,4 +22,115 @@ export default class ZModelWorkspaceManager extends DefaultWorkspaceManager { const stdlib = this.langiumDocuments.getOrCreateDocument(stdLibUri); _collector(stdlib); } + + override async initializeWorkspace( + folders: WorkspaceFolder[], + cancelToken = CancellationToken.None + ): Promise { + const fileExtensions = this.serviceRegistry.all.flatMap((e) => e.LanguageMetaData.fileExtensions); + const documents: LangiumDocument[] = []; + const collector = (document: LangiumDocument) => { + documents.push(document); + if (!this.langiumDocuments.hasDocument(document.uri)) { + this.langiumDocuments.addDocument(document); + } + }; + + await this.loadAdditionalDocuments(folders, collector); + await Promise.all( + folders + .map((wf) => [wf, this.getRootFolder(wf)] as [WorkspaceFolder, URI]) + .map(async (entry) => this.traverseFolder(...entry, fileExtensions, collector)) + ); + + // find plugin models + documents.forEach((doc) => { + const parsed = doc.parseResult.value as Model; + parsed.declarations.forEach((decl) => { + if (isPlugin(decl)) { + const providerField = decl.fields.find((f) => f.name === 'provider'); + if (providerField) { + const provider = getLiteral(providerField.value); + if (provider) { + this.pluginModels.add(provider); + } + } + } + }); + }); + + if (this.pluginModels.size > 0) { + console.log(`Used plugin documents: ${Array.from(this.pluginModels)}`); + + // the loaded plugin models would be removed from the set + const unLoadedPluginModels = new Set(this.pluginModels); + + await Promise.all( + folders + .map((wf) => [wf, this.getRootFolder(wf)] as [WorkspaceFolder, URI]) + .map(async (entry) => this.loadPluginModels(...entry, unLoadedPluginModels, collector)) + ); + + if (unLoadedPluginModels.size > 0) { + console.warn(`The following plugin documents could not be loaded: ${Array.from(unLoadedPluginModels)}`); + } + } + + // Only after creating all documents do we check whether we need to cancel the initialization + // The document builder will later pick up on all unprocessed documents + await interruptAndCheck(cancelToken); + await this.documentBuilder.build(documents, undefined, cancelToken); + } + + protected async loadPluginModels( + workspaceFolder: WorkspaceFolder, + folderPath: URI, + pluginModels: Set, + collector: (document: LangiumDocument) => void + ): Promise { + const content = await ( + await this.fileSystemProvider.readDirectory(folderPath) + ).sort((a, b) => { + // make sure the node_moudules folder is always the first one to be checked + // so it could be early exited if the plugin is found + if (a.isDirectory && b.isDirectory) { + const aName = Utils.basename(a.uri); + if (aName === 'node_modules') { + return -1; + } else { + return 1; + } + } else { + return 0; + } + }); + + for (const entry of content) { + if (entry.isDirectory) { + const name = Utils.basename(entry.uri); + if (name === 'node_modules') { + for (const plugin of Array.from(pluginModels)) { + const path = Utils.joinPath(entry.uri, plugin, PLUGIN_MODULE_NAME); + try { + this.fileSystemProvider.readFileSync(path); + const document = this.langiumDocuments.getOrCreateDocument(path); + collector(document); + console.log(`Adding plugin document from ${path.path}`); + + pluginModels.delete(plugin); + // early exit if all plugins are loaded + if (pluginModels.size === 0) { + return; + } + } catch { + // no-op. The module might be found in another node_modules folder + // will show the warning message eventually if not found + } + } + } else { + await this.loadPluginModels(workspaceFolder, entry.uri, pluginModels, collector); + } + } + } + } }