MIR Gen - Prevent returned values from being dropped

3 files changed, 164 insertions, 18 deletions

diff --git a/src/mir/from_hir.cpp b/src/mir/from_hir.cpp
index f014961a..d30140e4 100644
--- a/src/mir/from_hir.cpp
+++ b/src/mir/from_hir.cpp
@@ -337,6 +337,7 @@ namespace {
                 ::MIR::RValue   res;
                 bool diverged = false;
                 
+                auto block_res = m_builder.new_temporary(node.m_res_type);
                 auto scope = m_builder.new_scope_var(node.span());
                 
                 for(unsigned int i = 0; i < node.m_nodes.size()-1; i ++)
@@ -347,6 +348,7 @@ namespace {
                     auto stmt_scope = m_builder.new_scope_temp(sp);
                     this->visit_node_ptr(subnode);
                     if( m_builder.has_result() ) {
+                        // TODO: Drop.
                         m_builder.get_result(sp);
                     }
                     
@@ -371,8 +373,11 @@ namespace {
                     if( m_builder.has_result() || m_builder.block_active() ) {
                         ASSERT_BUG(sp, m_builder.block_active(), "Result yielded, but no active block");
                         ASSERT_BUG(sp, m_builder.has_result(), "Active block but no result yeilded");
-                        // PROBLEM:
+                        // PROBLEM: This can drop the result before we want to use it.
+                        
                         res = m_builder.get_result(sp);
+                        m_builder.raise_variables(sp, res);
+                        
                         m_builder.terminate_scope(sp, mv$(stmt_scope));
                         res_valid = true;
                     }
diff --git a/src/mir/from_hir.hpp b/src/mir/from_hir.hpp
index 3fd933fc..4aa04442 100644
--- a/src/mir/from_hir.hpp
+++ b/src/mir/from_hir.hpp
@@ -159,6 +159,10 @@ public:
     // Mark a value as initialised (used for Call, because it has to be done after the panic block is populated)
     void mark_value_assigned(const Span& sp, const ::MIR::LValue& val);
     
+    // Moves control of temporaries up to the next scope
+    void raise_variables(const Span& sp, const ::MIR::LValue& val);
+    void raise_variables(const Span& sp, const ::MIR::RValue& rval);
+    
     void set_cur_block(unsigned int new_block);
     ::MIR::BasicBlockId pause_cur_block();
     void end_block(::MIR::Terminator term);
diff --git a/src/mir/mir_builder.cpp b/src/mir/mir_builder.cpp
index 86dde1ac..5ee97a1d 100644
--- a/src/mir/mir_builder.cpp
+++ b/src/mir/mir_builder.cpp
@@ -336,6 +336,145 @@ void MirBuilder::mark_value_assigned(const Span& sp, const ::MIR::LValue& dst)
     )
 }
 
+void MirBuilder::raise_variables(const Span& sp, const ::MIR::LValue& val)
+{
+    TRACE_FUNCTION_F(val);
+    TU_MATCH_DEF(::MIR::LValue, (val), (e),
+    (
+        ),
+    (Variable,
+        auto idx = e;
+        auto scope_it = m_scope_stack.rbegin();
+        while( scope_it != m_scope_stack.rend() )
+        {
+            auto& scope_def = m_scopes.at(*scope_it);
+            
+            TU_IFLET( ScopeType, scope_def.data, Variables, e,
+                auto tmp_it = ::std::find( e.vars.begin(), e.vars.end(), idx );
+                if( tmp_it != e.vars.end() )
+                {
+                    e.vars.erase( tmp_it );
+                    DEBUG("Move variable " << idx << " from " << *scope_it);
+                    break ;
+                }
+            )
+            ++scope_it;
+        }
+        if( scope_it == m_scope_stack.rend() )
+        {
+            // Temporary wasn't defined in a visible scope?
+            return ;
+        }
+        ++scope_it;
+        
+        while( scope_it != m_scope_stack.rend() )
+        {
+            auto& scope_def = m_scopes.at(*scope_it);
+            
+            TU_IFLET( ScopeType, scope_def.data, Variables, e,
+                e.vars.push_back( idx );
+                DEBUG("- to " << *scope_it);
+                return ;
+            )
+            ++scope_it;
+        }
+        
+        DEBUG("- top");
+        ),
+    (Temporary,
+        auto idx = e.idx;
+        auto scope_it = m_scope_stack.rbegin();
+        while( scope_it != m_scope_stack.rend() )
+        {
+            auto& scope_def = m_scopes.at(*scope_it);
+            
+            TU_IFLET( ScopeType, scope_def.data, Temporaries, e,
+                auto tmp_it = ::std::find( e.temporaries.begin(), e.temporaries.end(), idx );
+                if( tmp_it != e.temporaries.end() )
+                {
+                    e.temporaries.erase( tmp_it );
+                    DEBUG("Move temporary " << idx << " from " << *scope_it);
+                    break ;
+                }
+            )
+            ++scope_it;
+        }
+        if( scope_it == m_scope_stack.rend() )
+        {
+            // Temporary wasn't defined in a visible scope?
+            return ;
+        }
+        ++scope_it;
+        
+        while( scope_it != m_scope_stack.rend() )
+        {
+            auto& scope_def = m_scopes.at(*scope_it);
+            
+            TU_IFLET( ScopeType, scope_def.data, Temporaries, e,
+                e.temporaries.push_back( idx );
+                DEBUG("- to " << *scope_it);
+                return ;
+            )
+            ++scope_it;
+        }
+        
+        DEBUG("- top");
+        )
+    )
+}
+void MirBuilder::raise_variables(const Span& sp, const ::MIR::RValue& rval)
+{
+    TU_MATCHA( (rval), (e),
+    (Use,
+        this->raise_variables(sp, e);
+        ),
+    (Constant,
+        ),
+    (SizedArray,
+        this->raise_variables(sp, e.val);
+        ),
+    (Borrow,
+        // TODO: Wait, is this valid?
+        this->raise_variables(sp, e.val);
+        ),
+    (Cast,
+        this->raise_variables(sp, e.val);
+        ),
+    (BinOp,
+        this->raise_variables(sp, e.val_l);
+        this->raise_variables(sp, e.val_r);
+        ),
+    (UniOp,
+        this->raise_variables(sp, e.val);
+        ),
+    (DstMeta,
+        this->raise_variables(sp, e.val);
+        ),
+    (DstPtr,
+        this->raise_variables(sp, e.val);
+        ),
+    (MakeDst,
+        this->raise_variables(sp, e.ptr_val);
+        this->raise_variables(sp, e.meta_val);
+        ),
+    (Tuple,
+        for(const auto& val : e.vals)
+            this->raise_variables(sp, val);
+        ),
+    (Array,
+        for(const auto& val : e.vals)
+            this->raise_variables(sp, val);
+        ),
+    (Variant,
+        this->raise_variables(sp, e.val);
+        ),
+    (Struct,
+        for(const auto& val : e.vals)
+            this->raise_variables(sp, val);
+        )
+    )
+}
+
 void MirBuilder::set_cur_block(unsigned int new_block)
 {
     ASSERT_BUG(Span(), !m_block_active, "Updating block when previous is active");
@@ -943,19 +1082,18 @@ VarState MirBuilder::get_temp_state(const Span& sp, unsigned int idx) const
     for( auto scope_idx : ::reverse(m_scope_stack) )
     {
         const auto& scope_def = m_scopes.at(scope_idx);
-        TU_MATCH_DEF( ScopeType, (scope_def.data), (e),
-        (
-            ),
-        (Temporaries,
+        if( scope_def.data.is_Temporaries() )
+        {
+            const auto& e = scope_def.data.as_Temporaries();
             auto it = ::std::find(e.temporaries.begin(), e.temporaries.end(), idx);
             if( it != e.temporaries.end() ) {
                 break ;
             }
-            ),
-        (Split,
+        }
+        else if( scope_def.data.is_Split() )
+        {
             // TODO: Does split account for temps? It should.
-            )
-        )
+        }
     }
     
     ASSERT_BUG(sp, idx < m_temporary_states.size(), "Temporary " << idx << " out of range for state table");
@@ -966,19 +1104,18 @@ void MirBuilder::set_temp_state(const Span& sp, unsigned int idx, VarState state
     for( auto scope_idx : ::reverse(m_scope_stack) )
     {
         auto& scope_def = m_scopes.at(scope_idx);
-        TU_MATCH_DEF( ScopeType, (scope_def.data), (e),
-        (
-            ),
-        (Temporaries,
+        if( scope_def.data.is_Temporaries() )
+        {
+            const auto& e = scope_def.data.as_Temporaries();
             auto it = ::std::find(e.temporaries.begin(), e.temporaries.end(), idx);
             if( it != e.temporaries.end() ) {
-                break;
+                break ;
             }
-            ),
-        (Split,
+        }
+        else if( scope_def.data.is_Split() )
+        {
             // TODO: Does split account for temps? It should.
-            )
-        )
+        }
     }
     
     ASSERT_BUG(sp, idx < m_temporary_states.size(), "Temporary " << idx << " out of range for state table");