Reuse state during common subexpression elimination.

4 files changed, 94 insertions, 73 deletions

diff --git a/Assembly.cpp b/Assembly.cpp
index 9530ded4..abcd4451 100644
--- a/Assembly.cpp
+++ b/Assembly.cpp
@@ -311,54 +311,45 @@ Assembly& Assembly::optimise(bool _enable)
 		copt << toString(*this);
 		count = 0;
 
-		//@todo CFG interface should be a generator, that returns an item and a pointer to a
-		// knownstate, which has to replace the current state if it is not null.
-		// Feed these items to the CSE, but also store them and replace the stored version
-		// if the items generated by the CSE are shorter. (or even use less gas?)
-		copt << "Performing control flow analysis...";
+		copt << "Performing optimisation...";
 		{
 			ControlFlowGraph cfg(m_items);
-			AssemblyItems optItems;
+			AssemblyItems optimisedItems;
 			for (BasicBlock const& block: cfg.optimisedBlocks())
-				copy(m_items.begin() + block.begin, m_items.begin() + block.end,
-					 back_inserter(optItems));
-			if (optItems.size() < m_items.size())
 			{
-				copt << "Old size: " << m_items.size() << ", new size: " << optItems.size();
-				m_items = move(optItems);
-				count++;
-			}
-		}
-
-		copt << "Performing common subexpression elimination...";
-		for (auto iter = m_items.begin(); iter != m_items.end();)
-		{
-			//@todo use only a single state / expression classes instance.
-			KnownState state(make_shared<ExpressionClasses>());
-			CommonSubexpressionEliminator eliminator(state);
-			auto orig = iter;
-			iter = eliminator.feedItems(iter, m_items.end());
-			AssemblyItems optItems;
-			bool shouldReplace = false;
-			try
-			{
-				optItems = eliminator.getOptimizedItems();
-				shouldReplace = (optItems.size() < size_t(iter - orig));
-			}
-			catch (StackTooDeepException const&)
-			{
-				// This might happen if the opcode reconstruction is not as efficient
-				// as the hand-crafted code.
-			}
-
-			if (shouldReplace)
-			{
-				copt << "Old size: " << (iter - orig) << ", new size: " << optItems.size();
-				count++;
-				for (auto moveIter = optItems.begin(); moveIter != optItems.end(); ++orig, ++moveIter)
-					*orig = move(*moveIter);
-				iter = m_items.erase(orig, iter);
+				assertThrow(!!block.startState, OptimizerException, "");
+				CommonSubexpressionEliminator eliminator(*block.startState);
+				auto iter = m_items.begin() + block.begin;
+				auto const end = m_items.begin() + block.end;
+				while (iter < end)
+				{
+					auto orig = iter;
+					iter = eliminator.feedItems(iter, end);
+					bool shouldReplace = false;
+					AssemblyItems optimisedChunk;
+					try
+					{
+						optimisedChunk = eliminator.getOptimizedItems();
+						shouldReplace = (optimisedChunk.size() < size_t(iter - orig));
+					}
+					catch (StackTooDeepException const&)
+					{
+						// This might happen if the opcode reconstruction is not as efficient
+						// as the hand-crafted code.
+					}
+
+					if (shouldReplace)
+					{
+						copt << "Old size: " << (iter - orig) << ", new size: " << optimisedChunk.size();
+						count++;
+						optimisedItems += optimisedChunk;
+					}
+					else
+						copy(orig, iter, back_inserter(optimisedItems));
+				}
 			}
+			if (optimisedItems.size() < m_items.size())
+				m_items = move(optimisedItems);
 		}
 	}
 
@@ -461,7 +452,8 @@ bytes Assembly::assemble() const
 	for (auto const& i: tagRef)
 	{
 		bytesRef r(ret.data() + i.first, bytesPerTag);
-		toBigEndian(tagPos[i.second], r);
+		//@todo in the failure case, we could use the position of the invalid jumpdest
+		toBigEndian(i.second < tagPos.size() ? tagPos[i.second] : (1 << (8 * bytesPerTag)) - 1, r);
 	}
 
 	if (!m_data.empty())
