fix asm remapper to actually work properly, add to tests

unimined · Jun 7, 2024 · 53119d1 · 53119d1
1 parent b7a7ba9
commit 53119d1
Show file tree

Hide file tree

Showing 11 changed files with 201 additions and 9 deletions.
diff --git a/expect-platform-test/build.gradle b/expect-platform-test/build.gradle
@@ -49,25 +49,41 @@ dependencies {
 tasks.register("aExpectPlatform", ExpectPlatformFiles) {
     platformName = "a"
     inputCollection = sourceSets.main.output
+
+    remap = [
+        "xyz/wagyourtail/unimined/expect/annotation/Environment": "xyz/wagyourtail/ept/a/Env",
+        "xyz/wagyourtail/unimined/expect/annotation/Environment\$EnvType": "xyz/wagyourtail/ept/a/Env\$EnvType",
+        "xyz/wagyourtail/unimined/expect/annotation/Environment\$EnvType.COMBINED": "JOINED",
+    ]
 }
 
 tasks.register("bExpectPlatform", ExpectPlatformFiles) {
     platformName = "b"
     inputCollection = sourceSets.main.output
+
+    remap = [
+            "xyz/wagyourtail/unimined/expect/annotation/Environment": "xyz/wagyourtail/ept/b/OnlyIn",
+            "xyz/wagyourtail/unimined/expect/annotation/Environment.value": "env",
+            "xyz/wagyourtail/unimined/expect/annotation/Environment\$EnvType": "xyz/wagyourtail/ept/b/OnlyIn\$Type",
+    ]
 }
 
 tasks.register("cExpectPlatform", ExpectPlatformFiles) {
     platformName = "c"
     inputCollection = sourceSets.main.output
+
+    remap = [
+            "xyz/wagyourtail/unimined/expect/annotation/Environment": "xyz/wagyourtail/ept/c/Environment",
+            "xyz/wagyourtail/unimined/expect/annotation/Environment.value": "type",
+            "xyz/wagyourtail/unimined/expect/annotation/Environment\$EnvType": "xyz/wagyourtail/ept/c/Environment\$EnvType",
+    ]
 }
 
 tasks.register('runA', JavaExec) {
     dependsOn(tasks.aExpectPlatform)
     classpath = sourceSets.a.runtimeClasspath + tasks.aExpectPlatform.outputCollection
     mainClass = 'xyz.wagyourtail.ept.Main'
     group = 'ept'
-
-    System.out.println("classpath: " + classpath.files)
 }
 
 tasks.register('runB', JavaExec) {
@@ -89,44 +105,74 @@ tasks.register('runAgentA', JavaExec) {
     mainClass = 'xyz.wagyourtail.ept.Main'
     group = 'ept'
 
-    expectPlatform.insertAgent(delegate, "a")
+    expectPlatform.insertAgent(delegate, "a", [
+                    "xyz/wagyourtail/unimined/expect/annotation/Environment": "xyz/wagyourtail/ept/a/Env",
+                    "xyz/wagyourtail/unimined/expect/annotation/Environment\$EnvType": "xyz/wagyourtail/ept/a/Env\$EnvType",
+                    "xyz/wagyourtail/unimined/expect/annotation/Environment\$EnvType.COMBINED": "JOINED",
+            ])
 }
 
 tasks.register('runAgentB', JavaExec) {
     classpath = sourceSets.b.runtimeClasspath + sourceSets.main.runtimeClasspath
     mainClass = 'xyz.wagyourtail.ept.Main'
     group = 'ept'
 
-    expectPlatform.insertAgent(delegate, "b")
+    expectPlatform.insertAgent(delegate, "b", [
+                    "xyz/wagyourtail/unimined/expect/annotation/Environment": "xyz/wagyourtail/ept/b/OnlyIn",
+                    "xyz/wagyourtail/unimined/expect/annotation/Environment.value": "env",
+                    "xyz/wagyourtail/unimined/expect/annotation/Environment\$EnvType": "xyz/wagyourtail/ept/b/OnlyIn\$Type",
+            ])
 }
 
 tasks.register('runAgentC', JavaExec) {
     classpath = sourceSets.c.runtimeClasspath + sourceSets.main.runtimeClasspath
     mainClass = 'xyz.wagyourtail.ept.Main'
     group = 'ept'
 
-    expectPlatform.insertAgent(delegate, "c")
+    expectPlatform.insertAgent(delegate, "c", [
+                    "xyz/wagyourtail/unimined/expect/annotation/Environment": "xyz/wagyourtail/ept/c/Environment",
+                    "xyz/wagyourtail/expect/unimined/annotation/Environment.value": "type",
+                    "xyz/wagyourtail/expect/unimined/annotation/Environment\$EnvType": "xyz/wagyourtail/ept/c/Environment\$EnvType",
+            ])
 }
 
 tasks.register('jarA', ExpectPlatformJar) {
     platformName = "a"
     inputFiles = sourceSets.main.output
     from sourceSets.a.output
     archiveFileName = "a.jar"
+
+    remap = [
+            "xyz/wagyourtail/unimined/expect/annotation/Environment": "xyz/wagyourtail/ept/a/Env",
+            "xyz/wagyourtail/unimined/expect/annotation/Environment\$EnvType": "xyz/wagyourtail/ept/a/Env\$EnvType",
+            "xyz/wagyourtail/unimined/expect/annotation/Environment\$EnvType.COMBINED": "JOINED",
+    ]
 }
 
 tasks.register('jarB', ExpectPlatformJar) {
     platformName = "b"
     inputFiles = sourceSets.main.output
     from sourceSets.b.output
     archiveFileName = "b.jar"
+
+    remap = [
+            "xyz/wagyourtail/unimined/expect/annotation/Environment": "xyz/wagyourtail/ept/b/OnlyIn",
+            "xyz/wagyourtail/unimined/expect/annotation/Environment.value": "env",
+            "xyz/wagyourtail/unimined/expect/annotation/Environment\$EnvType": "xyz/wagyourtail/ept/b/OnlyIn\$Type",
+    ]
 }
 
 tasks.register('jarC', ExpectPlatformJar) {
     platformName = "c"
     inputFiles = sourceSets.main.output
     from sourceSets.c.output
     archiveFileName = "c.jar"
+
+    remap = [
+            "xyz/wagyourtail/expect/unimined/annotation/Environment": "xyz/wagyourtail/ept/c/Environment",
+            "xyz/wagyourtail/expect/unimined/annotation/Environment.value": "type",
+            "xyz/wagyourtail/expect/unimined/annotation/Environment\$EnvType": "xyz/wagyourtail/ept/c/Environment\$EnvType",
+    ]
 }
 
 assemble.dependsOn(tasks.jarA)

diff --git a/expect-platform-test/src/a/java/xyz/wagyourtail/ept/a/Env.java b/expect-platform-test/src/a/java/xyz/wagyourtail/ept/a/Env.java
@@ -0,0 +1,19 @@
+package xyz.wagyourtail.ept.a;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ElementType.TYPE, ElementType.METHOD, ElementType.FIELD})
+public @interface Env {
+
+    EnvType value();
+
+    enum EnvType {
+        CLIENT,
+        SERVER,
+        JOINED
+    }
+}
diff --git a/expect-platform-test/src/a/java/xyz/wagyourtail/ept/a/MainImpl.java b/expect-platform-test/src/a/java/xyz/wagyourtail/ept/a/MainImpl.java
@@ -11,4 +11,14 @@ public static String platformTest2(String name) {
         return "Goodbye a! " + name;
     }
 
