Fix for rematerializer.

2 files changed, 16 insertions, 0 deletions

diff --git a/libjulia/optimiser/Rematerialiser.cpp b/libjulia/optimiser/Rematerialiser.cpp
index 09a9bb90..bf7d7d16 100644
--- a/libjulia/optimiser/Rematerialiser.cpp
+++ b/libjulia/optimiser/Rematerialiser.cpp
@@ -136,6 +136,10 @@ void Rematerialiser::handleAssignment(set<string> const& _variables, Expression*
 		else
 			m_substitutions.erase(name);
 	}
+	else
+		for (auto const& name: _variables)
+			m_substitutions.erase(name);
+
 	// Disallow substitutions that use a variable that will be reassigned by this assignment.
 	for (auto const& name: _variables)
 		for (auto const& ref: m_referencedBy[name])
diff --git a/test/libjulia/Rematerialiser.cpp b/test/libjulia/Rematerialiser.cpp
index 020f0020..d6cc0dae 100644
--- a/test/libjulia/Rematerialiser.cpp
+++ b/test/libjulia/Rematerialiser.cpp
@@ -118,6 +118,18 @@ BOOST_AUTO_TEST_CASE(branches_for)
 	);
 }
 
+BOOST_AUTO_TEST_CASE(branches_for_declared_in_init)
+{
+	CHECK(
+		"{ let b := 0 for { let a := 1 pop(a) } a { pop(a) } { b := 1 pop(a) } }",
+		"{ let b := 0 for { let a := 1 pop(1) } 1 { pop(1) } { b := 1 pop(1) } }"
+	);
+	CHECK(
+		"{ let b := 0 for { let a := 1 pop(a) } lt(a, 0) { pop(a) a := add(a, 3) } { b := 1 pop(a) } }",
+		"{ let b := 0 for { let a := 1 pop(1) } lt(a, 0) { pop(a) a := add(a, 3) } { b := 1 pop(a) } }"
+	);
+}
+
 BOOST_AUTO_TEST_CASE(reassignment)
 {
 	CHECK(