From 6037c4f26c4cf5d1b7353086ac39f348030aada8 Mon Sep 17 00:00:00 2001
From: Daniel Kirchner <daniel@ekpyron.org>
Date: Mon, 9 Aug 2021 12:49:02 +0200
Subject: [PATCH] Add function to FunctionReturnLabelSlot and assert that
 functions are identical in comparison operators.

---
 libyul/backends/evm/ControlFlowGraph.h       | 16 ++++++++++++++--
 libyul/backends/evm/StackLayoutGenerator.cpp |  2 +-
 test/libyul/StackLayoutGeneratorTest.cpp     |  2 +-
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/libyul/backends/evm/ControlFlowGraph.h b/libyul/backends/evm/ControlFlowGraph.h
index 60bec498b..476b515b6 100644
--- a/libyul/backends/evm/ControlFlowGraph.h
+++ b/libyul/backends/evm/ControlFlowGraph.h
@@ -24,6 +24,7 @@
 #include <libyul/AST.h>
 #include <libyul/AsmAnalysisInfo.h>
 #include <libyul/Dialect.h>
+#include <libyul/Exceptions.h>
 #include <libyul/Scope.h>
 
 #include <functional>
@@ -54,8 +55,19 @@ struct FunctionCallReturnLabelSlot
 /// the function.
 struct FunctionReturnLabelSlot
 {
-	bool operator==(FunctionReturnLabelSlot const&) const { return true; }
-	bool operator<(FunctionReturnLabelSlot const&) const { return false; }
+	std::reference_wrapper<Scope::Function const> function;
+	bool operator==(FunctionReturnLabelSlot const& _rhs) const
+	{
+		// There can never be return label slots of different functions on stack simultaneously.
+		yulAssert(&function.get() == &_rhs.function.get(), "");
+		return true;
+	}
+	bool operator<(FunctionReturnLabelSlot const& _rhs) const
+	{
+		// There can never be return label slots of different functions on stack simultaneously.
+		yulAssert(&function.get() == &_rhs.function.get(), "");
+		return false;
+	}
 	static constexpr bool canBeFreelyGenerated = false;
 };
 /// A slot containing the current value of a particular variable.
diff --git a/libyul/backends/evm/StackLayoutGenerator.cpp b/libyul/backends/evm/StackLayoutGenerator.cpp
index 03b02bc0b..d625db877 100644
--- a/libyul/backends/evm/StackLayoutGenerator.cpp
+++ b/libyul/backends/evm/StackLayoutGenerator.cpp
@@ -323,7 +323,7 @@ optional<Stack> StackLayoutGenerator::getExitLayoutOrStageDependencies(
 			Stack stack = _functionReturn.info->returnVariables | ranges::views::transform([](auto const& _varSlot){
 				return StackSlot{_varSlot};
 			}) | ranges::to<Stack>;
-			stack.emplace_back(FunctionReturnLabelSlot{});
+			stack.emplace_back(FunctionReturnLabelSlot{_functionReturn.info->function});
 			return stack;
 		},
 		[&](CFG::BasicBlock::Terminated const&) -> std::optional<Stack>
diff --git a/test/libyul/StackLayoutGeneratorTest.cpp b/test/libyul/StackLayoutGeneratorTest.cpp
index e6a8e1b81..6dbed8f67 100644
--- a/test/libyul/StackLayoutGeneratorTest.cpp
+++ b/test/libyul/StackLayoutGeneratorTest.cpp
@@ -117,7 +117,7 @@ public:
 			m_stream << joinHumanReadable(_info.returnVariables | ranges::views::transform(variableSlotToString));
 		}
 		m_stream << "\\l\\\n";
-		Stack functionEntryStack = {FunctionReturnLabelSlot{}};
+		Stack functionEntryStack = {FunctionReturnLabelSlot{_info.function}};
 		functionEntryStack += _info.parameters | ranges::views::reverse;
 		m_stream << stackToString(functionEntryStack) << "\"];\n";
 		m_stream << "FunctionEntry_" << _info.function.name.str() << " -> Block" << getBlockId(*_info.entry) << ";\n";