Refactor sampler node export process

19 files changed, 490 insertions, 264 deletions

diff --git a/spark-common/src/main/java/me/lucko/spark/common/command/modules/SamplerModule.java b/spark-common/src/main/java/me/lucko/spark/common/command/modules/SamplerModule.java
index 334e416..cc7d56b 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/command/modules/SamplerModule.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/command/modules/SamplerModule.java
@@ -37,12 +37,11 @@ import me.lucko.spark.common.sampler.SamplerMode;
 import me.lucko.spark.common.sampler.ThreadDumper;
 import me.lucko.spark.common.sampler.ThreadGrouper;
 import me.lucko.spark.common.sampler.async.AsyncSampler;
-import me.lucko.spark.common.sampler.node.MergeMode;
+import me.lucko.spark.common.sampler.java.MergeStrategy;
 import me.lucko.spark.common.sampler.source.ClassSourceLookup;
 import me.lucko.spark.common.tick.TickHook;
 import me.lucko.spark.common.util.FormatUtil;
 import me.lucko.spark.common.util.MediaTypes;
-import me.lucko.spark.common.util.MethodDisambiguator;
 import me.lucko.spark.common.ws.ViewerSocket;
 import me.lucko.spark.proto.SparkSamplerProtos;
 import net.kyori.adventure.text.Component;
@@ -507,12 +506,7 @@ public class SamplerModule implements CommandModule {
         return new Sampler.ExportProps()
                 .creator(resp.senderData())
                 .comment(Iterables.getFirst(arguments.stringFlag("comment"), null))
-                .mergeMode(() -> {
-                    MethodDisambiguator methodDisambiguator = new MethodDisambiguator(platform.createClassFinder());
-                    return arguments.boolFlag("separate-parent-calls")
-                            ? MergeMode.separateParentCalls(methodDisambiguator)
-                            : MergeMode.sameMethod(methodDisambiguator);
-                })
+                .mergeStrategy(arguments.boolFlag("separate-parent-calls") ? MergeStrategy.SEPARATE_PARENT_CALLS : MergeStrategy.SAME_METHOD)
                 .classSourceLookup(() -> ClassSourceLookup.create(platform));
     }
 
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 20e7973..d76b1a1 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
@@ -25,8 +25,8 @@ import me.lucko.spark.common.command.sender.CommandSender;
 import me.lucko.spark.common.monitor.memory.GarbageCollectorStatistics;
 import me.lucko.spark.common.platform.SparkMetadata;
 import me.lucko.spark.common.sampler.aggregator.DataAggregator;
-import me.lucko.spark.common.sampler.node.MergeMode;
 import me.lucko.spark.common.sampler.node.ThreadNode;
+import me.lucko.spark.common.sampler.node.exporter.NodeExporter;
 import me.lucko.spark.common.sampler.source.ClassSourceLookup;
 import me.lucko.spark.common.sampler.window.ProtoTimeEncoder;
 import me.lucko.spark.common.sampler.window.WindowStatisticsCollector;
