Don’t create new shells all the time!

[Sone.git] / src / main / java / net / pterodactylus / sone / data / ShellCache.java

diff --git a/src/main/java/net/pterodactylus/sone/data/ShellCache.java b/src/main/java/net/pterodactylus/sone/data/ShellCache.java
index 9da9a06..7e28285 100644 (file)
--- a/src/main/java/net/pterodactylus/sone/data/ShellCache.java
+++ b/src/main/java/net/pterodactylus/sone/data/ShellCache.java
@@ -17,6 +17,7 @@
 
 package net.pterodactylus.sone.data;
 
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -31,7 +32,27 @@ import java.util.Map;
 public class ShellCache<T> {
 
        /** The object cache. */
-       private final Map<String, T> cache = new HashMap<String, T>();
+       private final Map<String, T> objectCache = new HashMap<String, T>();
+
+       /** The shell cache. */
+       private final Map<String, Shell<T>> shellCache = new HashMap<String, Shell<T>>();
+
+       /** The shell creator. */
+       private final ShellCreator<T> shellCreator;
+
+       /**
+        * Creates a new shell cache.
+        *
+        * @param shellCreator
+        *            The creator for new shells
+        */
+       public ShellCache(ShellCreator<T> shellCreator) {
+               this.shellCreator = shellCreator;
+       }
+
+       //
+       // ACCESSORS
+       //
 
        /**
         * Stores the given object in this cache. If the given object is not a
@@ -42,11 +63,20 @@ public class ShellCache<T> {
         *            The ID of the object
         * @param object
         *            The object to store
+        * @return The real object, or the shell if there is no real object yet
         */
-       public void put(String id, T object) {
-               if (!(object instanceof Shell<?>) || !cache.containsKey(id)) {
-                       cache.put(id, object);
+       @SuppressWarnings("unchecked")
+       public T put(String id, T object) {
+               if (!(object instanceof Shell<?>)) {
+                       objectCache.put(id, object);
+                       shellCache.remove(id);
+                       return object;
                }
+               if (objectCache.containsKey(id)) {
+                       return objectCache.get(id);
+               }
+               shellCache.put(id, (Shell<T>) object);
+               return object;
        }
 
        /**
@@ -58,7 +88,23 @@ public class ShellCache<T> {
         *         with the given ID
         */
        public T get(String id) {
-               return cache.get(id);
+               if (!objectCache.containsKey(id)) {
+                       if (!shellCache.containsKey(id)) {
+                               Shell<T> shell = shellCreator.createShell(id);
+                               shellCache.put(id, shell);
+                       }
+                       return shellCache.get(id).getShelled();
+               }
+               return objectCache.get(id);
+       }
+
+       /**
+        * Returns all cached shells.
+        *
+        * @return All cached shells
+        */
+       public Collection<Shell<T>> getShells() {
+               return shellCache.values();
        }
 
 }