MIR Gen - Reset drop flags on entry to a loop (instead of after use)

3 files changed, 30 insertions, 17 deletions

diff --git a/src/mir/from_hir.cpp b/src/mir/from_hir.cpp
index bc9b04a7..786c2243 100644
--- a/src/mir/from_hir.cpp
+++ b/src/mir/from_hir.cpp
@@ -548,8 +548,8 @@ namespace {
         void visit(::HIR::ExprNode_Loop& node) override
         {
             TRACE_FUNCTION_FR("_Loop", "_Loop");
-            auto loop_body_scope = m_builder.new_scope_loop(node.span());
             auto loop_block = m_builder.new_bb_linked();
+            auto loop_body_scope = m_builder.new_scope_loop(node.span());
             auto loop_next = m_builder.new_bb_unlinked();
 
             auto loop_tmp_scope = m_builder.new_scope_temp(node.span());
@@ -2269,8 +2269,14 @@ namespace {
         root_node.visit( ev );
     }
 
+    // NOTE: Can't clean up yet, as consteval isn't done
+    //MIR_Cleanup(resolve, path, fcn, args, ptr->m_res_type);
     MIR_Validate(resolve, path, fcn, args, ptr->m_res_type);
 
+    if( getenv("MRUSTC_VALIDATE_FULL_EARLY") ) {
+        MIR_Validate_Full(resolve, path, fcn, args, ptr->m_res_type);
+    }
+
     return ::MIR::FunctionPointer(new ::MIR::Function(mv$(fcn)));
 }
 
diff --git a/src/mir/from_hir.hpp b/src/mir/from_hir.hpp
index b11a380a..c4be91a2 100644
--- a/src/mir/from_hir.hpp
+++ b/src/mir/from_hir.hpp
@@ -103,6 +103,9 @@ TAGGED_UNION(ScopeType, Owning,
         ::std::map<unsigned int,VarState>    changed_args;
         bool exit_state_valid;
         SplitEnd    exit_state;
+        // TODO: Any drop flags allocated in the loop must be re-initialised at the start of the loop (or before a loopback)
+        ::MIR::BasicBlockId   entry_bb;
+        ::std::vector<unsigned> drop_flags;
         }),
     // State which should end up with no mutation of variable states
     (Freeze, struct {
diff --git a/src/mir/mir_builder.cpp b/src/mir/mir_builder.cpp
index 5dcb2ed2..97942ba2 100644
--- a/src/mir/mir_builder.cpp
+++ b/src/mir/mir_builder.cpp
@@ -338,12 +338,6 @@ void MirBuilder::push_stmt_drop(const Span& sp, ::MIR::LValue val, unsigned int
     }
 
     this->push_stmt(sp, ::MIR::Statement::make_Drop({ ::MIR::eDropKind::DEEP, mv$(val), flag }));
-
-    if( flag != ~0u )
-    {
-        // Reset flag value back to default.
-        push_stmt_set_dropflag_val(sp, flag, m_output.drop_flags.at(flag));
-    }
 }
 void MirBuilder::push_stmt_drop_shallow(const Span& sp, ::MIR::LValue val, unsigned int flag/*=~0u*/)
 {
@@ -353,12 +347,6 @@ void MirBuilder::push_stmt_drop_shallow(const Span& sp, ::MIR::LValue val, unsig
     // TODO: Ensure that the type is a Box?
 
     this->push_stmt(sp, ::MIR::Statement::make_Drop({ ::MIR::eDropKind::SHALLOW, mv$(val), flag }));
-
-    if( flag != ~0u )
-    {
-        // Reset flag value back to default.
-        push_stmt_set_dropflag_val(sp, flag, m_output.drop_flags.at(flag));
-    }
 }
 void MirBuilder::push_stmt_asm(const Span& sp, ::MIR::Statement::Data_Asm data)
 {
@@ -697,6 +685,14 @@ unsigned int MirBuilder::new_drop_flag(bool default_state)
 {
     auto rv = m_output.drop_flags.size();
     m_output.drop_flags.push_back(default_state);
+    for(size_t i = m_scope_stack.size(); i --;)
+    {
+        if( auto* e = m_scopes.at(m_scope_stack[i]).data.opt_Loop() )
+        {
+            e->drop_flags.push_back(rv);
+            break;
+        }
+    }
     DEBUG("(" << default_state << ") = " << rv);
     return rv;
 }
@@ -741,6 +737,7 @@ ScopeHandle MirBuilder::new_scope_loop(const Span& sp)
 {
     unsigned int idx = m_scopes.size();
     m_scopes.push_back( ScopeDef {sp, ScopeType::make_Loop({})} );
+    m_scopes.back().data.as_Loop().entry_bb = m_current_block;
     m_scope_stack.push_back( idx );
     DEBUG("START (loop) scope " << idx);
     return ScopeHandle { *this, idx };
@@ -1549,13 +1546,20 @@ void MirBuilder::complete_scope(ScopeDef& sd)
     };
 
     // No macro for better debug output.
-    if( sd.data.is_Loop() )
+    if( auto* e = sd.data.opt_Loop() )
     {
-        auto& e = sd.data.as_Loop();
         TRACE_FUNCTION_F("Loop");
-        if( e.exit_state_valid )
+        if( e->exit_state_valid )
+        {
+            H::apply_end_state(sd.span, *this, e->exit_state);
+        }
+
+        // Insert sets of drop flags to the first block (at the start of that block)
+        auto& stmts = m_output.blocks.at(e->entry_bb).statements;
+        for(auto idx : e->drop_flags)
         {
-            H::apply_end_state(sd.span, *this, e.exit_state);
+            DEBUG("Reset df$" << idx);
+            stmts.insert( stmts.begin(), ::MIR::Statement::make_SetDropFlag({ idx, m_output.drop_flags.at(idx), ~0u }) );
         }
     }
     else if( sd.data.is_Split() )