Added tests and review suggestions

5 files changed, 176 insertions, 22 deletions

diff --git a/libsolidity/codegen/CompilerUtils.cpp b/libsolidity/codegen/CompilerUtils.cpp
index 2f45765a..92c45227 100644
--- a/libsolidity/codegen/CompilerUtils.cpp
+++ b/libsolidity/codegen/CompilerUtils.cpp
@@ -1170,6 +1170,14 @@ void CompilerUtils::popStackSlots(size_t _amount)
 		m_context << Instruction::POP;
 }
 
+void CompilerUtils::popAndJump(unsigned _toHeight, eth::AssemblyItem const& _jumpTo)
+{
+	unsigned amount = m_context.stackHeight() - _toHeight;
+	popStackSlots(amount);
+	m_context.appendJumpTo(_jumpTo);
+	m_context.adjustStackOffset(amount);
+}
+
 unsigned CompilerUtils::sizeOnStack(vector<shared_ptr<Type const>> const& _variableTypes)
 {
 	unsigned size = 0;
diff --git a/libsolidity/codegen/CompilerUtils.h b/libsolidity/codegen/CompilerUtils.h
index 0ff3ad7c..26df4765 100644
--- a/libsolidity/codegen/CompilerUtils.h
+++ b/libsolidity/codegen/CompilerUtils.h
@@ -241,6 +241,10 @@ public:
 	void popStackElement(Type const& _type);
 	/// Removes element from the top of the stack _amount times.
 	void popStackSlots(size_t _amount);
+	/// Pops slots from the stack such that its height is _toHeight.
+	/// Adds jump to _jumpTo.
+	/// Readjusts the stack offset to the original value.
+	void popAndJump(unsigned _toHeight, eth::AssemblyItem const& _jumpTo);
 
 	template <class T>
 	static unsigned sizeOnStack(std::vector<T> const& _variables);
diff --git a/libsolidity/codegen/ContractCompiler.cpp b/libsolidity/codegen/ContractCompiler.cpp
index def4a878..474d2b68 100644
--- a/libsolidity/codegen/ContractCompiler.cpp
+++ b/libsolidity/codegen/ContractCompiler.cpp
@@ -446,6 +446,7 @@ bool ContractCompiler::visit(FunctionDefinition const& _function)
 		if (auto c = m_context.nextConstructor(dynamic_cast<ContractDefinition const&>(*_function.scope())))
 			appendBaseConstructor(*c);
 
+	solAssert(m_returnTags.empty(), "");
 	m_breakTags.clear();
 	m_continueTags.clear();
 	m_currentFunction = &_function;
@@ -666,8 +667,6 @@ bool ContractCompiler::visit(WhileStatement const& _whileStatement)
 	eth::AssemblyItem loopEnd = m_context.newTag();
 	m_breakTags.push_back({loopEnd, m_context.stackHeight()});
 
-	visitLoop(&_whileStatement);
-
 	m_context << loopStart;
 
 	if (_whileStatement.isDoWhile())
@@ -710,7 +709,7 @@ bool ContractCompiler::visit(ForStatement const& _forStatement)
 	eth::AssemblyItem loopEnd = m_context.newTag();
 	eth::AssemblyItem loopNext = m_context.newTag();
 
-	visitLoop(&_forStatement);
+	storeStackHeight(&_forStatement);
 
 	if (_forStatement.initializationExpression())
 		_forStatement.initializationExpression()->accept(*this);
@@ -754,7 +753,7 @@ bool ContractCompiler::visit(Continue const& _continueStatement)
 {
 	CompilerContext::LocationSetter locationSetter(m_context, _continueStatement);
 	solAssert(!m_continueTags.empty(), "");
-	popAndJump(m_context.stackHeight() - m_continueTags.back().second, m_continueTags.back().first);
+	CompilerUtils(m_context).popAndJump(m_continueTags.back().second, m_continueTags.back().first);
 	return false;
 }
 
@@ -762,7 +761,7 @@ bool ContractCompiler::visit(Break const& _breakStatement)
 {
 	CompilerContext::LocationSetter locationSetter(m_context, _breakStatement);
 	solAssert(!m_breakTags.empty(), "");
-	popAndJump(m_context.stackHeight() - m_breakTags.back().second, m_breakTags.back().first);
+	CompilerUtils(m_context).popAndJump(m_breakTags.back().second, m_breakTags.back().first);
 	return false;
 }
 
@@ -789,7 +788,7 @@ bool ContractCompiler::visit(Return const& _return)
 			CompilerUtils(m_context).moveToStackVariable(*retVariable);
 	}
 
