5 files changed, 78 insertions, 24 deletions
diff --git a/libyul/optimiser/README.md b/libyul/optimiser/README.md
index faef818b..c2575179 100644
--- a/libyul/optimiser/README.md
+++ b/libyul/optimiser/README.md
@@ -135,6 +135,8 @@ If there are two assignments to a variable where the first one is a movable expr
 and the variable is not used between the two assignments (and the second is not inside
 a loop or conditional, the first one is not inside), the first assignment is removed.
 
+This step also removes movable expression statements.
+
 
 ## Function Unifier
 
diff --git a/libyul/optimiser/RedundantAssignEliminator.cpp b/libyul/optimiser/RedundantAssignEliminator.cpp
index 775b7673..b7217074 100644
--- a/libyul/optimiser/RedundantAssignEliminator.cpp
+++ b/libyul/optimiser/RedundantAssignEliminator.cpp
@@ -96,9 +96,15 @@ void RedundantAssignEliminator::operator()(FunctionDefinition const& _functionDe
 	(*this)(_functionDefinition.body);
 
 	for (auto const& param: _functionDefinition.parameters)
+	{
 		changeUndecidedTo(param.name, State::Unused);
+		finalize(param.name);
+	}
 	for (auto const& retParam: _functionDefinition.returnVariables)
+	{
 		changeUndecidedTo(retParam.name, State::Used);
+		finalize(retParam.name);
+	}
 }
 
 void RedundantAssignEliminator::operator()(ForLoop const& _forLoop)
@@ -150,30 +156,48 @@ void RedundantAssignEliminator::run(Block& _ast)
 	RedundantAssignEliminator rae;
 	rae(_ast);
 
-	std::set<Assignment const*> assignmentsToRemove;
-	for (auto const& variables: rae.m_assignments)
-		for (auto const& assignment: variables.second)
-		{
-			assertThrow(assignment.second != State::Undecided, OptimizerException, "");
-			if (assignment.second == State::Unused && MovableChecker{*assignment.first->value}.movable())
-				assignmentsToRemove.emplace(assignment.first);
-		}
-
-	AssignmentRemover remover{assignmentsToRemove};
+	AssignmentRemover remover{rae.m_assignmentsToRemove};
 	remover(_ast);
 }
 
-void RedundantAssignEliminator::join(RedundantAssignEliminator& _other)
+template <class K, class V, class F>
+void joinMap(std::map<K, V>& _a, std::map<K, V>&& _b, F _conflictSolver)
 {
-	for (auto& var: _other.m_assignments)
-		if (m_assignments.count(var.first))
+	// TODO Perhaps it is better to just create a sorted list
+	// and then use insert(begin, end)
+
+	auto ita = _a.begin();
+	auto aend = _a.end();
+	auto itb = _b.begin();
+	auto bend = _b.end();
+
+	for (; itb != bend; ++ita)
+	{
+		if (ita == aend)
+			ita = _a.insert(ita, std::move(*itb++));
+		else if (ita->first < itb->first)
+			continue;
+		else if (itb->first < ita->first)
+			ita = _a.insert(ita, std::move(*itb++));
+		else
 		{
-			map<Assignment const*, State>& assignmentsHere = m_assignments[var.first];
-			for (auto& assignment: var.second)
-				assignmentsHere[assignment.first].join(assignment.second);
+			_conflictSolver(ita->second, std::move(itb->second));
+			++itb;
 		}
-		else
-			m_assignments[var.first] = std::move(var.second);
+	}
+}
+
+void RedundantAssignEliminator::join(RedundantAssignEliminator& _other)
+{
+	m_assignmentsToRemove.insert(begin(_other.m_assignmentsToRemove), end(_other.m_assignmentsToRemove));
+
+	joinMap(m_assignments, std::move(_other.m_assignments), [](
+		map<Assignment const*, State>& _assignmentHere,
+		map<Assignment const*, State>&& _assignmentThere
+	)
+	{
+		return joinMap(_assignmentHere, std::move(_assignmentThere), State::join);
+	});
 }
 
 void RedundantAssignEliminator::changeUndecidedTo(YulString _variable, RedundantAssignEliminator::State _newState)
@@ -183,6 +207,19 @@ void RedundantAssignEliminator::changeUndecidedTo(YulString _variable, Redundant
 			assignment.second = _newState;
 }
 
