3 files changed, 68 insertions, 2 deletions

diff --git a/libjulia/optimiser/Rematerialiser.cpp b/libjulia/optimiser/Rematerialiser.cpp
index ca5f94b0..09a9bb90 100644
--- a/libjulia/optimiser/Rematerialiser.cpp
+++ b/libjulia/optimiser/Rematerialiser.cpp
@@ -29,6 +29,8 @@
 
 #include <libdevcore/CommonData.h>
 
+#include <boost/range/adaptor/reversed.hpp>
+
 using namespace std;
 using namespace dev;
 using namespace dev::julia;
@@ -46,6 +48,7 @@ void Rematerialiser::operator()(VariableDeclaration& _varDecl)
 	set<string> names;
 	for (auto const& var: _varDecl.variables)
 		names.insert(var.name);
+	m_variableScopes.back().first += names;
 	handleAssignment(names, _varDecl.value.get());
 }
 
@@ -74,9 +77,23 @@ void Rematerialiser::operator()(Switch& _switch)
 	handleAssignment(assignedVariables, nullptr);
 }
 
+void Rematerialiser::operator()(FunctionDefinition& _fun)
+{
+	m_variableScopes.push_back(make_pair(set<string>(), true));
+	for (auto const& parameter: _fun.parameters)
+		m_variableScopes.back().first.insert(parameter.name);
+	for (auto const& var: _fun.returnVariables)
+		m_variableScopes.back().first.insert(var.name);
+	ASTModifier::operator()(_fun);
+	m_variableScopes.pop_back();
+}
+
 void Rematerialiser::operator()(ForLoop& _for)
 {
-	(*this)(_for.pre);
+	// Special scope handling of the pre block.
+	m_variableScopes.push_back(make_pair(set<string>(), false));
+	for (auto& statement: _for.pre.statements)
+		visit(statement);
 
 	Assignments ass;
 	ass(_for.body);
@@ -88,6 +105,17 @@ void Rematerialiser::operator()(ForLoop& _for)
 	(*this)(_for.post);
 
 	handleAssignment(ass.names(), nullptr);
+
+	m_variableScopes.pop_back();
+}
+
+void Rematerialiser::operator()(Block& _block)
+{
+	size_t numScopes = m_variableScopes.size();
+	m_variableScopes.push_back(make_pair(set<string>(), false));
+	ASTModifier::operator()(_block);
+	m_variableScopes.pop_back();
+	solAssert(numScopes == m_variableScopes.size(), "");
 }
 
 void Rematerialiser::handleAssignment(set<string> const& _variables, Expression* _value)
@@ -128,6 +156,18 @@ void Rematerialiser::handleAssignment(set<string> const& _variables, Expression*
 	}
 }
 
+bool Rematerialiser::inScope(string const& _variableName) const
+{
+	for (auto const& scope: m_variableScopes | boost::adaptors::reversed)
+	{
+		if (scope.first.count(_variableName))
+			return true;
+		if (scope.second)
+			return false;
+	}
+	return false;
+}
+
 void Rematerialiser::visit(Expression& _e)
 {
 	if (_e.type() == typeid(Identifier))
@@ -136,7 +176,15 @@ void Rematerialiser::visit(Expression& _e)
 		if (m_substitutions.count(identifier.name))
 		{
 			string name = identifier.name;
-			_e = (ASTCopier{}).translate(*m_substitutions.at(name));
+			bool doSubstitute = true;
+			for (auto const& ref: m_references[name])
+				if (!inScope(ref))
+				{
+					doSubstitute = false;
+					break;
+				}
+			if (doSubstitute)
+				_e = (ASTCopier{}).translate(*m_substitutions.at(name));
 		}
 	}
 	ASTModifier::visit(_e);
diff --git a/libjulia/optimiser/Rematerialiser.h b/libjulia/optimiser/Rematerialiser.h
index f2a3102f..1accc3f6 100644
--- a/libjulia/optimiser/Rematerialiser.h
+++ b/libjulia/optimiser/Rematerialiser.h
@@ -44,20 +44,29 @@ public:
 	virtual void operator()(VariableDeclaration& _varDecl) override;
 	virtual void operator()(If& _if) override;
 	virtual void operator()(Switch& _switch) override;
+	virtual void operator()(FunctionDefinition&) override;
 	virtual void operator()(ForLoop&) override;
+	virtual void operator()(Block& _block) override;
 
 protected:
+	using ASTModifier::visit;
 	virtual void visit(Expression& _e) override;
 
 private:
 	void handleAssignment(std::set<std::string> const& _names, Expression* _value);
 
+	/// Returns true iff the variable is in scope.
+	bool inScope(std::string const& _variableName) const;
+
 	/// Substitutions to be performed, if possible.
 	std::map<std::string, Expression const*> m_substitutions;
 	/// m_references[a].contains(b) <=> the current expression assigned to a references b
 	std::map<std::string, std::set<std::string>> m_references;
 	/// m_referencedBy[b].contains(a) <=> the current expression assigned to a references b
 	std::map<std::string, std::set<std::string>> m_referencedBy;
+	/// List of scopes, where each scope is a set of variables and a bool that tells
+	/// whether it is a function body (true) or not.
+	std::vector<std::pair<std::set<std::string>, bool>> m_variableScopes;
 };
 
 }
diff --git a/test/libjulia/Rematerialiser.cpp b/test/libjulia/Rematerialiser.cpp
index ae254ad9..020f0020 100644
--- a/test/libjulia/Rematerialiser.cpp
+++ b/test/libjulia/Rematerialiser.cpp
@@ -126,4 +126,13 @@ BOOST_AUTO_TEST_CASE(reassignment)
 	);
 }
 
+BOOST_AUTO_TEST_CASE(do_not_move_out_of_scope)
+{
+	// Cannot replace by `let b := x` by `let b := a` since a is out of scope.
+	CHECK(
+		"{ let x { let a := sload(0) x := a } let b := x }",
+		"{ let x { let a := sload(0) x := a } let b := x }"
+	);
+}
+
 BOOST_AUTO_TEST_SUITE_END()