+    public static void environmentCheck(Env env) {
+        if (env.value() == Env.EnvType.CLIENT) {
+            System.out.println("Client only");
+        } else if (env.value() == Env.EnvType.SERVER) {
+            System.out.println("Server only");
+        } else {
+            System.out.println("Joined");
+        }
+    }
+
 }
diff --git a/expect-platform-test/src/b/java/xyz/wagyourtail/ept/b/MainImpl.java b/expect-platform-test/src/b/java/xyz/wagyourtail/ept/b/MainImpl.java
@@ -9,4 +9,14 @@ public static String platformTest(String name) {
     public static String platformTest2(String name) {
         return "Goodbye b! " + name;
     }
+
+    public static void environmentCheck(OnlyIn onlyIn) {
+        if (onlyIn.env() == OnlyIn.Type.CLIENT) {
+            System.out.println("Client only");
+        } else if (onlyIn.env() == OnlyIn.Type.SERVER) {
+            System.out.println("Server only");
+        } else {
+            System.out.println("Combined");
+        }
+    }
 }
diff --git a/expect-platform-test/src/b/java/xyz/wagyourtail/ept/b/OnlyIn.java b/expect-platform-test/src/b/java/xyz/wagyourtail/ept/b/OnlyIn.java
@@ -0,0 +1,20 @@
+package xyz.wagyourtail.ept.b;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ElementType.TYPE, ElementType.METHOD, ElementType.FIELD})
+public @interface OnlyIn {
+
+    Type env();
+
+    enum Type {
+        CLIENT,
+        SERVER,
+        COMBINED
+    }
+
+}
diff --git a/expect-platform-test/src/c/java/xyz/wagyourtail/ept/c/Environment.java b/expect-platform-test/src/c/java/xyz/wagyourtail/ept/c/Environment.java
@@ -0,0 +1,19 @@
+package xyz.wagyourtail.ept.c;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@Retention(RetentionPolicy.RUNTIME)
+@Target({ElementType.TYPE, ElementType.METHOD, ElementType.FIELD})
+public @interface Environment {
+
+    EnvType type();
+
+    enum EnvType {
+        CLIENT,
+        SERVER,
+        COMBINED
+    }
+}
diff --git a/expect-platform-test/src/c/java/xyz/wagyourtail/ept/c/MainImpl.java b/expect-platform-test/src/c/java/xyz/wagyourtail/ept/c/MainImpl.java
@@ -6,4 +6,14 @@ public static String platformTest2(String name) {
         return "Goodbye c! " + name;
     }
 
