MIR Opt - Partial dead temporary eliminaton

5 files changed, 429 insertions, 17 deletions

diff --git a/src/mir/check.cpp b/src/mir/check.cpp
index e867c822..7d6b431a 100644
--- a/src/mir/check.cpp
+++ b/src/mir/check.cpp
@@ -599,6 +599,15 @@ void MIR_Validate(const StaticTraitResolve& resolve, const ::HIR::ItemPath& path
                 // Check that the condition is an enum
                 ),
             (Call,
+                if( e.fcn.is_Value() )
+                {
+                    ::HIR::TypeRef  tmp;
+                    const auto& ty = state.get_lvalue_type(tmp, e.fcn.as_Value());
+                    if( ! ty.m_data.is_Function() )
+                    {
+                        //MIR_BUG(state, "Call Fcn::Value with non-function type - " << ty);
+                    }
+                }
                 // Typecheck arguments and return value
                 )
             )
diff --git a/src/mir/helpers.cpp b/src/mir/helpers.cpp
index 027c68b7..a0c8df07 100644
--- a/src/mir/helpers.cpp
+++ b/src/mir/helpers.cpp
@@ -24,7 +24,8 @@ void ::MIR::TypeResolve::print_msg(const char* tag, ::std::function<void(::std::
     os << ": ";
     cb(os);
     os << ::std::endl;
-    throw CheckFailure {};
+    abort();
+    //throw CheckFailure {};
 }
 
 const ::MIR::BasicBlock& ::MIR::TypeResolve::get_block(::MIR::BasicBlockId id) const
@@ -57,12 +58,15 @@ const ::HIR::TypeRef& ::MIR::TypeResolve::get_lvalue_type(::HIR::TypeRef& tmp, c
 {
     TU_MATCH(::MIR::LValue, (val), (e),
     (Variable,
+        MIR_ASSERT(*this, e < m_fcn.named_variables.size(), val << " out of range (" << m_fcn.named_variables.size() << ")");
         return m_fcn.named_variables.at(e);
         ),
     (Temporary,
+        MIR_ASSERT(*this, e.idx < m_fcn.temporaries.size(), val << " out of range (" << m_fcn.temporaries.size() << ")");
         return m_fcn.temporaries.at(e.idx);
         ),
     (Argument,
+        MIR_ASSERT(*this, e.idx < m_args.size(), val << " out of range (" << m_args.size() << ")");
         return m_args.at(e.idx).second;
         ),
     (Static,
diff --git a/src/mir/mir.cpp b/src/mir/mir.cpp
index f7a5f3e8..c878950f 100644
--- a/src/mir/mir.cpp
+++ b/src/mir/mir.cpp
@@ -84,6 +84,53 @@ namespace MIR {
         )
         return os;
     }
+    bool operator==(const LValue& a, const LValue& b)
+    {
+        if( a.tag() != b.tag() )
+            return false;
+        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 true;
+            ),
+        (Field,
+            if( *ea.val != *eb.val )
+                return false;
+            if( ea.field_index != eb.field_index )
+                return false;
+            return true;
+            ),
+        (Deref,
+            return *ea.val == *eb.val;
+            ),
+        (Index,
+            if( *ea.val != *eb.val )
+                return false;
+            if( *ea.idx != *eb.idx )
+                return false;
+            return true;
+            ),
+        (Downcast,
+            if( *ea.val != *eb.val )
+                return false;
+            if( ea.variant_index != eb.variant_index )
+                return false;
+            return true;
+            )
+        )
+        throw "";
+    }
 
     ::std::ostream& operator<<(::std::ostream& os, const RValue& x)
     {
diff --git a/src/mir/mir.hpp b/src/mir/mir.hpp
index 6bd371b5..93bd7244 100644
--- a/src/mir/mir.hpp
+++ b/src/mir/mir.hpp
@@ -58,6 +58,10 @@ TAGGED_UNION_EX(LValue, (), Variable, (
     )
     );
 extern ::std::ostream& operator<<(::std::ostream& os, const LValue& x);
+extern bool operator==(const LValue& a, const LValue& b);
+static inline bool operator!=(const LValue& a, const LValue& b) {
+    return !(a == b);
+}
 
 enum class eBinOp
 {
diff --git a/src/mir/optimise.cpp b/src/mir/optimise.cpp
index bb9536b0..d1f9c7fa 100644
--- a/src/mir/optimise.cpp
+++ b/src/mir/optimise.cpp
@@ -37,6 +37,159 @@ namespace {
             return rv;
         }
     }
+
+    bool visit_mir_lvalue_mut(::MIR::LValue& lv, bool is_write, ::std::function<bool(::MIR::LValue& , bool)> cb)
+    {
+        if( cb(lv, is_write) )
+            return true;
+        TU_MATCHA( (lv), (e),
+        (Variable,
+            ),
+        (Argument,
+            ),
+        (Temporary,
+            ),
+        (Static,
+            ),
+        (Return,
+            ),
+        (Field,
+            return visit_mir_lvalue_mut(*e.val, is_write, cb);
+            ),
+        (Deref,
+            return visit_mir_lvalue_mut(*e.val, is_write, cb);
+            ),
+        (Index,
+            if( visit_mir_lvalue_mut(*e.val, is_write, cb) )
+                return true;
+            return visit_mir_lvalue_mut(*e.idx, false, cb);
+            ),
+        (Downcast,
+            return visit_mir_lvalue_mut(*e.val, is_write, cb);
+            )
+        )
+        return false;
+    }
+
+    bool visit_mir_lvalues_mut(::MIR::Statement& stmt, ::std::function<bool(::MIR::LValue& , bool)> cb)
+    {
+        bool rv = false;
+        TU_MATCHA( (stmt), (e),
+        (Assign,
+            TU_MATCHA( (e.src), (se),
+            (Use,
+                visit_mir_lvalue_mut(se, false, cb);
+                ),
+            (Constant,
+                ),
+            (SizedArray,
+                visit_mir_lvalue_mut(se.val, false, cb);
+                ),
+            (Borrow,
+                visit_mir_lvalue_mut(se.val, (se.type != ::HIR::BorrowType::Shared), cb);
+                ),
+            (Cast,
+                visit_mir_lvalue_mut(se.val, false, cb);
+                ),
+            (BinOp,
+                visit_mir_lvalue_mut(se.val_l, false, cb);
+                visit_mir_lvalue_mut(se.val_r, false, cb);
+                ),
+            (UniOp,
+                visit_mir_lvalue_mut(se.val, false, cb);
+                ),
+            (DstMeta,
+                visit_mir_lvalue_mut(se.val, false, cb);
+                ),
+            (DstPtr,
+                visit_mir_lvalue_mut(se.val, false, cb);
+                ),
+            (MakeDst,
+                // TODO: Would prefer a flag to indicate "move"
+                visit_mir_lvalue_mut(se.ptr_val, false, cb);
+                visit_mir_lvalue_mut(se.meta_val, false, cb);
+                ),
+            (Tuple,
+                for(auto& v : se.vals)
+                    visit_mir_lvalue_mut(v, false, cb);
+                ),
+            (Array,
+                for(auto& v : se.vals)
+                    visit_mir_lvalue_mut(v, false, cb);
+                ),
+            (Variant,
+                visit_mir_lvalue_mut(se.val, false, cb);
+                ),
+            (Struct,
+                for(auto& v : se.vals)
+                    visit_mir_lvalue_mut(v, false, cb);
+                )
+            )
+            rv |= visit_mir_lvalue_mut(e.dst, true, cb);
+            ),
+        (Asm,
+            for(auto& v : e.inputs)
+                rv |= visit_mir_lvalue_mut(v.second, false, cb);
+            for(auto& v : e.outputs)
+                rv |= visit_mir_lvalue_mut(v.second, true, cb);
+            ),
+        (Drop,
+            rv |= visit_mir_lvalue_mut(e.slot, false, cb);
+            )
+        )
+        return rv;
+    }
+    bool visit_mir_lvalues(const ::MIR::Statement& stmt, ::std::function<bool(const ::MIR::LValue& , bool)> cb)
+    {
+        return visit_mir_lvalues_mut(const_cast<::MIR::Statement&>(stmt), [&](auto& lv, bool im){ return cb(lv, im); });
+    }
+
+    void visit_mir_lvalues_mut(::MIR::TypeResolve& state, ::MIR::Function& fcn, ::std::function<bool(::MIR::LValue& , bool)> cb)
+    {
+        for(unsigned int block_idx = 0; block_idx < fcn.blocks.size(); block_idx ++)
+        {
+            auto& block = fcn.blocks[block_idx];
+            if( block.terminator.tag() == ::MIR::Terminator::TAGDEAD )
+                continue ;
+            for(auto& stmt : block.statements)
+            {
+                state.set_cur_stmt(block_idx, (&stmt - &block.statements.front()));
+                visit_mir_lvalues_mut(stmt, cb);
+            }
+            state.set_cur_stmt_term(block_idx);
+            TU_MATCHA( (block.terminator), (e),
+            (Incomplete,
+                ),
+            (Return,
+                ),
+            (Diverge,
+                ),
+            (Goto,
+                ),
+            (Panic,
+                ),
+            (If,
+                visit_mir_lvalue_mut(e.cond, false, cb);
+                ),
+            (Switch,
+                visit_mir_lvalue_mut(e.val, false, cb);
+                ),
+            (Call,
+                if( e.fcn.is_Value() ) {
+                    visit_mir_lvalue_mut(e.fcn.as_Value(), false, cb);
+                }
+                for(auto& v : e.args)
+                    visit_mir_lvalue_mut(v, false, cb);
+                visit_mir_lvalue_mut(e.ret_val, true, cb);
+                )
+            )
+        }
+    }
+    void visit_mir_lvalues(::MIR::TypeResolve& state, const ::MIR::Function& fcn, ::std::function<bool(const ::MIR::LValue& , bool)> cb)
+    {
+        visit_mir_lvalues_mut(state, const_cast<::MIR::Function&>(fcn), [&](auto& lv, bool im){ return cb(lv, im); });
+    }
+
 }
 
 void MIR_Optimise(const StaticTraitResolve& resolve, const ::HIR::ItemPath& path, ::MIR::Function& fcn, const ::HIR::Function::args_t& args, const ::HIR::TypeRef& ret_type)
@@ -118,10 +271,11 @@ void MIR_Optimise(const StaticTraitResolve& resolve, const ::HIR::ItemPath& path
                 DEBUG("Append bb " << tgt << " to bb" << i);
 
                 assert( &fcn.blocks[tgt] != &block );
+                auto src_block = mv$(fcn.blocks[tgt]);
 
-                for(auto& stmt : fcn.blocks[tgt].statements)
+                for(auto& stmt : src_block.statements)
                     block.statements.push_back( mv$(stmt) );
-                block.terminator = mv$( fcn.blocks[tgt].terminator );
+                block.terminator = mv$( src_block.terminator );
             }
             i ++;
         }
@@ -136,10 +290,166 @@ void MIR_Optimise(const StaticTraitResolve& resolve, const ::HIR::ItemPath& path
     // - Can use little heristics like a Call pointing to an assignment of its RV
     // - Count the read/write count of a variable, if it's 1,1 then this optimisation is correct.
     // - If the count is read=*,write=1 and the write is of an argument, replace with the argument.
+    struct ValUse {
+        unsigned int    read = 0;
+        unsigned int    write = 0;
+    };
+    struct {
+        ::std::vector<ValUse> var_uses;
+        ::std::vector<ValUse> tmp_uses;
+
+        void use_lvalue(const ::MIR::LValue& lv, bool is_write) {
+            TU_MATCHA( (lv), (e),
+            (Variable,
+                auto& vu = var_uses[e];
+                if( is_write )
+                    vu.write += 1;
+                else
+                    vu.read += 1;
+                ),
+            (Argument,
+                ),
+            (Temporary,
+                auto& vu = tmp_uses[e.idx];
+                if( is_write )
+                    vu.write += 1;
+                else
+                    vu.read += 1;
+                ),
+            (Static,
+                ),
+            (Return,
+                ),
+            (Field,
+                use_lvalue(*e.val, is_write);
+                ),
+            (Deref,
+                use_lvalue(*e.val, is_write);
+                ),
+            (Index,
+                use_lvalue(*e.val, is_write);
+                use_lvalue(*e.idx, false);
+                ),
+            (Downcast,
+                use_lvalue(*e.val, is_write);
+                )
+            )
+        }
+        void read_lvalue(const ::MIR::LValue& lv) {
+            use_lvalue(lv, false);
+        }
+        void write_lvalue(const ::MIR::LValue& lv) {
+            use_lvalue(lv, true);
+        }
+    } val_uses = {
+        ::std::vector<ValUse>(fcn.named_variables.size()),
+        ::std::vector<ValUse>(fcn.temporaries.size())
+        };
+    visit_mir_lvalues(state, fcn, [&](const auto& lv, bool im){ val_uses.use_lvalue(lv, im); return false; });
+
+    // Rules:
+    // - Find an assignment `tmp = Use(...)` where the temporary is only written and read once
+    // - Stop on any conditional terminator
+    // - Any lvalues in the source lvalue must not be mutated between the source assignment and the usage.
+    //  > This includes mutation, borrowing, or moving.
+    {
+        ::std::map< unsigned int, ::MIR::LValue>    replacements;
+        for(const auto& block : fcn.blocks)
+        {
+            if( block.terminator.tag() == ::MIR::Terminator::TAGDEAD )
+                continue ;
+            //FOREACH_IDX(stmt_idx, block.statements)
+            for(unsigned int stmt_idx = 0; stmt_idx < block.statements.size(); stmt_idx ++)
+            {
+                const auto& stmt = block.statements[stmt_idx];
+                // > Assignment
+                if( ! stmt.is_Assign() )
+                    continue ;
+                const auto& e = stmt.as_Assign();
+                // > Of a temporary from with a RValue::Use
+                if( !( e.dst.is_Temporary() && e.src.is_Use() ) )
+                    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 ) )
+                    continue ;
+                // Get the root value(s) of the source
+                const auto& src = e.src.as_Use();
+
+                // TODO: Handle more complex values. (but don't bother for really complex values?)
+                if( !src.is_Temporary() || !src.is_Variable() )
+                    continue ;
+
+                // Eligable for replacement
+                // Find where this value is used
+                // - Stop on a conditional block terminator
+                // - Stop if any value mentioned in the source is mutated/invalidated
+                //bool stop = false;
+                bool found = false;
+                for(unsigned int si2 = stmt_idx+1; si2 < block.statements.size(); si2 ++)
+                {
+                    // Usage found.
+                    if( visit_mir_lvalues(block.statements[si2], [&](const auto& lv, bool ){ return lv == e.dst; }) )
+                    {
+                        found = true;
+                        //stop = true;
+                        break;
+                    }
+
+                    // Determine if source is mutated.
+                    // > Assume that any mutating access of the root value counts (over-cautious)
+                    if( visit_mir_lvalues(block.statements[si2], [&](const auto& lv, bool is_write){ return lv == src && is_write; }) )
+                    {
+                        //stop = true;
+                        break;
+                    }
+                }
+                // Schedule a replacement in a future pass
+                if( found )
+                {
+                    replacements.insert( ::std::make_pair(idx, e.src.as_Use().clone()) );
+                }
+            }
+        }
+
+        // Apply replacements
+        unsigned int replaced = 0;
+        while( replaced < replacements.size() )
+        {
+            visit_mir_lvalues_mut(state, fcn, [&](auto& lv, bool is_write){
+                if( lv.is_Temporary() && !is_write )
+                {
+                    auto idx = lv.as_Temporary().idx;
+                    auto it = replacements.find(idx);
+                    if( it != replacements.end() )
+                    {
+                        MIR_ASSERT(state, it->second.tag() != ::MIR::LValue::TAGDEAD, "Replacement of  " << lv << " fired twice");
+                        lv = ::std::move(it->second);
+                        replaced += 1;
+                    }
+                }
+                return false;
+                });
+        }
+        // Remove assignments of replaced values
+        for(auto& block : fcn.blocks)
+        {
+            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 )
+                    it = block.statements.erase(it);
+                else
+                    ++it;
+            }
+        }
+    }
 
     // GC pass on blocks and variables
     // - Find unused blocks, then delete and rewrite all references.
     {
+        ::std::vector<bool> used_temps( fcn.temporaries.size() );
         ::std::vector<bool> visited( fcn.blocks.size() );
         ::std::vector< ::MIR::BasicBlockId> to_visit;
         to_visit.push_back( 0 );
@@ -147,8 +457,16 @@ void MIR_Optimise(const StaticTraitResolve& resolve, const ::HIR::ItemPath& path
         {
             auto bb = to_visit.back(); to_visit.pop_back();
             visited[bb] = true;
-
             const auto& block = fcn.blocks[bb];
+
+            for(const auto& stmt : block.statements)
+            {
+                TU_IFLET( ::MIR::Statement, stmt, Assign, e,
+                    if( e.dst.is_Temporary() )
+                        used_temps[e.dst.as_Temporary().idx] = true;
+                )
+            }
+
             TU_MATCHA( (block.terminator), (e),
             (Incomplete,
                 ),
@@ -178,19 +496,24 @@ void MIR_Optimise(const StaticTraitResolve& resolve, const ::HIR::ItemPath& path
                     to_visit.push_back(e.ret_block);
                 if( !visited[e.panic_block] )
                     to_visit.push_back(e.panic_block);
+                if( e.ret_val.is_Temporary() )
+                    used_temps[e.ret_val.as_Temporary().idx] = true;
                 )
             )
         }
 
-        ::std::vector<unsigned int> rewrite_table;
+        ::std::vector<unsigned int> block_rewrite_table;
         for(unsigned int i = 0, j = 0; i < fcn.blocks.size(); i ++)
         {
-            if( visited[i] ) {
-                rewrite_table.push_back(j++);
-            }
-            else {
-                rewrite_table.push_back(~0u);
-            }
+            block_rewrite_table.push_back( visited[i] ? j ++ : ~0u );
+        }
+        ::std::vector<unsigned int> temp_rewrite_table;
+        unsigned int n_temp = fcn.temporaries.size();
+        for(unsigned int i = 0, j = 0; i < n_temp; i ++)
+        {
+            if( !used_temps[i] )
+                fcn.temporaries.erase(fcn.temporaries.begin() + j);
+            temp_rewrite_table.push_back( used_temps[i] ? j ++ : ~0u );
         }
 
         auto it = fcn.blocks.begin();
@@ -204,6 +527,23 @@ void MIR_Optimise(const StaticTraitResolve& resolve, const ::HIR::ItemPath& path
             }
             else
             {
+                auto lvalue_cb = [&](auto& lv, bool ) {
+                    if( lv.is_Temporary() ) {
+                        auto& e = lv.as_Temporary();
+                        MIR_ASSERT(state, e.idx < temp_rewrite_table.size(), "Temporary out of range - " << lv);
+                        MIR_ASSERT(state, temp_rewrite_table.at(e.idx) != ~0u, "LValue " << lv << " incorrectly marked as unused");
+                        e.idx = temp_rewrite_table.at(e.idx);
+                    }
+                    return false;
+                    };
+                unsigned int stmt_idx = 0;
+                for(auto& stmt : it->statements)
+                {
+                    state.set_cur_stmt(i, stmt_idx);
+                    visit_mir_lvalues_mut(stmt, lvalue_cb);
+                    stmt_idx ++;
+                }
+                state.set_cur_stmt_term(i);
                 // Rewrite and advance
                 TU_MATCHA( (it->terminator), (e),
                 (Incomplete,
@@ -213,21 +553,29 @@ void MIR_Optimise(const StaticTraitResolve& resolve, const ::HIR::ItemPath& path
                 (Diverge,
                     ),
                 (Goto,
-                    e = rewrite_table[e];
+                    e = block_rewrite_table[e];
                     ),
                 (Panic,
                     ),
                 (If,
-                    e.bb0 = rewrite_table[e.bb0];
-                    e.bb1 = rewrite_table[e.bb1];
+                    visit_mir_lvalue_mut(e.cond, false, lvalue_cb);
+                    e.bb0 = block_rewrite_table[e.bb0];
+                    e.bb1 = block_rewrite_table[e.bb1];
                     ),
                 (Switch,
+                    visit_mir_lvalue_mut(e.val, false, lvalue_cb);
                     for(auto& target : e.targets)
-                        target = rewrite_table[target];
+                        target = block_rewrite_table[target];
                     ),
                 (Call,
-                    e.ret_block = rewrite_table[e.ret_block];
-                    e.panic_block = rewrite_table[e.panic_block];
+                    if( e.fcn.is_Value() ) {
+                        visit_mir_lvalue_mut(e.fcn.as_Value(), false, lvalue_cb);
+                    }
+                    for(auto& v : e.args)
+                        visit_mir_lvalue_mut(v, false, lvalue_cb);
+                    visit_mir_lvalue_mut(e.ret_val, true, lvalue_cb);
+                    e.ret_block   = block_rewrite_table[e.ret_block];
+                    e.panic_block = block_rewrite_table[e.panic_block];
                     )
                 )