MIR Optimise - Allow assignment elimination with variables

3 files changed, 66 insertions, 17 deletions

diff --git a/src/mir/mir.cpp b/src/mir/mir.cpp
index c878950f..d5aae815 100644
--- a/src/mir/mir.cpp
+++ b/src/mir/mir.cpp
@@ -84,6 +84,49 @@ namespace MIR {
         )
         return os;
     }
+    bool operator<(const LValue& a, const LValue& b)
+    {
+        if( a.tag() != b.tag() )
+            return a.tag() < b.tag();
+        TU_MATCHA( (a, b), (ea, eb),
+        (Variable,
+            return ea < eb;
+            ),
+        (Temporary,
+            return ea.idx < eb.idx;
+            ),
+        (Argument,
+            return ea.idx < eb.idx;
+            ),
+        (Static,
+            return ea < eb;
+            ),
+        (Return,
+            return false;
+            ),
+        (Field,
+            if( *ea.val != *eb.val )
+                return *ea.val < *eb.val;
+            if( ea.field_index != eb.field_index )
+                return ea.field_index < eb.field_index;
+            return true;
+            ),
+        (Deref,
+            return *ea.val < *eb.val;
+            ),
+        (Index,
+            if( *ea.val != *eb.val )
+                return *ea.val < *eb.val;
+            return *ea.idx < *eb.idx;
+            ),
+        (Downcast,
+            if( *ea.val != *eb.val )
+                return *ea.val < *eb.val;
+            return ea.variant_index < eb.variant_index;
+            )
+        )
+        throw "";
+    }
     bool operator==(const LValue& a, const LValue& b)
     {
         if( a.tag() != b.tag() )
diff --git a/src/mir/mir.hpp b/src/mir/mir.hpp
index 93bd7244..cacaea99 100644
--- a/src/mir/mir.hpp
+++ b/src/mir/mir.hpp
@@ -58,6 +58,7 @@ TAGGED_UNION_EX(LValue, (), Variable, (
     )
     );
 extern ::std::ostream& operator<<(::std::ostream& os, const LValue& x);
+extern bool operator<(const LValue& a, const LValue& b);
 extern bool operator==(const LValue& a, const LValue& b);
 static inline bool operator!=(const LValue& a, const LValue& b) {
     return !(a == b);
diff --git a/src/mir/optimise.cpp b/src/mir/optimise.cpp
index 7f6cfce1..dac45de5 100644
--- a/src/mir/optimise.cpp
+++ b/src/mir/optimise.cpp
@@ -422,7 +422,7 @@ void MIR_Optimise(const StaticTraitResolve& resolve, const ::HIR::ItemPath& path
             }
         }
         // 2. Assignments (forward propagate)
-        ::std::map< unsigned int, ::MIR::LValue>    replacements;
+        ::std::map< ::MIR::LValue, ::MIR::LValue>    replacements;
         for(const auto& block : fcn.blocks)
         {
             if( block.terminator.tag() == ::MIR::Terminator::TAGDEAD )
@@ -436,20 +436,27 @@ void MIR_Optimise(const StaticTraitResolve& resolve, const ::HIR::ItemPath& path
                     continue ;
                 const auto& e = stmt.as_Assign();
                 // > Of a temporary from with a RValue::Use
-                // TODO: Variables too.
-                // TODO: Allow any rvalue type (if the usage is a RValue::Use)
-                if( !( e.dst.is_Temporary() && e.src.is_Use() ) )
+                // TODO: Variables too (can eliminate arguments)
+                if( e.dst.is_Temporary() )
+                {
+                    const auto& vu = val_uses.tmp_uses[e.dst.as_Temporary().idx];
+                    // > Where the temporary is written once and read once
+                    if( !( vu.read == 1 && vu.write == 1 ) )
+                        continue ;
+                }
+                else
+                {
                     continue ;
-                auto idx = e.dst.as_Temporary().idx;
-                const auto& vu = val_uses.tmp_uses[idx];
-                // > Where the temporary is written once and read once
-                if( !( vu.read == 1 && vu.write == 1 ) )
+                }
+                // TODO: Allow any rvalue type (if the usage is a RValue::Use)
+                // - Requires fiddling in `replacements
+                if( ! e.src.is_Use() )
                     continue ;
 
                 // Get the root value(s) of the source
                 // TODO: Handle more complex values. (but don't bother for really complex values?)
                 const auto& src = e.src.as_Use();
-                if( !( src.is_Temporary() || src.is_Variable() ) )
+                if( !( src.is_Temporary() || src.is_Variable() || src.is_Argument() ) )
                     continue ;
                 bool src_is_lvalue = true;
 
@@ -535,7 +542,7 @@ void MIR_Optimise(const StaticTraitResolve& resolve, const ::HIR::ItemPath& path
                 if( found )
                 {
                     DEBUG("> Replace " << e.dst << " with " << e.src.as_Use());
-                    replacements.insert( ::std::make_pair(idx, e.src.as_Use().clone()) );
+                    replacements.insert( ::std::make_pair(e.dst.clone(), e.src.as_Use().clone()) );
                 }
                 else
                 {
@@ -550,10 +557,9 @@ void MIR_Optimise(const StaticTraitResolve& resolve, const ::HIR::ItemPath& path
             for(auto& r : replacements)
             {
                 visit_mir_lvalue_mut(r.second, false, [&](auto& lv, bool is_write) {
-                    if( lv.is_Temporary() && !is_write )
+                    if( !is_write )
                     {
-                        auto idx = lv.as_Temporary().idx;
-                        auto it = replacements.find(idx);
+                        auto it = replacements.find(lv);
                         if( it != replacements.end() )
                         {
                             lv = it->second.clone();
@@ -573,10 +579,9 @@ void MIR_Optimise(const StaticTraitResolve& resolve, const ::HIR::ItemPath& path
         {
             auto old_replaced = replaced;
             visit_mir_lvalues_mut(state, fcn, [&](auto& lv, bool is_write){
-                if( lv.is_Temporary() && !is_write )
+                if( !is_write )
                 {
-                    auto idx = lv.as_Temporary().idx;
-                    auto it = replacements.find(idx);
+                    auto it = replacements.find(lv);
                     if( it != replacements.end() )
                     {
                         MIR_ASSERT(state, it->second.tag() != ::MIR::LValue::TAGDEAD, "Replacement of  " << lv << " fired twice");
@@ -594,7 +599,7 @@ void MIR_Optimise(const StaticTraitResolve& resolve, const ::HIR::ItemPath& path
             for(auto it = block.statements.begin(); it != block.statements.end(); )
             {
                 // If the statement was an assign of a replaced temporary, remove it.
-                if( it->is_Assign() && it->as_Assign().dst.is_Temporary() && replacements.count( it->as_Assign().dst.as_Temporary().idx ) > 0 )
+                if( it->is_Assign() && replacements.count( it->as_Assign().dst ) > 0 )
                     it = block.statements.erase(it);
                 else
                     ++it;