HIR Typecheck - Expand associated types returned from methods

6 files changed, 246 insertions, 32 deletions

diff --git a/src/hir/expr.hpp b/src/hir/expr.hpp
index 65cf3735..e5de1eff 100644
--- a/src/hir/expr.hpp
+++ b/src/hir/expr.hpp
@@ -299,6 +299,9 @@ struct ExprNode_CallPath:
     ::HIR::Path m_path;
     ::std::vector<ExprNodeP> m_args;
     
+    // - Cache for typeck
+    ::std::vector< ::HIR::TypeRef>  m_arg_types;
+    
     ExprNode_CallPath(Span sp, ::HIR::Path path, ::std::vector< ::HIR::ExprNodeP> args):
         ExprNode(mv$(sp)),
         m_path( mv$(path) ),
@@ -331,6 +334,7 @@ struct ExprNode_CallMethod:
     
     // - Set during typeck to the real path to the method
     ::HIR::Path m_method_path;
+    ::std::vector< ::HIR::TypeRef>  m_arg_types;
 
     ExprNode_CallMethod(Span sp, ::HIR::ExprNodeP val, ::std::string method_name, ::HIR::PathParams params, ::std::vector< ::HIR::ExprNodeP> args):
         ExprNode( mv$(sp) ),
diff --git a/src/hir/hir.cpp b/src/hir/hir.cpp
index 66f09379..12dee93f 100644
--- a/src/hir/hir.cpp
+++ b/src/hir/hir.cpp
@@ -30,10 +30,14 @@ namespace HIR {
 }
 
 namespace {
-    bool matches_type_int(const ::HIR::GenericParams& params, const ::HIR::TypeRef& left,  const ::HIR::TypeRef& right)
+    bool matches_type_int(const ::HIR::GenericParams& params, const ::HIR::TypeRef& left,  const ::HIR::TypeRef& right_in, ::HIR::t_cb_resolve_type ty_res)
     {
         //DEBUG("left = " << left << ", right = " << right);
         assert(! left.m_data.is_Infer() );
+        const auto& right = (right_in.m_data.is_Infer() ? ty_res(right_in) : right_in);
+        
+        // TODO: What indicates what out of ty_res?
+        
         if( right.m_data.is_Infer() ) {
             // TODO: Why is this false? A _ type could match anything
             return false;
@@ -76,7 +80,7 @@ namespace {
                     }
                     for( unsigned int i = 0; i < pre.m_params.m_types.size(); i ++ )
                     {
-                        if( ! matches_type_int(params, ple.m_params.m_types[i], pre.m_params.m_types[i]) )
+                        if( ! matches_type_int(params, ple.m_params.m_types[i], pre.m_params.m_types[i], ty_res) )
                             return false;
                     }
                 }
@@ -92,32 +96,32 @@ namespace {
             return false;
             ),
         (Array,
-            if( ! matches_type_int(params, *le.inner, *re.inner) )
+            if( ! matches_type_int(params, *le.inner, *re.inner, ty_res) )
                 return false;
             if( le.size_val != re.size_val )
                 return false;
             return true;
             ),
         (Slice,
-            return matches_type_int(params, *le.inner, *re.inner);
+            return matches_type_int(params, *le.inner, *re.inner, ty_res);
             ),
         (Tuple,
             if( le.size() != re.size() )
                 return false;
             for( unsigned int i = 0; i < le.size(); i ++ )
-                if( !matches_type_int(params, le[i], re[i]) )
+                if( !matches_type_int(params, le[i], re[i], ty_res) )
                     return false;
             return true;
             ),
         (Borrow,
             if( le.type != re.type )
                 return false;
-            return matches_type_int(params, *le.inner, *re.inner);
+            return matches_type_int(params, *le.inner, *re.inner, ty_res);
             ),
         (Pointer,
             if( le.is_mut != re.is_mut )
                 return false;
-            return matches_type_int(params, *le.inner, *re.inner);
+            return matches_type_int(params, *le.inner, *re.inner, ty_res);
             ),
         (Function,
             DEBUG("TODO: Compare " << left << " and " << right);
@@ -128,17 +132,17 @@ namespace {
     }
 }
 
-bool ::HIR::TraitImpl::matches_type(const ::HIR::TypeRef& type) const
+bool ::HIR::TraitImpl::matches_type(const ::HIR::TypeRef& type, ::HIR::t_cb_resolve_type ty_res) const
 {
-    return matches_type_int(m_params, m_type, type);
+    return matches_type_int(m_params, m_type, type, ty_res);
 }
-bool ::HIR::TypeImpl::matches_type(const ::HIR::TypeRef& type) const
+bool ::HIR::TypeImpl::matches_type(const ::HIR::TypeRef& type, ::HIR::t_cb_resolve_type ty_res) const
 {
-    return matches_type_int(m_params, m_type, type);
+    return matches_type_int(m_params, m_type, type, ty_res);
 }
-bool ::HIR::MarkerImpl::matches_type(const ::HIR::TypeRef& type) const
+bool ::HIR::MarkerImpl::matches_type(const ::HIR::TypeRef& type, ::HIR::t_cb_resolve_type ty_res) const
 {
-    return matches_type_int(m_params, m_type, type);
+    return matches_type_int(m_params, m_type, type, ty_res);
 }
 
 
@@ -254,3 +258,18 @@ const ::HIR::Function& ::HIR::Crate::get_function_by_path(const Span& sp, const
         BUG(sp, "Enum path " << path << " didn't point to an enum");
     }
 }
+
+const 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
+{
+    auto its = this->m_trait_impls.equal_range( trait );
+    for( auto it = its.first; it != its.second; ++ it )
+    {
+        const auto& impl = it->second;
+        if( impl.matches_type(type, ty_res) ) {
+            if( callback(impl) ) {
+                return true;
+            }
+        }
+    }
+    return false;
+}
diff --git a/src/hir/hir.hpp b/src/hir/hir.hpp
index 0b2847e9..73e7b1af 100644
--- a/src/hir/hir.hpp
+++ b/src/hir/hir.hpp
@@ -202,6 +202,10 @@ TAGGED_UNION(ValueItem, Import,
     (StructConstructor, struct { ::HIR::SimplePath ty; })
     );
 
+// --------------------------------------------------------------------
+
+typedef ::std::function<const ::HIR::TypeRef&(const ::HIR::TypeRef&)> t_cb_resolve_type;
+
 class TypeImpl
 {
 public:
@@ -210,7 +214,10 @@ public:
     
     ::std::map< ::std::string, ::HIR::Function> m_methods;
     
-    bool matches_type(const ::HIR::TypeRef& tr) const;
+    bool matches_type(const ::HIR::TypeRef& tr, t_cb_resolve_type ty_res) const;
+    bool matches_type(const ::HIR::TypeRef& tr) const {
+        return matches_type(tr, [](const auto& x)->const auto&{ return x; });
+    }
 };
 
 class TraitImpl
@@ -224,7 +231,10 @@ public:
     ::std::map< ::std::string, ::HIR::ExprPtr> m_constants;
     ::std::map< ::std::string, ::HIR::TypeRef> m_types;
     
-    bool matches_type(const ::HIR::TypeRef& tr) const;
+    bool matches_type(const ::HIR::TypeRef& tr, t_cb_resolve_type ty_res) const;
+    bool matches_type(const ::HIR::TypeRef& tr) const {
+        return matches_type(tr, [](const auto& x)->const auto&{ return x; });
+    }
 };
 
 class MarkerImpl
