HIR - Sort trait impls too

7 files changed, 174 insertions, 115 deletions

diff --git a/src/hir/deserialise.cpp b/src/hir/deserialise.cpp
index 9d5b6cd2..f2794f29 100644
--- a/src/hir/deserialise.cpp
+++ b/src/hir/deserialise.cpp
@@ -839,6 +839,14 @@
     template<> DEF_D( ::MacroRulesPtr, return d.deserialise_macrorulesptr(); )
     template<> DEF_D( unsigned int, return static_cast<unsigned int>(d.deserialise_count()); )
 
+    template<typename T>
+    DEF_D( ::HIR::Crate::ImplGroup<T>,
+        ::HIR::Crate::ImplGroup<T>  rv;
+        rv.named = d.deserialise_pathmap< ::std::vector<T>>();
+        rv.non_named = d.deserialise_vec<T>();
+        rv.generic = d.deserialise_vec<T>();
+        return rv;
+        )
     template<> DEF_D( ::HIR::ExternLibrary, return d.deserialise_extlib(); )
 
     ::HIR::LifetimeDef HirDeserialiser::deserialise_lifetimedef()
@@ -1325,12 +1333,9 @@
         rv.m_crate_name = this->m_crate_name;
         rv.m_root_module = deserialise_module();
 
-        rv.m_named_type_impls = deserialise_pathmap< ::std::vector<::HIR::TypeImpl>>();
-        rv.m_primitive_type_impls = deserialise_vec< ::HIR::TypeImpl>();
-        rv.m_generic_type_impls = deserialise_vec< ::HIR::TypeImpl>();
-
-        rv.m_trait_impls = deserialise_pathmap< ::std::vector<::HIR::TraitImpl>>();
-        rv.m_marker_impls = deserialise_pathmap< ::std::vector<::HIR::MarkerImpl>>();
+        rv.m_type_impls = D< ::HIR::Crate::ImplGroup<::HIR::TypeImpl> >::des(*this);
+        rv.m_trait_impls = deserialise_pathmap< ::HIR::Crate::ImplGroup<::HIR::TraitImpl>>();
+        rv.m_marker_impls = deserialise_pathmap< ::HIR::Crate::ImplGroup<::HIR::MarkerImpl>>();
 
         rv.m_exported_macros = deserialise_istrumap< ::MacroRulesPtr>();
         rv.m_lang_items = deserialise_strumap< ::HIR::SimplePath>();
diff --git a/src/hir/from_ast.cpp b/src/hir/from_ast.cpp
index cb8438bf..0ad59d56 100644
--- a/src/hir/from_ast.cpp
+++ b/src/hir/from_ast.cpp
@@ -1676,7 +1676,8 @@ void LowerHIR_Module_Impls(const ::AST::Module& ast_mod,  ::HIR::Crate& hir_crat
                     )
                 }
 
