From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001
From: Martijn Muijsers <martijnmuijsers@live.nl>
Date: Wed, 30 Nov 2022 21:15:33 +0100
Subject: [PATCH] Avoid Class#isAssignableFrom call in ClassInstanceMultiMap

License: MIT (https://opensource.org/licenses/MIT)
Gale - https://galemc.org

This patch is based on the following mixins:
"com/ishland/vmp/mixins/general/collections/MixinTypeFilterableList.java"
By: ishland <ishlandmc@yeah.net>
As part of: VMP (https://github.com/RelativityMC/VMP-fabric)
Licensed under: MIT (https://opensource.org/licenses/MIT)

"net/caffeinemc/mods/lithium/mixin/collections/entity_filtering/ClassInstanceMultiMapMixin.java"
By: Angeline <jellysquid3@users.noreply.github.com>
As part of: Lithium (https://github.com/CaffeineMC/lithium-fabric)
Licensed under: LGPL-3.0 (https://www.gnu.org/licenses/lgpl-3.0.html)

diff --git a/net/minecraft/util/ClassInstanceMultiMap.java b/net/minecraft/util/ClassInstanceMultiMap.java
index 4cd1b8ce60f7befca63a5b6396c0781ae3583dcd..c79aaf21566fe84fd45c8666230c6c5f4fe1db27 100644
--- a/net/minecraft/util/ClassInstanceMultiMap.java
+++ b/net/minecraft/util/ClassInstanceMultiMap.java
@@ -56,18 +56,41 @@ public class ClassInstanceMultiMap<T> extends AbstractCollection<T> {
     }
 
     public <S> Collection<S> find(Class<S> type) {
-        if (!this.baseClass.isAssignableFrom(type)) {
-            throw new IllegalArgumentException("Don't know how to search for " + type);
-        } else {
-            List<? extends T> list = this.byClass
-                .computeIfAbsent(type, clazz -> this.allInstances.stream().filter(clazz::isInstance).collect(Util.toMutableList()));
-            return (Collection<S>)Collections.unmodifiableCollection(list);
+        // Gale start - Lithium/VMP - avoid Class#isAssignableFrom call in ClassInstanceMultiMap
+        /*
+        Only perform the slow Class#isAssignableFrom(Class) if a list doesn't exist for the type, otherwise
+        we can assume it's already valid. The slow-path code is moved to a separate method to help the JVM inline this.
+        */
+        Collection<T> collection = this.byClass.get(type);
+
+        if (collection == null) {
+            collection = this.createAllOfType(type);
         }
+
+        return (Collection<S>) collection;
+    }
+
+    private <S> Collection<T> createAllOfType(Class<S> type) {
+        return this.byClass.computeIfAbsent(
+            type,
+            clazz -> {
+                it.unimi.dsi.fastutil.objects.ObjectArrayList<T> ts = new it.unimi.dsi.fastutil.objects.ObjectArrayList<>(this.allInstances.size());
+
+                for (T _allElement : ((it.unimi.dsi.fastutil.objects.ObjectArrayList<T>) this.allInstances).elements()) {
+                    if (clazz.isInstance(_allElement)) {
+                        ts.add(_allElement);
+                    }
+                }
+
+                return ts;
+            }
+        );
+        // Gale end - Lithium/VMP - avoid Class#isAssignableFrom call in ClassInstanceMultiMap
     }
 
     @Override
     public Iterator<T> iterator() {
-        return (Iterator<T>)(this.allInstances.isEmpty() ? Collections.emptyIterator() : Iterators.unmodifiableIterator(this.allInstances.iterator()));
+        return this.allInstances.isEmpty() ? Collections.emptyIterator() : Iterators.unmodifiableIterator(this.allInstances.iterator());
     }
 
     public List<T> getAllInstances() {