Add chain builder and adjust visitor factory

PaperMC · Feb 13, 2024 · 6b85c9d · 6b85c9d
1 parent 433ac83
commit 6b85c9d
Show file tree

Hide file tree

Showing 7 changed files with 148 additions and 94 deletions.
diff --git a/...ection-rewriter/rewriter/src/main/java/io/papermc/reflectionrewriter/DefineClassRule.java b/...ection-rewriter/rewriter/src/main/java/io/papermc/reflectionrewriter/DefineClassRule.java
@@ -40,13 +40,14 @@ private DefineClassRule(final String proxyClassName, final boolean assumeClassLo
     // extend (S)CL. However since the MethodHandles.Lookup portion always needs to run, the actual benefit would
     // be beyond minute (if not actually worse).
     @Override
-    public @Nullable Rewrite rewrite(final ClassProcessingContext context,
-       final boolean invokeDynamic,
-       final int opcode,
-       final String owner,
-       final String name,
-       MethodTypeDesc descriptor,
-       final boolean isInterface
+    public @Nullable Rewrite rewrite(
+        final ClassProcessingContext context,
+        final boolean invokeDynamic,
+        final int opcode,
+        final String owner,
+        final String name,
+        MethodTypeDesc descriptor,
+        final boolean isInterface
     ) {
         if (!name.equals("defineClass") || isStatic(opcode, invokeDynamic)) {
             return null;

diff --git a/src/main/java/io/papermc/asm/RewriteRuleVisitorFactory.java b/src/main/java/io/papermc/asm/RewriteRuleVisitorFactory.java
@@ -0,0 +1,96 @@
+package io.papermc.asm;
+
+import io.papermc.asm.rules.RewriteRule;
+import java.util.function.Supplier;
+import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
+import org.checkerframework.checker.nullness.qual.NonNull;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.checkerframework.framework.qual.DefaultQualifier;
+import org.objectweb.asm.ClassVisitor;
+
+@DefaultQualifier(NonNull.class)
+public final class RewriteRuleVisitorFactory {
+
+    private final int api;
+    private final Supplier<RewriteRule> ruleFactory;
+    private volatile @MonotonicNonNull RewriteRule rule;
+    private final ClassInfoProvider classInfoProvider;
+
+    public RewriteRuleVisitorFactory(
+        final int api,
+        final Supplier<RewriteRule> ruleFactory,
+        final ClassInfoProvider classInfoProvider
+    ) {
+        this.api = api;
+        this.ruleFactory = ruleFactory;
+        this.classInfoProvider = classInfoProvider;
+    }
+
+    public RewriteRuleVisitorFactory(
+        final int api,
+        final RewriteRule rule,
+        final ClassInfoProvider classInfoProvider
+    ) {
+        this(api, () -> rule, classInfoProvider);
+    }
+
+    public ClassVisitor createVisitor(final ClassVisitor parent) {
+        final MutableProcessingContext context = new MutableProcessingContext();
+        final ClassVisitor ruleVisitor = this.rule().createVisitor(this.api, parent, context);
+        return new ContextFillerVisitor(this.api, ruleVisitor, context);
+    }
+
+    private RewriteRule rule() {
+        @Nullable RewriteRule rule = this.rule;
+        if (rule != null) {
+            return rule;
+        }
+
+        synchronized (this) {
+            rule = this.rule;
+            if (rule == null) {
+                rule = this.ruleFactory.get();
+                this.rule = rule;
+            }
+        }
+
+        return rule;
+    }
+
+    private final class MutableProcessingContext implements ClassProcessingContext {
+        private @MonotonicNonNull String name;
+        private @Nullable String superName;
+
+        @Override
+        public ClassInfoProvider classInfoProvider() {
+            return RewriteRuleVisitorFactory.this.classInfoProvider;
+        }
+
+        @Override
+        public String processingClassName() {
+            return this.name;
+        }
+
+        @Override
+        public @Nullable String processingClassSuperClassName() {
+            return this.superName;
+        }
+    }
+
+    private static final class ContextFillerVisitor extends ClassVisitor {
+
+        private final MutableProcessingContext context;
+
+        private ContextFillerVisitor(final int api, final ClassVisitor classVisitor, final MutableProcessingContext context) {
+            super(api, classVisitor);
+            this.context = context;
+        }
+
+        @Override
+        public void visit(final int version, final int access, final String name, final @Nullable String signature, final @Nullable String superName, final String @Nullable [] interfaces) {
+            super.visit(version, access, name, signature, superName, interfaces);
+            this.context.name = name;
+            this.context.superName = superName;
+        }
+    }
+}
diff --git a/src/main/java/io/papermc/asm/RewriteRulesVisitorFactory.java b/src/main/java/io/papermc/asm/RewriteRulesVisitorFactory.java
diff --git a/src/main/java/io/papermc/asm/rules/RewriteRule.java b/src/main/java/io/papermc/asm/rules/RewriteRule.java
@@ -2,25 +2,28 @@
 
 import io.papermc.asm.ClassProcessingContext;
 import io.papermc.asm.rules.builder.RuleFactory;
+import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Set;
 import java.util.function.Consumer;
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.commons.GeneratorAdapter;
 
+@FunctionalInterface
 public interface RewriteRule {
 
     RewriteRule EMPTY = (api, parent, context) -> new ClassVisitor(api, parent) {};
 
     @SafeVarargs
-    static RewriteRule forOwner(final Class<?> owner, final Consumer<? super RuleFactory> firstFactoryConsumer, final Consumer<? super RuleFactory> ...factoryConsumers) {
+    static RewriteRule forOwner(final Class<?> owner, final Consumer<? super RuleFactory> firstFactoryConsumer, final Consumer<? super RuleFactory>... factoryConsumers) {
         return forOwners(Collections.singleton(owner), firstFactoryConsumer, factoryConsumers);
     }
 
     @SafeVarargs
-    static RewriteRule forOwners(final Set<Class<?>> owners, final Consumer<? super RuleFactory> firstFactoryConsumer, final Consumer<? super RuleFactory> ...factoryConsumers) {
+    static RewriteRule forOwners(final Set<Class<?>> owners, final Consumer<? super RuleFactory> firstFactoryConsumer, final Consumer<? super RuleFactory>... factoryConsumers) {
         final RuleFactory factory = RuleFactory.create(owners);
         firstFactoryConsumer.accept(factory);
         for (final Consumer<? super RuleFactory> factoryConsumer : factoryConsumers) {
@@ -33,11 +36,15 @@ static RewriteRule chain(final RewriteRule... rules) {
         return chain(Arrays.asList(rules));
     }
 
-    static RewriteRule chain(final List<? extends RewriteRule> rules) {
+    static RewriteRule chain(final Collection<? extends RewriteRule> rules) {
         return new Chain(List.copyOf(rules));
     }
 
-    ClassVisitor createVisitor(int api, ClassVisitor parent, final ClassProcessingContext context);
+    static ChainBuilder chain() {
+        return new ChainBuilder();
+    }
+
+    ClassVisitor createVisitor(int api, ClassVisitor parent, ClassProcessingContext context);
 
     default void generateMethods(final MethodGeneratorFactory methodGeneratorFactory) {
     }
@@ -47,12 +54,9 @@ interface MethodGeneratorFactory {
         GeneratorAdapter create(int access, String name, String descriptor);
     }
 
-    final class Chain implements RewriteRule {
-
-        private final List<RewriteRule> rules;
-
-        private Chain(final List<RewriteRule> rules) {
-            this.rules = rules;
+    record Chain(List<RewriteRule> rules) implements RewriteRule {
+        public Chain(final List<RewriteRule> rules) {
+            this.rules = List.copyOf(rules);
         }
 
         @Override
@@ -69,4 +73,29 @@ public void generateMethods(final MethodGeneratorFactory methodGeneratorFactory)
             this.rules.forEach(rule -> rule.generateMethods(methodGeneratorFactory));
         }
     }
+
+    final class ChainBuilder {
+        private final List<RewriteRule> rules = new ArrayList<>();
+
+        private ChainBuilder() {
+        }
+
+        public ChainBuilder then(final RewriteRule rule) {
+            this.rules.add(rule);
+            return this;
+        }
+
+        public ChainBuilder then(final Collection<? extends RewriteRule> rules) {
+            this.rules.addAll(rules);
+            return this;
+        }
+
+        public ChainBuilder then(final RewriteRule... rules) {
+            return this.then(Arrays.asList(rules));
+        }
+
+        public RewriteRule build() {
+            return new Chain(this.rules);
+        }
+    }
 }
diff --git a/src/main/java/io/papermc/asm/rules/builder/matcher/FieldMatcher.java b/src/main/java/io/papermc/asm/rules/builder/matcher/FieldMatcher.java
@@ -9,9 +9,9 @@
 
 public interface FieldMatcher {
 
-    boolean matchesName(final String name);
+    boolean matchesName(String name);
 
-    boolean matches(final String name, final String descriptor);
+    boolean matches(String name, String descriptor);
 
     static Builder builder() {
         return new Builder();

diff --git a/src/main/java/io/papermc/asm/rules/builder/matcher/MethodMatcher.java b/src/main/java/io/papermc/asm/rules/builder/matcher/MethodMatcher.java
@@ -10,9 +10,9 @@
 
 public interface MethodMatcher {
 
-    boolean matchesName(final String name);
+    boolean matchesName(String name);
 
-    boolean matches(final String name, final String descriptor);
+    boolean matches(String name, String descriptor);
 
     static Builder builder() {
         return new Builder();

diff --git a/src/main/java/io/papermc/asm/rules/generate/GeneratedMethodHolder.java b/src/main/java/io/papermc/asm/rules/generate/GeneratedMethodHolder.java
@@ -14,5 +14,5 @@ default ClassDesc staticRedirectOwner() {
         return this.generatedMethodOwner();
     }
 
-    void generateMethod(final Map.Entry<Executable, ? extends MethodTypeDesc> pair, final RewriteRule.MethodGeneratorFactory factory);
+    void generateMethod(Map.Entry<Executable, ? extends MethodTypeDesc> pair, RewriteRule.MethodGeneratorFactory factory);
 }