Added assertion and tests suggestions

3 files changed, 17 insertions, 3 deletions

diff --git a/libsolidity/codegen/CompilerContext.cpp b/libsolidity/codegen/CompilerContext.cpp
index 51f76399..3b1b4ec0 100644
--- a/libsolidity/codegen/CompilerContext.cpp
+++ b/libsolidity/codegen/CompilerContext.cpp
@@ -128,6 +128,8 @@ void CompilerContext::addVariable(VariableDeclaration const& _declaration,
 {
 	solAssert(m_asm->deposit() >= 0 && unsigned(m_asm->deposit()) >= _offsetToCurrent, "");
 	unsigned sizeOnStack = _declaration.annotation().type->sizeOnStack();
+	// Variables should not have stack size other than [1, 2],
+	// but that might change when new types are introduced.
 	solAssert(sizeOnStack == 1 || sizeOnStack == 2, "");
 	m_localVariables[&_declaration].push_back(unsigned(m_asm->deposit()) - _offsetToCurrent);
 }
@@ -146,6 +148,7 @@ void CompilerContext::removeVariablesAboveStackHeight(unsigned _stackHeight)
 	for (auto _var: m_localVariables)
 	{
 		solAssert(!_var.second.empty(), "");
+		solAssert(_var.second.back() <= stackHeight(), "");
 		if (_var.second.back() >= _stackHeight)
 			toRemove.push_back(_var.first);
 	}
@@ -153,6 +156,11 @@ void CompilerContext::removeVariablesAboveStackHeight(unsigned _stackHeight)
 		removeVariable(*_var);
 }
 
+unsigned CompilerContext::numberOfLocalVariables() const
+{
+	return m_localVariables.size();
+}
+
 eth::Assembly const& CompilerContext::compiledContract(const ContractDefinition& _contract) const
 {
 	auto ret = m_compiledContracts.find(&_contract);
diff --git a/libsolidity/codegen/CompilerContext.h b/libsolidity/codegen/CompilerContext.h
index f3934615..3f357821 100644
--- a/libsolidity/codegen/CompilerContext.h
+++ b/libsolidity/codegen/CompilerContext.h
@@ -74,6 +74,8 @@ public:
 	void removeVariable(Declaration const& _declaration);
 	/// Removes all local variables currently allocated above _stackHeight.
 	void removeVariablesAboveStackHeight(unsigned _stackHeight);
+	/// Returns the number of currently allocated local variables.
+	unsigned numberOfLocalVariables() const;
 
 	void setCompiledContracts(std::map<ContractDefinition const*, eth::Assembly const*> const& _contracts) { m_compiledContracts = _contracts; }
 	eth::Assembly const& compiledContract(ContractDefinition const& _contract) const;
diff --git a/libsolidity/codegen/ContractCompiler.cpp b/libsolidity/codegen/ContractCompiler.cpp
index 207e0af7..93f698bc 100644
--- a/libsolidity/codegen/ContractCompiler.cpp
+++ b/libsolidity/codegen/ContractCompiler.cpp
@@ -498,8 +498,12 @@ bool ContractCompiler::visit(FunctionDefinition const& _function)
 	m_context.adjustStackOffset(-(int)c_returnValuesSize);
 
 	/// The constructor and the fallback function doesn't to jump out.
-	if (!_function.isConstructor() && !_function.isFallback())
-		m_context.appendJump(eth::AssemblyItem::JumpType::OutOfFunction);
+	if (!_function.isConstructor())
+	{
+		solAssert(m_context.numberOfLocalVariables() == 0, "");
+		if (!_function.isFallback())
+			m_context.appendJump(eth::AssemblyItem::JumpType::OutOfFunction);
+	}
 
 	return false;
 }
@@ -999,10 +1003,10 @@ eth::AssemblyPointer ContractCompiler::cloneRuntime() const
 void ContractCompiler::popScopedVariables(ASTNode const* _node)
 {
 	unsigned blockHeight = m_scopeStackHeight.at(m_modifierDepth).at(_node);
+	m_context.removeVariablesAboveStackHeight(blockHeight);
 	solAssert(m_context.stackHeight() >= blockHeight, "");
 	unsigned stackDiff = m_context.stackHeight() - blockHeight;
 	CompilerUtils(m_context).popStackSlots(stackDiff);
-	m_context.removeVariablesAboveStackHeight(blockHeight);
 	m_scopeStackHeight[m_modifierDepth].erase(_node);
 	if (m_scopeStackHeight[m_modifierDepth].size() == 0)
 		m_scopeStackHeight.erase(m_modifierDepth);