1 files changed, 66 insertions, 26 deletions

diff --git a/KnownState.cpp b/KnownState.cpp
index 5a70a74f..b84e656a 100644
--- a/KnownState.cpp
+++ b/KnownState.cpp
@@ -162,29 +162,41 @@ KnownState::StoreOperation KnownState::feedItem(AssemblyItem const& _item, bool
 
 /// Helper function for KnownState::reduceToCommonKnowledge, removes everything from
 /// _this which is not in or not equal to the value in _other.
-template <class _Mapping, class _KeyType> void intersect(
-	_Mapping& _this,
-	_Mapping const& _other,
-	function<_KeyType(_KeyType)> const& _keyTrans = [](_KeyType _k) { return _k; }
-)
+template <class _Mapping> void intersect(_Mapping& _this, _Mapping const& _other)
 {
 	for (auto it = _this.begin(); it != _this.end();)
-		if (_other.count(_keyTrans(it->first)) && _other.at(_keyTrans(it->first)) == it->second)
+		if (_other.count(it->first) && _other.at(it->first) == it->second)
 			++it;
 		else
 			it = _this.erase(it);
 }
 
-template <class _Mapping> void intersect(_Mapping& _this, _Mapping const& _other)
-{
-	intersect<_Mapping, ExpressionClasses::Id>(_this, _other, [](ExpressionClasses::Id _k) { return _k; });
-}
-
 void KnownState::reduceToCommonKnowledge(KnownState const& _other)
 {
 	int stackDiff = m_stackHeight - _other.m_stackHeight;
-	function<int(int)> stackKeyTransform = [=](int _key) -> int { return _key - stackDiff; };
-	intersect(m_stackElements, _other.m_stackElements, stackKeyTransform);
+	for (auto it = m_stackElements.begin(); it != m_stackElements.end();)
+		if (_other.m_stackElements.count(it->first - stackDiff))
+		{
+			Id other = _other.m_stackElements.at(it->first - stackDiff);
+			if (it->second == other)
+				++it;
+			else
+			{
+				set<u256> theseTags = tagsInExpression(it->second);
+				set<u256> otherTags = tagsInExpression(other);
+				if (!theseTags.empty() && !otherTags.empty())
+				{
+					theseTags.insert(otherTags.begin(), otherTags.end());
+					it->second = tagUnion(theseTags);
+					++it;
+				}
+				else
+					it = m_stackElements.erase(it);
+			}
+		}
+		else
+			it = m_stackElements.erase(it);
+
 	// Use the smaller stack height. Essential to terminate in case of loops.
 	if (m_stackHeight > _other.m_stackHeight)
 	{
@@ -201,10 +213,15 @@ void KnownState::reduceToCommonKnowledge(KnownState const& _other)
 
 bool KnownState::operator==(const KnownState& _other) const
 {
-	return m_storageContent == _other.m_storageContent &&
-		m_memoryContent == _other.m_memoryContent &&
-		m_stackHeight == _other.m_stackHeight &&
-		m_stackElements == _other.m_stackElements;
+	if (m_storageContent != _other.m_storageContent || m_memoryContent != _other.m_memoryContent)
+		return false;
+	int stackDiff = m_stackHeight - _other.m_stackHeight;
+	auto thisIt = m_stackElements.cbegin();
+	auto otherIt = _other.m_stackElements.cbegin();
+	for (; thisIt != m_stackElements.cend() && otherIt != _other.m_stackElements.cend(); ++thisIt, ++otherIt)
+		if (thisIt->first - stackDiff != otherIt->first || thisIt->second != otherIt->second)
+			return false;
+	return (thisIt == m_stackElements.cend() && otherIt == _other.m_stackElements.cend());
 }
 
 ExpressionClasses::Id KnownState::stackElement(int _stackHeight, SourceLocation const& _location)
@@ -212,18 +229,17 @@ ExpressionClasses::Id KnownState::stackElement(int _stackHeight, SourceLocation
 	if (m_stackElements.count(_stackHeight))
 		return m_stackElements.at(_stackHeight);
 	// Stack element not found (not assigned yet), create new unknown equivalence class.
-	//@todo check that we do not infer incorrect equivalences when the stack is cleared partially
-	//in between.
-	return m_stackElements[_stackHeight] = initialStackElement(_stackHeight, _location);
+	return m_stackElements[_stackHeight] =
+			m_expressionClasses->find(AssemblyItem(UndefinedItem, _stackHeight, _location));
 }
 
-ExpressionClasses::Id KnownState::initialStackElement(
-	int _stackHeight,
-	SourceLocation const& _location
-)
+void KnownState::clearTagUnions()
 {
-	// This is a special assembly item that refers to elements pre-existing on the initial stack.
-	return m_expressionClasses->find(AssemblyItem(UndefinedItem, u256(_stackHeight), _location));
+	for (auto it = m_stackElements.begin(); it != m_stackElements.end();)
+		if (m_tagUnions.left.count(it->second))
+			it = m_stackElements.erase(it);
+		else
+			++it;
 }
 
 void KnownState::setStackElement(int _stackHeight, Id _class)
@@ -352,3 +368,27 @@ KnownState::Id KnownState::applySha3(
 	return m_knownSha3Hashes[arguments] = v;
 }
 
+set<u256> KnownState::tagsInExpression(KnownState::Id _expressionId)
+{
+	if (m_tagUnions.left.count(_expressionId))
+		return m_tagUnions.left.at(_expressionId);
+	// Might be a tag, then return the set of itself.
+	ExpressionClasses::Expression expr = m_expressionClasses->representative(_expressionId);
+	if (expr.item && expr.item->type() == PushTag)
+		return set<u256>({expr.item->data()});
+	else
+		return set<u256>();
+}
+
+KnownState::Id KnownState::tagUnion(set<u256> _tags)
+{
+	if (m_tagUnions.right.count(_tags))
+		return m_tagUnions.right.at(_tags);
+	else
+	{
+		Id id = m_expressionClasses->newClass(SourceLocation());
+		m_tagUnions.right.insert(make_pair(_tags, id));
+		return id;
+	}
+}
+