@@ -42,6 +42,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.function.Function;
 import java.util.function.Supplier;
 
 /**
@@ -198,7 +199,7 @@ public abstract class AbstractSampler implements Sampler {
         proto.setMetadata(metadata);
     }
 
-    protected void writeDataToProto(SamplerData.Builder proto, DataAggregator dataAggregator, MergeMode mergeMode, ClassSourceLookup classSourceLookup, Supplier<ClassFinder> classFinderSupplier) {
+    protected void writeDataToProto(SamplerData.Builder proto, DataAggregator dataAggregator, Function<ProtoTimeEncoder, NodeExporter> nodeExporterFunction, ClassSourceLookup classSourceLookup, Supplier<ClassFinder> classFinderSupplier) {
         List<ThreadNode> data = dataAggregator.exportData();
         data.sort(Comparator.comparing(ThreadNode::getThreadLabel));
 
@@ -213,8 +214,10 @@ public abstract class AbstractSampler implements Sampler {
         this.windowStatisticsCollector.ensureHasStatisticsForAllWindows(timeWindows);
         proto.putAllTimeWindowStatistics(this.windowStatisticsCollector.export());
 
+        NodeExporter exporter = nodeExporterFunction.apply(timeEncoder);
+
         for (ThreadNode entry : data) {
-            proto.addThreads(entry.toProto(mergeMode, timeEncoder));
+            proto.addThreads(exporter.export(entry));
             classSourceVisitor.visit(entry);
         }
 
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/Sampler.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/Sampler.java
index bb74cd2..5aca704 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/Sampler.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/Sampler.java
@@ -22,7 +22,7 @@ package me.lucko.spark.common.sampler;
 
 import me.lucko.spark.common.SparkPlatform;
 import me.lucko.spark.common.command.sender.CommandSender;
-import me.lucko.spark.common.sampler.node.MergeMode;
+import me.lucko.spark.common.sampler.java.MergeStrategy;
 import me.lucko.spark.common.sampler.source.ClassSourceLookup;
 import me.lucko.spark.common.ws.ViewerSocket;
 import me.lucko.spark.proto.SparkProtos;
@@ -111,7 +111,7 @@ public interface Sampler {
     final class ExportProps {
         private CommandSender.Data creator;
         private String comment;
-        private Supplier<MergeMode> mergeMode;
+        private MergeStrategy mergeStrategy;
         private Supplier<ClassSourceLookup> classSourceLookup;
         private SocketChannelInfo channelInfo;
 
@@ -126,8 +126,8 @@ public interface Sampler {
             return this.comment;
         }
 
-        public Supplier<MergeMode> mergeMode() {
-            return this.mergeMode;
+        public MergeStrategy mergeStrategy() {
+            return this.mergeStrategy;
         }
 
         public Supplier<ClassSourceLookup> classSourceLookup() {
@@ -148,8 +148,8 @@ public interface Sampler {
             return this;
         }
 
-        public ExportProps mergeMode(Supplier<MergeMode> mergeMode) {
-            this.mergeMode = mergeMode;
+        public ExportProps mergeStrategy(MergeStrategy mergeStrategy) {
+            this.mergeStrategy = mergeStrategy;
             return this;
         }
 
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncDataAggregator.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncDataAggregator.java
index b9a80e0..484493a 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncDataAggregator.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncDataAggregator.java
@@ -33,7 +33,7 @@ public class AsyncDataAggregator extends AbstractDataAggregator {
 
     /** A describer for async-profiler stack trace elements. */
     private static final StackTraceNode.Describer<AsyncStackTraceElement> STACK_TRACE_DESCRIBER = (element, parent) ->
