HIR Expand - Fix closure conversion

3 files changed, 93 insertions, 8 deletions

diff --git a/src/hir/expr.cpp b/src/hir/expr.cpp
index 4fb79470..a7a5c31b 100644
--- a/src/hir/expr.cpp
+++ b/src/hir/expr.cpp
@@ -23,6 +23,7 @@ DEF_VISIT(ExprNode_Return, node,
     visit_node_ptr(node.m_value);
 )
 DEF_VISIT(ExprNode_Let, node,
+    visit_pattern(node.m_pattern);
     visit_type(node.m_type);
     if( node.m_value ) {
         visit_node_ptr(node.m_value);
@@ -36,6 +37,8 @@ DEF_VISIT(ExprNode_Match, node,
     visit_node_ptr(node.m_value);
     for(auto& arm : node.m_arms)
     {
+        for(auto& pat : arm.m_patterns)
+            visit_pattern(pat);
         if( arm.m_cond )
             visit_node_ptr(arm.m_cond);
         visit_node_ptr(arm.m_code);
@@ -147,8 +150,10 @@ DEF_VISIT(ExprNode_ArraySized, node,
 )
 
 DEF_VISIT(ExprNode_Closure, node,
-    for(auto& arg : node.m_args)
+    for(auto& arg : node.m_args) {
+        visit_pattern(arg.first);
         visit_type(arg.second);
+    }
     visit_type(node.m_return);
     if(node.m_code)
     {
@@ -159,6 +164,67 @@ DEF_VISIT(ExprNode_Closure, node,
 #undef DEF_VISIT
 
 // TODO: Merge this with the stuff in ::HIR::Visitor
+void ::HIR::ExprVisitorDef::visit_pattern(::HIR::Pattern& pat)
+{
+    TU_MATCH(::HIR::Pattern::Data, (pat.m_data), (e),
+    (Any,
+        ),
+    (Box,
+        this->visit_pattern(*e.sub);
+        ),
+    (Ref,
+        this->visit_pattern(*e.sub);
+        ),
+    (Tuple,
+        for(auto& subpat : e.sub_patterns)
+            this->visit_pattern(subpat);
+        ),
+    (SplitTuple,
+        for(auto& subpat : e.leading)
+            this->visit_pattern(subpat);
+        for(auto& subpat : e.trailing)
+            this->visit_pattern(subpat);
+        ),
+    (StructValue,
+        // Nothing.
+        ),
+    (StructTuple,
+        for(auto& subpat : e.sub_patterns)
+            this->visit_pattern(subpat);
+        ),
+    (Struct,
+        for(auto& fld_pat : e.sub_patterns)
+        {
+            this->visit_pattern(fld_pat.second);
+        }
+        ),
+    // Refutable
+    (Value,
+        ),
+    (Range,
+        ),
+    (EnumValue,
+        ),
+    (EnumTuple,
+        for(auto& subpat : e.sub_patterns)
+            this->visit_pattern(subpat);
+        ),
+    (EnumStruct,
+        for(auto& fld_pat : e.sub_patterns)
+            this->visit_pattern(fld_pat.second);
+        ),
+    (Slice,
+        for(auto& subpat : e.sub_patterns)
+            this->visit_pattern(subpat);
+        ),
+    (SplitSlice,
+        for(auto& subpat : e.leading)
+            this->visit_pattern(subpat);
+        for(auto& subpat : e.trailing)
+            this->visit_pattern(subpat);
+        )
+    )
+}
 void ::HIR::ExprVisitorDef::visit_type(::HIR::TypeRef& ty)
 {
     TU_MATCH(::HIR::TypeRef::Data, (ty.m_data), (e),
diff --git a/src/hir/expr.hpp b/src/hir/expr.hpp
index ae9ce7ed..4b385c99 100644
--- a/src/hir/expr.hpp
+++ b/src/hir/expr.hpp
@@ -810,6 +810,7 @@ public:
     NV(ExprNode_Closure);
     #undef NV
     
+    virtual void visit_pattern(::HIR::Pattern& pat);
     virtual void visit_type(::HIR::TypeRef& ty);
     virtual void visit_trait_path(::HIR::TraitPath& p);
     virtual void visit_path_params(::HIR::PathParams& ty);
diff --git a/src/hir_expand/closures.cpp b/src/hir_expand/closures.cpp
index 74255849..03b5e8f2 100644
--- a/src/hir_expand/closures.cpp
+++ b/src/hir_expand/closures.cpp
@@ -65,6 +65,25 @@ namespace {
         {
         }
         
+        void visit_pattern(::HIR::Pattern& pat) override {
+            if( pat.m_binding.is_valid() ) {
+                auto binding_it = ::std::find(m_local_vars.begin(), m_local_vars.end(), pat.m_binding.m_slot);
+                if( binding_it != m_local_vars.end() ) {
+                    // NOTE: Offset of 1 is for `self` (`args` is destructured)
+                    pat.m_binding.m_slot = 1 + binding_it - m_local_vars.begin();
+                }
+                else {
+                    BUG(Span(), "Pattern binds to non-local");
+                }
+            }
+            
+            TU_IFLET(::HIR::Pattern::Data, (pat.m_data), SplitSlice, e,
+                TODO(Span(), "Fixup binding in split slice");
+            )
+            
+            ::HIR::ExprVisitorDef::visit_pattern(pat);
+        }
+        
         void visit_type(::HIR::TypeRef& ty) override {
             TU_IFLET(::HIR::TypeRef::Data, ty.m_data, Generic, e,
                 auto n = m_monomorph_cb(ty).clone();
@@ -98,8 +117,8 @@ namespace {
             // 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 2 is for `self` and `args`
-                node.m_slot = 2 + binding_it - m_local_vars.begin();
+                // NOTE: Offset of 1 is for `self` (`args` is destructured)
+                node.m_slot = 1 + binding_it - m_local_vars.begin();
                 return ;
             }
             
@@ -366,9 +385,9 @@ namespace {
             
             DEBUG("--- Mutate inner code");
             // 2. Iterate over the nodes and rewrite variable accesses to either renumbered locals, or field accesses
-            // - TODO: Monomorphise all referenced types within this.
             ExprVisitor_Mutate    ev { node.m_res_type, ent.local_vars, ent.node.m_var_captures, monomorph_cb };
             ev.visit_node_ptr( node.m_code );
+            // NOTE: `ev` is used down in `Args` to convert the argument destructuring pattern
             
             // - Types of local variables
             DEBUG("--- Build locals and captures");
@@ -404,6 +423,7 @@ namespace {
             ::std::vector< ::HIR::TypeRef>  args_ty_inner;
             for(const auto& arg : node.m_args) {
                 args_pat_inner.push_back( arg.first.clone() );
+                ev.visit_pattern( args_pat_inner.back() );
                 args_ty_inner.push_back( monomorphise_type_with(sp, arg.second, monomorph_cb) );
             }
             ::HIR::TypeRef  args_ty { mv$(args_ty_inner) };
@@ -516,14 +536,12 @@ namespace {
             }
         }
         
-        void visit(::HIR::ExprNode_Let& node) override
+        void visit_pattern(::HIR::Pattern& pat) override
         {
             if( !m_closure_stack.empty() )
             {
-                add_closure_def_from_pattern(node.span(), node.m_pattern);
+                add_closure_def_from_pattern(Span(), pat);
             }
-            
-            ::HIR::ExprVisitorDef::visit(node);
         }
         void visit(::HIR::ExprNode_Variable& node) override
         {