diff --git a/CommonSubexpressionEliminator.cpp b/CommonSubexpressionEliminator.cpp
index 4b85eba4..5beb7966 100644
--- a/CommonSubexpressionEliminator.cpp
+++ b/CommonSubexpressionEliminator.cpp
@@ -45,16 +45,22 @@ vector<AssemblyItem> CommonSubexpressionEliminator::getOptimizedItems()
 	for (int height = minHeight; height <= m_state.stackHeight(); ++height)
 		targetStackContents[height] = m_state.stackElement(height, SourceLocation());
 
-	// Debug info:
-	//stream(cout, initialStackContents, targetStackContents);
-
 	AssemblyItems items = CSECodeGenerator(m_state.expressionClasses(), m_storeOperations).generateCode(
 		m_initialState.stackHeight(),
 		initialStackContents,
 		targetStackContents
 	);
 	if (m_breakingItem)
+	{
 		items.push_back(*m_breakingItem);
+		m_state.feedItem(*m_breakingItem);
+	}
+
+	// cleanup
+	m_initialState = m_state;
+	m_breakingItem = nullptr;
+	m_storeOperations.clear();
+
 	return items;
 }
 
@@ -113,6 +119,7 @@ AssemblyItems CSECodeGenerator::generateCode(
 {
 	m_stackHeight = _initialStackHeight;
 	m_stack = _initialStack;
+	m_targetStack = _targetStackContents;
 	for (auto const& item: m_stack)
 		if (!m_classPositions.count(item.second))
 			m_classPositions[item.second] = item.first;
@@ -122,7 +129,7 @@ AssemblyItems CSECodeGenerator::generateCode(
 	// generate the dependency graph starting from final storage and memory writes and target stack contents
 	for (auto const& p: m_storeOperations)
 		addDependencies(p.second.back().expression);
-	for (auto const& targetItem: _targetStackContents)
+	for (auto const& targetItem: m_targetStack)
 	{
 		m_finalClasses.insert(targetItem.second);
 		addDependencies(targetItem.second);
@@ -141,8 +148,10 @@ AssemblyItems CSECodeGenerator::generateCode(
 			generateClassElement(seqAndId.second, true);
 
 	// generate the target stack elements
-	for (auto const& targetItem: _targetStackContents)
+	for (auto const& targetItem: m_targetStack)
 	{
+		if (m_stack.count(targetItem.first) && m_stack.at(targetItem.first) == targetItem.second)
+			continue; // already there
 		int position = generateClassElement(targetItem.second);
 		assertThrow(position != c_invalidPosition, OptimizerException, "");
 		if (position == targetItem.first)
@@ -164,21 +173,24 @@ AssemblyItems CSECodeGenerator::generateCode(
 
 	// check validity
 	int finalHeight = 0;
-	if (!_targetStackContents.empty())
+	if (!m_targetStack.empty())
 		// have target stack, so its height should be the final height
-		finalHeight = (--_targetStackContents.end())->first;
+		finalHeight = (--m_targetStack.end())->first;
 	else if (!_initialStack.empty())
 		// no target stack, only erase the initial stack
 		finalHeight = _initialStack.begin()->first - 1;
 	else
 		// neither initial no target stack, no change in height
-		finalHeight = 0;
+		finalHeight = _initialStackHeight;
 	assertThrow(finalHeight == m_stackHeight, OptimizerException, "Incorrect final stack height.");
+
 	return m_generatedItems;
 }
 
 void CSECodeGenerator::addDependencies(Id _c)
 {
+	if (m_classPositions.count(_c))
+		return; // it is already on the stack
 	if (m_neededBy.count(_c))
 		return; // we already computed the dependencies for _c
 	ExpressionClasses::Expression expr = m_expressionClasses.representative(_c);
@@ -340,7 +352,7 @@ int CSECodeGenerator::generateClassElement(Id _c, bool _allowSequenced)
 		// this will not append a swap but remove the one that is already there
 		appendOrRemoveSwap(m_stackHeight - 1, location);
 	for (auto arg: arguments)
-		if (canBeRemoved(arg, _c))
+		if (m_classPositions[arg] != c_invalidPosition && canBeRemoved(arg, _c))
 			m_classPositions[arg] = c_invalidPosition;
 	for (size_t i = 0; i < arguments.size(); ++i)
 		m_stack.erase(m_stackHeight - i);
@@ -371,13 +383,22 @@ int CSECodeGenerator::classElementPosition(Id _id) const
 	return m_classPositions.at(_id);
 }
 
-bool CSECodeGenerator::canBeRemoved(Id _element, Id _result)
+bool CSECodeGenerator::canBeRemoved(Id _element, Id _result, int _fromPosition)
 {
-	// Returns false if _element is finally needed or is needed by a class that has not been
-	// computed yet. Note that m_classPositions also includes classes that were deleted in the meantime.
-	if (m_finalClasses.count(_element))
-		return false;
+	// Default for _fromPosition is the canonical position of the element.
+	if (_fromPosition == c_invalidPosition)
+		_fromPosition = classElementPosition(_element);
 
+	bool isCopy = _fromPosition != classElementPosition(_element);
+	if (m_finalClasses.count(_element))
+		// It is part of the target stack. It can be removed if it is a copy that is not in the target position.
+		return isCopy && (!m_targetStack.count(_fromPosition) || m_targetStack[_fromPosition] != _element);
+	else if (isCopy)
+		// It is only a copy, can be removed.
+		return true;
+
+	// Can be removed unless it is needed by a class that has not been computed yet.
+	// Note that m_classPositions also includes classes that were deleted in the meantime.
 	auto range = m_neededBy.equal_range(_element);
 	for (auto it = range.first; it != range.second; ++it)
 		if (it->second != _result && !m_classPositions.count(it->second))
@@ -391,7 +412,7 @@ bool CSECodeGenerator::removeStackTopIfPossible()
 		return false;
 	assertThrow(m_stack.count(m_stackHeight) > 0, OptimizerException, "");
 	Id top = m_stack[m_stackHeight];
-	if (!canBeRemoved(top))
+	if (!canBeRemoved(top, Id(-1), m_stackHeight))
 		return false;
 	m_generatedItems.push_back(AssemblyItem(Instruction::POP));
 	m_stack.erase(m_stackHeight);
diff --git a/CommonSubexpressionEliminator.h b/CommonSubexpressionEliminator.h
index 6e1ba40b..2a9a3125 100644
--- a/CommonSubexpressionEliminator.h
+++ b/CommonSubexpressionEliminator.h
@@ -71,13 +71,6 @@ public:
 	/// @returns the resulting items after optimization.
 	AssemblyItems getOptimizedItems();
 
-	/// Streams debugging information to @a _out.
-	std::ostream& stream(
-		std::ostream& _out,
-		std::map<int, Id> _initialStack = std::map<int, Id>(),
-		std::map<int, Id> _targetStack = std::map<int, Id>()
-	) const;
-
 private:
 	/// Feeds the item into the system for analysis.
 	void feedItem(AssemblyItem const& _item, bool _copyItem = false);
@@ -134,8 +127,9 @@ private:
 	/// @note throws an exception if it is not on the stack.
 	int classElementPosition(Id _id) const;
 
-	/// @returns true if @a _element can be removed - in general or, if given, while computing @a _result.
-	bool canBeRemoved(Id _element, Id _result = Id(-1));
+	/// @returns true if the copy of @a _element can be removed from stack position _fromPosition
+	/// - in general or, if given, while computing @a _result.
+	bool canBeRemoved(Id _element, Id _result = Id(-1), int _fromPosition = c_invalidPosition);
 
 	/// Appends code to remove the topmost stack element if it can be removed.
 	bool removeStackTopIfPossible();
@@ -167,6 +161,7 @@ private:
 	std::map<std::pair<StoreOperation::Target, Id>, StoreOperations> m_storeOperations;
 	/// The set of equivalence classes that should be present on the stack at the end.
 	std::set<Id> m_finalClasses;
+	std::map<int, Id> m_targetStack;
 };
 
 template <class _AssemblyItemIterator>
@@ -175,6 +170,7 @@ _AssemblyItemIterator CommonSubexpressionEliminator::feedItems(
 	_AssemblyItemIterator _end
 )
 {
+	assertThrow(!m_breakingItem, OptimizerException, "Invalid use of CommonSubexpressionEliminator.");
 	for (; _iterator != _end && !SemanticInformation::breaksCSEAnalysisBlock(*_iterator); ++_iterator)
 		feedItem(*_iterator);
 	if (_iterator != _end)
diff --git a/ControlFlowGraph.cpp b/ControlFlowGraph.cpp
index 2e28317a..7ed56ff1 100644
--- a/ControlFlowGraph.cpp
+++ b/ControlFlowGraph.cpp
@@ -142,7 +142,7 @@ void ControlFlowGraph::removeUnusedBlocks()
 		BasicBlock const& block = m_blocks.at(blocksToProcess.back());
 		blocksToProcess.pop_back();
 		for (BlockId tag: block.pushedTags)
-			if (!neededBlocks.count(tag))
+			if (!neededBlocks.count(tag) && m_blocks.count(tag))
 			{
 				neededBlocks.insert(tag);
 				blocksToProcess.push_back(tag);
@@ -191,12 +191,12 @@ void ControlFlowGraph::setPrevLinks()
 		if (push.type() != PushTag)
 			continue;
 		BlockId nextId(push.data());
-		if (m_blocks.at(nextId).prev)
+		if (m_blocks.count(nextId) && m_blocks.at(nextId).prev)
 			continue;
 		bool hasLoop = false;
-		for (BlockId id = nextId; id && !hasLoop; id = m_blocks.at(id).next)
+		for (BlockId id = nextId; id && m_blocks.count(id) && !hasLoop; id = m_blocks.at(id).next)
 			hasLoop = (id == blockId);
-		if (hasLoop)
+		if (hasLoop || !m_blocks.count(nextId))
 			continue;
 
 		m_blocks[nextId].prev = blockId;
@@ -225,6 +225,8 @@ void ControlFlowGraph::gatherKnowledge()
 	{
 		//@todo we might have to do something like incrementing the sequence number for each JUMPDEST
 		assertThrow(!!workQueue.back().first, OptimizerException, "");
+		if (!m_blocks.count(workQueue.back().first))
+			continue; // too bad, we do not know the tag, probably an invalid jump
 		BasicBlock& block = m_blocks.at(workQueue.back().first);
 		KnownStatePointer state = workQueue.back().second;
 		workQueue.pop_back();
@@ -281,6 +283,15 @@ void ControlFlowGraph::gatherKnowledge()
 		)
 			workQueue.push_back(make_pair(block.next, state->copy()));
 	}
+
+	// Remove all blocks we never visited here. This might happen because a tag is pushed but
+	// never used for a JUMP.
+	// Note that this invalidates some contents of pushedTags
+	for (auto it = m_blocks.begin(); it != m_blocks.end();)
+		if (!it->second.startState)
+			m_blocks.erase(it++);
+		else
+			it++;
 }
 
 BasicBlocks ControlFlowGraph::rebuildCode()
@@ -288,7 +299,8 @@ BasicBlocks ControlFlowGraph::rebuildCode()
 	map<BlockId, unsigned> pushes;
 	for (auto& idAndBlock: m_blocks)
 		for (BlockId ref: idAndBlock.second.pushedTags)
-			pushes[ref]++;
+			if (m_blocks.count(ref))
+				pushes[ref]++;
 
 	set<BlockId> blocksToAdd;
 	for (auto it: m_blocks)