From 49fc0ba47fac077a2fe965c44fd382a1ab02089c Mon Sep 17 00:00:00 2001
From: Luck <git@lucko.me>
Date: Sun, 14 Nov 2021 19:10:46 +0000
Subject: Improve Velocity ClassSourceLookup

---
 .../spark/velocity/VelocityClassSourceLookup.java    | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

(limited to 'spark-velocity/src/main')

diff --git a/spark-velocity/src/main/java/me/lucko/spark/velocity/VelocityClassSourceLookup.java b/spark-velocity/src/main/java/me/lucko/spark/velocity/VelocityClassSourceLookup.java
index 4989bf4..bcb8176 100644
--- a/spark-velocity/src/main/java/me/lucko/spark/velocity/VelocityClassSourceLookup.java
+++ b/spark-velocity/src/main/java/me/lucko/spark/velocity/VelocityClassSourceLookup.java
@@ -27,6 +27,9 @@ import me.lucko.spark.common.util.ClassSourceLookup;
 
 import org.checkerframework.checker.nullness.qual.Nullable;
 
+import java.util.HashMap;
+import java.util.Map;
+
 public class VelocityClassSourceLookup extends ClassSourceLookup.ByClassLoader {
     private static final Class<?> PLUGIN_CLASS_LOADER;
 
@@ -38,21 +41,22 @@ public class VelocityClassSourceLookup extends ClassSourceLookup.ByClassLoader {
         }
     }
 
-    private final PluginManager pluginManager;
+    private final Map<ClassLoader, String> classLoadersToPlugin;
 
     public VelocityClassSourceLookup(PluginManager pluginManager) {
-        this.pluginManager = pluginManager;
+        this.classLoadersToPlugin = new HashMap<>();
+        for (PluginContainer plugin : pluginManager.getPlugins()) {
+            plugin.getInstance().ifPresent(instance -> {
+                String id = plugin.getDescription().getName().orElseGet(() -> plugin.getDescription().getId());
+                this.classLoadersToPlugin.put(instance.getClass().getClassLoader(), id);
+            });
+        }
     }
 
     @Override
     public @Nullable String identify(ClassLoader loader) {
         if (PLUGIN_CLASS_LOADER.isInstance(loader)) {
-            for (PluginContainer plugin : this.pluginManager.getPlugins()) {
-                Object instance = plugin.getInstance().orElse(null);
-                if (instance != null && instance.getClass().getClassLoader() == loader) {
-                    return plugin.getDescription().getName().orElseGet(() -> plugin.getDescription().getId());
-                }
-            }
+            return this.classLoadersToPlugin.get(loader);
         }
         return null;
     }
-- 
cgit