From 4ae1176157b66b2d1634919d2fdeba85359104af Mon Sep 17 00:00:00 2001
From: md_5 <git@md-5.net>
Date: Mon, 12 Jan 2015 20:32:03 +1100
Subject: [PATCH] Prevent unbounded / double cache growth.

---
 src/main/java/net/md_5/specialsource/Jar.java          | 10 +---------
 src/main/java/net/md_5/specialsource/JarComparer.java  |  6 +++++-
 .../java/net/md_5/specialsource/repo/CachingRepo.java  |  8 ++++----
 3 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/src/main/java/net/md_5/specialsource/Jar.java b/src/main/java/net/md_5/specialsource/Jar.java
index f852d7a..74d86a0 100644
--- a/src/main/java/net/md_5/specialsource/Jar.java
+++ b/src/main/java/net/md_5/specialsource/Jar.java
@@ -65,7 +65,6 @@ public class Jar {
     private final String filename;
     private final LinkedHashMap<String, JarFile> jarForResource;
     private final Set<String> contains = new HashSet<String>();
-    private final Map<String, ClassNode> classes = new HashMap<String, ClassNode>();
 
     /**
      * Check if this jar contains the given class. Takes the internal name of a
@@ -124,12 +123,6 @@ public class Jar {
      * @return
      */
     public ClassNode getNode(String clazz) {
-        // Lets try a cache hit
-        ClassNode cache = classes.get(clazz);
-        if (cache != null) {
-            return cache;
-        }
-
         // No luck, so lets try read it
         try {
             InputStream is = getClass(clazz);
@@ -138,8 +131,7 @@ public class Jar {
                 ClassReader cr = new ClassReader(is);
                 ClassNode node = new ClassNode();
                 cr.accept(node, 0);
-                // Add it to the cache 
-                classes.put(clazz, node);
+
                 return node;
             }
         } catch (IOException ex) {
diff --git a/src/main/java/net/md_5/specialsource/JarComparer.java b/src/main/java/net/md_5/specialsource/JarComparer.java
index 3e28478..f92c685 100644
--- a/src/main/java/net/md_5/specialsource/JarComparer.java
+++ b/src/main/java/net/md_5/specialsource/JarComparer.java
@@ -33,6 +33,8 @@ import java.util.Collection;
 import java.util.Map;
 import net.md_5.specialsource.provider.InheritanceProvider;
 import net.md_5.specialsource.provider.JarProvider;
+import net.md_5.specialsource.repo.ClassRepo;
+import net.md_5.specialsource.repo.JarRepo;
 import net.md_5.specialsource.util.NoDupeList;
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.FieldVisitor;
@@ -47,6 +49,7 @@ public class JarComparer extends ClassVisitor {
 
     private final MethodReferenceFinder methodVisitor = new MethodReferenceFinder();
     public final Jar jar;
+    private final ClassRepo jarRepo;
     private final InheritanceProvider inheritance;
     private String myName;
     public int iterDepth;
@@ -70,6 +73,7 @@ public class JarComparer extends ClassVisitor {
     public JarComparer(Jar jar) {
         super(Opcodes.ASM5);
         this.jar = jar;
+        this.jarRepo = new JarRepo(jar);
         this.inheritance = new JarProvider(jar);
     }
 
@@ -116,7 +120,7 @@ public class JarComparer extends ClassVisitor {
 
         String newParent = null;
 
-        ClassNode n = jar.getNode(currentParent);
+        ClassNode n = jarRepo.findClass(currentParent);
         if (n == null) {
             return newParent;
         }
diff --git a/src/main/java/net/md_5/specialsource/repo/CachingRepo.java b/src/main/java/net/md_5/specialsource/repo/CachingRepo.java
index a994955..8c57998 100644
--- a/src/main/java/net/md_5/specialsource/repo/CachingRepo.java
+++ b/src/main/java/net/md_5/specialsource/repo/CachingRepo.java
@@ -28,17 +28,17 @@
  */
 package net.md_5.specialsource.repo;
 
-import java.util.Map;
-import java.util.WeakHashMap;
+import com.google.common.cache.Cache;
+import com.google.common.cache.CacheBuilder;
 import org.objectweb.asm.tree.ClassNode;
 
 public abstract class CachingRepo implements ClassRepo {
 
-    private final Map<String, ClassNode> cache = new WeakHashMap<String, ClassNode>();
+    private final Cache<String, ClassNode> cache = CacheBuilder.newBuilder().maximumSize(4096).build();
 
     @Override
     public final ClassNode findClass(String internalName) {
-        ClassNode fromCache = cache.get(internalName);
+        ClassNode fromCache = cache.getIfPresent(internalName);
         if (fromCache != null) {
             return fromCache;
         }