Add interface for getting managed files

Project loading can now depend on a single interface that provides
access to managed files, rather than ServerState itself.

src/main/java/software/amazon/smithy/lsp/ManagedFiles.java

@@ -0,0 +1,16 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.lsp;
+
+import software.amazon.smithy.lsp.document.Document;
+
+public interface ManagedFiles {
+    /**
+     * @param uri Uri of the document to get
+     * @return The document if found and it is managed, otherwise {@code null}
+     */
+    Document getManagedDocument(String uri);
+}

src/main/java/software/amazon/smithy/lsp/ServerState.java

@@ -31,7 +31,7 @@
 /**
  * Keeps track of the state of the server.
  */
-public final class ServerState {
+public final class ServerState implements ManagedFiles {
     private static final Logger LOGGER = Logger.getLogger(ServerState.class.getName());
 
     private final Map<String, Project> projects;
@@ -73,10 +73,7 @@ public Project findProjectByRoot(String root) {
         return projects.get(root);
     }
 
-    /**
-     * @param uri Uri of the document to get
-     * @return The document if found and it is managed, otherwise {@code null}
-     */
+    @Override
     public Document getManagedDocument(String uri) {
         if (managedUris.contains(uri)) {
             ProjectAndFile projectAndFile = findProjectAndFile(uri);
@@ -88,19 +85,6 @@ public Document getManagedDocument(String uri) {
         return null;
     }
 
-    /**
-     * @param path The path of the document to get
-     * @return The document if found and it is managed, otherwise {@code null}
-     */
-    public Document getManagedDocument(Path path) {
-        if (managedUris.isEmpty()) {
-            return null;
-        }
-
-        String uri = LspAdapter.toUri(path.toString());
-        return getManagedDocument(uri);
-    }
-
     ProjectAndFile findProjectAndFile(String uri) {
         String path = LspAdapter.toPath(uri);
 

src/main/java/software/amazon/smithy/lsp/project/ProjectConfigLoader.java

@@ -15,8 +15,9 @@
 import java.util.Map;
 import java.util.logging.Logger;
 import software.amazon.smithy.build.model.SmithyBuildConfig;
-import software.amazon.smithy.lsp.ServerState;
+import software.amazon.smithy.lsp.ManagedFiles;
 import software.amazon.smithy.lsp.document.Document;
+import software.amazon.smithy.lsp.protocol.LspAdapter;
 import software.amazon.smithy.lsp.util.Result;
 import software.amazon.smithy.model.node.Node;
 import software.amazon.smithy.model.node.NodeMapper;
@@ -87,11 +88,7 @@ public final class ProjectConfigLoader {
     private ProjectConfigLoader() {
     }
 
-    static Result<ProjectConfig, List<Exception>> loadFromRoot(Path workspaceRoot) {
-        return loadFromRoot(workspaceRoot, new ServerState());
-    }
-
-    static Result<ProjectConfig, List<Exception>> loadFromRoot(Path workspaceRoot, ServerState state) {
+    static Result<ProjectConfig, List<Exception>> loadFromRoot(Path workspaceRoot, ManagedFiles managedFiles) {
         SmithyBuildConfig.Builder builder = SmithyBuildConfig.builder();
         List<Exception> exceptions = new ArrayList<>();
         Map<String, BuildFile> buildFiles = new HashMap<>();
@@ -100,7 +97,7 @@ static Result<ProjectConfig, List<Exception>> loadFromRoot(Path workspaceRoot, S
         if (Files.isRegularFile(smithyBuildPath)) {
             LOGGER.info("Loading smithy-build.json from " + smithyBuildPath);
             Result<SmithyBuildConfig, Exception> result = Result.ofFallible(() -> {
-                BuildFile buildFile = addBuildFile(buildFiles, smithyBuildPath, state);
+                BuildFile buildFile = addBuildFile(buildFiles, smithyBuildPath, managedFiles);
                 return SmithyBuildConfig.fromNode(
                         Node.parseJsonWithComments(buildFile.document().copyText(), buildFile.path()));
             });
@@ -116,7 +113,7 @@ static Result<ProjectConfig, List<Exception>> loadFromRoot(Path workspaceRoot, S
             Path extPath = workspaceRoot.resolve(ext);
             if (Files.isRegularFile(extPath)) {
                 Result<SmithyBuildExtensions, Exception> result = Result.ofFallible(() -> {
-                    BuildFile buildFile = addBuildFile(buildFiles, extPath, state);
+                    BuildFile buildFile = addBuildFile(buildFiles, extPath, managedFiles);
                     return loadSmithyBuildExtensions(buildFile);
                 });
                 result.get().ifPresent(extensionsBuilder::merge);
@@ -129,7 +126,7 @@ static Result<ProjectConfig, List<Exception>> loadFromRoot(Path workspaceRoot, S
         if (Files.isRegularFile(smithyProjectPath)) {
             LOGGER.info("Loading .smithy-project.json from " + smithyProjectPath);
             Result<ProjectConfig.Builder, Exception> result = Result.ofFallible(() -> {
-                BuildFile buildFile = addBuildFile(buildFiles, smithyProjectPath, state);
+                BuildFile buildFile = addBuildFile(buildFiles, smithyProjectPath, managedFiles);
                 return ProjectConfig.Builder.load(buildFile);
             });
             if (result.isOk()) {
@@ -154,14 +151,16 @@ static Result<ProjectConfig, List<Exception>> loadFromRoot(Path workspaceRoot, S
         return Result.ok(finalConfigBuilder.build());
     }
 
-    private static BuildFile addBuildFile(Map<String, BuildFile> buildFiles, Path path, ServerState state) {
-        Document managed = state.getManagedDocument(path);
+    private static BuildFile addBuildFile(Map<String, BuildFile> buildFiles, Path path, ManagedFiles managedFiles) {
+        String pathString = path.toString();
+        String uri = LspAdapter.toUri(pathString);
+        Document managed = managedFiles.getManagedDocument(uri);
         BuildFile buildFile;
         if (managed != null) {
-            buildFile = new BuildFile(path.toString(), managed);
+            buildFile = new BuildFile(pathString, managed);
         } else {
             Document document = Document.of(IoUtils.readUtf8File(path));
-            buildFile = new BuildFile(path.toString(), document);
+            buildFile = new BuildFile(pathString, document);
         }
         buildFiles.put(buildFile.path(), buildFile);
         return buildFile;

src/main/java/software/amazon/smithy/lsp/project/ProjectLoader.java

@@ -22,7 +22,7 @@
 import java.util.function.Supplier;
 import java.util.logging.Logger;
 import java.util.stream.Stream;
-import software.amazon.smithy.lsp.ServerState;
+import software.amazon.smithy.lsp.ManagedFiles;
 import software.amazon.smithy.lsp.document.Document;
 import software.amazon.smithy.lsp.protocol.LspAdapter;
 import software.amazon.smithy.lsp.util.Result;
@@ -47,7 +47,7 @@ private ProjectLoader() {
     /**
      * Loads a detachedProjects (single-file) {@link Project} with the given file.
      *
-     * <p>Unlike {@link #load(Path, ServerState)}, this method isn't
+     * <p>Unlike {@link #load(Path, ManagedFiles)}, this method isn't
      * fallible since it doesn't do any IO that we would want to recover an
      * error from.
      *
@@ -106,11 +106,11 @@ public static Project loadDetached(String uri, String text) {
      * reason about how the project was structured.
      *
      * @param root Path of the project root
-     * @param state Server's current state
+     * @param managedFiles Files managed by the server
      * @return Result of loading the project
      */
-    public static Result<Project, List<Exception>> load(Path root, ServerState state) {
-        Result<ProjectConfig, List<Exception>> configResult = ProjectConfigLoader.loadFromRoot(root, state);
+    public static Result<Project, List<Exception>> load(Path root, ManagedFiles managedFiles) {
+        Result<ProjectConfig, List<Exception>> configResult = ProjectConfigLoader.loadFromRoot(root, managedFiles);
         if (configResult.isErr()) {
             return Result.err(configResult.unwrapErr());
         }
@@ -129,7 +129,7 @@ public static Result<Project, List<Exception>> load(Path root, ServerState state
 
         LoadModelResult result;
         try {
-            result = doLoad(state::getManagedDocument, dependencies, allSmithyFilePaths);
+            result = doLoad(managedFiles, dependencies, allSmithyFilePaths);
         } catch (Exception e) {
             return Result.err(Collections.singletonList(e));
         }
@@ -154,10 +154,6 @@ private record LoadModelResult(
     ) {
     }
 
-    private interface ManagedFiles {
-        Document getManagedDocument(String uri);
-    }
-
     private static LoadModelResult doLoad(
             ManagedFiles managedFiles,
             List<Path> dependencies,
@@ -178,7 +174,7 @@ private static LoadModelResult doLoad(
         ValidatedResult<Model> modelResult = loadModel(managedFiles, allSmithyFilePaths, assembler, smithyFiles);
 
         Project.RebuildIndex rebuildIndex = Project.RebuildIndex.create(modelResult);
-        addExtraSmithyFiles(managedFiles, rebuildIndex.filesToDefinedShapes().keySet(), smithyFiles);
+        addDependencySmithyFiles(managedFiles, rebuildIndex.filesToDefinedShapes().keySet(), smithyFiles);
 
         return new LoadModelResult(
                 assemblerFactory,
@@ -207,7 +203,8 @@ private static ValidatedResult<Model> loadModel(
         return assembler.assemble();
     }
 
-    private static void addExtraSmithyFiles(
+    // Smithy files in jars were loaded by the model assembler via model discovery, so we need to collect those.
+    private static void addDependencySmithyFiles(
             ManagedFiles managedFiles,
             Set<String> loadedSmithyFilePaths,
             Map<String, SmithyFile> smithyFiles

src/test/java/software/amazon/smithy/lsp/project/ProjectConfigLoaderTest.java

@@ -18,14 +18,15 @@
 import org.junit.jupiter.api.Test;
 import software.amazon.smithy.build.model.MavenConfig;
 import software.amazon.smithy.build.model.MavenRepository;
+import software.amazon.smithy.lsp.ServerState;
 import software.amazon.smithy.lsp.util.Result;
 
 public class ProjectConfigLoaderTest {
     @Test
     public void loadsConfigWithEnvVariable() {
         System.setProperty("FOO", "bar");
         Path root = toPath(getClass().getResource("env-config"));
-        Result<ProjectConfig, List<Exception>> result = ProjectConfigLoader.loadFromRoot(root);
+        Result<ProjectConfig, List<Exception>> result = load(root);
 
         assertThat(result.isOk(), is(true));
         ProjectConfig config = result.unwrap();
@@ -41,7 +42,7 @@ public void loadsConfigWithEnvVariable() {
     @Test
     public void loadsLegacyConfig() {
         Path root = toPath(getClass().getResource("legacy-config"));
-        Result<ProjectConfig, List<Exception>> result = ProjectConfigLoader.loadFromRoot(root);
+        Result<ProjectConfig, List<Exception>> result = load(root);
 
         assertThat(result.isOk(), is(true));
         ProjectConfig config = result.unwrap();
@@ -56,7 +57,7 @@ public void loadsLegacyConfig() {
     @Test
     public void prefersNonLegacyConfig() {
         Path root = toPath(getClass().getResource("legacy-config-with-conflicts"));
-        Result<ProjectConfig, List<Exception>> result = ProjectConfigLoader.loadFromRoot(root);
+        Result<ProjectConfig, List<Exception>> result = load(root);
 
         assertThat(result.isOk(), is(true));
         ProjectConfig config = result.unwrap();
@@ -71,10 +72,14 @@ public void prefersNonLegacyConfig() {
     @Test
     public void mergesBuildExts() {
         Path root = toPath(getClass().getResource("build-exts"));
-        Result<ProjectConfig, List<Exception>> result = ProjectConfigLoader.loadFromRoot(root);
+        Result<ProjectConfig, List<Exception>> result = load(root);
 
         assertThat(result.isOk(), is(true));
         ProjectConfig config = result.unwrap();
         assertThat(config.imports(), containsInAnyOrder(containsString("main.smithy"), containsString("other.smithy")));
     }
+
+    private static Result<ProjectConfig, List<Exception>> load(Path root) {
+        return ProjectConfigLoader.loadFromRoot(root, new ServerState());
+    }
 }