HIR Expand ErasedType - Replace non-local erased types in return types

3 files changed, 155 insertions, 48 deletions

diff --git a/src/hir/item_path.hpp b/src/hir/item_path.hpp
index 9b579b72..8ed140af 100644
--- a/src/hir/item_path.hpp
+++ b/src/hir/item_path.hpp
@@ -11,13 +11,13 @@ namespace HIR {
 
 class ItemPath
 {
+public:
     const ItemPath* parent = nullptr;
     const ::HIR::TypeRef* ty = nullptr;
     const ::HIR::SimplePath* trait = nullptr;
     const ::HIR::PathParams* trait_params = nullptr;
     const char* name = nullptr;
     
-public:
     ItemPath() {}
     ItemPath(const ItemPath& p, const char* n):
         parent(&p),
@@ -70,6 +70,24 @@ public:
         return ItemPath(*this, name.c_str());
     }
     
+    bool operator==(const ::HIR::SimplePath& sp) const {
+        if( sp.m_crate_name != "" )  return false;
+        
+        unsigned int i = sp.m_components.size();
+        const auto* n = this;
+        while( n && i -- )
+        {
+            if( !n->name )
+                return false;
+            if( n->name != sp.m_components[i] )
+                return false;
+            n = n->parent;
+        }
+        if( i > 0 || n->name )
+            return false;
+        return true;
+    }
+    
     friend ::std::ostream& operator<<(::std::ostream& os, const ItemPath& x) {
         if( x.parent ) {
             os << *x.parent;
diff --git a/src/hir_expand/closures.cpp b/src/hir_expand/closures.cpp
index ccd293af..2523eb68 100644
--- a/src/hir_expand/closures.cpp
+++ b/src/hir_expand/closures.cpp
@@ -511,6 +511,28 @@ namespace {
                 }
                 };
             auto monomorph = [&](const auto& ty){ return monomorphise_type_with(sp, ty, monomorph_cb); };
+            auto cb_replace = [&](const auto& tpl, auto& rv) {
+                if( tpl.m_data.is_Infer() ) {
+                    BUG(sp, "");
+                }
+                else if( tpl.m_data.is_Generic() ) {
+                    rv = monomorph_cb(tpl).clone();
+                    return true;
+                }
+                //else if( tpl.m_data.is_ErasedType() ) {
+                //    const auto& e = tpl.m_data.as_ErasedType();
+                //    
+                //    // TODO: Share code with 
+                //    TODO(sp, "Repalce ErasedType with origin " << e.m_origin << " #" << e.m_index);
+                //    //ASSERT_BUG(sp, e.m_index < fcn_ptr->m_code.m_erased_types.size(), "");
+                //    //const auto& erased_type_replacement = fcn_ptr->m_code.m_erased_types.at(e.m_index);
+                //    //rv = monomorphise_type_with(sp, erased_type_replacement,  monomorph_cb, false);
+                //    //return true;
+                //}
+                else {
+                    return false;
+                }
+                };
             
             // - Clone the bounds (from both levels)
             auto monomorph_bound = [&](const ::HIR::GenericBound& b)->auto {
@@ -612,14 +634,15 @@ namespace {
             // - Args
             ::std::vector< ::HIR::Pattern>  args_pat_inner;
             ::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) );
+                args_ty_inner.push_back( clone_ty_with(sp, arg.second, cb_replace) );
             }
             ::HIR::TypeRef  args_ty { mv$(args_ty_inner) };
             ::HIR::Pattern  args_pat { {}, ::HIR::Pattern::Data::make_Tuple({ mv$(args_pat_inner) }) };
-            ::HIR::TypeRef  ret_type = monomorphise_type_with(sp, node.m_return, monomorph_cb);
+            ::HIR::TypeRef  ret_type = clone_ty_with(sp, node.m_return, cb_replace);
             
             DEBUG("args_ty = " << args_ty << ", ret_type = " << ret_type);
             
@@ -634,6 +657,7 @@ namespace {
                 DEBUG("-- Fixing types in signature");
                 fixup.visit_type( args_ty );
                 fixup.visit_type( ret_type );
+                // TODO: Replace erased types too
             }
             
             // ---
diff --git a/src/hir_expand/erased_types.cpp b/src/hir_expand/erased_types.cpp
index 93bef966..5cbb9726 100644
--- a/src/hir_expand/erased_types.cpp
+++ b/src/hir_expand/erased_types.cpp
@@ -11,6 +11,55 @@
 #include <algorithm>
 #include "main_bindings.hpp"
 
+const ::HIR::Function& HIR_Expand_ErasedType_GetFunction(const Span& sp, const StaticTraitResolve& resolve, const ::HIR::Path& origin_path, t_cb_generic& monomorph_cb, ::HIR::PathParams& impl_params)
+{
+    const ::HIR::Function*  fcn_ptr = nullptr;
+    TU_MATCHA( (origin_path.m_data), (pe),
+    (UfcsUnknown,
+        BUG(Span(), "UfcsUnknown in ErasedType - " << origin_path);
+        ),
+    (Generic,
+        monomorph_cb = monomorphise_type_get_cb(sp, nullptr, nullptr, &pe.m_params);
+        fcn_ptr = &resolve.m_crate.get_function_by_path(sp, pe.m_path);
+        ),
+    (UfcsKnown,
+        // NOTE: This isn't possible yet (will it be? or will it expand to an associated type?)
+        TODO(sp, "Replace ErasedType - " << origin_path << " with source (UfcsKnown)");
+        ),
+    (UfcsInherent,
+        // 1. Find correct impl block for the path
+        const ::HIR::TypeImpl* impl_ptr = nullptr;
+        resolve.m_crate.find_type_impls(*pe.type, [&](const auto& ty)->const auto& { return ty; },
+            [&](const auto& impl) {
+                DEBUG("- impl" << impl.m_params.fmt_args() << " " << impl.m_type);
+                auto it = impl.m_methods.find(pe.item);
+                if( it == impl.m_methods.end() )
+                    return false;
+                fcn_ptr = &it->second.data;
+                impl_ptr = &impl;
+                return true;
+            });
+        ASSERT_BUG(sp, fcn_ptr, "Failed to locate function " << origin_path);
+        assert(impl_ptr);
+
+        // 2. Obtain monomorph_cb (including impl params)
+        impl_params.m_types.resize(impl_ptr->m_params.m_types.size());
+        impl_ptr->m_type .match_test_generics(sp, *pe.type, [](const auto& x)->const auto&{return x;}, [&](auto idx, const auto& ty) {
+            assert( idx < impl_params.m_types.size() );
+            impl_params.m_types[idx] = ty.clone();
+            return ::HIR::Compare::Equal;
+            });
+        for(const auto& t : impl_params.m_types)
+            if( t == ::HIR::TypeRef() )
+                TODO(sp, "Handle ErasedType where an impl parameter comes from a bound - " << origin_path);
+        
+        monomorph_cb = monomorphise_type_get_cb(sp, &*pe.type, &impl_params, &pe.params);
+        )
+    )
+    assert(fcn_ptr);
+    return *fcn_ptr;
+}
+
 namespace {
     
     
@@ -53,51 +102,7 @@ namespace {
                 
                 ::HIR::PathParams   impl_params;    // cache.
                 t_cb_generic    monomorph_cb;
-                const ::HIR::Function*  fcn_ptr = nullptr;
-                TU_MATCHA( (e.m_origin.m_data), (pe),
-                (UfcsUnknown,
-                    BUG(Span(), "UfcsUnknown in ErasedType - " << ty);
-                    ),
-                (Generic,
-                    monomorph_cb = monomorphise_type_get_cb(sp, nullptr, nullptr, &pe.m_params);
-                    fcn_ptr = &m_resolve.m_crate.get_function_by_path(sp, pe.m_path);
-                    ),
-                (UfcsKnown,
-                    // NOTE: This isn't possible yet (will it be? or will it expand to an associated type?)
-                    TODO(sp, "Replace ErasedType - " << ty << " with source (UfcsKnown)");
-                    ),
-                (UfcsInherent,
-                    // 1. Find correct impl block for the path
-                    const ::HIR::TypeImpl* impl_ptr = nullptr;
-                    m_resolve.m_crate.find_type_impls(*pe.type, [&](const auto& ty)->const auto& { return ty; },
-                        [&](const auto& impl) {
-                            DEBUG("- impl" << impl.m_params.fmt_args() << " " << impl.m_type);
-                            auto it = impl.m_methods.find(pe.item);
-                            if( it == impl.m_methods.end() )
-                                return false;
-                            fcn_ptr = &it->second.data;
-                            impl_ptr = &impl;
-                            return true;
-                        });
-                    ASSERT_BUG(sp, fcn_ptr, "Failed to locate function " << e.m_origin);
-                    assert(impl_ptr);
-
-                    // 2. Obtain monomorph_cb (including impl params)
-                    impl_params.m_types.resize(impl_ptr->m_params.m_types.size());
-                    impl_ptr->m_type .match_test_generics(sp, *pe.type, [](const auto& x)->const auto&{return x;}, [&](auto idx, const auto& ty) {
-                        assert( idx < impl_params.m_types.size() );
-                        impl_params.m_types[idx] = ty.clone();
-                        return ::HIR::Compare::Equal;
-                        });
-                    for(const auto& t : impl_params.m_types)
-                        if( t == ::HIR::TypeRef() )
-                            TODO(sp, "Handle ErasedType where an impl parameter comes from a bound - " << e.m_origin);
-                    
-                    monomorph_cb = monomorphise_type_get_cb(sp, &*pe.type, &impl_params, &pe.params);
-                    )
-                )
-                assert(fcn_ptr);
-                const auto& fcn = *fcn_ptr;
+                const auto& fcn = HIR_Expand_ErasedType_GetFunction(sp, m_resolve, e.m_origin, monomorph_cb, impl_params);
                 const auto& erased_types = fcn.m_code.m_erased_types;
                 
                 ASSERT_BUG(sp, e.m_index < erased_types.size(), "Erased type index out of range for " << e.m_origin << " - " << e.m_index << " >= " << erased_types.size());
@@ -120,6 +125,7 @@ namespace {
         public ::HIR::Visitor
     {
         StaticTraitResolve  m_resolve;
+        ::HIR::ItemPath m_fcn_path;
     public:
         OuterVisitor(const ::HIR::Crate& crate):
             m_resolve(crate)
@@ -133,6 +139,65 @@ namespace {
                 ev.visit_root( exp );
             }
         }
+        
+        void visit_function(::HIR::ItemPath p, ::HIR::Function& fcn) override
+        {
+            m_fcn_path = p;
+            ::HIR::Visitor::visit_function(p, fcn);
+            m_fcn_path = ::HIR::ItemPath();
+        }
+        
+        void visit_type(::HIR::TypeRef& ty) override
+        {
+            static const Span   sp;
+            if( ty.m_data.is_ErasedType() )
+            {
+                ASSERT_BUG(sp, m_fcn_path.get_name(), "");
+                
+                const auto& e = ty.m_data.as_ErasedType();
+                
+                TU_MATCHA( (e.m_origin.m_data), (pe),
+                (Generic,
+                    if( m_fcn_path == pe.m_path ) {
+                        ::HIR::Visitor::visit_type(ty);
+                        return ;
+                    }
+                    ),
+                (UfcsUnknown,
+                    BUG(sp, "UfcsUnknown unexpected");
+                    ),
+                (UfcsKnown,
+                    BUG(sp, "UfcsKnown not supported");
+                    ),
+                (UfcsInherent,
+                    if( m_fcn_path.parent && m_fcn_path.parent->ty && !m_fcn_path.parent->trait && *m_fcn_path.parent->ty == *pe.type && m_fcn_path.name == pe.item ) {
+                        ::HIR::Visitor::visit_type(ty);
+                        return ;
+                    }
+                    )
+                )
+                
+                TRACE_FUNCTION_FR(ty, ty);
+                
+                ::HIR::PathParams   impl_params;
+                t_cb_generic    monomorph_cb;
+                const auto& fcn = HIR_Expand_ErasedType_GetFunction(sp, m_resolve, e.m_origin, monomorph_cb, impl_params);
+                const auto& erased_types = fcn.m_code.m_erased_types;
+                
+                ASSERT_BUG(sp, e.m_index < erased_types.size(), "Erased type index out of range for " << e.m_origin << " - " << e.m_index << " >= " << erased_types.size());
+                const auto& tpl = erased_types[e.m_index];
+                
+                auto new_ty = monomorphise_type_with(sp, tpl, monomorph_cb);
+                DEBUG("> " << ty << " => " << new_ty);
+                ty = mv$(new_ty);
+                // Recurse (TODO: Cleanly prevent infinite recursion - TRACE_FUNCTION does crude prevention)
+                visit_type(ty);
+            }
+            else
+            {
+                ::HIR::Visitor::visit_type(ty);
+            }
+        }
     };
 }