2 files changed, 46 insertions, 26 deletions

diff --git a/libjulia/backends/evm/EVMCodeTransform.cpp b/libjulia/backends/evm/EVMCodeTransform.cpp
index cd6fd276..638ff4d9 100644
--- a/libjulia/backends/evm/EVMCodeTransform.cpp
+++ b/libjulia/backends/evm/EVMCodeTransform.cpp
@@ -91,8 +91,7 @@ void CodeTransform::operator()(Label const& _label)
 	solAssert(m_scope, "");
 	solAssert(m_scope->identifiers.count(_label.name), "");
 	Scope::Label& label = boost::get<Scope::Label>(m_scope->identifiers.at(_label.name));
-	assignLabelIdIfUnset(label.id);
-	m_assembly.appendLabel(*label.id);
+	m_assembly.appendLabel(labelID(label));
 	checkStackHeight(&_label);
 }
 
@@ -120,12 +119,11 @@ void CodeTransform::operator()(FunctionCall const& _call)
 	for (auto const& arg: _call.arguments | boost::adaptors::reversed)
 		visitExpression(arg);
 	m_assembly.setSourceLocation(_call.location);
-	assignLabelIdIfUnset(function->id);
 	if (m_evm15)
-		m_assembly.appendJumpsub(*function->id, function->arguments.size(), function->returns.size());
+		m_assembly.appendJumpsub(functionEntryID(*function), function->arguments.size(), function->returns.size());
 	else
 	{
-		m_assembly.appendJumpTo(*function->id, function->returns.size() - function->arguments.size() - 1);
+		m_assembly.appendJumpTo(functionEntryID(*function), function->returns.size() - function->arguments.size() - 1);
 		m_assembly.appendLabel(returnLabel);
 		m_stackAdjustment--;
 	}
@@ -181,8 +179,7 @@ void CodeTransform::operator()(assembly::Identifier const& _identifier)
 		},
 		[=](Scope::Label& _label)
 		{
-			assignLabelIdIfUnset(_label.id);
-			m_assembly.appendLabelReference(*_label.id);
+			m_assembly.appendLabelReference(labelID(_label));
 		},
 		[=](Scope::Function&)
 		{
@@ -282,7 +279,6 @@ void CodeTransform::operator()(FunctionDefinition const& _function)
 	solAssert(m_scope, "");
 	solAssert(m_scope->identifiers.count(_function.name), "");
 	Scope::Function& function = boost::get<Scope::Function>(m_scope->identifiers.at(_function.name));
-	assignLabelIdIfUnset(function.id);
 
 	int const localStackAdjustment = m_evm15 ? 0 : 1;
 	int height = localStackAdjustment;
@@ -303,12 +299,12 @@ void CodeTransform::operator()(FunctionDefinition const& _function)
 	if (m_evm15)
 	{
 		m_assembly.appendJumpTo(afterFunction, -stackHeightBefore);
-		m_assembly.appendBeginsub(*function.id, _function.arguments.size());
+		m_assembly.appendBeginsub(functionEntryID(function), _function.arguments.size());
 	}
 	else
 	{
 		m_assembly.appendJumpTo(afterFunction, -stackHeightBefore + height);
-		m_assembly.appendLabel(*function.id);
+		m_assembly.appendLabel(functionEntryID(function));
 	}
 	m_stackAdjustment += localStackAdjustment;
 
@@ -321,7 +317,7 @@ void CodeTransform::operator()(FunctionDefinition const& _function)
 		m_assembly.appendConstant(u256(0));
 	}
 
-	CodeTransform(m_assembly, m_info, m_evm15, m_identifierAccess, localStackAdjustment)
+	CodeTransform(m_assembly, m_info, m_evm15, m_identifierAccess, localStackAdjustment, m_context)
 		.run(_function.body);
 
 	{
@@ -367,7 +363,7 @@ void CodeTransform::operator()(FunctionDefinition const& _function)
 
 void CodeTransform::operator()(Block const& _block)
 {
-	CodeTransform(m_assembly, m_info, m_evm15, m_identifierAccess, m_stackAdjustment).run(_block);
+	CodeTransform(m_assembly, m_info, m_evm15, m_identifierAccess, m_stackAdjustment, m_context).run(_block);
 }
 
 AbstractAssembly::LabelID CodeTransform::labelFromIdentifier(Identifier const& _identifier)
@@ -377,8 +373,7 @@ AbstractAssembly::LabelID CodeTransform::labelFromIdentifier(Identifier const& _
 		[=](Scope::Variable&) { solAssert(false, "Expected label"); },
 		[&](Scope::Label& _label)
 		{
-			assignLabelIdIfUnset(_label.id);
-			label = *_label.id;
+			label = labelID(_label);
 		},
 		[=](Scope::Function&) { solAssert(false, "Expected label"); }
 	)))