-	popAndJump(m_context.stackHeight() - m_returnTags.back().second, m_returnTags.back().first);
+	CompilerUtils(m_context).popAndJump(m_returnTags.back().second, m_returnTags.back().first);
 	return false;
 }
 
@@ -872,7 +871,7 @@ bool ContractCompiler::visit(PlaceholderStatement const& _placeholderStatement)
 
 bool ContractCompiler::visit(Block const& _block)
 {
-	m_scopeStackHeight[m_modifierDepth][&_block] = m_context.stackHeight();
+	storeStackHeight(&_block);
 	return true;
 }
 
@@ -999,7 +998,7 @@ eth::AssemblyPointer ContractCompiler::cloneRuntime() const
 
 void ContractCompiler::popScopedVariables(ASTNode const* _node)
 {
-	unsigned blockHeight = m_scopeStackHeight[m_modifierDepth][_node];
+	unsigned blockHeight = m_scopeStackHeight.at(m_modifierDepth).at(_node);
 	unsigned stackDiff = m_context.stackHeight() - blockHeight;
 	CompilerUtils(m_context).popStackSlots(stackDiff);
 	m_context.removeVariablesAboveStackHeight(blockHeight);
@@ -1008,14 +1007,7 @@ void ContractCompiler::popScopedVariables(ASTNode const* _node)
 		m_scopeStackHeight.erase(m_modifierDepth);
 }
 
-void ContractCompiler::visitLoop(BreakableStatement const* _loop)
-{
-	m_scopeStackHeight[m_modifierDepth][_loop] = m_context.stackHeight();
-}
-
-void ContractCompiler::popAndJump(unsigned _amount, eth::AssemblyItem const& _jumpTo)
+void ContractCompiler::storeStackHeight(ASTNode const* _node)
 {
-	CompilerUtils(m_context).popStackSlots(_amount);
-	m_context.appendJumpTo(_jumpTo);
-	m_context.adjustStackOffset(_amount);
+	m_scopeStackHeight[m_modifierDepth][_node] = m_context.stackHeight();
 }
diff --git a/libsolidity/codegen/ContractCompiler.h b/libsolidity/codegen/ContractCompiler.h
index 72f46495..8516ec2c 100644
--- a/libsolidity/codegen/ContractCompiler.h
+++ b/libsolidity/codegen/ContractCompiler.h
@@ -128,12 +128,8 @@ private:
 	/// Frees the variables of a certain scope (to be used when leaving).
 	void popScopedVariables(ASTNode const* _node);
 
-	/// Pops _amount slots from the stack and jumps to _jumpTo.
-	/// Also readjusts the stack offset to the original value.
-	void popAndJump(unsigned _amount, eth::AssemblyItem const& _jumpTo);
-
 	/// Sets the stack height for the visited loop.
-	void visitLoop(BreakableStatement const* _loop);
+	void storeStackHeight(ASTNode const* _node);
 
 	bool const m_optimise;
 	/// Pointer to the runtime compiler in case this is a creation compiler.
diff --git a/test/libsolidity/SolidityEndToEndTest.cpp b/test/libsolidity/SolidityEndToEndTest.cpp
index fbc6a9c6..52fdfed7 100644
--- a/test/libsolidity/SolidityEndToEndTest.cpp
+++ b/test/libsolidity/SolidityEndToEndTest.cpp
@@ -523,6 +523,57 @@ BOOST_AUTO_TEST_CASE(do_while_loop_continue)
 	ABI_CHECK(callContractFunction("f()"), encodeArgs(42));
 }
 
