From 6ae09ee0cc9b667fb6327fcbeaa173602119b1e2 Mon Sep 17 00:00:00 2001
From: Alex Beregszaszi <alex@rtfs.hu>
Date: Thu, 22 Apr 2021 15:08:51 +0100
Subject: [PATCH] Add AssemblyStack.assembleEVM which returns both deploy and
 runtime assemblies

---
 libyul/AssemblyStack.cpp | 59 +++++++++++++++++++++++++---------------
 libyul/AssemblyStack.h   | 18 +++++++++++-
 2 files changed, 54 insertions(+), 23 deletions(-)

diff --git a/libyul/AssemblyStack.cpp b/libyul/AssemblyStack.cpp
index 55682cc95..121bc48e4 100644
--- a/libyul/AssemblyStack.cpp
+++ b/libyul/AssemblyStack.cpp
@@ -216,7 +216,41 @@ MachineAssemblyObject AssemblyStack::assemble(Machine _machine) const
 	return MachineAssemblyObject();
 }
 
-std::pair<MachineAssemblyObject, MachineAssemblyObject> AssemblyStack::assembleWithDeployed(optional<string_view> _deployName) const
+std::pair<MachineAssemblyObject, MachineAssemblyObject>
+AssemblyStack::assembleWithDeployed(optional<string_view> _deployName) const
+{
+	auto [creationAssembly, deployedAssembly] = assembleEVMWithDeployed(_deployName);
+	yulAssert(creationAssembly, "");
+
+	MachineAssemblyObject creationObject;
+	creationObject.bytecode = make_shared<evmasm::LinkerObject>(creationAssembly->assemble());
+	yulAssert(creationObject.bytecode->immutableReferences.empty(), "Leftover immutables.");
+	creationObject.assembly = creationAssembly->assemblyString();
+	creationObject.sourceMappings = make_unique<string>(
+		evmasm::AssemblyItem::computeSourceMapping(
+			creationAssembly->items(),
+			{{scanner().charStream() ? scanner().charStream()->name() : "", 0}}
+		)
+	);
+
+	MachineAssemblyObject deployedObject;
+	if (deployedAssembly)
+	{
+		deployedObject.bytecode = make_shared<evmasm::LinkerObject>(deployedAssembly->assemble());
+		deployedObject.assembly = deployedAssembly->assemblyString();
+		deployedObject.sourceMappings = make_unique<string>(
+			evmasm::AssemblyItem::computeSourceMapping(
+				deployedAssembly->items(),
+				{{scanner().charStream() ? scanner().charStream()->name() : "", 0}}
+			)
+		);
+	}
+
+	return {std::move(creationObject), std::move(deployedObject)};
+}
+
+std::pair<std::shared_ptr<evmasm::Assembly>, std::shared_ptr<evmasm::Assembly>>
+AssemblyStack::assembleEVMWithDeployed(optional<string_view> _deployName) const
 {
 	yulAssert(m_analysisSuccessful, "");
 	yulAssert(m_parserResult, "");
@@ -227,18 +261,6 @@ std::pair<MachineAssemblyObject, MachineAssemblyObject> AssemblyStack::assembleW
 	EthAssemblyAdapter adapter(assembly);
 	compileEVM(adapter, m_optimiserSettings.optimizeStackAllocation);
 
-	MachineAssemblyObject creationObject;
-	creationObject.bytecode = make_shared<evmasm::LinkerObject>(assembly.assemble());
-	yulAssert(creationObject.bytecode->immutableReferences.empty(), "Leftover immutables.");
-	creationObject.assembly = assembly.assemblyString();
-	creationObject.sourceMappings = make_unique<string>(
-		evmasm::AssemblyItem::computeSourceMapping(
-			assembly.items(),
-			{{scanner().charStream() ? scanner().charStream()->name() : "", 0}}
-		)
-	);
-
-	MachineAssemblyObject deployedObject;
 	optional<size_t> subIndex;
 
 	// Pick matching assembly if name was given
@@ -260,17 +282,10 @@ std::pair<MachineAssemblyObject, MachineAssemblyObject> AssemblyStack::assembleW
 	if (subIndex.has_value())
 	{
 		evmasm::Assembly& runtimeAssembly = assembly.sub(*subIndex);
-		deployedObject.bytecode = make_shared<evmasm::LinkerObject>(runtimeAssembly.assemble());
-		deployedObject.assembly = runtimeAssembly.assemblyString();
-		deployedObject.sourceMappings = make_unique<string>(
-			evmasm::AssemblyItem::computeSourceMapping(
-				runtimeAssembly.items(),
-				{{scanner().charStream() ? scanner().charStream()->name() : "", 0}}
-			)
-		);
+		return {make_shared<evmasm::Assembly>(assembly), make_shared<evmasm::Assembly>(runtimeAssembly)};
 	}
 
-	return {std::move(creationObject), std::move(deployedObject)};
+	return {make_shared<evmasm::Assembly>(assembly), {}};
 }
 
 string AssemblyStack::print() const
diff --git a/libyul/AssemblyStack.h b/libyul/AssemblyStack.h
index e0186b041..925d0f87d 100644
--- a/libyul/AssemblyStack.h
+++ b/libyul/AssemblyStack.h
@@ -35,6 +35,11 @@
 #include <memory>
 #include <string>
 
+namespace solidity::evmasm
+{
+class Assembly;
+}
+
 namespace solidity::langutil
 {
 class Scanner;
@@ -93,7 +98,18 @@ public:
 	/// In addition to the value returned by @a assemble, returns
 	/// a second object that is the runtime code.
 	/// Only available for EVM.
-	std::pair<MachineAssemblyObject, MachineAssemblyObject> assembleWithDeployed(std::optional<std::string_view> _deployeName = {}) const;
+	std::pair<MachineAssemblyObject, MachineAssemblyObject>
+	assembleWithDeployed(
+		std::optional<std::string_view> _deployName = {}
+	) const;
+
+	/// Run the assembly step (should only be called after parseAndAnalyze).
+	/// Similar to @a assemblyWithDeployed, but returns EVM assembly objects.
+	/// Only available for EVM.
+	std::pair<std::shared_ptr<evmasm::Assembly>, std::shared_ptr<evmasm::Assembly>>
+	assembleEVMWithDeployed(
+		std::optional<std::string_view> _deployName = {}
+	) const;
 
 	/// @returns the errors generated during parsing, analysis (and potentially assembly).
 	langutil::ErrorList const& errors() const { return m_errors; }