@@ -388,6 +383,20 @@ AbstractAssembly::LabelID CodeTransform::labelFromIdentifier(Identifier const& _
 	return label;
 }
 
+AbstractAssembly::LabelID CodeTransform::labelID(Scope::Label const& _label)
+{
+	if (!m_context->labelIDs.count(&_label))
+		m_context->labelIDs[&_label] = m_assembly.newLabelId();
+	return m_context->labelIDs[&_label];
+}
+
+AbstractAssembly::LabelID CodeTransform::functionEntryID(Scope::Function const& _function)
+{
+	if (!m_context->functionEntryIDs.count(&_function))
+		m_context->functionEntryIDs[&_function] = m_assembly.newLabelId();
+	return m_context->functionEntryIDs[&_function];
+}
+
 void CodeTransform::visitExpression(Statement const& _expression)
 {
 	int height = m_assembly.stackHeight();
@@ -446,9 +455,3 @@ void CodeTransform::checkStackHeight(void const* _astElement)
 		to_string(m_assembly.stackHeight() - m_stackAdjustment)
 	);
 }
-
-void CodeTransform::assignLabelIdIfUnset(boost::optional<AbstractAssembly::LabelID>& _labelId)
-{
-	if (!_labelId)
-		_labelId.reset(m_assembly.newLabelId());
-}
diff --git a/libjulia/backends/evm/EVMCodeTransform.h b/libjulia/backends/evm/EVMCodeTransform.h
index f65948dc..57781c91 100644
--- a/libjulia/backends/evm/EVMCodeTransform.h
+++ b/libjulia/backends/evm/EVMCodeTransform.h
@@ -64,7 +64,14 @@ public:
 		solidity::assembly::AsmAnalysisInfo& _analysisInfo,
 		bool _evm15 = false,
 		ExternalIdentifierAccess const& _identifierAccess = ExternalIdentifierAccess()
-	): CodeTransform(_assembly, _analysisInfo, _evm15, _identifierAccess, _assembly.stackHeight())
+	): CodeTransform(
+		_assembly,
+		_analysisInfo,
+		_evm15,
+		_identifierAccess,
+		_assembly.stackHeight(),
+		std::make_shared<Context>()
+	)
 	{
 	}
 
@@ -72,18 +79,26 @@ public:
 	void run(solidity::assembly::Block const& _block);
 
 protected:
+	struct Context
+	{
+		std::map<solidity::assembly::Scope::Label const*, AbstractAssembly::LabelID> labelIDs;
+		std::map<solidity::assembly::Scope::Function const*, AbstractAssembly::LabelID> functionEntryIDs;
+	};
+
 	CodeTransform(
 		julia::AbstractAssembly& _assembly,
 		solidity::assembly::AsmAnalysisInfo& _analysisInfo,
 		bool _evm15,
 		ExternalIdentifierAccess const& _identifierAccess,
-		int _stackAdjustment
+		int _stackAdjustment,
+		std::shared_ptr<Context> _context
 	):
 		m_assembly(_assembly),
 		m_info(_analysisInfo),
 		m_evm15(_evm15),
 		m_identifierAccess(_identifierAccess),
-		m_stackAdjustment(_stackAdjustment)
+		m_stackAdjustment(_stackAdjustment),
+		m_context(_context)
 	{}
 
 public:
@@ -102,6 +117,10 @@ public:
 
 private:
 	AbstractAssembly::LabelID labelFromIdentifier(solidity::assembly::Identifier const& _identifier);
+	/// @returns the label ID corresponding to the given label, allocating a new one if
+	/// necessary.
+	AbstractAssembly::LabelID labelID(solidity::assembly::Scope::Label const& _label);
+	AbstractAssembly::LabelID functionEntryID(solidity::assembly::Scope::Function const& _function);
 	/// Generates code for an expression that is supposed to return a single value.
 	void visitExpression(solidity::assembly::Statement const& _expression);
 
@@ -116,9 +135,6 @@ private:
 
 	void checkStackHeight(void const* _astElement);
 
-	/// Assigns the label's or function's id to a value taken from eth::Assembly if it has not yet been set.
-	void assignLabelIdIfUnset(boost::optional<AbstractAssembly::LabelID>& _labelId);
-
 	julia::AbstractAssembly& m_assembly;
 	solidity::assembly::AsmAnalysisInfo& m_info;
 	solidity::assembly::Scope* m_scope = nullptr;
@@ -129,6 +145,7 @@ private:
 	/// for inline assembly and different stack heights depending on the EVM backend used
 	/// (EVM 1.0 or 1.5).
 	int m_stackAdjustment = 0;
+	std::shared_ptr<Context> m_context;
 };
 
 }