+BOOST_AUTO_TEST_CASE(do_while_loop_multiple_local_vars)
+{
+	char const* sourceCode = R"(
+		contract test {
+			function f(uint x) public pure returns(uint r) {
+				uint i = 0;
+				do
+				{
+					uint z = x * 2;
+					if (z < 4) break;
+					else {
+						uint k = z + 1;
+						if (k < 8) {
+							x++;
+							continue;
+						}
+					}
+					if (z > 12) return 0;
+					x++;
+					i++;
+				} while (true);
+				return 42;
+			}
+		}
+	)";
+	compileAndRun(sourceCode);
+
+	auto do_while = [](u256 n) -> u256
+	{
+		u256 i = 0;
+		do
+		{
+			u256 z = n * 2;
+			if (z < 4) break;
+			else {
+				u256 k = z + 1;
+				if (k < 8) {
+					n++;
+					continue;
+				}
+			}
+			if (z > 12) return 0;
+			n++;
+			i++;
+		} while (true);
+		return 42;
+	};
+
+	testContractAgainstCppOnRange("f(uint256)", do_while, 0, 12);
+}
+
 BOOST_AUTO_TEST_CASE(nested_loops)
 {
 	// tests that break and continue statements in nested loops jump to the correct place
@@ -574,6 +625,109 @@ BOOST_AUTO_TEST_CASE(nested_loops)
 	testContractAgainstCppOnRange("f(uint256)", nested_loops_cpp, 0, 12);
 }
 
+BOOST_AUTO_TEST_CASE(nested_loops_multiple_local_vars)
+{
+	// tests that break and continue statements in nested loops jump to the correct place
+	char const* sourceCode = R"(
+		contract test {
+			function f(uint x) returns(uint y) {
+				while (x > 0) {
+					uint z = x + 10;
+					uint k = z + 1;
+					if (k > 20) break;
+					if (k > 15) {
+						x--;
+						continue;
+					}
+					while (k > 10) {
+						uint m = k - 1;
+						if (m == 10) return x;
+						return k;
+					}
+					x--;
+					break;
+				}
+				return x;
+			}
+		}
+	)";
+	compileAndRun(sourceCode);
+
+	auto nested_loops_cpp = [](u256 n) -> u256
+	{
+		while (n > 0)
+		{
+			u256 z = n + 10;
+			u256 k = z + 1;
+			if (k > 20) break;
+			if (k > 15) {
+				n--;
+				continue;
+			}
+			while (k > 10)
+			{
+				u256 m = k - 1;
+				if (m == 10) return n;
+				return k;
+			}
+			n--;
+			break;
+		}
+
+		return n;
+	};
+
+	testContractAgainstCppOnRange("f(uint256)", nested_loops_cpp, 0, 12);
+}
+
+BOOST_AUTO_TEST_CASE(for_loop_multiple_local_vars)
+{
+	char const* sourceCode = R"(
+		contract test {
+			function f(uint x) public pure returns(uint r) {
+				for (uint i = 0; i < 12; i++)
+				{
+					uint z = x + 1;
+					if (z < 4) break;
+					else {
+						uint k = z * 2;
+						if (i + k < 10) {
+							x++;
+							continue;
+						}
+					}
+					if (z > 8) return 0;
+					x++;
+				}
+				return 42;
+			}
+		}
+	)";
+	compileAndRun(sourceCode);
+
+	auto for_loop = [](u256 n) -> u256
+	{
+		for (u256 i = 0; i < 12; i++)
+		{
+			u256 z = n + 1;
+			if (z < 4) break;
+			else {
+				u256 k = z * 2;
+				if (i + k < 10) {
+					n++;
+					continue;
+				}
+			}
+			if (z > 8) return 0;
+			n++;
+		}
+		return 42;
+	};
+
+	testContractAgainstCppOnRange("f(uint256)", for_loop, 0, 12);
+}
+
+
 BOOST_AUTO_TEST_CASE(for_loop)
 {
 	char const* sourceCode = R"(