Unwrap checkcasts for blockstate & loot contexts

codebook-lvt/src/main/java/io/papermc/codebook/lvt/InstructionUnwrapper.java

@@ -0,0 +1,134 @@
+/*
+ * codebook is a remapper utility for the PaperMC project.
+ *
+ * Copyright (c) 2023 Kyle Wood (DenWav)
+ *                    Contributors
+ *
+ * This library is free software; you can redistribute it and/or
+ * modify it under the terms of the GNU Lesser General Public
+ * License as published by the Free Software Foundation;
+ * version 3 only, no later versions.
+ *
+ * This library is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+ * Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public
+ * License along with this library; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301
+ * USA
+ */
+
+package io.papermc.codebook.lvt;
+
+import com.google.inject.Injector;
+import io.papermc.codebook.report.ReportType;
+import io.papermc.codebook.report.Reports;
+import io.papermc.codebook.report.type.CheckCastWraps;
+import java.util.Set;
+import java.util.stream.Collectors;
+import org.checkerframework.checker.nullness.qual.Nullable;
+import org.objectweb.asm.Opcodes;
+import org.objectweb.asm.tree.AbstractInsnNode;
+import org.objectweb.asm.tree.MethodInsnNode;
+import org.objectweb.asm.tree.VarInsnNode;
+
+public class InstructionUnwrapper {
+
+    private final Reports reports;
+    private final Injector reportsInjector;
+
+    private static final MethodMatcher BOX_METHODS = new MethodMatcher(Set.of(
+            new Method("java/lang/Byte", "byteValue", "()B"),
+            new Method("java/lang/Short", "shortValue", "()S"),
+            new Method("java/lang/Integer", "intValue", "()I"),
+            new Method("java/lang/Long", "longValue", "()J"),
+            new Method("java/lang/Float", "floatValue", "()F"),
+            new Method("java/lang/Double", "doubleValue", "()D"),
+            new Method("java/lang/Boolean", "booleanValue", "()Z"),
+            new Method("java/lang/Character", "charValue", "()C")));
+
+    private static final MethodMatcher UNWRAP_AFTER_CAST = new MethodMatcher(Set.of(
+            new Method(
+                    "net/minecraft/world/level/block/state/BlockState",
+                    "getValue",
+                    "(Lnet/minecraft/world/level/block/state/properties/Property;)Ljava/lang/Comparable;"),
+            new Method(
+                    "net/minecraft/world/level/storage/loot/LootContext",
+                    "getParamOrNull",
+                    "(Lnet/minecraft/world/level/storage/loot/parameters/LootContextParam;)Ljava/lang/Object;"),
+            new Method(
+                    "net/minecraft/world/level/storage/loot/LootParams$Builder",
+                    "getOptionalParameter",
+                    "(Lnet/minecraft/world/level/storage/loot/parameters/LootContextParam;)Ljava/lang/Object;")));
+
+    public InstructionUnwrapper(final Reports reports, final Injector reportsInjector) {
+        this.reports = reports;
+        this.reportsInjector = reportsInjector;
+    }
+
+    public @Nullable AbstractInsnNode unwrapFromAssignment(final VarInsnNode assignment) {
+        @Nullable AbstractInsnNode prev = assignment.getPrevious();
+        if (prev == null) {
+            return null;
+        }
+
+        // unwrap unboxing methods and the subsequent checkcast to the boxed type
+        if (prev.getOpcode() == Opcodes.INVOKEVIRTUAL && BOX_METHODS.matches(prev)) {
+            prev = prev.getPrevious();
+            if (prev != null && prev.getOpcode() == Opcodes.CHECKCAST) {
+                prev = prev.getPrevious();
+            }
+        }
+        if (prev == null) {
+            return null;
+        }
+
+        if (prev.getOpcode() == Opcodes.CHECKCAST) {
+            final AbstractInsnNode tempPrev = prev.getPrevious();
+            if (tempPrev.getOpcode() == Opcodes.INVOKEVIRTUAL
+                    || tempPrev.getOpcode() == Opcodes.INVOKEINTERFACE
+                    || tempPrev.getOpcode() == Opcodes.INVOKESTATIC) {
+                final MethodInsnNode methodInsn = (MethodInsnNode) tempPrev;
+                if (UNWRAP_AFTER_CAST.matches(methodInsn)) {
+                    prev = methodInsn;
+                } else {
+                    if (this.reports.shouldGenerate(ReportType.CHECK_CAST_WRAPS)) {
+                        this.reportsInjector.getInstance(CheckCastWraps.class).report(methodInsn);
+                    }
+                    return null;
+                }
+            }
+        }
+
+        return prev;
+    }
+
+    private record MethodMatcher(Set<Method> methods, Set<String> methodNames) {
+
+        private MethodMatcher(final Set<Method> methods) {
+            this(methods, methods.stream().map(Method::name).collect(Collectors.toUnmodifiableSet()));
+        }
+
+        boolean matches(final AbstractInsnNode insn) {
+            return insn instanceof final MethodInsnNode methodInsnNode
+                    && this.methodNames.contains(methodInsnNode.name)
+                    && this.methods.stream().anyMatch(m -> m.matches(methodInsnNode));
+        }
+    }
+
+    private record Method(String owner, String name, String desc, boolean itf) {
+
+        private Method(final String owner, final String name, final String desc) {
+            this(owner, name, desc, false);
+        }
+
+        boolean matches(final MethodInsnNode insn) {
+            return this.owner.equals(insn.owner)
+                    && this.name.equals(insn.name)
+                    && this.desc.equals(insn.desc)
+                    && this.itf == insn.itf;
+        }
+    }
+}

