4 files changed, 94 insertions, 17 deletions

diff --git a/libevmasm/Assembly.cpp b/libevmasm/Assembly.cpp
index 0247593b..b7859c1f 100644
--- a/libevmasm/Assembly.cpp
+++ b/libevmasm/Assembly.cpp
@@ -94,7 +94,10 @@ unsigned Assembly::bytesRequired(unsigned subTagSize) const
 	}
 }
 
-string Assembly::locationFromSources(StringMap const& _sourceCodes, SourceLocation const& _location) const
+namespace
+{
+
+string locationFromSources(StringMap const& _sourceCodes, SourceLocation const& _location)
 {
 	if (_location.isEmpty() || _sourceCodes.empty() || _location.start >= _location.end || _location.start < 0)
 		return "";
@@ -115,27 +118,92 @@ string Assembly::locationFromSources(StringMap const& _sourceCodes, SourceLocati
 	return cut;
 }
 
-ostream& Assembly::streamAsm(ostream& _out, string const& _prefix, StringMap const& _sourceCodes) const
+class Functionalizer
 {
-	for (size_t i = 0; i < m_items.size(); ++i)
+public:
+	Functionalizer (ostream& _out, string const& _prefix, StringMap const& _sourceCodes):
+		m_out(_out), m_prefix(_prefix), m_sourceCodes(_sourceCodes)
+	{}
+
+	void feed(AssemblyItem const& _item)
 	{
-		AssemblyItem const& item = m_items[i];
-		if (!item.location().isEmpty() && (i == 0 || m_items[i - 1].location() != item.location()))
+		if (!_item.location().isEmpty() && _item.location() != m_location)
 		{
-			_out << _prefix << "    /*";
-			if (item.location().sourceName)
-				_out << " \"" + *item.location().sourceName + "\"";
-			if (!item.location().isEmpty())
-				_out << ":" << to_string(item.location().start) + ":" + to_string(item.location().end);
-			_out << " */" << endl;
+			flush();
+			printLocation();
+			m_location = _item.location();
 		}
-		_out << _prefix << (item.type() == Tag ? "" : "  ") << item.toAssemblyText() << endl;
+		if (!(
+			_item.canBeFunctional() &&
+			_item.returnValues() <= 1 &&
+			_item.arguments() <= int(m_pending.size())
+		))
+		{
+			flush();
+			m_out << m_prefix << (_item.type() == Tag ? "" : "  ") << _item.toAssemblyText() << endl;
+			return;
+		}
+		string expression = _item.toAssemblyText();
+		if (_item.arguments() > 0)
+		{
+			expression += "(";
+			for (int i = 0; i < _item.arguments(); ++i)
+			{
+				expression += m_pending.back();
+				m_pending.pop_back();
+				if (i + 1 < _item.arguments())
+					expression += ", ";
+			}
+			expression += ")";
+		}
+
+		m_pending.push_back(expression);
+		if (_item.returnValues() != 1)
+			flush();
+	}
+
+	void flush()
+	{
+		for (string const& expression: m_pending)
+			m_out << m_prefix << "  " << expression << endl;
+		m_pending.clear();
+	}
+
+	void printLocation()
+	{
+		if (!m_location.sourceName && m_location.isEmpty())
+			return;
+		m_out << m_prefix << "    /*";
+		if (m_location.sourceName)
+			m_out << " \"" + *m_location.sourceName + "\"";
+		if (!m_location.isEmpty())
+			m_out << ":" << to_string(m_location.start) + ":" + to_string(m_location.end);
+		m_out << "  " << locationFromSources(m_sourceCodes, m_location);
+		m_out << " */" << endl;
 	}
 
+private:
+	strings m_pending;
+	SourceLocation m_location;
+
+	ostream& m_out;
+	string const& m_prefix;
+	StringMap const& m_sourceCodes;
+};
+
+}
+
+ostream& Assembly::streamAsm(ostream& _out, string const& _prefix, StringMap const& _sourceCodes) const
+{
+	Functionalizer f(_out, _prefix, _sourceCodes);
+
+	for (auto const& i: m_items)
+		f.feed(i);
+	f.flush();
+
 	if (!m_data.empty() || !m_subs.empty())
 	{
 		_out << _prefix << "stop" << endl;
-		Json::Value data;
 		for (auto const& i: m_data)
 			assertThrow(u256(i.first) < m_subs.size(), AssemblyException, "Data not yet implemented.");
 
diff --git a/libevmasm/Assembly.h b/libevmasm/Assembly.h
index 9e7f9f7b..528c9e74 100644
--- a/libevmasm/Assembly.h
+++ b/libevmasm/Assembly.h
@@ -118,7 +118,6 @@ protected:
 	/// returns the replaced tags.
 	std::map<u256, u256> optimiseInternal(bool _enable, bool _isCreation, size_t _runs);
 
-	std::string locationFromSources(StringMap const& _sourceCodes, SourceLocation const& _location) const;
 	void donePath() { if (m_totalDeposit != INT_MAX && m_totalDeposit != m_deposit) BOOST_THROW_EXCEPTION(InvalidDeposit()); }
 	unsigned bytesRequired(unsigned subTagSize) const;
 
diff --git a/libevmasm/AssemblyItem.cpp b/libevmasm/AssemblyItem.cpp
index 6c7d5425..26d9fded 100644
--- a/libevmasm/AssemblyItem.cpp
+++ b/libevmasm/AssemblyItem.cpp
@@ -76,12 +76,20 @@ unsigned AssemblyItem::bytesRequired(unsigned _addressLength) const
 	BOOST_THROW_EXCEPTION(InvalidOpcode());
 }
 
-int AssemblyItem::deposit() const
+int AssemblyItem::arguments() const
+{
+	if (type() == Operation)
+		return instructionInfo(instruction()).args;
+	else
+		return 0;
+}
+
+int AssemblyItem::returnValues() const
 {
 	switch (m_type)
 	{
 	case Operation:
-		return instructionInfo(instruction()).ret - instructionInfo(instruction()).args;
+		return instructionInfo(instruction()).ret;
 	case Push:
 	case PushString:
 	case PushTag:
diff --git a/libevmasm/AssemblyItem.h b/libevmasm/AssemblyItem.h
index 002b3c87..464368fb 100644
--- a/libevmasm/AssemblyItem.h
+++ b/libevmasm/AssemblyItem.h
@@ -116,7 +116,9 @@ public:
 	/// @returns an upper bound for the number of bytes required by this item, assuming that
 	/// the value of a jump tag takes @a _addressLength bytes.
 	unsigned bytesRequired(unsigned _addressLength) const;
-	int deposit() const;
+	int arguments() const;
+	int returnValues() const;
+	int deposit() const { return returnValues() - arguments(); }
 
 	/// @returns true if the assembly item can be used in a functional context.
 	bool canBeFunctional() const;