Fix problems with statement blocks.

3 files changed, 48 insertions, 21 deletions

diff --git a/libsolidity/ast/AST.h b/libsolidity/ast/AST.h
index 1a204dca..69186cb7 100644
--- a/libsolidity/ast/AST.h
+++ b/libsolidity/ast/AST.h
@@ -818,6 +818,8 @@ public:
 	virtual void accept(ASTVisitor& _visitor) override;
 	virtual void accept(ASTConstVisitor& _visitor) const override;
 
+	std::vector<ASTPointer<Statement>> const& statements() const { return m_statements; }
+
 private:
 	std::vector<ASTPointer<Statement>> m_statements;
 };
diff --git a/libsolidity/formal/Why3Translator.cpp b/libsolidity/formal/Why3Translator.cpp
index 944b1e7b..14b324e0 100644
--- a/libsolidity/formal/Why3Translator.cpp
+++ b/libsolidity/formal/Why3Translator.cpp
@@ -113,6 +113,8 @@ void Why3Translator::addLine(string const& _line)
 
 void Why3Translator::add(string const& _str)
 {
+	if (m_currentLine.empty())
+		m_indentationAtLineStart = m_indentation;
 	m_currentLine += _str;
 }
 
@@ -120,7 +122,7 @@ void Why3Translator::newLine()
 {
 	if (!m_currentLine.empty())
 	{
-		for (size_t i = 0; i < m_indentation; ++i)
+		for (size_t i = 0; i < m_indentationAtLineStart; ++i)
 			m_result.push_back('\t');
 		m_result += m_currentLine;
 		m_result.push_back('\n');
@@ -130,6 +132,7 @@ void Why3Translator::newLine()
 
 void Why3Translator::unindent()
 {
+	newLine();
 	solAssert(m_indentation > 0, "");
 	m_indentation--;
 }
@@ -177,7 +180,6 @@ bool Why3Translator::visit(FunctionDefinition const& _function)
 		add(" (arg_" + param->name() + ": " + paramType + ")");
 	}
 	add(":");
-	newLine();
 
 	indent();
 	indent();
@@ -191,7 +193,7 @@ bool Why3Translator::visit(FunctionDefinition const& _function)
 			retString += ", ";
 		retString += paramType;
 	}
-	addLine(retString + ")");
+	add(retString + ")");
 	unindent();
 
 	addSourceFromDocStrings(_function.annotation());
@@ -218,6 +220,7 @@ bool Why3Translator::visit(FunctionDefinition const& _function)
 	addLine("try");
 
 	_function.body().accept(*this);
+	add(";");
 	addLine("raise Ret");
 
 	string retVals;
@@ -238,9 +241,18 @@ bool Why3Translator::visit(FunctionDefinition const& _function)
 bool Why3Translator::visit(Block const& _node)
 {
 	addSourceFromDocStrings(_node.annotation());
-	addLine("begin");
+	add("begin");
 	indent();
-	return true;
+	for (size_t i = 0; i < _node.statements().size(); ++i)
+	{
+		_node.statements()[i]->accept(*this);
+		if (!m_currentLine.empty() && i != _node.statements().size() - 1)
+			add(";");
+		newLine();
+	}
+	unindent();
+	add("end");
+	return false;
 }
 
 bool Why3Translator::visit(IfStatement const& _node)
@@ -250,12 +262,12 @@ bool Why3Translator::visit(IfStatement const& _node)
 	add("if ");
 	_node.condition().accept(*this);
 	add(" then");
-	newLine();
-	_node.trueStatement().accept(*this);
+	visitIndentedUnlessBlock(_node.trueStatement());
 	if (_node.falseStatement())
 	{
-		addLine("else");
-		_node.falseStatement()->accept(*this);
+		newLine();
+		add("else");
+		visitIndentedUnlessBlock(*_node.falseStatement());
 	}
 	return false;
 }
