Even more complex source set classloader setup, this allows me to use MC/FC classes in my ImmutableMap directly

common/build.gradle

@@ -1,17 +1,26 @@
 apply plugin: 'java'
 
 sourceSets {
-    googleextension {
+    // Only defines classes that override the package-private Google methods as public, to allow access from other
+    // packages. Loaded on the app classloader
+    googleaccess {
         java
     }
+    // Defines custom implementations of ImmutableMap (and related classes). These classes live on the transforming/knot
+    // classloader, but need to be loaded after googleaccess
+    googleimpl {
+        java {
+            compileClasspath += googleaccess.output + main.compileClasspath + main.output
+        }
+    }
     main {
         java {
-            runtimeClasspath += googleextension.output
+            runtimeClasspath += googleimpl.output
         }
     }
     test {
         java {
-            runtimeClasspath += googleextension.output
+            runtimeClasspath += googleaccess.output + googleimpl.output
         }
     }
 }
@@ -23,8 +32,8 @@ dependencies {
     modImplementation "net.fabricmc:fabric-loader:${rootProject.fabric_loader_version}"
     testImplementation(platform('org.junit:junit-bom:5.7.1'))
     testImplementation('org.junit.jupiter:junit-jupiter')
-    googleextensionCompileOnly('com.google.guava:guava:21.0')
-    googleextensionCompileOnly('org.jetbrains:annotations:19.0.0')
+    googleaccessCompileOnly('com.google.guava:guava:21.0')
+    googleaccessCompileOnly('org.jetbrains:annotations:19.0.0')
 }
 
 architectury {
@@ -33,7 +42,11 @@ architectury {
 }
 
 jar {
-    from(sourceSets.googleextension.output)
+    // Move googleaccess to a directory that does not match the package to prevent accidental loads
+    into('googleaccess') {
+        from sourceSets.googleaccess.output
+    }
+    from(sourceSets.googleimpl.output)
 }
 
 test {

common/src/googleaccess/java/com/google/common/collect/FerriteCoreEntrySetAccess.java

@@ -0,0 +1,14 @@
+package com.google.common.collect;
+
+import java.util.Map;
+
+/**
+ * Same as {@link FerriteCoreImmutableMapAccess}
+ */
+public abstract class FerriteCoreEntrySetAccess<K, V> extends ImmutableSet<Map.Entry<K, V>> {
+
+    public FerriteCoreEntrySetAccess() {}
+
+    @Override
+    public abstract boolean isPartialView();
+}

common/src/googleaccess/java/com/google/common/collect/FerriteCoreImmutableMapAccess.java

@@ -0,0 +1,17 @@
+package com.google.common.collect;
+
+import java.util.Map;
+
+/**
+ * Redeclares the package-private members of ImmutableMap as public, so it can be extended in other packages
+ */
+public abstract class FerriteCoreImmutableMapAccess<K, V> extends ImmutableMap<K, V> {
+
+    public FerriteCoreImmutableMapAccess() {}
+
+    @Override
+    public abstract ImmutableSet<Map.Entry<K, V>> createEntrySet();
+
+    @Override
+    public abstract boolean isPartialView();
+}

common/src/googleextension/java/com/google/common/collect/FerriteCoreEntrySet.java

@@ -1,43 +0,0 @@
-package com.google.common.collect;
-
-import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
-
-import java.util.Map;
-
-public class FerriteCoreEntrySet<K> extends ImmutableSet<Map.Entry<K, Comparable<?>>> {
-    private final Object viewedState;
-
-    public FerriteCoreEntrySet(Object viewedState) {
-        this.viewedState = viewedState;
-    }
-
-    @Override
-    @NotNull
-    public UnmodifiableIterator<Map.Entry<K, Comparable<?>>> iterator() {
-        return new FerriteCoreIterator<>(
-                i -> (Map.Entry<K, Comparable<?>>) FerriteCoreImmutableMap.entryByStateAndIndex.apply(viewedState, i),
-                size()
-        );
-    }
-
-    @Override
-    public int size() {
-        return FerriteCoreImmutableMap.numProperties.applyAsInt(viewedState);
-    }
-
-    @Override
-    public boolean contains(@Nullable Object object) {
-        if (!(object instanceof Map.Entry)) {
-            return false;
-        }
-        Map.Entry<?, ?> entry = (Map.Entry<?, ?>) object;
-        Object valueInMap = FerriteCoreImmutableMap.getByStateAndKey.apply(viewedState, entry.getKey());
-        return valueInMap != null && valueInMap.equals(((Map.Entry<?, ?>) object).getValue());
-    }
-
-    @Override
-    boolean isPartialView() {
-        return false;
-    }
-}

common/src/googleextension/java/com/google/common/collect/FerriteCoreImmutableMap.java

@@ -1,50 +0,0 @@
-package com.google.common.collect;
-
-import org.jetbrains.annotations.NotNull;
-import org.jetbrains.annotations.Nullable;
-
-import java.util.Map;
-import java.util.function.BiFunction;
-import java.util.function.ToIntFunction;
-
-public class FerriteCoreImmutableMap<K> extends ImmutableMap<K, Comparable<?>> {
-    // This is a very inconvenient "handle" on a FastMap, but
-    // a) we need classloader separation
-    // b) by keeping the functions static we can keep the size of the object down
-    public static ToIntFunction<Object> numProperties;
-    public static BiFunction<Object, Object, Comparable<?>> getByStateAndKey;
-    public static BiFunction<Object, Integer, Entry<?, Comparable<?>>> entryByStateAndIndex;
-
-    // Actually a FastMapStateHolder, but classloader separation…
-    private final Object viewedState;
-
-    public FerriteCoreImmutableMap(Object viewedState) {
-        this.viewedState = viewedState;
-    }
-
-    @Override
-    public int size() {
-        return numProperties.applyAsInt(viewedState);
-    }
-
-    @Override
-    public Comparable<?> get(@Nullable Object key) {
-        return getByStateAndKey.apply(viewedState, key);
-    }
-
-    @Override
-    ImmutableSet<Map.Entry<K, Comparable<?>>> createEntrySet() {
-        return new FerriteCoreEntrySet<>(viewedState);
-    }
-
-    @Override
-    @NotNull
-    public ImmutableSet<Entry<K, Comparable<?>>> entrySet() {
-        return new FerriteCoreEntrySet<>(viewedState);
-    }
-
-    @Override
-    boolean isPartialView() {
-        return false;
-    }
-}

common/src/googleextension/java/com/google/common/collect/FerriteCoreIterator.java

@@ -1,28 +0,0 @@
-package com.google.common.collect;
-
-import java.util.Map;
-import java.util.function.IntFunction;
-
-public class FerriteCoreIterator<K, V> extends UnmodifiableIterator<Map.Entry<K, V>> {
-    private final IntFunction<Map.Entry<K, V>> getIth;
-    private final int length;
-
-    private int currentIndex;
-
-    public FerriteCoreIterator(IntFunction<Map.Entry<K, V>> getIth, int length) {
-        this.getIth = getIth;
-        this.length = length;
-    }
-
-    @Override
-    public boolean hasNext() {
-        return currentIndex < length;
-    }
-
-    @Override
-    public Map.Entry<K, V> next() {
-        Map.Entry<K, V> next = getIth.apply(currentIndex);
-        ++currentIndex;
-        return next;
-    }
-}

common/src/googleimpl/java/malte0811/ferritecore/fastmap/immutable/FastMapEntryEntrySet.java

@@ -0,0 +1,44 @@
+package malte0811.ferritecore.fastmap.immutable;
+
+import com.google.common.collect.FerriteCoreEntrySetAccess;
+import com.google.common.collect.UnmodifiableIterator;
+import malte0811.ferritecore.ducks.FastMapStateHolder;
+import net.minecraft.state.Property;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+import java.util.Map;
+
+public class FastMapEntryEntrySet extends FerriteCoreEntrySetAccess<Property<?>, Comparable<?>> {
+    private final FastMapStateHolder<?> viewedState;
+
+    public FastMapEntryEntrySet(FastMapStateHolder<?> viewedState) {
+        this.viewedState = viewedState;
+    }
+
+    @Override
+    @NotNull
+    public UnmodifiableIterator<Map.Entry<Property<?>, Comparable<?>>> iterator() {
+        return new FastMapEntryIterator(viewedState);
+    }
+
+    @Override
+    public int size() {
+        return viewedState.getStateMap().numProperties();
+    }
+
+    @Override
+    public boolean contains(@Nullable Object object) {
+        if (!(object instanceof Map.Entry)) {
+            return false;
+        }
+        Map.Entry<?, ?> entry = (Map.Entry<?, ?>) object;
+        Comparable<?> valueInMap = viewedState.getStateMap().getValue(viewedState.getStateIndex(), entry.getKey());
+        return valueInMap != null && valueInMap.equals(((Map.Entry<?, ?>) object).getValue());
+    }
+
+    @Override
+    public boolean isPartialView() {
+        return false;
+    }
+}

common/src/googleimpl/java/malte0811/ferritecore/fastmap/immutable/FastMapEntryImmutableMap.java

@@ -0,0 +1,40 @@
+package malte0811.ferritecore.fastmap.immutable;
+
+import com.google.common.collect.FerriteCoreImmutableMapAccess;
+import com.google.common.collect.ImmutableSet;
+import malte0811.ferritecore.ducks.FastMapStateHolder;
+import net.minecraft.state.Property;
+import org.jetbrains.annotations.Nullable;
+
+public class FastMapEntryImmutableMap extends FerriteCoreImmutableMapAccess<Property<?>, Comparable<?>> {
+    private final FastMapStateHolder<?> viewedState;
+
+    public FastMapEntryImmutableMap(FastMapStateHolder<?> viewedState) {
+        this.viewedState = viewedState;
+    }
+
+    @Override
+    public int size() {
+        return viewedState.getStateMap().numProperties();
+    }
+
+    @Override
+    public Comparable<?> get(@Nullable Object key) {
+        return viewedState.getStateMap().getValue(viewedState.getStateIndex(), key);
+    }
+
+    @Override
+    public ImmutableSet<Entry<Property<?>, Comparable<?>>> createEntrySet() {
+        return new FastMapEntryEntrySet(viewedState);
+    }
+
+    @Override
+    public ImmutableSet<Entry<Property<?>, Comparable<?>>> entrySet() {
+        return new FastMapEntryEntrySet(viewedState);
+    }
+
+    @Override
+    public boolean isPartialView() {
+        return false;
+    }
+}

common/src/googleimpl/java/malte0811/ferritecore/fastmap/immutable/FastMapEntryIterator.java

@@ -0,0 +1,30 @@
+package malte0811.ferritecore.fastmap.immutable;
+
+import com.google.common.collect.UnmodifiableIterator;
+import malte0811.ferritecore.ducks.FastMapStateHolder;
+import net.minecraft.state.Property;
+
+import java.util.Map;
+
+public class FastMapEntryIterator extends UnmodifiableIterator<Map.Entry<Property<?>, Comparable<?>>> {
+    private final FastMapStateHolder<?> viewedState;
+    private int currentIndex = 0;
+
+    public FastMapEntryIterator(FastMapStateHolder<?> viewedState) {
+        this.viewedState = viewedState;
+    }
+
+    @Override
+    public boolean hasNext() {
+        return currentIndex < viewedState.getStateMap().numProperties();
+    }
+
+    @Override
+    public Map.Entry<Property<?>, Comparable<?>> next() {
+        Map.Entry<Property<?>, Comparable<?>> next = viewedState.getStateMap().getEntry(
+                currentIndex, viewedState.getStateIndex()
+        );
+        ++currentIndex;
+        return next;
+    }
+}

common/src/main/java/malte0811/ferritecore/classloading/FastImmutableMapDefiner.java

@@ -13,9 +13,6 @@ import java.lang.invoke.MethodHandle;
 import java.lang.invoke.MethodHandles;
 import java.lang.invoke.MethodType;
 import java.lang.reflect.Method;
-import java.util.Map;
-import java.util.function.BiFunction;
-import java.util.function.ToIntFunction;
 
 /**
  * Helper to define classes in the com.google.common.collect package without issues due to jar signing and classloaders
@@ -50,37 +47,21 @@ public class FastImmutableMapDefiner {
             }
         }
     });
+
     /**
-     * Creates a MethodHandle for the constructor of {@link com.google.common.collect.FerriteCoreImmutableMap} which
-     * takes one argument, which has to be an instance FastMapStateHolder. The Lazy also sets up the callbacks used by
-     * the map to get data out of the StateHolder
+     * Creates a MethodHandle for the constructor of FastMapEntryImmutableMap which takes one argument, which has to be
+     * an instance FastMapStateHolder. This also handles the necessary classloader acrobatics.
      */
     private static final LazyValue<MethodHandle> MAKE_IMMUTABLE_FAST_MAP = new LazyValue<>(() -> {
         try {
-            define("com.google.common.collect.FerriteCoreIterator");
-            define("com.google.common.collect.FerriteCoreEntrySet");
-            Class<?> map = define("com.google.common.collect.FerriteCoreImmutableMap");
-            map.getField("numProperties").set(
-                    null, (ToIntFunction<Object>) o -> ((FastMapStateHolder<?>) o).getStateMap().numProperties()
-            );
-            map.getField("getByStateAndKey").set(
-                    null, (BiFunction<Object, Object, Comparable<?>>) (o, key) -> {
-                        FastMapStateHolder<?> stateHolder = (FastMapStateHolder<?>) o;
-                        if (key instanceof Property<?>) {
-                            return stateHolder.getStateMap().getValue(stateHolder.getStateIndex(), (Property<?>) key);
-                        } else {
-                            return null;
-                        }
-                    }
-            );
-            map.getField("entryByStateAndIndex").set(
-                    null, (BiFunction<Object, Integer, Map.Entry<?, Comparable<?>>>) (o, key) -> {
-                        FastMapStateHolder<?> stateHolder = (FastMapStateHolder<?>) o;
-                        return stateHolder.getStateMap().getEntry(key, stateHolder.getStateIndex());
-                    }
-            );
+            // Load these in the app classloader!
+            defineInAppClassloader("com.google.common.collect.FerriteCoreEntrySetAccess");
+            defineInAppClassloader("com.google.common.collect.FerriteCoreImmutableMapAccess");
+            // This lives in the transforming classloader, but must not be loaded before the previous classes are in
+            // the app classloader!
+            Class<?> map = Class.forName("malte0811.ferritecore.fastmap.immutable.FastMapEntryImmutableMap");
             MethodHandles.Lookup lookup = MethodHandles.lookup();
-            return lookup.findConstructor(map, MethodType.methodType(void.class, Object.class));
+            return lookup.findConstructor(map, MethodType.methodType(void.class, FastMapStateHolder.class));
         } catch (Exception x) {
             throw new RuntimeException(x);
         }
@@ -96,14 +77,15 @@ public class FastImmutableMapDefiner {
         }
     }
 
-    private static Class<?> define(String name) throws Exception {
+    private static void defineInAppClassloader(String name) throws Exception {
         InputStream byteInput = FastImmutableMapDefiner.class.getResourceAsStream(
-                '/' + name.replace('.', '/') + ".class"
+                "/googleaccess/" + name.replace('.', '/') + ".class"
         );
         byte[] classBytes = new byte[byteInput.available()];
         final int bytesRead = byteInput.read(classBytes);
         Preconditions.checkState(bytesRead == classBytes.length);
-        return DEFINE_CLASS.getValue().define(classBytes, name);
+        Class<?> loaded = DEFINE_CLASS.getValue().define(classBytes, name);
+        Preconditions.checkState(loaded.getClassLoader() == ImmutableMap.class.getClassLoader());
     }
 
     private interface Definer {

common/src/main/java/malte0811/ferritecore/fastmap/FastMap.java

@@ -98,6 +98,15 @@ public class FastMap<Value> {
         return propId.getValue(stateIndex);
     }
 
+    @Nullable
+    public Comparable<?> getValue(int stateIndex, Object key) {
+        if (key instanceof Property<?>) {
+            return getValue(stateIndex, (Property<?>) key);
+        } else {
+            return null;
+        }
+    }
+
     /**
      * Returns the given property and its value in the given state
      *