Display source info for mixin injected methods (#249)

Co-authored-by: Luck <git@lucko.me>

3 files changed, 246 insertions, 25 deletions

diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/AbstractSampler.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/AbstractSampler.java
index 1c217db..3cfef0b 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/AbstractSampler.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/AbstractSampler.java
@@ -164,8 +164,16 @@ public abstract class AbstractSampler implements Sampler {
             classSourceVisitor.visit(entry);
         }
 
-        if (classSourceVisitor.hasMappings()) {
-            proto.putAllClassSources(classSourceVisitor.getMapping());
+        if (classSourceVisitor.hasClassSourceMappings()) {
+            proto.putAllClassSources(classSourceVisitor.getClassSourceMapping());
+        }
+
+        if (classSourceVisitor.hasMethodSourceMappings()) {
+            proto.putAllMethodSources(classSourceVisitor.getMethodSourceMapping());
+        }
+
+        if (classSourceVisitor.hasLineSourceMappings()) {
+            proto.putAllLineSources(classSourceVisitor.getLineSourceMapping());
         }
     }
 }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/util/ClassSourceLookup.java b/spark-common/src/main/java/me/lucko/spark/common/util/ClassSourceLookup.java
index bd9ec37..668f31a 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/util/ClassSourceLookup.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/util/ClassSourceLookup.java
@@ -38,9 +38,11 @@ import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.Function;
+import java.util.stream.Collectors;
 
 /**
- * A function which defines the source of given {@link Class}es.
+ * A function which defines the source of given {@link Class}es or (Mixin) method calls.
  */
 public interface ClassSourceLookup {
 
@@ -53,6 +55,26 @@ public interface ClassSourceLookup {
     @Nullable String identify(Class<?> clazz) throws Exception;
 
     /**
+     * Identify the given method call.
+     *
+     * @param methodCall the method call info
+     * @return the source of the method call
+     */
+    default @Nullable String identify(MethodCall methodCall) throws Exception {
+        return null;
+    }
+
+    /**
+     * Identify the given method call.
+     *
+     * @param methodCall the method call info
+     * @return the source of the method call
+     */
+    default @Nullable String identify(MethodCallByLine methodCall) throws Exception {
+        return null;
+    }
+
+    /**
      * A no-operation {@link ClassSourceLookup}.
      */
     ClassSourceLookup NO_OP = new ClassSourceLookup() {
@@ -156,9 +178,17 @@ public interface ClassSourceLookup {
     interface Visitor {
         void visit(ThreadNode node);
 
-        boolean hasMappings();
+        boolean hasClassSourceMappings();
+
+        Map<String, String> getClassSourceMapping();
+
+        boolean hasMethodSourceMappings();
+
+        Map<String, String> getMethodSourceMapping();
+
+        boolean hasLineSourceMappings();
 
-        Map<String, String> getMapping();
+        Map<String, String> getLineSourceMapping();
     }
 
     static Visitor createVisitor(ClassSourceLookup lookup) {
@@ -177,25 +207,46 @@ public interface ClassSourceLookup {
         }
 
         @Override
-        public boolean hasMappings() {
+        public boolean hasClassSourceMappings() {
+            return false;
+        }
+
+        @Override
+        public Map<String, String> getClassSourceMapping() {
+            return Collections.emptyMap();
+        }
+
+        @Override
+        public boolean hasMethodSourceMappings() {
+            return false;
+        }
+
+        @Override
+        public Map<String, String> getMethodSourceMapping() {
+            return Collections.emptyMap();
+        }
+
+        @Override
+        public boolean hasLineSourceMappings() {
             return false;
         }
 
         @Override
-        public Map<String, String> getMapping() {
+        public Map<String, String> getLineSourceMapping() {
             return Collections.emptyMap();
         }
     }
 
     /**
-     * Visitor which scans {@link StackTraceNode}s and accumulates class identities.
+     * Visitor which scans {@link StackTraceNode}s and accumulates class/method call identities.
      */
     class VisitorImpl implements Visitor {
         private final ClassSourceLookup lookup;
         private final ClassFinder classFinder = new ClassFinder();
 
-        // class name --> identifier (plugin name)
-        private final Map<String, String> map = new HashMap<>();
+        private final SourcesMap<String> classSources = new SourcesMap<>(Function.identity());
+        private final SourcesMap<MethodCall> methodSources = new SourcesMap<>(MethodCall::toString);
+        private final SourcesMap<MethodCallByLine> lineSources = new SourcesMap<>(MethodCallByLine::toString);
 
         VisitorImpl(ClassSourceLookup lookup) {
             this.lookup = lookup;
@@ -208,34 +259,194 @@ public interface ClassSourceLookup {
             }
         }
 
+        private void visitStackNode(StackTraceNode node) {
+            this.classSources.computeIfAbsent(
+                    node.getClassName(),
+                    className -> {
+                        Class<?> clazz = this.classFinder.findClass(className);
+                        if (clazz == null) {
+                            return null;
+                        }
+                        return this.lookup.identify(clazz);
+                    });
+
+            if (node.getMethodDescription() != null) {
+                MethodCall methodCall = new MethodCall(node.getClassName(), node.getMethodName(), node.getMethodDescription());
+                this.methodSources.computeIfAbsent(methodCall, this.lookup::identify);
+            } else {
+                MethodCallByLine methodCall = new MethodCallByLine(node.getClassName(), node.getMethodName(), node.getLineNumber());
+                this.lineSources.computeIfAbsent(methodCall, this.lookup::identify);
+            }
+
+            // recursively
+            for (StackTraceNode child : node.getChildren()) {
+                visitStackNode(child);
+            }
+        }
+
         @Override
-        public boolean hasMappings() {
-            return !this.map.isEmpty();
+        public boolean hasClassSourceMappings() {
+            return this.classSources.hasMappings();
         }
 
         @Override
-        public Map<String, String> getMapping() {
-            this.map.values().removeIf(Objects::isNull);
-            return this.map;
+        public Map<String, String> getClassSourceMapping() {
+            return this.classSources.export();
         }
 
-        private void visitStackNode(StackTraceNode node) {
-            String className = node.getClassName();
-            if (!this.map.containsKey(className)) {
+        @Override
+        public boolean hasMethodSourceMappings() {
+            return this.methodSources.hasMappings();
+        }
+
+        @Override
+        public Map<String, String> getMethodSourceMapping() {
+            return this.methodSources.export();
+        }
+
+        @Override
+        public boolean hasLineSourceMappings() {
+            return this.lineSources.hasMappings();
+        }
+
+        @Override
+        public Map<String, String> getLineSourceMapping() {
+            return this.lineSources.export();
+        }
+    }
+
+    final class SourcesMap<T> {
+        // <key> --> identifier (plugin name)
+        private final Map<T, String> map = new HashMap<>();
+        private final Function<? super T, String> keyToStringFunction;
+
+        private SourcesMap(Function<? super T, String> keyToStringFunction) {
+            this.keyToStringFunction = keyToStringFunction;
+        }
+
+        public void computeIfAbsent(T key, ComputeSourceFunction<T> function) {
+            if (!this.map.containsKey(key)) {
                 try {
-                    Class<?> clazz = this.classFinder.findClass(className);
-                    Objects.requireNonNull(clazz);
-                    this.map.put(className, this.lookup.identify(clazz));
+                    this.map.put(key, function.compute(key));
                 } catch (Throwable e) {
-                    this.map.put(className, null);
+                    this.map.put(key, null);
                 }
             }
+        }
 
-            // recursively
-            for (StackTraceNode child : node.getChildren()) {
-                visitStackNode(child);
+        public boolean hasMappings() {
+            this.map.values().removeIf(Objects::isNull);
+            return !this.map.isEmpty();
+        }
+
+        public Map<String, String> export() {
+            this.map.values().removeIf(Objects::isNull);
+            if (this.keyToStringFunction.equals(Function.identity())) {
+                //noinspection unchecked
+                return (Map<String, String>) this.map;
+            } else {
+                return this.map.entrySet().stream().collect(Collectors.toMap(
+                        e -> this.keyToStringFunction.apply(e.getKey()),
+                        Map.Entry::getValue
+                ));
             }
         }
+
+        private interface ComputeSourceFunction<T> {
+            String compute(T key) throws Exception;
+        }
+    }
+
+    /**
+     * Encapsulates information about a given method call using the name + method description.
+     */
+    final class MethodCall {
+        private final String className;
+        private final String methodName;
+        private final String methodDescriptor;
+
+        public MethodCall(String className, String methodName, String methodDescriptor) {
+            this.className = className;
+            this.methodName = methodName;
+            this.methodDescriptor = methodDescriptor;
+        }
+
+        public String getClassName() {
+            return this.className;
+        }
+
+        public String getMethodName() {
+            return this.methodName;
+        }
+
+        public String getMethodDescriptor() {
+            return this.methodDescriptor;
+        }
+
+        @Override
+        public String toString() {
+            return this.className + ";" + this.methodName + ";" + this.methodDescriptor;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (!(o instanceof MethodCall)) return false;
+            MethodCall that = (MethodCall) o;
+            return this.className.equals(that.className) &&
+                    this.methodName.equals(that.methodName) &&
+                    this.methodDescriptor.equals(that.methodDescriptor);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(this.className, this.methodName, this.methodDescriptor);
+        }
+    }
+
+    /**
+     * Encapsulates information about a given method call using the name + line number.
+     */
+    final class MethodCallByLine {
+        private final String className;
+        private final String methodName;
+        private final int lineNumber;
+
+        public MethodCallByLine(String className, String methodName, int lineNumber) {
+            this.className = className;
+            this.methodName = methodName;
+            this.lineNumber = lineNumber;
+        }
+
+        public String getClassName() {
+            return this.className;
+        }
+
+        public String getMethodName() {
+            return this.methodName;
+        }
+
+        public int getLineNumber() {
+            return this.lineNumber;
+        }
+
+        @Override
+        public String toString() {
+            return this.className + ";" + this.lineNumber;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (!(o instanceof MethodCallByLine)) return false;
+            MethodCallByLine that = (MethodCallByLine) o;
+            return this.lineNumber == that.lineNumber && this.className.equals(that.className);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hash(this.className, this.lineNumber);
+        }
     }
 
 }
diff --git a/spark-common/src/main/proto/spark/spark_sampler.proto b/spark-common/src/main/proto/spark/spark_sampler.proto
index 8d9512a..f670ddf 100644
--- a/spark-common/src/main/proto/spark/spark_sampler.proto
+++ b/spark-common/src/main/proto/spark/spark_sampler.proto
@@ -11,6 +11,8 @@ message SamplerData {
   SamplerMetadata metadata = 1;
   repeated ThreadNode threads = 2;
   map<string, string> class_sources = 3; // optional
+  map<string, string> method_sources = 4; // optional
+  map<string, string> line_sources = 5; // optional
 }
 
 message SamplerMetadata {