@@ -266,10 +278,10 @@ bool Why3Translator::visit(WhileStatement const& _node)
 
 	add("while ");
 	_node.condition().accept(*this);
-	add(" do");
 	newLine();
-	_node.body().accept(*this);
-	addLine("done;");
+	add("do");
+	visitIndentedUnlessBlock(_node.body());
+	add("done");
 	return false;
 }
 
@@ -286,14 +298,12 @@ bool Why3Translator::visit(Return const& _node)
 			error(_node, "Directly returning tuples not supported. Rather assign to return variable.");
 			return false;
 		}
-		newLine();
 		add("begin _" + params.front()->name() + " := ");
 		_node.expression()->accept(*this);
 		add("; raise Ret end");
-		newLine();
 	}
 	else
-		addLine("raise Ret;");
+		add("raise Ret");
 	return false;
 }
 
@@ -310,8 +320,6 @@ bool Why3Translator::visit(VariableDeclarationStatement const& _node)
 	{
 		add("_" + _node.declarations().front()->name() + " := ");
 		_node.initialValue()->accept(*this);
-		add(";");
-		newLine();
 	}
 	return false;
 }
@@ -517,7 +525,21 @@ bool Why3Translator::visit(Literal const& _literal)
 	return false;
 }
 
-void Why3Translator::addSourceFromDocStrings(const DocumentedAnnotation& _annotation)
+void Why3Translator::visitIndentedUnlessBlock(Statement const& _statement)
+{
+	bool isBlock = !!dynamic_cast<Block const*>(&_statement);
+	if (isBlock)
+		newLine();
+	else
+		indent();
+	_statement.accept(*this);
+	if (isBlock)
+		newLine();
+	else
+		unindent();
+}
+
+void Why3Translator::addSourceFromDocStrings(DocumentedAnnotation const& _annotation)
 {
 	auto why3Range = _annotation.docTags.equal_range("why3");
 	for (auto i = why3Range.first; i != why3Range.second; ++i)
diff --git a/libsolidity/formal/Why3Translator.h b/libsolidity/formal/Why3Translator.h
index 21ead977..1aa46424 100644
--- a/libsolidity/formal/Why3Translator.h
+++ b/libsolidity/formal/Why3Translator.h
@@ -64,7 +64,7 @@ private:
 	/// if the type is not supported.
 	std::string toFormalType(Type const& _type) const;
 
-	void indent() { m_indentation++; }
+	void indent() { newLine(); m_indentation++; }
 	void unindent();
 	void addLine(std::string const& _line);
 	void add(std::string const& _str);
@@ -74,13 +74,11 @@ private:
 	virtual bool visit(ContractDefinition const& _contract) override;
 	virtual bool visit(FunctionDefinition const& _function) override;
 	virtual bool visit(Block const&) override;
-	virtual void endVisit(Block const&) override { unindent(); addLine("end;"); }
 	virtual bool visit(IfStatement const& _node) override;
 	virtual bool visit(WhileStatement const& _node) override;
 	virtual bool visit(Return const& _node) override;
 	virtual bool visit(VariableDeclarationStatement const& _node) override;
 	virtual bool visit(ExpressionStatement const&) override;
-	virtual void endVisit(ExpressionStatement const&) override { add(";"); newLine(); }
 	virtual bool visit(Assignment const& _node) override;
 	virtual bool visit(TupleExpression const& _node) override;
 	virtual void endVisit(TupleExpression const&) override { add(")"); }
@@ -98,8 +96,13 @@ private:
 		return false;
 	}
 
+	/// Visits the givin statement and indents it unless it is a block
+	/// (which does its own indentation).
+	void visitIndentedUnlessBlock(Statement const& _statement);
+
 	void addSourceFromDocStrings(DocumentedAnnotation const& _annotation);
 
+	size_t m_indentationAtLineStart = 0;
 	size_t m_indentation = 0;
 	std::string m_currentLine;
 	/// True if we have already seen a contract. For now, only a single contract