@@ -235,7 +245,10 @@ public:
     bool    is_positive;
     ::HIR::TypeRef  m_type;
     
-    bool matches_type(const ::HIR::TypeRef& tr) const;
+    bool matches_type(const ::HIR::TypeRef& tr, t_cb_resolve_type ty_res) const;
+    bool matches_type(const ::HIR::TypeRef& tr) const {
+        return matches_type(tr, [](const auto& x)->const auto&{ return x; });
+    }
 };
 
 class Crate
@@ -261,6 +274,8 @@ public:
     
     const ::HIR::ValueItem& get_valitem_by_path(const Span& sp, const ::HIR::SimplePath& path) const;
     const ::HIR::Function& get_function_by_path(const Span& sp, const ::HIR::SimplePath& path) const;
+    
+    const bool find_trait_impls(const ::HIR::SimplePath& path, const ::HIR::TypeRef& type, t_cb_resolve_type ty_res, ::std::function<bool(const ::HIR::TraitImpl&)> callback) const;
 };
 
 }   // namespace HIR
diff --git a/src/hir/type.cpp b/src/hir/type.cpp
index a5f04df5..2d87ffb9 100644
--- a/src/hir/type.cpp
+++ b/src/hir/type.cpp
@@ -184,6 +184,55 @@ bool ::HIR::TypeRef::operator==(const ::HIR::TypeRef& x) const
     )
     throw "";
 }
