Fix LVT index calculation for long and double parameters (#7)

api/src/main/java/net/neoforged/jst/api/PsiHelper.java

@@ -1,10 +1,12 @@
 package net.neoforged.jst.api;
 
+import com.intellij.lang.jvm.types.JvmPrimitiveTypeKind;
 import com.intellij.psi.PsiClass;
-import com.intellij.psi.PsiLambdaExpression;
 import com.intellij.psi.PsiMethod;
 import com.intellij.psi.PsiModifier;
 import com.intellij.psi.PsiParameter;
+import com.intellij.psi.PsiParameterListOwner;
+import com.intellij.psi.PsiPrimitiveType;
 import com.intellij.psi.PsiTypes;
 import com.intellij.psi.SyntaxTraverser;
 import com.intellij.psi.util.CachedValueProvider;
@@ -114,28 +116,43 @@ private static boolean isNonStaticInnerClassConstructor(PsiMethod method) {
         return false;
     }
 
-    public static int getBinaryIndex(PsiParameter psiParameter, int index) {
-        var declarationScope = psiParameter.getDeclarationScope();
-        if (declarationScope instanceof PsiMethod psiMethod) {
+    /**
+     * Gets the local variable table indices of the parameters for the given method
+     * or lambda expression
+     */
+    public static int[] getParameterLvtIndices(PsiParameterListOwner methodOrLambda) {
+
+        // Account for hidden parameters before the first actual parameter
+        int currentIndex = 0;
+        if (methodOrLambda instanceof PsiMethod psiMethod) {
             if (!psiMethod.hasModifierProperty(PsiModifier.STATIC)) {
-                index++; // this pointer
+                currentIndex++; // this pointer
             }
 
             // Try to account for hidden parameters only present in bytecode since the
             // mapping data refers to parameters using those indices
             if (isEnumConstructor(psiMethod)) {
-                index += 2;
+                currentIndex += 2;
             } else if (isNonStaticInnerClassConstructor(psiMethod)) {
-                index += 1;
+                currentIndex += 1;
             }
+        }
 
-            return index;
-        } else if (declarationScope instanceof PsiLambdaExpression psiLambda) {
-            // Naming lambdas doesn't really work
-            return index;
-        } else {
-            return -1;
+        var parameters = methodOrLambda.getParameterList().getParameters();
+        var lvti = new int[parameters.length];
+        for (int i = 0; i < lvti.length; i++) {
+            lvti[i] = currentIndex++;
+            // double and long use 2 slots in the LVT
+            if (parameters[i].getType() instanceof PsiPrimitiveType primitiveType) {
+                var kind = primitiveType.getKind();
+                if (kind == JvmPrimitiveTypeKind.LONG || kind == JvmPrimitiveTypeKind.DOUBLE) {
+                    currentIndex++;
+                }
+            }
         }
+
+        return lvti;
+
     }
 
     public static boolean isRecordConstructor(PsiMethod psiMethod) {

cli/build.gradle

@@ -22,6 +22,7 @@ dependencies {
 
     testImplementation platform("org.junit:junit-bom:$junit_version")
     testImplementation 'org.junit.jupiter:junit-jupiter'
+    testImplementation "org.assertj:assertj-core:$assertj_version"
 }
 
 test {

cli/src/test/java/net/neoforged/jst/cli/PsiHelperTest.java

@@ -14,6 +14,7 @@
 
 import java.io.IOException;
 
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 
 class PsiHelperTest {
@@ -66,13 +67,12 @@ void testSignature() {
         }
 
         @Test
-        void testMethodParameterIndex() {
-            var firstParam = ctor.getParameterList().getParameter(0);
-            int index = PsiHelper.getBinaryIndex(firstParam, 0);
-            // Binary parameters are:
+        void testMethodParameterIndices() {
+            // LVT of method should be:
             // 0) this
             // 1) first method parameter
-            assertEquals(1, index);
+            assertThat(PsiHelper.getParameterLvtIndices(ctor))
+                    .containsExactly(1);
         }
     }
 
@@ -95,12 +95,11 @@ void testSignature() {
         }
 
         @Test
-        void testMethodParameterIndex() {
-            var firstParam = ctor.getParameterList().getParameter(0);
-            int index = PsiHelper.getBinaryIndex(firstParam, 0);
+        void testMethodParameterIndices() {
             // Binary parameters are:
             // 0) first method parameter
-            assertEquals(0, index);
+            assertThat(PsiHelper.getParameterLvtIndices(ctor))
+                    .containsExactly(0);
         }
     }
 
@@ -124,15 +123,14 @@ void testSignature() {
         }
 
         @Test
-        void testMethodParameterIndex() {
-            var firstParam = ctor.getParameterList().getParameter(0);
-            int index = PsiHelper.getBinaryIndex(firstParam, 0);
+        void testMethodParameterIndices() {
             // Binary parameters are:
             // 0) this
             // 1) enum literal name
             // 2) enum literal ordinal
             // 3) first method parameter
-            assertEquals(3, index);
+            assertThat(PsiHelper.getParameterLvtIndices(ctor))
+                    .containsExactly(3);
         }
     }
 
@@ -157,14 +155,13 @@ void testSignature() {
         }
 
         @Test
-        void testMethodParameterIndex() {
-            var firstParam = ctor.getParameterList().getParameter(0);
-            int index = PsiHelper.getBinaryIndex(firstParam, 0);
+        void testMethodParameterIndices() {
             // Binary parameters are:
             // 0) this
             // 1) outer class pointer
             // 2) first method parameter
-            assertEquals(2, index);
+            assertThat(PsiHelper.getParameterLvtIndices(ctor))
+                    .containsExactly(2);
         }
     }
 
@@ -189,16 +186,29 @@ void testSignature() {
         }
 
         @Test
-        void testMethodParameterIndex() {
-            var firstParam = ctor.getParameterList().getParameter(0);
-            int index = PsiHelper.getBinaryIndex(firstParam, 0);
+        void testMethodParameterIndices() {
             // Binary parameters are:
             // 0) this
             // 1) first method parameter
-            assertEquals(1, index);
+            assertThat(PsiHelper.getParameterLvtIndices(ctor))
+                    .containsExactly(1);
         }
     }
 
+    @Test
+    void testLvtIndicesForPrimitiveTypes() {
+        var m = parseSingleMethod("""
+                class Outer {
+                    static void m(byte p1, short p2, int p3, long p4, float p5, double p6, boolean p7) {
+                    }
+                }
+                """);
+
+        assertEquals("(BSIJFDZ)V", PsiHelper.getBinaryMethodSignature(m));
+        assertThat(PsiHelper.getParameterLvtIndices(m))
+                .containsExactly(0, 1, 2, 3, 5, 6, 8);
+    }
+
     private PsiMethod parseSingleMethod(@Language("JAVA") String javaCode) {
         return parseSingleElement(javaCode, PsiMethod.class);
     }

parchment/src/main/java/net/neoforged/jst/parchment/GatherReplacementsVisitor.java

@@ -79,7 +79,8 @@ public void visitElement(@NotNull PsiElement element) {
                 Map<String, String> renamedParameters = new HashMap<>();
                 List<String> parameterOrder = new ArrayList<>();
 
-                PsiParameter[] parameters = psiMethod.getParameterList().getParameters();
+                var parameters = psiMethod.getParameterList().getParameters();
+                var parametersLvtIndices = PsiHelper.getParameterLvtIndices(psiMethod);
                 boolean hadReplacements = false;
                 for (int i = 0; i < parameters.length; i++) {
                     var psiParameter = parameters[i];
@@ -90,7 +91,7 @@ public void visitElement(@NotNull PsiElement element) {
 
                     // Parchment stores parameter indices based on the index of the parameter in the actual compiled method
                     // to account for synthetic parameter not found in the source-code, we must adjust the index accordingly.
-                    var jvmIndex = PsiHelper.getBinaryIndex(psiParameter, i);
+                    var jvmIndex = parametersLvtIndices[i];
 
                     var paramData = methodData.getParameter(jvmIndex);
                     // Optionally replace the parameter name, but skip record constructors, since those could have

tests/data/param_indices/expected/TestClass.java

@@ -2,6 +2,7 @@ public class TestClass {
     public TestClass(int mapped) {}
     public void instanceMethod(int mapped) {}
     public static void staticMethod(int mapped) {}
+    public static void staticMethodWithLongAndDouble(int mappedInt, long mappedLong, double mappedDouble, float mappedFloat) {}
     public class InnerClass {
         public InnerClass(int mapped) {}
         public void instanceMethod(int mapped) {}

tests/data/param_indices/parchment.json

@@ -48,6 +48,28 @@
               "name": "mapped"
             }
           ]
+        },
+        {
+          "name": "staticMethodWithLongAndDouble",
+          "descriptor": "(IJDF)V",
+          "parameters": [
+            {
+              "index": 0,
+              "name": "mappedInt"
+            },
+            {
+              "index": 1,
+              "name": "mappedLong"
+            },
+            {
+              "index": 3,
+              "name": "mappedDouble"
+            },
+            {
+              "index": 5,
+              "name": "mappedFloat"
+            }
+          ]
         }
       ]
     },

tests/data/param_indices/source/TestClass.java

@@ -2,6 +2,7 @@ public class TestClass {
     public TestClass(int p) {}
     public void instanceMethod(int p) {}
     public static void staticMethod(int p) {}
+    public static void staticMethodWithLongAndDouble(int p1, long p2, double p3, float p4) {}
     public class InnerClass {
         public InnerClass(int p) {}
         public void instanceMethod(int p) {}