From 36136e878e245575b6aed4f20f833a2dca5a4e0e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=87=91=E6=88=9F?= <jinji.lf@alibaba-inc.com>
Date: Tue, 27 Oct 2020 11:54:00 +0800
Subject: [PATCH] fit kotlin companion object as static method

---
 .../agent/handler/SourceClassHandler.java       | 17 +++++++++++++----
 .../alibaba/testable/agent/util/ClassUtil.java  | 16 ++++++++++++++++
 .../testable/agent/util/ClassUtilTest.java      | 12 ++++++++++++
 3 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/testable-agent/src/main/java/com/alibaba/testable/agent/handler/SourceClassHandler.java b/testable-agent/src/main/java/com/alibaba/testable/agent/handler/SourceClassHandler.java
index 7873432..875b149 100644
--- a/testable-agent/src/main/java/com/alibaba/testable/agent/handler/SourceClassHandler.java
+++ b/testable-agent/src/main/java/com/alibaba/testable/agent/handler/SourceClassHandler.java
@@ -84,7 +84,8 @@ public class SourceClassHandler extends BaseClassHandler {
 
     private String getMemberInjectMethodName(List<MethodInfo> memberInjectMethodList, MethodInsnNode node) {
         for (MethodInfo m : memberInjectMethodList) {
-            if (m.getClazz().equals(node.owner) && m.getName().equals(node.name) && m.getDesc().equals(node.desc)) {
+            String nodeOwner = ClassUtil.fitCompanionClassName(node.owner);
+            if (m.getClazz().equals(nodeOwner) && m.getName().equals(node.name) && m.getDesc().equals(node.desc)) {
                 return m.getMockName();
             }
         }
@@ -166,16 +167,24 @@ public class SourceClassHandler extends BaseClassHandler {
         String testClassName = ClassUtil.getTestClassName(cn.name);
         mn.instructions.insertBefore(instructions[start], new FieldInsnNode(GETSTATIC, testClassName,
             ConstPool.TESTABLE_INJECT_REF, ClassUtil.toByteCodeClassName(testClassName)));
-        if (Opcodes.INVOKESTATIC == opcode) {
-            // append a null value if it was a static invoke
+        if (Opcodes.INVOKESTATIC == opcode || isCompanionMethod(ownerClass, opcode)) {
+            // append a null value if it was a static invoke or in kotlin companion class
             mn.instructions.insertBefore(instructions[start], new InsnNode(ACONST_NULL));
+            if (ClassUtil.isCompanionClassName(ownerClass)) {
+                // for kotlin companion class, remove the byte code of reference to "companion" static field
+                mn.instructions.remove(instructions[end - 1]);
+            }
         }
         mn.instructions.insertBefore(instructions[end], new MethodInsnNode(INVOKEVIRTUAL, testClassName,
-            substitutionMethod, addFirstParameter(method.desc, ownerClass), false));
+            substitutionMethod, addFirstParameter(method.desc, ClassUtil.fitCompanionClassName(ownerClass)), false));
         mn.instructions.remove(instructions[end]);
         return mn.instructions.toArray();
     }
 
+    private boolean isCompanionMethod(String ownerClass, int opcode) {
+        return Opcodes.INVOKEVIRTUAL == opcode && ClassUtil.isCompanionClassName(ownerClass);
+    }
+
     private String addFirstParameter(String desc, String ownerClass) {
         return "(" + ClassUtil.toByteCodeClassName(ownerClass) + desc.substring(1);
     }
diff --git a/testable-agent/src/main/java/com/alibaba/testable/agent/util/ClassUtil.java b/testable-agent/src/main/java/com/alibaba/testable/agent/util/ClassUtil.java
index 9f00b9d..f3bf44e 100644
--- a/testable-agent/src/main/java/com/alibaba/testable/agent/util/ClassUtil.java
+++ b/testable-agent/src/main/java/com/alibaba/testable/agent/util/ClassUtil.java
@@ -72,6 +72,22 @@ public class ClassUtil {
         return false;
     }
 
+    /**
+     * fit kotlin companion class name to original name
+     * @param name a class name (which could be a companion class)
+     */
+    public static boolean isCompanionClassName(String name) {
+        return name.endsWith("$Companion");
+    }
+
+    /**
+     * fit kotlin companion class name to original name
+     * @param name a class name (which could be a companion class)
+     */
+    public static String fitCompanionClassName(String name) {
+        return name.replaceAll("\\$Companion$", "");
+    }
+
     /**
      * get test class name from source class name
      * @param sourceClassName source class name
diff --git a/testable-agent/src/test/java/com/alibaba/testable/agent/util/ClassUtilTest.java b/testable-agent/src/test/java/com/alibaba/testable/agent/util/ClassUtilTest.java
index 7cd9456..c561c0a 100644
--- a/testable-agent/src/test/java/com/alibaba/testable/agent/util/ClassUtilTest.java
+++ b/testable-agent/src/test/java/com/alibaba/testable/agent/util/ClassUtilTest.java
@@ -40,5 +40,17 @@ class ClassUtilTest {
         assertEquals("Ljava/lang/String;", ClassUtil.toByteCodeClassName("java.lang.String"));
     }
 
+    @Test
+    void should_able_to_fit_companion_class_name() {
+        assertEquals("com/intellij/rt/debugger/agent/CaptureAgent$ParamKeyProvider",
+            ClassUtil.fitCompanionClassName("com/intellij/rt/debugger/agent/CaptureAgent$ParamKeyProvider"));
+        assertEquals("com/alibaba/testable/demo/BlackBox",
+            ClassUtil.fitCompanionClassName("com/alibaba/testable/demo/BlackBox"));
+        assertEquals("com/alibaba/testable/demo/BlackBox$Companion",
+            ClassUtil.fitCompanionClassName("com/alibaba/testable/demo/BlackBox$Companion$Companion"));
+        assertEquals("com/alibaba/testable/demo/BlackBox",
+            ClassUtil.fitCompanionClassName("com/alibaba/testable/demo/BlackBox$Companion"));
+    }
+
 }