+    public static void environmentCheck(Environment env) {
+        if (env.type() == Environment.EnvType.CLIENT) {
+            System.out.println("Client only");
+        } else if (env.type() == Environment.EnvType.SERVER) {
+            System.out.println("Server only");
+        } else {
+            System.out.println("Combined");
+        }
+    }
+
 }
diff --git a/expect-platform-test/src/main/java/xyz/wagyourtail/ept/Main.java b/expect-platform-test/src/main/java/xyz/wagyourtail/ept/Main.java
@@ -1,5 +1,6 @@
 package xyz.wagyourtail.ept;
 
+import xyz.wagyourtail.unimined.expect.annotation.Environment;
 import xyz.wagyourtail.unimined.expect.annotation.ExpectPlatform;
 import xyz.wagyourtail.unimined.expect.annotation.PlatformOnly;
 import xyz.wagyourtail.unimined.expect.Target;
@@ -18,6 +19,8 @@ public static void main(String[] args) {
         try {
             Main.class.getDeclaredMethod("platformOnlyTest");
             System.out.println("platformOnlyTest exists");
+
+            environmentCheck(Main.class.getDeclaredMethod("clientTest").getAnnotation(Environment.class));
         } catch (NoSuchMethodException e) {
             System.out.println("platformOnlyTest does not exist");
         }
@@ -42,4 +45,13 @@ public static String platformTest2(String name) {
     @PlatformOnly({"a", "b"})
     public static void platformOnlyTest() {
     }
+
+    @Environment(Environment.EnvType.COMBINED)
+    public static void clientTest() {
+    }
+
+    @ExpectPlatform
+    public static void environmentCheck(Environment env) {
+        throw new AssertionError();
+    }
 }
diff --git a/src/agent/java/xyz/wagyourtail/unimined/expect/ExpectPlatformAgent.java b/src/agent/java/xyz/wagyourtail/unimined/expect/ExpectPlatformAgent.java
@@ -14,11 +14,11 @@ public class ExpectPlatformAgent {
     private static final String platform = System.getProperty(EXPECT_PLATFORM);
     private static final String remap = System.getProperty(REMAP);
 
+    private static final TransformPlatform transformPlatform = new TransformPlatform(platform, remap);
 
     public static void premain(String args, Instrumentation inst) {
-        System.out.println("[ExpectPlatformAgent] Starting agent");
         System.out.println("[ExpectPlatformAgent] Platform: " + platform);
-        System.out.println("[ExpectPlatformAgent] Remap: " + remap);
+        System.out.println("[ExpectPlatformAgent] Remap: " + transformPlatform.getRemap());
         if (platform == null) {
             throw new IllegalStateException("-D" + EXPECT_PLATFORM + " not set");
         }
@@ -30,7 +30,6 @@ public static void agentmain(String args, Instrumentation inst) {
     }
 
     public static class ExpectPlatformTransformer implements ClassFileTransformer {
-        TransformPlatform transformPlatform = new TransformPlatform(platform, remap);
 
         @Override
         public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) {

diff --git a/src/shared/java/xyz/wagyourtail/unimined/expect/BetterClassRemapper.java b/src/shared/java/xyz/wagyourtail/unimined/expect/BetterClassRemapper.java
@@ -0,0 +1,43 @@
+package xyz.wagyourtail.unimined.expect;
+
+import org.objectweb.asm.*;
+import org.objectweb.asm.commons.*;
+
+public class BetterClassRemapper extends ClassRemapper {
+
+    public BetterClassRemapper(ClassVisitor classVisitor, Remapper remapper) {
+        super(classVisitor, remapper);
+    }
+
+    @Override
+    protected AnnotationVisitor createAnnotationRemapper(String descriptor, AnnotationVisitor annotationVisitor) {
+        return new AnnotationRemapper(api, Type.getType(descriptor).getInternalName(), annotationVisitor, remapper) {
+            @Override
+            public void visitEnum(String name, String descriptor, String value) {
+                super.visitEnum(name, descriptor, remapper.mapFieldName(Type.getType(descriptor).getInternalName(), value, descriptor));
+            }
+        };
+    }
+
+    @Override
+    protected MethodVisitor createMethodRemapper(MethodVisitor methodVisitor) {
+        return new MethodRemapper(api, methodVisitor, remapper) {
+
+            @Override
+            protected AnnotationVisitor createAnnotationRemapper(String descriptor, AnnotationVisitor annotationVisitor) {
+                return BetterClassRemapper.this.createAnnotationRemapper(descriptor, annotationVisitor);
+            }
+
+        };
+    }
+
+    @Override
+    protected FieldVisitor createFieldRemapper(FieldVisitor fieldVisitor) {
+        return new FieldRemapper(api, fieldVisitor, remapper) {
+            @Override
+            public AnnotationVisitor visitAnnotation(String descriptor, boolean visible) {
+                return BetterClassRemapper.this.createAnnotationRemapper(descriptor, super.visitAnnotation(descriptor, visible));
+            }
+        };
+    }
+}
diff --git a/src/shared/java/xyz/wagyourtail/unimined/expect/TransformPlatform.java b/src/shared/java/xyz/wagyourtail/unimined/expect/TransformPlatform.java
@@ -28,6 +28,10 @@ public TransformPlatform(String platformName, Map<String, String> map) {
         remap.putAll(map);
     }
 
+    public Map<String, String> getRemap() {
+        return remap;
+    }
+
     public void transform(Path inputRoot, Path outputRoot) throws IOException {
         try (Stream<Path> files = Files.walk(inputRoot)) {
             files.parallel().forEach(path -> {
@@ -60,7 +64,7 @@ public ClassNode transform(ClassNode classNode) {
         Map<MethodNode, AnnotationNode> platformOnly = new HashMap<>();
 
         ClassNode target = new ClassNode();
-        ClassRemapper remapper = new ClassRemapper(target, new SimpleRemapper(remap));
+        ClassRemapper remapper = new BetterClassRemapper(target, new SimpleRemapper(remap));
         classNode.accept(remapper);
 
         classNode = target;