Cache list-classpath

2 files changed, 46 insertions, 0 deletions

diff --git a/src/main/java/org/javacs/LruCache.java b/src/main/java/org/javacs/LruCache.java
new file mode 100644
index 0000000..5f82388
--- /dev/null
+++ b/src/main/java/org/javacs/LruCache.java
@@ -0,0 +1,33 @@
+package org.javacs;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.ArrayBlockingQueue;
+import java.util.function.Function;
+
+class LruCache<K, V> {
+    private final ArrayBlockingQueue<K> fifo;
+    private final Map<K, V> map;
+    private final Function<K, V> loader;
+
+    LruCache(int capacity, Function<K, V> loader) {
+        this.fifo = new ArrayBlockingQueue<>(capacity);
+        this.map = new HashMap<K, V>(capacity);
+        this.loader = loader;
+    }
+
+    public V get(K key) {
+        // If we already have key in cache, return it
+        if (map.containsKey(key)) return map.get(key);
+        // If we need to make room for another entry, remove the oldest entry
+        if (fifo.remainingCapacity() == 0) {
+            var evict = fifo.remove();
+            map.remove(evict);
+        }
+        // Add key to map
+        map.put(key, loader.apply(key));
+        fifo.add(key);
+
+        return map.get(key);
+    }
+}
diff --git a/src/main/java/org/javacs/SourceFileManager.java b/src/main/java/org/javacs/SourceFileManager.java
index b7e701e..dfef972 100644
--- a/src/main/java/org/javacs/SourceFileManager.java
+++ b/src/main/java/org/javacs/SourceFileManager.java
@@ -22,17 +22,30 @@ class SourceFileManager extends ForwardingJavaFileManager<StandardJavaFileManage
         LOG.warning(error.getMessage(null));
     }
 
+    private final LruCache<String, Iterable<JavaFileObject>> cacheClassPath = new LruCache<>(1000, this::listClassPath);
+
     @Override
     public Iterable<JavaFileObject> list(
             Location location, String packageName, Set<JavaFileObject.Kind> kinds, boolean recurse) throws IOException {
         if (location == StandardLocation.SOURCE_PATH) {
             var stream = FileStore.list(packageName).stream().map(this::asJavaFileObject);
             return stream::iterator;
+        } else if (location == StandardLocation.CLASS_PATH) {
+            // Listing large class paths is expensive
+            return cacheClassPath.get(packageName);
         } else {
             return super.list(location, packageName, kinds, recurse);
         }
     }
 
+    private Iterable<JavaFileObject> listClassPath(String packageName) {
+        try {
+            return super.list(StandardLocation.CLASS_PATH, packageName, Set.of(JavaFileObject.Kind.values()), false);
+        } catch (IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     private JavaFileObject asJavaFileObject(Path file) {
         return new SourceFileObject(file);
     }