HIR Closures - Rework to correctly support borrow types

1 files changed, 113 insertions, 40 deletions

diff --git a/src/hir_expand/closures.cpp b/src/hir_expand/closures.cpp
index 6530c3d2..60b85c73 100644
--- a/src/hir_expand/closures.cpp
+++ b/src/hir_expand/closures.cpp
@@ -50,14 +50,20 @@ namespace {
     {
         const ::HIR::TypeRef&   m_closure_type;
         const ::std::vector<unsigned int>&  m_local_vars;
-        const ::std::vector<unsigned int>&  m_captures;
+        const ::std::vector< ::std::pair<unsigned int, ::HIR::ValueUsage> >&  m_captures;
         
         typedef ::std::function< const ::HIR::TypeRef&(const ::HIR::TypeRef&)>  t_monomorph_cb;
         t_monomorph_cb m_monomorph_cb;
         
         ::HIR::ExprNodeP    m_replacement;
     public:
-        ExprVisitor_Mutate(const ::HIR::TypeRef& closure_type, const ::std::vector<unsigned int>& local_vars, const ::std::vector<unsigned int>& captures, t_monomorph_cb mcb):
+        ExprVisitor_Mutate(
+            const ::HIR::TypeRef& closure_type,
+            const ::std::vector<unsigned int>& local_vars,
+            const ::std::vector< ::std::pair<unsigned int, ::HIR::ValueUsage>>& captures,
+            t_monomorph_cb mcb
+            )
+            :
             m_closure_type(closure_type),
             m_local_vars(local_vars),
             m_captures(captures),
@@ -111,26 +117,42 @@ namespace {
         {
             // Do nothing, inner closures should just be value references now
             assert( ! node.m_code );
+            
+            for(auto& subnode : node.m_captures)
+            {
+                visit_node_ptr(subnode);
+            }
         }
         void visit(::HIR::ExprNode_Variable& node) override
         {
             // 1. Is it a closure-local?
-            auto binding_it = ::std::find(m_local_vars.begin(), m_local_vars.end(), node.m_slot);
-            if( binding_it != m_local_vars.end() ) {
-                // NOTE: Offset of 1 is for `self` (`args` is destructured)
-                node.m_slot = 1 + binding_it - m_local_vars.begin();
-                return ;
+            {
+                auto binding_it = ::std::find(m_local_vars.begin(), m_local_vars.end(), node.m_slot);
+                if( binding_it != m_local_vars.end() ) {
+                    // NOTE: Offset of 1 is for `self` (`args` is destructured)
+                    node.m_slot = 1 + binding_it - m_local_vars.begin();
+                    return ;
+                }
             }
             
             // 2. Is it a capture?
-            binding_it = ::std::find(m_captures.begin(), m_captures.end(), node.m_slot);
-            if( binding_it != m_captures.end() ) {
-                m_replacement = NEWNODE(node.m_res_type.clone(), Field, node.span(),
-                    get_self(node.span()),
-                    FMT(binding_it - m_captures.begin())
-                    );
-                m_replacement->m_usage = node.m_usage;
-                return ;
+            {
+                auto binding_it = ::std::find_if(m_captures.begin(), m_captures.end(), [&](const auto& x){return x.first == node.m_slot;});
+                if( binding_it != m_captures.end() )
+                {
+                    m_replacement = NEWNODE(node.m_res_type.clone(), Field, node.span(),
+                        get_self(node.span()),
+                        FMT(binding_it - m_captures.begin())
+                        );
+                    if( binding_it->second != ::HIR::ValueUsage::Move ) {
+                        auto bt = (binding_it->second == ::HIR::ValueUsage::Mutate ? ::HIR::BorrowType::Unique : ::HIR::BorrowType::Shared);
+                        
+                        m_replacement->m_res_type = ::HIR::TypeRef::new_borrow( bt, mv$(m_replacement->m_res_type) );
+                        m_replacement = NEWNODE(node.m_res_type.clone(), Deref, node.span(),  mv$(m_replacement));
+                    }
+                    m_replacement->m_usage = node.m_usage;
+                    return ;
+                }
             }
             
             BUG(node.span(), "Encountered non-captured and unknown-origin variable - " << node.m_name << " #" << node.m_slot);
@@ -356,7 +378,8 @@ namespace {
         struct ClosureScope {
             ::HIR::ExprNode_Closure&    node;
             ::std::vector<unsigned int> local_vars;
-            // NOTE: Capture list is in the node
+            // - Lists captured variables to be stored in autogenerated struct (and how they're used, influencing the borrow type)
+            ::std::vector< ::std::pair<unsigned int, ::HIR::ValueUsage> > captured_vars;
             
             ClosureScope(::HIR::ExprNode_Closure& node):
                 node(node)
@@ -374,6 +397,7 @@ namespace {
         
         /// Stack of active closures
         ::std::vector<ClosureScope> m_closure_stack;
+        
     public:
         ExprVisitor_Extract(const StaticTraitResolve& resolve, const ::HIR::SimplePath& mod_path, ::std::vector< ::HIR::TypeRef>& var_types, out_impls_t& out_impls, const new_type_cb_t& new_type):
             m_resolve(resolve),
@@ -407,6 +431,16 @@ namespace {
             auto ent = mv$( m_closure_stack.back() );
             m_closure_stack.pop_back();
             
+            // --- Apply the capture set for this closure to the parent ---
+            if( m_closure_stack.size() > 0 )
+            {
+                for(const auto& cap : ent.captured_vars)
+                {
+                    mark_used_variable(node.span(), cap.first, cap.second);
+                }
+            }
+
+            
             // --- Extract and mutate code into a trait impl on the closure type ---
             
             // 1. Construct closure type (saving path/index in the node)
@@ -415,7 +449,7 @@ namespace {
             ::HIR::PathParams   impl_path_params;
             // - 0xFFFF "Self" -> 0 "Super"
             constructor_path_params.m_types.push_back( ::HIR::TypeRef("Self", 0xFFFF) );
-            params.m_types.push_back( ::HIR::TypeParamDef { "Super", {}, false } );  // TODO: Maybe Self is sized?
+            params.m_types.push_back( ::HIR::TypeParamDef { "Super", {}, false } );  // TODO: Determine if parent Self is Sized (or even present?)
             // - Top-level params come first
             unsigned ofs_impl = params.m_types.size();
             for(const auto& ty_def : m_resolve.impl_generics().m_types) {
@@ -480,7 +514,8 @@ namespace {
             
             DEBUG("--- Mutate inner code");
             // 2. Iterate over the nodes and rewrite variable accesses to either renumbered locals, or field accesses
-            ExprVisitor_Mutate    ev { node.m_res_type, ent.local_vars, ent.node.m_var_captures, monomorph_cb };
+            // - TODO: If this closure is a move closure, mutate `captured_vars` such that all captures are tagged with ValueUsage::Move
+            ExprVisitor_Mutate    ev { node.m_res_type, ent.local_vars, ent.captured_vars, monomorph_cb };
             ev.visit_node_ptr( node.m_code );
             // NOTE: `ev` is used down in `Args` to convert the argument destructuring pattern
             
@@ -492,11 +527,44 @@ namespace {
                 auto ty_mono = monomorphise_type_with(sp, m_variable_types.at(binding_idx).clone(), monomorph_cb);
                 local_types.push_back( mv$(ty_mono) );
             }
-            // - Types of captured variables (to be monomorphised)
+            // - Generate types of captures, and construct the actual capture values
+            //  > Capture types (with borrows and using closure's type params)
             ::std::vector< ::HIR::VisEnt< ::HIR::TypeRef> > capture_types;
-            for(const auto binding_idx : node.m_var_captures) {
-                // TODO: By-borrow captures (check the value usage type)
-                auto ty_mono = monomorphise_type_with(sp, m_variable_types.at(binding_idx).clone(), monomorph_cb);
+            //  > Capture value nodes
+            ::std::vector< ::HIR::ExprNodeP>    capture_nodes;
+            capture_types.reserve( ent.captured_vars.size() );
+            capture_nodes.reserve( ent.captured_vars.size() );
+            for(const auto binding : ent.captured_vars)
+            {
+                const auto binding_idx = binding.first;
+                const auto binding_type = binding.second;
+
+                const auto& cap_ty = m_variable_types.at(binding_idx);
+                auto ty_mono = monomorphise_type_with(sp, cap_ty, monomorph_cb);
+                
+                auto val_node = NEWNODE(cap_ty.clone(), Variable,  sp, "", binding_idx);
+                ::HIR::BorrowType   bt;
+                
+                // TODO: If move closure, all move.
+                switch(binding_type)
+                {
+                case ::HIR::ValueUsage::Unknown:
+                    BUG(sp, "ValueUsage::Unkown on " << binding_idx);
+                case ::HIR::ValueUsage::Borrow:
+                    bt = ::HIR::BorrowType::Shared;
+                    capture_nodes.push_back(NEWNODE( ::HIR::TypeRef::new_borrow(bt, cap_ty.clone()), Borrow,  sp, bt, mv$(val_node) ));
+                    ty_mono = ::HIR::TypeRef::new_borrow(bt, mv$(ty_mono));
+                    break;
+                case ::HIR::ValueUsage::Mutate:
+                    bt = ::HIR::BorrowType::Unique;
+                    capture_nodes.push_back(NEWNODE( ::HIR::TypeRef::new_borrow(bt, cap_ty.clone()), Borrow,  sp, bt, mv$(val_node) ));
+                    ty_mono = ::HIR::TypeRef::new_borrow(bt, mv$(ty_mono));
+                    break;
+                case ::HIR::ValueUsage::Move:
+                    capture_nodes.push_back( mv$(val_node) );
+                    break;
+                }
+                
                 // - Fix type to replace closure types with known paths
                 ExprVisitor_Fixup::fix_type(m_resolve.m_crate, ty_mono);
                 capture_types.push_back( ::HIR::VisEnt< ::HIR::TypeRef> { false, mv$(ty_mono) } );
@@ -512,6 +580,7 @@ namespace {
             
             // Mark the object pathname in the closure.
             node.m_obj_path = ::HIR::GenericPath( closure_struct_path, mv$(constructor_path_params) );
+            node.m_captures = mv$(capture_nodes);
             //node.m_res_type = ::HIR::TypeRef( node.m_obj_path.clone() );
             DEBUG("-- Object name: " << node.m_obj_path);
             ::HIR::TypeRef  closure_type = ::HIR::TypeRef( ::HIR::GenericPath(node.m_obj_path.m_path.clone(), mv$(impl_path_params)) );
@@ -539,8 +608,9 @@ namespace {
                 fixup.visit_root( body_code );
             }
             
+            // ---
             // 3. Create trait impls
-            
+            // ---
             ::HIR::PathParams   trait_params;
             trait_params.m_types.push_back( args_ty.clone() );
             switch(node.m_class)
@@ -773,23 +843,26 @@ namespace {
         }
         void mark_used_variable(const Span& sp, unsigned int slot, ::HIR::ValueUsage usage)
         {
-            //for(const auto& closure_rec : m_closure_stack)
-            //{
-            //    const auto& closure_defs = closure_rec.local_vars;
-                const auto& closure_defs = m_closure_stack.back().local_vars;
-                if( ::std::binary_search(closure_defs.begin(), closure_defs.end(), slot) ) {
-                    // Ignore, this is local to the current closure
-                    return ;
-                }
-            //}
+            const auto& closure_defs = m_closure_stack.back().local_vars;
+            if( ::std::binary_search(closure_defs.begin(), closure_defs.end(), slot) ) {
+                // Ignore, this is local to the current closure
+                return ;
+            }
 
+            if( usage == ::HIR::ValueUsage::Move && m_resolve.type_is_copy(sp, m_variable_types.at(slot)) ) {
+                usage = ::HIR::ValueUsage::Borrow;
+            }
+            
             assert(m_closure_stack.size() > 0 );
             auto& closure_rec = m_closure_stack.back();
             auto& closure = closure_rec.node;
             
-            auto it = ::std::lower_bound(closure.m_var_captures.begin(), closure.m_var_captures.end(), slot);
-            if( it == closure.m_var_captures.end() || *it != slot ) {
-                closure.m_var_captures.insert( it, slot );
+            auto it = ::std::lower_bound(closure_rec.captured_vars.begin(), closure_rec.captured_vars.end(), slot, [](const auto& a, const auto& b){ return a.first < b; });
+            if( it == closure_rec.captured_vars.end() || it->first != slot ) {
+                closure_rec.captured_vars.insert( it, ::std::make_pair(slot, usage) );
+            }
+            else {
+                it->second = ::std::max(it->second, usage);
             }
             DEBUG("Captured " << slot << " - " << m_variable_types.at(slot));
             
@@ -804,12 +877,12 @@ namespace {
                 closure.m_class = ::std::max(closure.m_class, ::HIR::ExprNode_Closure::Class::Mut);
                 break;
             case ::HIR::ValueUsage::Move:
-                if( m_resolve.type_is_copy( sp, m_variable_types.at(slot) ) ) {
-                    closure.m_class = ::std::max(closure.m_class, ::HIR::ExprNode_Closure::Class::Shared);
-                }
-                else {
+                //if( m_resolve.type_is_copy( sp, m_variable_types.at(slot) ) ) {
+                //    closure.m_class = ::std::max(closure.m_class, ::HIR::ExprNode_Closure::Class::Shared);
+                //}
+                //else {
                     closure.m_class = ::std::max(closure.m_class, ::HIR::ExprNode_Closure::Class::Once);
-                }
+                //}
                 break;
             }
         }