Handling Multiple Publicodes Models per Workspace

server/src/codeAction.ts

@@ -97,7 +97,7 @@ export function createRule(
   ctx: LSContext,
   { uri, ruleName, ruleNameToCreate, position }: CreateRuleParams,
 ) {
-  const ruleDefPos = getRuleDef(ctx, ruleName)?.defPos;
+  const ruleDefPos = getRuleDef(ctx, ruleName, fileURLToPath(uri))?.defPos;
 
   if (!ruleDefPos) {
     ctx.connection.console.error(

server/src/completion.ts

@@ -7,7 +7,7 @@ import {
   MarkupKind,
   ServerRequestHandler,
 } from "vscode-languageserver/node.js";
-import { DottedName, LSContext } from "./context";
+import { DottedName, LSContext, getModelFromFile } from "./context";
 import { RuleNode } from "publicodes";
 import { mechanisms } from "./completion-items/mechanisms";
 import { keywords } from "./completion-items/keywords";
@@ -92,8 +92,9 @@ export function completionHandler(
       return keywordsAndMechanismsCompletionItems;
     }
 
+    const model = getModelFromFile(ctx, filePath);
     const completionItems = getRuleCompletionItems(
-      ctx,
+      model?.parsedRules ?? {},
       refNodeAtCursorPosition,
       params.context?.triggerKind ?? CompletionTriggerKind.Invoked,
       fullRefName,
@@ -130,7 +131,7 @@ export function completionResolveHandler(_ctx: LSContext) {
  * @param currentRefNode The current node currenlty in completion which allows to filter only corresponding childs rules
  */
 const getRuleCompletionItems = (
-  ctx: LSContext,
+  parsedRules: Record<string, RuleNode>,
   currentRefNode: TSParser.SyntaxNode,
   triggerKind: CompletionTriggerKind,
   currRuleName: DottedName | undefined,
@@ -211,7 +212,7 @@ const getRuleCompletionItems = (
     return hasCommonNamespace;
   };
 
-  return Object.entries(ctx.parsedRules)
+  return Object.entries(parsedRules)
     .filter(([dottedName, _]) => {
       const splittedDottedName = dottedName.split(" . ");
 

server/src/context.ts

@@ -1,3 +1,5 @@
+import { existsSync } from "fs";
+import path, { dirname, join, normalize } from "path";
 import { TextDocument } from "vscode-languageserver-textdocument";
 import {
   Connection,
@@ -39,6 +41,15 @@ export type DottedName = string;
 // TODO: use the publicodes types
 export type RawPublicodes = Record<DottedName, any>;
 
+export type Model = {
+  root: FilePath;
+  ruleToFileNameMap: Map<DottedName, FilePath>;
+  rawPublicodesRules: RawPublicodes;
+  parsedRules: Record<string, any>;
+  engine: Engine<string>;
+  files: Set<FilePath>;
+};
+
 export type FileInfos = {
   // List of rules in the file extracted from the tree-sitter's CST
   ruleDefs: RuleDef[];
@@ -82,12 +93,11 @@ export type LSContext = {
   diagnosticsURI: Set<URI>;
   documentSettings: Map<string, Thenable<DocumentSettings>>;
   documents: TextDocuments<TextDocument>;
-  engine: Engine<string>;
+  fileToModel: Map<FilePath, FilePath>;
   fileInfos: Map<FilePath, FileInfos>;
   globalSettings: DocumentSettings;
-  parsedRules: Record<string, any>;
-  rawPublicodesRules: RawPublicodes;
-  ruleToFileNameMap: Map<DottedName, FilePath>;
+  models: Map<FilePath, Model>;
+  workspaceFolders: FilePath[];
 };
 
 /**
@@ -98,14 +108,188 @@ export type LSContext = {
 export function getRuleDef(
   ctx: LSContext,
   ruleName: string,
+  fromFilePath?: FilePath,
 ): RuleDef | undefined {
-  const fileName = ctx.ruleToFileNameMap.get(ruleName);
+  const fileName = getRuleFilePath(ctx, ruleName, fromFilePath);
   if (!fileName) {
-    ctx.connection.console.error(`[getRuleDefPos] file not found: ${ruleName}`);
+    ctx.connection.console.error(
+      `[getRuleDefPos] file not found: ${ruleName}`,
+    );
     return;
   }
 
   return ctx.fileInfos
     .get(fileName)
     ?.ruleDefs.find((rule) => rule.dottedName === ruleName);
 }
+
+export const PUBLICODES_CONFIG_FILES = [
+  ".publicodes.config.ts",
+  ".publicodes.config.js",
+];
+
+function findWorkspaceFolder(
+  ctx: LSContext,
+  filePath: FilePath,
+): FilePath | undefined {
+  const normalizedFile = normalize(filePath);
+
+  const matchingFolders = ctx.workspaceFolders.filter((folder) => {
+    const normalizedFolder = normalize(folder);
+
+    return (
+      normalizedFile === normalizedFolder ||
+      normalizedFile.startsWith(`${normalizedFolder}${path.sep}`)
+    );
+  });
+
+  return matchingFolders.sort((a, b) => b.length - a.length)[0];
+}
+
+export function findModelRoot(ctx: LSContext, filePath: FilePath): FilePath {
+  const workspaceRoot =
+    findWorkspaceFolder(ctx, filePath) ?? dirname(normalize(filePath));
+
+  let currentDir = dirname(normalize(filePath));
+
+  while (true) {
+    if (
+      PUBLICODES_CONFIG_FILES.some((configFile) =>
+        existsSync(join(currentDir, configFile)),
+      )
+    ) {
+      return currentDir;
+    }
+
+    if (existsSync(join(currentDir, "package.json"))) {
+      return currentDir;
+    }
+
+    if (currentDir === workspaceRoot) {
+      break;
+    }
+
+    const parent = dirname(currentDir);
+    if (parent === currentDir) {
+      break;
+    }
+    currentDir = parent;
+  }
+
+  return workspaceRoot;
+}
+
+export function getModelFromFile(
+  ctx: LSContext,
+  filePath: FilePath,
+): Model | undefined {
+  const modelRoot = ctx.fileToModel.get(filePath);
+
+  if (modelRoot == undefined) {
+    return undefined;
+  }
+
+  return ctx.models.get(modelRoot);
+}
+
+export function ensureModelForFile(
+  ctx: LSContext,
+  filePath: FilePath,
+): Model {
+  const modelRoot = findModelRoot(ctx, filePath);
+  const currentModelRoot = ctx.fileToModel.get(filePath);
+
+  if (currentModelRoot && currentModelRoot !== modelRoot) {
+    const currentModel = ctx.models.get(currentModelRoot);
+    currentModel?.files.delete(filePath);
+    currentModel?.ruleToFileNameMap.forEach((path, rule) => {
+      if (path === filePath) {
+        currentModel.ruleToFileNameMap.delete(rule);
+      }
+    });
+  }
+
+  let model = ctx.models.get(modelRoot);
+
+  if (model == undefined) {
+    model = {
+      root: modelRoot,
+      ruleToFileNameMap: new Map(),
+      rawPublicodesRules: {},
+      parsedRules: {},
+      engine: new Engine({}),
+      files: new Set(),
+    };
+    ctx.models.set(modelRoot, model);
+  }
+
+  model.files.add(filePath);
+  ctx.fileToModel.set(filePath, modelRoot);
+
+  return model;
+}
+
+export function clearRuleIndexForFile(ctx: LSContext, filePath: FilePath) {
+  const model = getModelFromFile(ctx, filePath);
+
+  if (model == undefined) {
+    return;
+  }
+
+  model.ruleToFileNameMap.forEach((path, rule) => {
+    if (path === filePath) {
+      model.ruleToFileNameMap.delete(rule);
+    }
+  });
+}
+
+export function removeFileFromModel(ctx: LSContext, filePath: FilePath) {
+  const modelRoot = ctx.fileToModel.get(filePath);
+  if (modelRoot == undefined) {
+    return;
+  }
+
+  const model = ctx.models.get(modelRoot);
+  if (model == undefined) {
+    return;
+  }
+
+  clearRuleIndexForFile(ctx, filePath);
+  model.files.delete(filePath);
+  ctx.fileToModel.delete(filePath);
+
+  if (model.files.size === 0) {
+    ctx.models.delete(modelRoot);
+  }
+}
+
+export function getRuleFilePath(
+  ctx: LSContext,
+  ruleName: string,
+  fromFilePath?: FilePath,
+): FilePath | undefined {
+  const model =
+    fromFilePath != undefined ? getModelFromFile(ctx, fromFilePath) : undefined;
+
+  if (model) {
+    return model.ruleToFileNameMap.get(ruleName);
+  }
+
+  let foundFile: FilePath | undefined = undefined;
+
+  ctx.models.forEach((model) => {
+    const file = model.ruleToFileNameMap.get(ruleName);
+
+    if (!file) {
+      return;
+    }
+
+    if (foundFile && file !== foundFile) {
+      foundFile = undefined;
+      return;
+    }
+    foundFile = file;
+  });
+
+  return foundFile;
+}

server/src/helpers.ts

@@ -1,6 +1,5 @@
-import TSParser from "tree-sitter";
 import { Range, URI } from "vscode-languageserver";
-import { LSContext, Position } from "./context";
+import { LSContext, Position, removeFileFromModel } from "./context";
 import { fileURLToPath } from "node:url";
 
 /**
@@ -43,11 +42,7 @@ export function deleteFileFromCtx(ctx: LSContext, uri: URI) {
 
   ctx.fileInfos.delete(path);
 
-  ctx.ruleToFileNameMap.forEach((path, rule) => {
-    if (path === fileURLToPath(uri)) {
-      ctx.ruleToFileNameMap.delete(rule);
-    }
-  });
+  removeFileFromModel(ctx, path);
 
   if (ctx.diagnostics.has(path)) {
     ctx.diagnostics.delete(path);

server/src/initialized.ts

@@ -2,7 +2,7 @@ import { DidChangeConfigurationNotification } from "vscode-languageserver/node.j
 import { LSContext } from "./context";
 import { parseDir } from "./parseRules";
 import validate from "./validate";
-import { pathToFileURL } from "node:url";
+import { fileURLToPath, pathToFileURL } from "node:url";
 
 export default function intializedHandler(ctx: LSContext) {
   return () => {
@@ -17,6 +17,9 @@ export default function intializedHandler(ctx: LSContext) {
     if (ctx.config.hasWorkspaceFolderCapability) {
       ctx.connection.workspace.getWorkspaceFolders().then((folders) => {
         if (folders) {
+          ctx.workspaceFolders = folders.map((folder) =>
+            fileURLToPath(folder.uri),
+          );
           folders.forEach((folder) => {
             parseDir(ctx, folder.uri);
           });

server/src/onDefinition.ts

@@ -4,7 +4,7 @@ import {
   HandlerResult,
   LocationLink,
 } from "vscode-languageserver";
-import { LSContext } from "./context";
+import { LSContext, getRuleFilePath } from "./context";
 import { getFullRefName } from "./treeSitter";
 import { fileURLToPath, pathToFileURL } from "url";
 
@@ -47,7 +47,11 @@ export default function (ctx: LSContext) {
         fileURLToPath(textDocument.uri),
         node,
       );
-      const filePath = ctx.ruleToFileNameMap.get(fullRefName);
+      const filePath = getRuleFilePath(
+        ctx,
+        fullRefName,
+        fileURLToPath(textDocument.uri),
+      );
       if (filePath == undefined) {
         return [];
       }

server/src/onHover.ts

@@ -1,5 +1,5 @@
 import { HandlerResult, Hover, HoverParams } from "vscode-languageserver";
-import { LSContext } from "./context";
+import { LSContext, getModelFromFile } from "./context";
 import { getFullRefName } from "./treeSitter";
 import { fileURLToPath } from "url";
 import { serializeEvaluation } from "publicodes";
@@ -32,14 +32,19 @@ export default function (ctx: LSContext) {
     switch (node.type) {
       case "name": {
         try {
+          const model = getModelFromFile(ctx, fileURLToPath(textDocument.uri));
+          if (!model) {
+            return;
+          }
+
           const fullRefName = getFullRefName(
             ctx,
             fileURLToPath(textDocument.uri),
             node,
           );
 
-          const rawRule = ctx.rawPublicodesRules[fullRefName];
-          const nodeValue = ctx.engine.evaluate(fullRefName);
+          const rawRule = model.rawPublicodesRules[fullRefName];
+          const nodeValue = model.engine.evaluate(fullRefName);
 
           // TODO: polish the hover message
           const value = `**${rawRule?.titre ?? fullRefName}** (${serializeEvaluation(nodeValue)})