-                hir_crate.m_trait_impls.insert( ::std::make_pair(mv$(trait_name), ::std::vector<::HIR::TraitImpl>()) ).first->second.push_back( ::HIR::TraitImpl {
+                // Sorted later on
+                hir_crate.m_trait_impls[mv$(trait_name)].generic.push_back( ::HIR::TraitImpl {
                     mv$(params),
                     mv$(trait_args),
                     mv$(type),
@@ -1696,7 +1697,7 @@ void LowerHIR_Module_Impls(const ::AST::Module& ast_mod,  ::HIR::Crate& hir_crat
             else
             {
                 auto type = LowerHIR_Type(impl.def().type());
-                hir_crate.m_marker_impls.insert( ::std::make_pair(mv$(trait_name), ::std::vector<::HIR::MarkerImpl>()) ).first->second.push_back(::HIR::MarkerImpl {
+                hir_crate.m_marker_impls[mv$(trait_name)].generic.push_back(::HIR::MarkerImpl {
                     mv$(params),
                     mv$(trait_args),
                     true,
@@ -1750,7 +1751,7 @@ void LowerHIR_Module_Impls(const ::AST::Module& ast_mod,  ::HIR::Crate& hir_crat
             }
 
             // Sorted later on
-            hir_crate.m_generic_type_impls.push_back( ::HIR::TypeImpl {
+            hir_crate.m_type_impls.generic.push_back( ::HIR::TypeImpl {
                 mv$(params),
                 mv$(type),
                 mv$(methods),
@@ -1771,7 +1772,8 @@ void LowerHIR_Module_Impls(const ::AST::Module& ast_mod,  ::HIR::Crate& hir_crat
         auto trait_name = mv$(trait.m_path);
         auto trait_args = mv$(trait.m_params);
 
-        hir_crate.m_marker_impls.insert( ::std::make_pair(mv$(trait_name), ::std::vector<::HIR::MarkerImpl>()) ).first->second.push_back(::HIR::MarkerImpl {
+        // Sorting done later
+        hir_crate.m_marker_impls[mv$(trait_name)].generic.push_back(::HIR::MarkerImpl {
             mv$(params),
             mv$(trait_args),
             false,
diff --git a/src/hir/hir.hpp b/src/hir/hir.hpp
index 12b874a7..2169c0c3 100644
--- a/src/hir/hir.hpp
+++ b/src/hir/hir.hpp
@@ -527,17 +527,46 @@ public:
 
     Module  m_root_module;
 
+    template<typename T>
+    struct ImplGroup
+    {
+        ::std::map<::HIR::SimplePath, ::std::vector<T>> named;
+        ::std::vector<T> non_named; // TODO: use a map of HIR::TypeRef::Data::Tag
+        ::std::vector<T> generic;
+
+        const ::std::vector<T>* get_list_for_type(const ::HIR::TypeRef& ty) const {
+            static ::std::vector<T> empty;
+            if( const auto* p = ty.get_sort_path() ) {
+                auto it = named.find(*p);
+                if( it != named.end() )
+                    return &it->second;
+                else
+                    return nullptr;
+            }
+            else {
+                // TODO: Sort these by type tag, use the `Primitive` group if `ty` is Infer
+                return &non_named;
+            }
+        }
+        ::std::vector<T>& get_list_for_type_mut(const ::HIR::TypeRef& ty) {
+            if( const auto* p = ty.get_sort_path() ) {
+                return named[*p];
+            }
+            else {
+                // TODO: Ivars match with core types
+                return non_named;
+            }
+        }
+    };
     /// Impl blocks on just a type, split into three groups
     // - Named type (sorted on the path)
     // - Primitive types
     // - Unsorted (generics, and everything before outer type resolution)
-    ::std::map<::HIR::SimplePath, ::std::vector<::HIR::TypeImpl>> m_named_type_impls;
-    ::std::vector< ::HIR::TypeImpl> m_primitive_type_impls;
-    ::std::vector< ::HIR::TypeImpl> m_generic_type_impls;
+    ImplGroup<::HIR::TypeImpl>  m_type_impls;
 
     /// Impl blocks
-    ::std::map< ::HIR::SimplePath, ::std::vector<::HIR::TraitImpl> > m_trait_impls;
-    ::std::map< ::HIR::SimplePath, ::std::vector<::HIR::MarkerImpl> > m_marker_impls;
+    ::std::map< ::HIR::SimplePath, ImplGroup<::HIR::TraitImpl> > m_trait_impls;
+    ::std::map< ::HIR::SimplePath, ImplGroup<::HIR::MarkerImpl> > m_marker_impls;
 
     /// Macros exported by this crate
     ::std::unordered_map< RcString, ::MacroRulesPtr >  m_exported_macros;
diff --git a/src/hir/hir_ops.cpp b/src/hir/hir_ops.cpp
index 9e867980..14533fb4 100644
--- a/src/hir/hir_ops.cpp
+++ b/src/hir/hir_ops.cpp
@@ -935,45 +935,60 @@ bool ::HIR::TraitImpl::overlaps_with(const Crate& crate, const ::HIR::TraitImpl&
     }
 }
 
-bool ::HIR::Crate::find_trait_impls(const ::HIR::SimplePath& trait, const ::HIR::TypeRef& type, t_cb_resolve_type ty_res, ::std::function<bool(const ::HIR::TraitImpl&)> callback) const
+namespace
 {
-    auto it = this->m_trait_impls.find( trait );
-    if( it != this->m_trait_impls.end() )
+    template<typename ImplType>
+    bool find_impls_list(const ::std::vector<ImplType>& impl_list, const ::HIR::TypeRef& type, ::HIR::t_cb_resolve_type ty_res, ::std::function<bool(const ImplType&)> callback)
     {
-        for(const auto& impl : it->second)
+        for(const auto& impl : impl_list)
         {
-            if( impl.matches_type(type, ty_res) ) {
-                if( callback(impl) ) {
+            if( impl.matches_type(type, ty_res) )
+            {
+                if( callback(impl) )
+                {
                     return true;
                 }
             }
         }
+        return false;
     }
-    for( const auto& ec : this->m_ext_crates )
-    {
-        if( ec.second.m_data->find_trait_impls(trait, type, ty_res, callback) ) {
-            return true;
-        }
-    }
-    return false;
 }
-bool ::HIR::Crate::find_auto_trait_impls(const ::HIR::SimplePath& trait, const ::HIR::TypeRef& type, t_cb_resolve_type ty_res, ::std::function<bool(const ::HIR::MarkerImpl&)> callback) const
+namespace
 {
-    auto it = this->m_marker_impls.find( trait );
-    if( it != this->m_marker_impls.end() )
+    bool find_trait_impls_int(
+            const ::HIR::Crate& crate, const ::HIR::SimplePath& trait, const ::HIR::TypeRef& type,
+            ::HIR::t_cb_resolve_type ty_res, ::std::function<bool(const ::HIR::TraitImpl&)> callback
+            )
     {
-        for(const auto& impl : it->second)
+        auto it = crate.m_trait_impls.find( trait );
+        if( it != crate.m_trait_impls.end() )
         {
-            if( impl.matches_type(type, ty_res) ) {
-                if( callback(impl) ) {
+            // 1. Find named impls (associated with named types)
+            if( const auto* impl_list = it->second.get_list_for_type(type) )
+            {
+                if( find_impls_list(*impl_list, type, ty_res, callback) )
                     return true;
-                }
             }
+
+            // 2. Search fully generic list.
+            if( find_impls_list(it->second.generic, type, ty_res, callback) )
+                return true;
         }
+
+        return false;
+    }
+}
+
+bool ::HIR::Crate::find_trait_impls(const ::HIR::SimplePath& trait, const ::HIR::TypeRef& type, t_cb_resolve_type ty_res, ::std::function<bool(const ::HIR::TraitImpl&)> callback) const
+{
+    if( find_trait_impls_int(*this, trait, type, ty_res, callback) )
+    {
+        return true;
     }
     for( const auto& ec : this->m_ext_crates )
     {
-        if( ec.second.m_data->find_auto_trait_impls(trait, type, ty_res, callback) ) {
+        if( find_trait_impls_int(*ec.second.m_data, trait, type, ty_res, callback) )
+        {
             return true;
         }
     }
@@ -981,46 +996,57 @@ bool ::HIR::Crate::find_auto_trait_impls(const ::HIR::SimplePath& trait, const :
 }
 namespace
 {
-    bool find_type_impls_list(const ::std::vector<HIR::TypeImpl>& impl_list, const ::HIR::TypeRef& type, ::HIR::t_cb_resolve_type ty_res, ::std::function<bool(const ::HIR::TypeImpl&)> callback)
+    bool find_auto_trait_impls_int(
+            const ::HIR::Crate& crate, const ::HIR::SimplePath& trait, const ::HIR::TypeRef& type,
+            ::HIR::t_cb_resolve_type ty_res, ::std::function<bool(const ::HIR::MarkerImpl&)> callback
+            )
     {
-        for(const auto& impl : impl_list)
+        auto it = crate.m_marker_impls.find( trait );
+        if( it != crate.m_marker_impls.end() )
         {
-            if( impl.matches_type(type, ty_res) )
+            // 1. Find named impls (associated with named types)
+            if( const auto* impl_list = it->second.get_list_for_type(type) )
             {
-                if( callback(impl) )
-                {
+                if( find_impls_list(*impl_list, type, ty_res, callback) )
                     return true;
-                }
             }
+
+            // 2. Search fully generic list.
+            if( find_impls_list(it->second.generic, type, ty_res, callback) )
+                return true;
         }
+
         return false;
     }
-    bool find_type_impls_int(const ::HIR::Crate& crate, const ::HIR::SimplePath* sort_path, const ::HIR::TypeRef& type, ::HIR::t_cb_resolve_type ty_res, ::std::function<bool(const ::HIR::TypeImpl&)> callback)
+}
+bool ::HIR::Crate::find_auto_trait_impls(const ::HIR::SimplePath& trait, const ::HIR::TypeRef& type, t_cb_resolve_type ty_res, ::std::function<bool(const ::HIR::MarkerImpl&)> callback) const
+{
+    if( find_auto_trait_impls_int(*this, trait, type, ty_res, callback) )
+    {
+        return true;
+    }
+    for( const auto& ec : this->m_ext_crates )
     {
-        // 1. Find named impls (associated with named types)
-        if( sort_path )
+        if( find_auto_trait_impls_int(*ec.second.m_data, trait, type, ty_res, callback) )
         {
-            // TODO: This fails if the type is marked as #[fundamental]
-            //if( sort_path->m_crate_name != crate.m_crate_name ) {
-            //    return false;
-            //}
-
-            auto impl_list_it = crate.m_named_type_impls.find(*sort_path);
-            if( impl_list_it != crate.m_named_type_impls.end() )
-            {
-                if( find_type_impls_list(impl_list_it->second, type, ty_res, callback) )
-                    return true;
-            }
+            return true;
         }
-        // - If this type has no associated path, look in the primitives list
-        else
+    }
+    return false;
+}
+namespace
+{
+    bool find_type_impls_int(const ::HIR::Crate& crate, const ::HIR::TypeRef& type, ::HIR::t_cb_resolve_type ty_res, ::std::function<bool(const ::HIR::TypeImpl&)> callback)
+    {
+        // 1. Find named impls (associated with named types)
+        if( const auto* impl_list = crate.m_type_impls.get_list_for_type(type) )
         {
-            if( find_type_impls_list(crate.m_primitive_type_impls, type, ty_res, callback) )
+            if( find_impls_list(*impl_list, type, ty_res, callback) )
                 return true;
         }
 
         // 2. Search fully generic list?
-        if( find_type_impls_list(crate.m_generic_type_impls, type, ty_res, callback) )
+        if( find_impls_list(crate.m_type_impls.generic, type, ty_res, callback) )
             return true;
 
         return false;
@@ -1028,40 +1054,15 @@ namespace
 }
 bool ::HIR::Crate::find_type_impls(const ::HIR::TypeRef& type, t_cb_resolve_type ty_res, ::std::function<bool(const ::HIR::TypeImpl&)> callback) const
 {
-    const ::HIR::SimplePath*    path = nullptr;
-    // - Generic paths get sorted
-    if( TU_TEST1(type.m_data, Path, .path.m_data.is_Generic()) )
-    {
-        path = &type.m_data.as_Path().path.m_data.as_Generic().m_path;
-    }
-    // - So do trait objects
-    else if( type.m_data.is_TraitObject() )
-    {
-        path = &type.m_data.as_TraitObject().m_trait.m_path.m_path;
-    }
-    else
-    {
-        // Keep as nullptr, will search primitive list
-    }
-
-    if( path )
-    {
-        DEBUG(type << ", path=" << *path);
-    }
-    else
-    {
-        DEBUG(type << ", no path");
-    }
-
     // > Current crate
-    if( find_type_impls_int(*this, path, type, ty_res, callback) )
+    if( find_type_impls_int(*this, type, ty_res, callback) )
     {
         return true;
     }
     for( const auto& ec : this->m_ext_crates )
     {
         //DEBUG("- " << ec.first);
-        if( find_type_impls_int(*ec.second.m_data, path, type, ty_res, callback) )
+        if( find_type_impls_int(*ec.second.m_data, type, ty_res, callback) )
         {
             return true;
         }
diff --git a/src/hir/serialise.cpp b/src/hir/serialise.cpp
index 83e6d016..7725d49f 100644
--- a/src/hir/serialise.cpp
+++ b/src/hir/serialise.cpp
@@ -308,17 +308,20 @@
             serialise(pm.path);
             serialise_vec(pm.attributes);
         }
+        template<typename T>
+        void serialise(const ::HIR::Crate::ImplGroup<T>& ig)
+        {
+            serialise_pathmap(ig.named);
+            serialise_vec(ig.non_named);
+            serialise_vec(ig.generic);
+        }
 
         void serialise_crate(const ::HIR::Crate& crate)
         {
             m_out.write_string(crate.m_crate_name);
             serialise_module(crate.m_root_module);
 
-            //std::map<HIR::SimplePath,std::vector<TypeImpl>>
-            serialise_pathmap(crate.m_named_type_impls);
-            serialise_vec(crate.m_primitive_type_impls);
-            serialise_vec(crate.m_generic_type_impls);
-
+            serialise(crate.m_type_impls);
             serialise_pathmap(crate.m_trait_impls);
             serialise_pathmap(crate.m_marker_impls);
 
diff --git a/src/hir/type.hpp b/src/hir/type.hpp
index 7a5a0a1c..13300bbc 100644
--- a/src/hir/type.hpp
+++ b/src/hir/type.hpp
@@ -328,6 +328,24 @@ public:
 
     // Compares this type with another, using `resolve_placeholder` to get replacements for generics/infers in `x`
     Compare compare_with_placeholders(const Span& sp, const ::HIR::TypeRef& x, t_cb_resolve_type resolve_placeholder) const;
+
+    const ::HIR::SimplePath* get_sort_path() const {
+        // - Generic paths get sorted
+        if( TU_TEST1(this->m_data, Path, .path.m_data.is_Generic()) )
+        {
+            return &this->m_data.as_Path().path.m_data.as_Generic().m_path;
+        }
+        // - So do trait objects
+        else if( this->m_data.is_TraitObject() )
+        {
+            return &this->m_data.as_TraitObject().m_trait.m_path.m_path;
+        }
+        else
+        {
+            // Keep as nullptr, will search primitive list
+            return nullptr;
+        }
+    }
 };
 
 extern ::std::ostream& operator<<(::std::ostream& os, const ::HIR::TypeRef& ty);
diff --git a/src/hir/visitor.cpp b/src/hir/visitor.cpp
index 8bf674b3..1cc54090 100644
--- a/src/hir/visitor.cpp
+++ b/src/hir/visitor.cpp
@@ -12,38 +12,39 @@
 {
 }
 
+namespace {
+    template<typename T>
+    void visit_impls(::HIR::Crate::ImplGroup<T>& g, ::std::function<void(T&)> cb) {
+        for( auto& impl_group : g.named )
+        {
+            for( auto& impl : impl_group.second )
+            {
+                cb(impl);
+            }
+        }
+        for( auto& impl : g.non_named )
+        {
+            cb(impl);
+        }
+        for( auto& impl : g.generic )
+        {
+            cb(impl);
+        }
+    }
+}
+
 void ::HIR::Visitor::visit_crate(::HIR::Crate& crate)
 {
     this->visit_module(::HIR::ItemPath(crate.m_crate_name), crate.m_root_module );
 
-    for( auto& ty_impl_group : crate.m_named_type_impls )
-    {
-        for( auto& ty_impl : ty_impl_group.second )
-        {
-            this->visit_type_impl(ty_impl);
-        }
-    }
-    for( auto& ty_impl : crate.m_primitive_type_impls )
-    {
-        this->visit_type_impl(ty_impl);
-    }
-    for( auto& ty_impl : crate.m_generic_type_impls )
-    {
-        this->visit_type_impl(ty_impl);
-    }
+    visit_impls<::HIR::TypeImpl>(crate.m_type_impls, [&](::HIR::TypeImpl& ty_impl){ this->visit_type_impl(ty_impl); });
     for( auto& impl_group : crate.m_trait_impls )
     {
-        for( auto& impl : impl_group.second )
-        {
-            this->visit_trait_impl(impl_group.first, impl);
-        }
+        visit_impls<::HIR::TraitImpl>(impl_group.second, [&](::HIR::TraitImpl& ty_impl){ this->visit_trait_impl(impl_group.first, ty_impl); });
     }
     for( auto& impl_group : crate.m_marker_impls )
     {
-        for( auto& impl : impl_group.second )
-        {
-            this->visit_marker_impl(impl_group.first, impl);
-        }
+        visit_impls<::HIR::MarkerImpl>(impl_group.second, [&](::HIR::MarkerImpl& ty_impl){ this->visit_marker_impl(impl_group.first, ty_impl); });
     }
 }