+void RedundantAssignEliminator::finalize(YulString _variable)
+{
+	for (auto& assignment: m_assignments[_variable])
+	{
+		assertThrow(assignment.second != State::Undecided, OptimizerException, "");
+		if (assignment.second == State{State::Unused} && MovableChecker{*assignment.first->value}.movable())
+			// TODO the only point where we actually need this
+			// to be a set is for the for loop
+			m_assignmentsToRemove.insert(assignment.first);
+	}
+	m_assignments.erase(_variable);
+}
+
 void AssignmentRemover::operator()(Block& _block)
 {
 	boost::range::remove_erase_if(_block.statements, [=](Statement const& _statement) -> bool {
diff --git a/libyul/optimiser/RedundantAssignEliminator.h b/libyul/optimiser/RedundantAssignEliminator.h
index 805a1f63..76106aae 100644
--- a/libyul/optimiser/RedundantAssignEliminator.h
+++ b/libyul/optimiser/RedundantAssignEliminator.h
@@ -127,10 +127,10 @@ private:
 		State(Value _value = Undecided): m_value(_value) {}
 		inline bool operator==(State _other) const { return m_value == _other.m_value; }
 		inline bool operator!=(State _other) const { return !operator==(_other); }
-		inline void join(State const& _other)
+		static inline void join(State& _a, State const& _b)
 		{
 			// Using "max" works here because of the order of the values in the enum.
-			m_value = Value(std::max(int(_other.m_value), int(m_value)));
+			_a.m_value =  Value(std::max(int(_a.m_value), int(_b.m_value)));
 		}
 	private:
 		Value m_value = Undecided;
@@ -149,8 +149,12 @@ private:
 		}
 		~BlockScope()
 		{
+			// This should actually store all declared variables
+			// into a different mapping
 			for (auto const& var: m_rae.m_declaredVariables)
 				m_rae.changeUndecidedTo(var, State::Unused);
+			for (auto const& var: m_rae.m_declaredVariables)
+				m_rae.finalize(var);
 			swap(m_rae.m_declaredVariables, m_outerDeclaredVariables);
 		}
 
@@ -164,10 +168,13 @@ private:
 	/// Will destroy @a _other.
 	void join(RedundantAssignEliminator& _other);
 	void changeUndecidedTo(YulString _variable, State _newState);
+	void finalize(YulString _variable);
 
 	std::set<YulString> m_declaredVariables;
 	// TODO check that this does not cause nondeterminism!
+	// This could also be a pseudo-map from state to assignment.
 	std::map<YulString, std::map<Assignment const*, State>> m_assignments;
+	std::set<Assignment const*> m_assignmentsToRemove;
 };
 
 class AssignmentRemover: public ASTModifier
diff --git a/libyul/optimiser/UnusedPruner.cpp b/libyul/optimiser/UnusedPruner.cpp
index a7b32873..71e86798 100644
--- a/libyul/optimiser/UnusedPruner.cpp
+++ b/libyul/optimiser/UnusedPruner.cpp
@@ -85,6 +85,16 @@ void UnusedPruner::operator()(Block& _block)
 					}};
 			}
 		}
+		else if (statement.type() == typeid(ExpressionStatement))
+		{
+			ExpressionStatement& exprStmt = boost::get<ExpressionStatement>(statement);
+			if (MovableChecker(exprStmt.expression).movable())
+			{
+				// pop(x) should be movable!
+				subtractReferences(ReferencesCounter::countReferences(exprStmt.expression));
+				statement = Block{std::move(exprStmt.location), {}};
+			}
+		}
 
 	removeEmptyBlocks(_block);
 
diff --git a/libyul/optimiser/UnusedPruner.h b/libyul/optimiser/UnusedPruner.h
index 2dd74940..b5aea3dd 100644
--- a/libyul/optimiser/UnusedPruner.h
+++ b/libyul/optimiser/UnusedPruner.h
@@ -32,10 +32,8 @@ namespace yul
 {
 
 /**
- * Optimisation stage that removes unused variables and functions.
- *
- * TODO: Also remove intermediate variable assignments from movable expressions
- * which are not referenced until after the next assignment to the same variable.
+ * Optimisation stage that removes unused variables and functions and also
+ * removes movable expression statements.
  *
  * Note that this does not remove circular references.
  *