codebook-lvt/src/main/java/io/papermc/codebook/lvt/LvtNamer.java

@@ -70,7 +70,8 @@ public LvtNamer(final HypoContext context, final MappingSet mappings, final Repo
         this.lvtTypeSuggester = new LvtTypeSuggester(context);
         this.reports = reports;
         this.reportsInjector = Guice.createInjector(reports);
-        this.lvtAssignSuggester = new RootLvtSuggester(context, this.lvtTypeSuggester, this.reportsInjector);
+        this.lvtAssignSuggester =
+                new RootLvtSuggester(context, this.lvtTypeSuggester, this.reports, this.reportsInjector);
     }
 
     public void processClass(final AsmClassData classData) throws IOException {

codebook-lvt/src/main/java/io/papermc/codebook/lvt/RootLvtSuggester.java

@@ -44,18 +44,20 @@
 import io.papermc.codebook.lvt.suggestion.SingleVerbSuggester;
 import io.papermc.codebook.lvt.suggestion.StringSuggester;
 import io.papermc.codebook.lvt.suggestion.VerbPrefixBooleanSuggester;
+import io.papermc.codebook.lvt.suggestion.context.AssignmentContext;
 import io.papermc.codebook.lvt.suggestion.context.ContainerContext;
+import io.papermc.codebook.lvt.suggestion.context.SuggesterContext;
 import io.papermc.codebook.lvt.suggestion.context.field.FieldCallContext;
 import io.papermc.codebook.lvt.suggestion.context.field.FieldInsnContext;
 import io.papermc.codebook.lvt.suggestion.context.method.MethodCallContext;
 import io.papermc.codebook.lvt.suggestion.context.method.MethodInsnContext;
 import io.papermc.codebook.lvt.suggestion.numbers.MthRandomSuggester;
 import io.papermc.codebook.lvt.suggestion.numbers.RandomSourceSuggester;
+import io.papermc.codebook.report.Reports;
 import io.papermc.codebook.report.type.MissingMethodLvtSuggestion;
 import java.io.IOException;
 import java.util.List;
 import java.util.Set;
-import java.util.stream.Collectors;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.objectweb.asm.Opcodes;
 import org.objectweb.asm.tree.AbstractInsnNode;
@@ -85,16 +87,21 @@ public final class RootLvtSuggester extends AbstractModule implements LvtSuggest
             GenericSuggester.class);
 
     private final HypoContext hypoContext;
-    private final LvtTypeSuggester lvtTypeSuggester;
+    final LvtTypeSuggester lvtTypeSuggester;
     private final Injector injector;
     private final List<? extends LvtSuggester> suggesters;
+    private final InstructionUnwrapper unwrapper;
 
     public RootLvtSuggester(
-            final HypoContext hypoContext, final LvtTypeSuggester lvtTypeSuggester, final Injector reports) {
+            final HypoContext hypoContext,
+            final LvtTypeSuggester lvtTypeSuggester,
+            final Reports reports,
+            final Injector reportsInjector) {
         this.hypoContext = hypoContext;
         this.lvtTypeSuggester = lvtTypeSuggester;
-        this.injector = reports.createChildInjector(this);
+        this.injector = reportsInjector.createChildInjector(this);
         this.suggesters = SUGGESTERS.stream().map(this.injector::getInstance).toList();
+        this.unwrapper = new InstructionUnwrapper(reports, reportsInjector);
     }
 
     @Override
@@ -139,7 +146,8 @@ public String suggestName(
         }
 
         if (assignmentNode != null) {
-            final @Nullable String suggestedName = this.suggestNameFromFirstAssignment(parent, assignmentNode);
+            final @Nullable String suggestedName = this.suggestNameFromFirstAssignment(
+                    ContainerContext.from(parent), new AssignmentContext(assignmentNode, lvt));
             if (suggestedName != null) {
                 return determineFinalName(suggestedName, scopedNames);
             }
@@ -172,50 +180,9 @@ public static String determineFinalName(final String suggestedName, final Set<St
         }
     }
 