+void ::HIR::TypeRef::match_generics(const Span& sp, const ::HIR::TypeRef& x, ::std::function<void(unsigned int, const ::HIR::TypeRef&)> callback) const
+{
+    if( m_data.is_Infer() ) {
+        BUG(sp, "");
+    }
+    if( m_data.is_Generic() ) {
+        callback(m_data.as_Generic().binding, x);
+        return ;
+    }
+    if( m_data.tag() != x.m_data.tag() ) {
+        BUG(sp, "");
+    }
+    TU_MATCH(::HIR::TypeRef::Data, (m_data, x.m_data), (te, xe),
+    (Infer, throw "";),
+    (Generic, throw "";),
+    (Primitive,
+        ),
+    (Diverge,
+        ),
+    (Path,
+        TODO(sp, "Path");
+        ),
+    (TraitObject,
+        TODO(sp, "TraitObject");
+        ),
+    (Array,
+        te.inner->match_generics( sp, *xe.inner, callback );
+        ),
+    (Slice,
+        te.inner->match_generics( sp, *xe.inner, callback );
+        ),
+    (Tuple,
+        if( te.size() != xe.size() ) {
+            BUG(sp, "");
+        }
+        for(unsigned int i = 0; i < te.size(); i ++ )
+            te[i].match_generics( sp, xe[i], callback );
+        ),
+    (Pointer,
+        te.inner->match_generics( sp, *xe.inner, callback );
+        ),
+    (Borrow,
+        te.inner->match_generics( sp, *xe.inner, callback );
+        ),
+    (Function,
+        TODO(sp, "Function");
+        )
+    )
+}
 
 namespace {
     ::HIR::TypeRef::TypePathBinding clone_binding(const ::HIR::TypeRef::TypePathBinding& x) {
diff --git a/src/hir/type.hpp b/src/hir/type.hpp
index 74bd25fe..bd61eea6 100644
--- a/src/hir/type.hpp
+++ b/src/hir/type.hpp
@@ -6,6 +6,7 @@
 #include <tagged_union.hpp>
 #include <hir/path.hpp>
 #include <hir/expr_ptr.hpp>
+#include <span.hpp>
 
 namespace HIR {
 
@@ -155,6 +156,11 @@ public:
     
     bool operator==(const ::HIR::TypeRef& x) const;
     bool operator!=(const ::HIR::TypeRef& x) const { return !(*this == x); }
+
+    
+    // Match generics in `this` with types from `x`
+    // Raises a bug against `sp` if there is a form mismatch or `this` has an infer
+    void match_generics(const Span& sp, const ::HIR::TypeRef& x, ::std::function<void(unsigned int, const ::HIR::TypeRef&)> callback) const;
 };
 
 extern ::std::ostream& operator<<(::std::ostream& os, const ::HIR::TypeRef& ty);
diff --git a/src/hir_typeck/expr.cpp b/src/hir_typeck/expr.cpp
index 4151e405..e1a357f3 100644
--- a/src/hir_typeck/expr.cpp
+++ b/src/hir_typeck/expr.cpp
@@ -1081,6 +1081,109 @@ namespace {
             }
         }
         
+        
+        ::HIR::TypeRef expand_associated_types(const Span& sp, ::HIR::TypeRef input) const
+        {
+            TU_MATCH(::HIR::TypeRef::Data, (input.m_data), (e),
+            (Infer,
+                ),
+            (Diverge,
+                ),
+            (Primitive,
+                ),
+            (Path,
+                TU_MATCH(::HIR::Path::Data, (e.path.m_data), (e2),
+                (Generic,
+                    ),
+                (UfcsInherent,
+                    TODO(sp, "Path - UfcsInherent - " << e.path);
+                    ),
+                (UfcsKnown,
+                    // HACK - Shortcut to prevent expensive search if the type is a parameter
+                    if( e2.type->m_data.is_Generic() ) {
+                        return input;
+                    }
+                    
+                    // Search for a matching trait impl
+                    const ::HIR::TraitImpl* impl_ptr = nullptr;
+                    bool rv = this->m_crate.find_trait_impls(e2.trait.m_path, *e2.type,
+                        [&](const auto& ty)->const auto& {
+                            if( ty.m_data.is_Infer() )
+                                return this->get_type(ty);
+                            else
+                                return ty;
+                        },
+                        [&](const auto& impl) {
+                            DEBUG("Found impl - " << e2.trait.m_path << impl.m_trait_args << " for " << impl.m_type);
+                            if( impl.m_trait_args.m_types.size() > 0 )
+                            {
+                                TODO(sp, "Check trait type parameters in expand_associated_types");
+                            }
+                            impl_ptr = &impl;
+                            return true;
+                        });
+                    if( !rv )
+                        break;
+                    
+                    // An impl was found:
+                    assert(impl_ptr);
+                    // - Populate the impl's type arguments
+                    ::std::vector< const ::HIR::TypeRef*>   impl_args;
+                    impl_args.resize( impl_ptr->m_params.m_types.size() );
+                    
+                    // 1. Match impl_ptr->m_type with *e2.type
+                    impl_ptr->m_type.match_generics(sp, *e2.type, [&](unsigned int slot, const ::HIR::TypeRef& ty) {
+                        DEBUG("Set " << slot << " = " << ty);
+                        if( slot >= impl_args.size() ) {
+                            BUG(sp, "Impl parameter out of range - " << slot);
+                        }
+                        impl_args.at(slot) = &ty;
+                        });
+                    // 2. Match impl_ptr->m_trait_args with e2.trait.m_params
+                    
+                    // - Monomorphise the output type
+                    auto new_type = monomorphise_type_with(sp, impl_ptr->m_types.at( e2.item ), [&](const auto& ty)->const auto& {
+                        const auto& ge = ty.m_data.as_Generic();
+                        assert(ge.binding < impl_args.size());
+                        return *impl_args[ge.binding];
+                        });
+                    DEBUG("Converted UfcsKnown - " << e.path << " = " << new_type);
+                    return new_type;
+                    ),
+                (UfcsUnknown,
+                    BUG(sp, "Encountered UfcsUnknown");
+                    )
+                )
+                ),
+            (Generic,
+                ),
+            (TraitObject,
+                // Recurse?
+                ),
+            (Array,
+                *e.inner = expand_associated_types(sp, mv$(*e.inner));
+                ),
+            (Slice,
+                *e.inner = expand_associated_types(sp, mv$(*e.inner));
+                ),
+            (Tuple,
+                for(auto& sub : e) {
+                    sub = expand_associated_types(sp, mv$(sub));
+                }
+                ),
+            (Borrow,
+                *e.inner = expand_associated_types(sp, mv$(*e.inner));
+                ),
+            (Pointer,
+                *e.inner = expand_associated_types(sp, mv$(*e.inner));
+                ),
+            (Function,
+                // Recurse?
+                )
+            )
+            return input;
+        }
+        
         /// Searches for a trait impl that matches the provided trait name and type
         bool find_trait_impls(const ::HIR::SimplePath& trait, const ::HIR::TypeRef& type,  ::std::function<bool(const ::HIR::PathParams&)> callback)
         {
@@ -1784,7 +1887,7 @@ namespace {
             
             TRACE_FUNCTION_F("path = " << path);
             unsigned int arg_ofs = (is_method ? 1 : 0);
-            // TODO: Construct method to get a reference to an item along with the params decoded out of the pat
+            // TODO: Construct method to get a reference to an item along with the params decoded out of the path?
             TU_MATCH(::HIR::Path::Data, (path.m_data), (e),
             (Generic,
                 const auto& fcn = this->context.m_crate.get_function_by_path(sp, e.m_path);
@@ -1813,19 +1916,19 @@ namespace {
                 
                 fcn_ptr = &fcn;
                 
-                const auto& path_params = e.params;
+                //const auto& path_params = e.params;
                 monomorph_cb = [&](const auto& gt)->const auto& {
                         const auto& ge = gt.m_data.as_Generic();
                         if( ge.binding == 0xFFFF ) {
                             return *e.type;
                         }
                         // TODO: Don't the function-level params use 256-511?
-                        else if( ge.binding < 256 ) {
-                            return path_params.m_types[ge.binding];
-                        }
-                        else {
-                        }
-                        TODO(sp, "Monomorphise for trait method");
+                        //else if( ge.binding < 256 ) {
+                        //    return path_params.m_types[ge.binding];
+                        //}
+                        //else {
+                        //}
+                        TODO(sp, "Monomorphise for trait method - " << ge.name << " " << ge.binding);
                     };
                 ),
             (UfcsUnknown,
@@ -1843,21 +1946,39 @@ namespace {
                 ERROR(sp, E0000, "Incorrect number of arguments to " << path);
             }
             
+            // TODO: Save this list (as it's static, since ivars are inserted)
+            // - Prevents need to do lookups on every cycle
+            ::std::vector< ::HIR::TypeRef>  arg_types;
+            for(const auto& arg : fcn.m_args) {
+                if( monomorphise_type_needed(arg.second) ) {
+                    arg_types.push_back( this->context.expand_associated_types(sp, monomorphise_type_with(sp, arg.second,  monomorph_cb)) );
+                }
+                else {
+                    arg_types.push_back( arg.second.clone() );
+                }
+            }
+            if( monomorphise_type_needed(fcn.m_return) ) {
+                arg_types.push_back( this->context.expand_associated_types(sp, monomorphise_type_with(sp, fcn.m_return,  monomorph_cb)) );
+            }
+            else {
+                arg_types.push_back( fcn.m_return.clone() );
+            }
+            
             // TODO: Avoid needing to monomorphise here
             // - Have two callbacks to apply_equality that are used to expand `Generic`s (cleared once used)
+            // - Problem: Converting known associated types into the concrete type
+            //  > Fixable when doing monomorphisation here (because the type is getting cloned anyway)
+            //  > Could also be handled in apply_eqality
             for( unsigned int i = arg_ofs; i < fcn.m_args.size(); i ++ )
             {
                 auto& arg_expr_ptr = args[i - arg_ofs];
-                const auto& arg_ty = fcn.m_args[i].second;
+                const auto& arg_ty = arg_types[i];
                 DEBUG("Arg " << i << ": " << arg_ty);
-                ::HIR::TypeRef  mono_type;
-                const auto& ty = (monomorphise_type_needed(arg_ty) ? (mono_type = monomorphise_type_with(sp, arg_ty,  monomorph_cb)) : arg_ty);
-                this->context.apply_equality(sp, ty, arg_expr_ptr->m_res_type,  &arg_expr_ptr);
+                this->context.apply_equality(sp, arg_ty, arg_expr_ptr->m_res_type,  &arg_expr_ptr);
             }
-            DEBUG("RV " << fcn.m_return);
-            ::HIR::TypeRef  mono_type;
-            const auto& ty = (monomorphise_type_needed(fcn.m_return) ? (mono_type = monomorphise_type_with(sp, fcn.m_return,  monomorph_cb)) : fcn.m_return);
-            this->context.apply_equality(sp, res_type, ty);
+            
+            DEBUG("RV " << arg_types.back());
+            this->context.apply_equality(sp, res_type, arg_types.back());
         }
         
         // - Call Path: Locate path and build return