From e6b5b7f16357ba971522d5b1c8785af175d5657c Mon Sep 17 00:00:00 2001
From: Matt <mcoley2@gmu.edu>
Date: Fri, 8 Nov 2024 00:23:18 -0500
Subject: [PATCH] Create a workspace-aware ClassWriter utility

---
 .../util/visitors/WorkspaceClassWriter.java   | 49 +++++++++++++++++++
 1 file changed, 49 insertions(+)
 create mode 100644 recaf-core/src/main/java/software/coley/recaf/util/visitors/WorkspaceClassWriter.java

diff --git a/recaf-core/src/main/java/software/coley/recaf/util/visitors/WorkspaceClassWriter.java b/recaf-core/src/main/java/software/coley/recaf/util/visitors/WorkspaceClassWriter.java
new file mode 100644
index 000000000..d68e902f9
--- /dev/null
+++ b/recaf-core/src/main/java/software/coley/recaf/util/visitors/WorkspaceClassWriter.java
@@ -0,0 +1,49 @@
+package software.coley.recaf.util.visitors;
+
+import jakarta.annotation.Nonnull;
+import jakarta.annotation.Nullable;
+import org.objectweb.asm.ClassReader;
+import org.objectweb.asm.ClassWriter;
+import software.coley.recaf.services.inheritance.InheritanceGraph;
+
+/**
+ * Class writer that pulls inheritance information from a workspace via {@link InheritanceGraph}.
+ *
+ * @author Matt Coley
+ */
+public class WorkspaceClassWriter extends ClassWriter {
+	private final InheritanceGraph graph;
+
+	/**
+	 * @param graph
+	 * 		Graph to pull inheritance relations from.
+	 * @param flags
+	 * 		Writer flags.
+	 */
+	public WorkspaceClassWriter(@Nonnull InheritanceGraph graph, int flags) {
+		this(graph, null, flags);
+	}
+
+	/**
+	 * @param graph
+	 * 		Graph to pull inheritance relations from.
+	 * @param reader
+	 * 		Reader to pre-populate the constant pool with. Speeds up writing process a bit.
+	 * @param flags
+	 * 		Writer flags.
+	 */
+	public WorkspaceClassWriter(@Nonnull InheritanceGraph graph, @Nullable ClassReader reader, int flags) {
+		super(reader, flags);
+		this.graph = graph;
+	}
+
+	@Override
+	protected String getCommonSuperClass(String type1, String type2) {
+		// Default assumption if a type isn't given.
+		if (type1 == null || type2 == null)
+			return "java/lang/Object";
+
+		// Find common parent in workspace
+		return graph.getCommon(type1, type2);
+	}
+}