-    private static final Set<BoxMethod> BOX_METHODS = Set.of(
-            new BoxMethod("java/lang/Byte", "byteValue", "()B"),
-            new BoxMethod("java/lang/Short", "shortValue", "()S"),
-            new BoxMethod("java/lang/Integer", "intValue", "()I"),
-            new BoxMethod("java/lang/Long", "longValue", "()J"),
-            new BoxMethod("java/lang/Float", "floatValue", "()F"),
-            new BoxMethod("java/lang/Double", "doubleValue", "()D"),
-            new BoxMethod("java/lang/Boolean", "booleanValue", "()Z"),
-            new BoxMethod("java/lang/Character", "charValue", "()C"));
-    private static final Set<String> BOX_METHOD_NAMES =
-            BOX_METHODS.stream().map(BoxMethod::name).collect(Collectors.toUnmodifiableSet());
-
-    private record BoxMethod(String owner, String name, String desc) {
-        boolean is(final MethodInsnNode node) {
-            return this.owner.equals(node.owner)
-                    && this.name.equals(node.name)
-                    && this.desc.equals(node.desc)
-                    && !node.itf;
-        }
-    }
-
-    private @Nullable AbstractInsnNode walkBack(final VarInsnNode assignmentNode) {
-        AbstractInsnNode prev = assignmentNode.getPrevious();
-        if (prev != null) {
-            final int op = prev.getOpcode();
-            if (op == Opcodes.INVOKEVIRTUAL) {
-                final MethodInsnNode methodInsnNode = (MethodInsnNode) prev;
-                if (BOX_METHOD_NAMES.contains(methodInsnNode.name)
-                        && BOX_METHODS.stream().anyMatch(bm -> bm.is(methodInsnNode))) {
-                    prev = prev.getPrevious();
-                    if (prev != null && prev.getOpcode() == Opcodes.CHECKCAST) {
-                        return prev.getPrevious();
-                    }
-                    return prev;
-                }
-            }
-            return prev;
-        }
-        return null;
-    }
-
-    private @Nullable String suggestNameFromFirstAssignment(final MethodData parent, final VarInsnNode varInsn)
-            throws IOException {
-        final @Nullable AbstractInsnNode prev = this.walkBack(varInsn);
+    private @Nullable String suggestNameFromFirstAssignment(
+            final ContainerContext container, final AssignmentContext assignment) throws IOException {
+        final @Nullable AbstractInsnNode prev = this.unwrapper.unwrapFromAssignment(assignment.assignmentNode());
         if (prev == null) {
             return null;
         }
@@ -236,36 +203,49 @@ boolean is(final MethodInsnNode node) {
             return null;
         }
 
-        return this.suggestFromMethod(
+        final @Nullable String suggestion = this.suggestFromMethod(
                 MethodCallContext.create(method),
                 MethodInsnContext.create(owner, methodInsnNode),
-                ContainerContext.from(parent));
+                container,
+                assignment,
+                new SuggesterContext(this.hypoContext, this.lvtTypeSuggester));
+        if (suggestion == null) {
+            this.injector
+                    .getInstance(MissingMethodLvtSuggestion.class)
+                    .reportMissingMethodLvtSuggestion(method, methodInsnNode);
+        }
+        return suggestion;
     }
 
     @Override
     public @Nullable String suggestFromMethod(
-            final MethodCallContext call, final MethodInsnContext insn, final ContainerContext container)
+            final MethodCallContext call,
+            final MethodInsnContext insn,
+            final ContainerContext container,
+            final AssignmentContext assignment,
+            final SuggesterContext suggester)
             throws IOException {
         @Nullable String suggestion;
         for (final LvtSuggester delegate : this.suggesters) {
-            suggestion = delegate.suggestFromMethod(call, insn, container);
+            suggestion = delegate.suggestFromMethod(call, insn, container, assignment, suggester);
             if (suggestion != null) {
                 return suggestion;
             }
         }
-        this.injector
-                .getInstance(MissingMethodLvtSuggestion.class)
-                .reportMissingMethodLvtSuggestion(call.data(), insn.node());
         return null;
     }
 
     @Override
     public @Nullable String suggestFromField(
-            final FieldCallContext call, final FieldInsnContext insn, final ContainerContext container)
+            final FieldCallContext call,
+            final FieldInsnContext insn,
+            final ContainerContext container,
+            final AssignmentContext assignment,
+            final SuggesterContext suggester)
             throws IOException {
         @Nullable String suggestion;
         for (final LvtSuggester delegate : this.suggesters) {
-            suggestion = delegate.suggestFromField(call, insn, container);
+            suggestion = delegate.suggestFromField(call, insn, container, assignment, suggester);
             if (suggestion != null) {
                 return suggestion;
             }