-            new StackTraceNode.Description(element.getClassName(), element.getMethodName(), element.getMethodDescription());
+            new StackTraceNode.AsyncDescription(element.getClassName(), element.getMethodName(), element.getMethodDescription());
 
     protected AsyncDataAggregator(ThreadGrouper threadGrouper) {
         super(threadGrouper);
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncNodeExporter.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncNodeExporter.java
new file mode 100644
index 0000000..ef68c46
--- /dev/null
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncNodeExporter.java
@@ -0,0 +1,63 @@
+/*
+ * This file is part of spark.
+ *
+ *  Copyright (c) lucko (Luck) <luck@lucko.me>
+ *  Copyright (c) contributors
+ *
+ *  This program is free software: you can redistribute it and/or modify
+ *  it under the terms of the GNU General Public License as published by
+ *  the Free Software Foundation, either version 3 of the License, or
+ *  (at your option) any later version.
+ *
+ *  This program 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 General Public License for more details.
+ *
+ *  You should have received a copy of the GNU General Public License
+ *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package me.lucko.spark.common.sampler.async;
+
+import me.lucko.spark.common.sampler.node.StackTraceNode;
+import me.lucko.spark.common.sampler.node.exporter.AbstractNodeExporter;
+import me.lucko.spark.common.sampler.window.ProtoTimeEncoder;
+import me.lucko.spark.proto.SparkSamplerProtos;
+
+import java.util.Collection;
+
+/**
+ * Node exporter for the {@link AsyncSampler}.
+ */
+public class AsyncNodeExporter extends AbstractNodeExporter {
+    public AsyncNodeExporter(ProtoTimeEncoder timeEncoder) {
+        super(timeEncoder);
+    }
+
+    @Override
+    protected SparkSamplerProtos.StackTraceNode export(StackTraceNode stackTraceNode, Iterable<Integer> childrenRefs) {
+        SparkSamplerProtos.StackTraceNode.Builder proto = SparkSamplerProtos.StackTraceNode.newBuilder()
+                .setClassName(stackTraceNode.getClassName())
+                .setMethodName(stackTraceNode.getMethodName());
+
+        double[] times = stackTraceNode.encodeTimesForProto(this.timeEncoder);
+        for (double time : times) {
+            proto.addTimes(time);
+        }
+
+        String methodDescription = stackTraceNode.getMethodDescription();
+        if (methodDescription != null) {
+            proto.setMethodDesc(methodDescription);
+        }
+
+        proto.addAllChildrenRefs(childrenRefs);
+
+        return proto.build();
+    }
+
+    @Override
+    protected Collection<StackTraceNode> exportChildren(Collection<StackTraceNode> children) {
+        return children;
+    }
+}
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncSampler.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncSampler.java
index 5350558..62af021 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncSampler.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncSampler.java
@@ -221,7 +221,7 @@ public class AsyncSampler extends AbstractSampler {
             proto.setChannelInfo(exportProps.channelInfo());
         }
         writeMetadataToProto(proto, platform, exportProps.creator(), exportProps.comment(), this.dataAggregator);
-        writeDataToProto(proto, this.dataAggregator, exportProps.mergeMode().get(), exportProps.classSourceLookup().get(), platform::createClassFinder);
+        writeDataToProto(proto, this.dataAggregator, AsyncNodeExporter::new, exportProps.classSourceLookup().get(), platform::createClassFinder);
         return proto.build();
     }
 
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaDataAggregator.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaDataAggregator.java
index b7b841e..5b6a470 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaDataAggregator.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaDataAggregator.java
@@ -39,7 +39,7 @@ public abstract class JavaDataAggregator extends AbstractDataAggregator {
     /** A describer for java.lang.StackTraceElement */
     private static final StackTraceNode.Describer<StackTraceElement> STACK_TRACE_DESCRIBER = (element, parent) -> {
         int parentLineNumber = parent == null ? StackTraceNode.NULL_LINE_NUMBER : parent.getLineNumber();
-        return new StackTraceNode.Description(element.getClassName(), element.getMethodName(), element.getLineNumber(), parentLineNumber);
+        return new StackTraceNode.JavaDescription(element.getClassName(), element.getMethodName(), element.getLineNumber(), parentLineNumber);
     };
 
     /** The worker pool for inserting stack nodes */
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaNodeExporter.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaNodeExporter.java
new file mode 100644
index 0000000..c110086
--- /dev/null
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaNodeExporter.java
@@ -0,0 +1,97 @@
+/*
+ * This file is part of spark.
+ *
+ *  Copyright (c) lucko (Luck) <luck@lucko.me>
+ *  Copyright (c) contributors
+ *
+ *  This program is free software: you can redistribute it and/or modify
+ *  it under the terms of the GNU General Public License as published by
+ *  the Free Software Foundation, either version 3 of the License, or
+ *  (at your option) any later version.
+ *
+ *  This program 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 General Public License for more details.
+ *
+ *  You should have received a copy of the GNU General Public License
+ *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package me.lucko.spark.common.sampler.java;
+
+import me.lucko.spark.common.sampler.node.StackTraceNode;
+import me.lucko.spark.common.sampler.node.exporter.AbstractNodeExporter;
+import me.lucko.spark.common.sampler.window.ProtoTimeEncoder;
+import me.lucko.spark.common.util.MethodDisambiguator;
+import me.lucko.spark.proto.SparkSamplerProtos;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+
+/**
+ * Node exporter for the {@link JavaSampler}.
+ */
+public class JavaNodeExporter extends AbstractNodeExporter {
+    private final MergeStrategy mergeStrategy;
+    private final MethodDisambiguator methodDisambiguator;
+
+    public JavaNodeExporter(ProtoTimeEncoder timeEncoder, MergeStrategy mergeStrategy, MethodDisambiguator methodDisambiguator) {
+        super(timeEncoder);
+        this.mergeStrategy = mergeStrategy;
+        this.methodDisambiguator = methodDisambiguator;
+    }
+
+    protected SparkSamplerProtos.StackTraceNode export(StackTraceNode stackTraceNode, Iterable<Integer> childrenRefs) {
+        SparkSamplerProtos.StackTraceNode.Builder proto = SparkSamplerProtos.StackTraceNode.newBuilder()
+                .setClassName(stackTraceNode.getClassName())
+                .setMethodName(stackTraceNode.getMethodName());
+
+        double[] times = stackTraceNode.encodeTimesForProto(this.timeEncoder);
+        for (double time : times) {
+            proto.addTimes(time);
+        }
+
+        int lineNumber = stackTraceNode.getLineNumber();
+        if (lineNumber >= 0) {
+            proto.setLineNumber(lineNumber);
+        }
+
+        if (this.mergeStrategy.separateParentCalls()) {
+            int parentLineNumber = stackTraceNode.getParentLineNumber();
+            if (parentLineNumber >= 0) {
+                proto.setParentLineNumber(parentLineNumber);
+            }
+        }
+
+        this.methodDisambiguator.disambiguate(stackTraceNode)
+                .map(MethodDisambiguator.MethodDescription::getDescription)
+                .ifPresent(proto::setMethodDesc);
+
+        proto.addAllChildrenRefs(childrenRefs);
+
+        return proto.build();
+    }
+
+    @Override
+    protected Collection<StackTraceNode> exportChildren(Collection<StackTraceNode> children) {
+        if (children.isEmpty()) {
+            return children;
+        }
+
+        List<StackTraceNode> list = new ArrayList<>(children.size());
+
+        outer:
+        for (StackTraceNode child : children) {
+            for (StackTraceNode other : list) {
+                if (this.mergeStrategy.shouldMerge(this.methodDisambiguator, other, child)) {
+                    other.merge(child);
+                    continue outer;
+                }
+            }
+            list.add(child);
+        }
+        return list;
+    }
+}
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaSampler.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaSampler.java
index e3ae73a..20f9383 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaSampler.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaSampler.java
@@ -28,6 +28,7 @@ import me.lucko.spark.common.sampler.SamplerSettings;
 import me.lucko.spark.common.sampler.window.ProfilingWindowUtils;
 import me.lucko.spark.common.sampler.window.WindowStatisticsCollector;
 import me.lucko.spark.common.tick.TickHook;
+import me.lucko.spark.common.util.MethodDisambiguator;
 import me.lucko.spark.common.util.SparkThreadFactory;
 import me.lucko.spark.common.ws.ViewerSocket;
 import me.lucko.spark.proto.SparkSamplerProtos.SamplerData;
@@ -192,8 +193,12 @@ public class JavaSampler extends AbstractSampler implements Runnable {
         if (exportProps.channelInfo() != null) {
             proto.setChannelInfo(exportProps.channelInfo());
         }
+
         writeMetadataToProto(proto, platform, exportProps.creator(), exportProps.comment(), this.dataAggregator);
-        writeDataToProto(proto, this.dataAggregator, exportProps.mergeMode().get(), exportProps.classSourceLookup().get(), platform::createClassFinder);
+
+        MethodDisambiguator methodDisambiguator = new MethodDisambiguator(platform.createClassFinder());
+        writeDataToProto(proto, this.dataAggregator, timeEncoder -> new JavaNodeExporter(timeEncoder, exportProps.mergeStrategy(), methodDisambiguator), exportProps.classSourceLookup().get(), platform::createClassFinder);
+
         return proto.build();
     }
 
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/node/MergeMode.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/java/MergeStrategy.java
index 18a0ed3..eac1c6f 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/node/MergeMode.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/java/MergeStrategy.java
@@ -18,37 +18,27 @@
  *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-package me.lucko.spark.common.sampler.node;
+package me.lucko.spark.common.sampler.java;
 
+import me.lucko.spark.common.sampler.node.StackTraceNode;
 import me.lucko.spark.common.util.MethodDisambiguator;
 
 import java.util.Objects;
 
 /**
- * Function to determine if {@link StackTraceNode}s should be merged.
+ * Strategy used to determine if {@link StackTraceNode}s should be merged.
  */
-public final class MergeMode {
+public enum MergeStrategy {
 
-    public static MergeMode sameMethod(MethodDisambiguator methodDisambiguator) {
-        return new MergeMode(methodDisambiguator, false);
-    }
-
-    public static MergeMode separateParentCalls(MethodDisambiguator methodDisambiguator) {
-        return new MergeMode(methodDisambiguator, true);
-    }
+    SAME_METHOD(false),
+    SEPARATE_PARENT_CALLS(true);
 
-    private final MethodDisambiguator methodDisambiguator;
     private final boolean separateParentCalls;
 
-    MergeMode(MethodDisambiguator methodDisambiguator, boolean separateParentCalls) {
-        this.methodDisambiguator = methodDisambiguator;
+    MergeStrategy(boolean separateParentCalls) {
         this.separateParentCalls = separateParentCalls;
     }
 
-    public MethodDisambiguator getMethodDisambiguator() {
-        return this.methodDisambiguator;
-    }
-
     public boolean separateParentCalls() {
         return this.separateParentCalls;
     }
@@ -56,11 +46,12 @@ public final class MergeMode {
     /**
      * Test if two stack trace nodes should be considered the same and merged.
      *
+     * @param disambiguator the method disambiguator
      * @param n1 the first node
      * @param n2 the second node
      * @return if the nodes should be merged
      */
-    public boolean shouldMerge(StackTraceNode n1, StackTraceNode n2) {
+    public boolean shouldMerge(MethodDisambiguator disambiguator, StackTraceNode n1, StackTraceNode n2) {
         // are the class names the same?
         if (!n1.getClassName().equals(n2.getClassName())) {
             return false;
@@ -77,8 +68,8 @@ public final class MergeMode {
         }
 
         // are the method descriptions the same? (is it the same method?)
-        String desc1 = this.methodDisambiguator.disambiguate(n1).map(MethodDisambiguator.MethodDescription::getDesc).orElse(null);
-        String desc2 = this.methodDisambiguator.disambiguate(n2).map(MethodDisambiguator.MethodDescription::getDesc).orElse(null);
+        String desc1 = disambiguator.disambiguate(n1).map(MethodDisambiguator.MethodDescription::getDescription).orElse(null);
+        String desc2 = disambiguator.disambiguate(n2).map(MethodDisambiguator.MethodDescription::getDescription).orElse(null);
 
         if (desc1 == null && desc2 == null) {
             return true;
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/node/AbstractNode.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/node/AbstractNode.java
index 163365c..d3b77b4 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/node/AbstractNode.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/node/AbstractNode.java
@@ -22,10 +22,7 @@ package me.lucko.spark.common.sampler.node;
 
 import me.lucko.spark.common.sampler.window.ProtoTimeEncoder;
 
-import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Collections;
-import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
@@ -86,7 +83,7 @@ public abstract class AbstractNode {
      *
      * @return the total times
      */
-    protected double[] encodeTimesForProto(ProtoTimeEncoder encoder) {
+    public double[] encodeTimesForProto(ProtoTimeEncoder encoder) {
         return encoder.encode(this.times);
     }
 
@@ -107,35 +104,11 @@ public abstract class AbstractNode {
      *
      * @param other the other node
      */
-    protected void merge(AbstractNode other) {
+    public void merge(AbstractNode other) {
         other.times.forEach((key, value) -> getTimeAccumulator(key).add(value.longValue()));
         for (Map.Entry<StackTraceNode.Description, StackTraceNode> child : other.children.entrySet()) {
             resolveChild(child.getKey()).merge(child.getValue());
         }
     }
 
-    protected List<StackTraceNode> exportChildren(MergeMode mergeMode) {
-        if (this.children.isEmpty()) {
-            return Collections.emptyList();
-        }
-
-        List<StackTraceNode> list = new ArrayList<>(this.children.size());
-
-        outer:
-        for (StackTraceNode child : this.children.values()) {
-            // attempt to find an existing node we can merge into
-            for (StackTraceNode other : list) {
-                if (mergeMode.shouldMerge(other, child)) {
-                    other.merge(child);
-                    continue outer;
-                }
-            }
-
-            // just add
-            list.add(child);
-        }
-
-        return list;
-    }
-
 }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/node/StackTraceNode.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/node/StackTraceNode.java
index fec7123..27cfa54 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/node/StackTraceNode.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/node/StackTraceNode.java
@@ -20,9 +20,6 @@
 
 package me.lucko.spark.common.sampler.node;
 
-import me.lucko.spark.common.sampler.window.ProtoTimeEncoder;
-import me.lucko.spark.common.util.MethodDisambiguator;
-import me.lucko.spark.proto.SparkSamplerProtos;
 import org.checkerframework.checker.nullness.qual.Nullable;
 
 import java.util.Objects;
@@ -45,58 +42,33 @@ public final class StackTraceNode extends AbstractNode {
     }
 
     public String getClassName() {
-        return this.description.className;
+        return this.description.className();
     }
 
     public String getMethodName() {
-        return this.description.methodName;
+        return this.description.methodName();
     }
 
     public String getMethodDescription() {
-        return this.description.methodDescription;
+        return this.description instanceof AsyncDescription
+                ? ((AsyncDescription) this.description).methodDescription()
+                : null;
     }
 
     public int getLineNumber() {
-        return this.description.lineNumber;
+        return this.description instanceof JavaDescription
+                ? ((JavaDescription) this.description).lineNumber()
+                : NULL_LINE_NUMBER;
     }
 
     public int getParentLineNumber() {
-        return this.description.parentLineNumber;
-    }
-
-    public SparkSamplerProtos.StackTraceNode toProto(MergeMode mergeMode, ProtoTimeEncoder timeEncoder, Iterable<Integer> childrenRefs) {
-        SparkSamplerProtos.StackTraceNode.Builder proto = SparkSamplerProtos.StackTraceNode.newBuilder()
-                .setClassName(this.description.className)
-                .setMethodName(this.description.methodName);
-
-        double[] times = encodeTimesForProto(timeEncoder);
-        for (double time : times) {
-            proto.addTimes(time);
-        }
-
-        if (this.description.lineNumber >= 0) {
-            proto.setLineNumber(this.description.lineNumber);
-        }
-
-        if (mergeMode.separateParentCalls() && this.description.parentLineNumber >= 0) {
-            proto.setParentLineNumber(this.description.parentLineNumber);
-        }
-
-        if (this.description.methodDescription != null) {
-            proto.setMethodDesc(this.description.methodDescription);
-        } else {
-            mergeMode.getMethodDisambiguator().disambiguate(this)
-                    .map(MethodDisambiguator.MethodDescription::getDesc)
-                    .ifPresent(proto::setMethodDesc);
-        }
-
-        proto.addAllChildrenRefs(childrenRefs);
-
-        return proto.build();
+        return this.description instanceof JavaDescription
+                ? ((JavaDescription) this.description).parentLineNumber()
+                : NULL_LINE_NUMBER;
     }
 
     /**
-     * Function to construct a {@link StackTraceNode.Description} from a stack trace element
+     * Function to construct a {@link Description} from a stack trace element
      * of type {@code T}.
      *
      * @param <T> the stack trace element type, e.g. {@link java.lang.StackTraceElement}
@@ -114,53 +86,101 @@ public final class StackTraceNode extends AbstractNode {
         Description describe(T element, @Nullable T parent);
     }
 
-    /**
-     * Encapsulates the attributes of a {@link StackTraceNode}.
-     */
-    public static final class Description {
+    public interface Description {
+        String className();
+
+        String methodName();
+    }
+
+    public static final class AsyncDescription implements Description {
         private final String className;
         private final String methodName;
-
-        // async-profiler
         private final String methodDescription;
 
-        // Java
+        private final int hash;
+
+        public AsyncDescription(String className, String methodName, String methodDescription) {
+            this.className = className;
+            this.methodName = methodName;
+            this.methodDescription = methodDescription;
+            this.hash = Objects.hash(this.className, this.methodName, this.methodDescription);
+        }
+
+        @Override
+        public String className() {
+            return this.className;
+        }
+
+        @Override
+        public String methodName() {
+            return this.methodName;
+        }
+
+        public String methodDescription() {
+            return this.methodDescription;
+        }
+
+        @Override
+        public boolean equals(Object o) {
+            if (this == o) return true;
+            if (o == null || getClass() != o.getClass()) return false;
+            AsyncDescription description = (AsyncDescription) o;
+            return this.hash == description.hash &&
+                    this.className.equals(description.className) &&
+                    this.methodName.equals(description.methodName) &&
+                    Objects.equals(this.methodDescription, description.methodDescription);
+        }
+
+        @Override
+        public int hashCode() {
+            return this.hash;
+        }
+    }
+
+    public static final class JavaDescription implements Description {
+        private final String className;
+        private final String methodName;
         private final int lineNumber;
         private final int parentLineNumber;
 
         private final int hash;
 
-        // Constructor used by the Java sampler
-        public Description(String className, String methodName, int lineNumber, int parentLineNumber) {
+        public JavaDescription(String className, String methodName, int lineNumber, int parentLineNumber) {
             this.className = className;
             this.methodName = methodName;
-            this.methodDescription = null;
             this.lineNumber = lineNumber;
             this.parentLineNumber = parentLineNumber;
             this.hash = Objects.hash(this.className, this.methodName, this.lineNumber, this.parentLineNumber);
         }
 
-        // Constructor used by the async-profiler sampler
-        public Description(String className, String methodName, String methodDescription) {
-            this.className = className;
-            this.methodName = methodName;
-            this.methodDescription = methodDescription;
-            this.lineNumber = StackTraceNode.NULL_LINE_NUMBER;
-            this.parentLineNumber = StackTraceNode.NULL_LINE_NUMBER;
-            this.hash = Objects.hash(this.className, this.methodName, this.methodDescription);
+        @Override
+        public String className() {
+            return this.className;
+        }
+
+        @Override
+        public String methodName() {
+            return this.methodName;
+        }
+
+        public int lineNumber() {
+            return this.lineNumber;
+        }
+
+        public int parentLineNumber() {
+            return this.parentLineNumber;
         }
 
         @Override
         public boolean equals(Object o) {
             if (this == o) return true;
             if (o == null || getClass() != o.getClass()) return false;
-            Description description = (Description) o;
+            JavaDescription description = (JavaDescription) o;
             return this.hash == description.hash &&
                     this.lineNumber == description.lineNumber &&
                     this.parentLineNumber == description.parentLineNumber &&
                     this.className.equals(description.className) &&
-                    this.methodName.equals(descripti