Optimize EnvPackageCache using a nested map

ballerina-platform · Jan 16, 2024 · 998b57c · 998b57c
1 parent a26ad33
commit 998b57c
Showing 1 changed file with 22 additions and 17 deletions.
diff --git a/...ang/src/main/java/io/ballerina/projects/internal/environment/EnvironmentPackageCache.java b/...ang/src/main/java/io/ballerina/projects/internal/environment/EnvironmentPackageCache.java
@@ -22,15 +22,20 @@
  */
 public class EnvironmentPackageCache implements WritablePackageCache {
 
-    private final Map<PackageId, Project> projects = new HashMap<>();
+    private final Map<PackageId, Project> projectsById = new HashMap<>();
+    private final Map<PackageOrg, Map<PackageName, Map<PackageVersion, Project>>>
+            projectsByOrgNameVersion = new HashMap<>();
 
     public void cache(Package pkg) {
-        projects.put(pkg.packageId(), pkg.project());
+        projectsById.put(pkg.packageId(), pkg.project());
+        projectsByOrgNameVersion.computeIfAbsent(pkg.packageOrg(), k -> new HashMap<>())
+                .computeIfAbsent(pkg.packageName(), k -> new HashMap<>())
+                .put(pkg.packageVersion(), pkg.project());
     }
 
     @Override
     public Optional<Package> getPackage(PackageId packageId) {
-        Project project = projects.get(packageId);
+        Project project = projectsById.get(packageId);
         if (project == null) {
             return Optional.empty();
         }
@@ -40,7 +45,7 @@ public Optional<Package> getPackage(PackageId packageId) {
 
     @Override
     public Package getPackageOrThrow(PackageId packageId) {
-        Project project = projects.get(packageId);
+        Project project = projectsById.get(packageId);
         if (project == null) {
             throw new IllegalStateException("Cannot find a Package for the given PackageId: " + packageId);
         }
@@ -51,23 +56,15 @@ public Package getPackageOrThrow(PackageId packageId) {
     public Optional<Package> getPackage(PackageOrg packageOrg,
                                         PackageName packageName,
                                         PackageVersion version) {
-        // Do we have a need to improve this logic?
-        for (Project project : projects.values()) {
-            PackageDescriptor pkgDesc = project.currentPackage().descriptor();
-            if (pkgDesc.org().equals(packageOrg) && pkgDesc.name().equals(packageName) &&
-                    pkgDesc.version().equals(version)) {
-                return Optional.of(project.currentPackage());
-            }
-        }
-        return Optional.empty();
+        return Optional.ofNullable(projectsByOrgNameVersion.getOrDefault(packageOrg, new HashMap<>())
+                .getOrDefault(packageName, new HashMap<>())
+                .get(version)).map(Project::currentPackage);
     }
 
     @Override
     public List<Package> getPackages(PackageOrg packageOrg, PackageName packageName) {
-        // Do we have a need to improve this logic?
-        // TODO Optimize this logic
         List<Package> foundList = new ArrayList<>();
-        for (Project project : projects.values()) {
+        for (Project project : projectsById.values()) {
             PackageManifest pkgDesc = project.currentPackage().manifest();
             if (pkgDesc.org().equals(packageOrg) &&
                     pkgDesc.name().equals(packageName)) {
@@ -79,6 +76,14 @@ public List<Package> getPackages(PackageOrg packageOrg, PackageName packageName)
 
     @Override
     public void removePackage(PackageId packageId) {
-        projects.remove(packageId);
+        Optional<Project> project = Optional.ofNullable(projectsById.get(packageId));
+        if (project.isEmpty()) {
+            return;
+        }
+        PackageDescriptor pkgDesc = project.get().currentPackage().descriptor();
+        projectsByOrgNameVersion.get(pkgDesc.org())
+                .get(pkgDesc.name())
+                .remove(pkgDesc.version());
+        projectsById.remove(packageId);
     }
 }