MIR+HIR Validation - Expand associated types in return type

2 files changed, 16 insertions, 27 deletions

diff --git a/src/hir_typeck/expr_check.cpp b/src/hir_typeck/expr_check.cpp
index 7af82bb1..f79f41c6 100644
--- a/src/hir_typeck/expr_check.cpp
+++ b/src/hir_typeck/expr_check.cpp
@@ -35,33 +35,21 @@ namespace {
             const auto& sp = node_ptr->span();
             node_ptr->visit(*this);
             
-            // TODO: If the return type contains ErasedType, then this check should instead check the structure and bounds.
-            // > Or just replace ErasedType-s with the known types and then equate
-            if( visit_ty_with(ret_type, [](const auto& ty) {
-                if( ty.m_data.is_ErasedType() )
+            // Monomorphise erased type
+            ::HIR::TypeRef  new_ret_type = clone_ty_with(sp, ret_type, [&](const auto& tpl, auto& rv) {
+                if( tpl.m_data.is_ErasedType() )
+                {
+                    const auto& e = tpl.m_data.as_ErasedType();
+                    ASSERT_BUG(sp, e.m_index < node_ptr.m_erased_types.size(), "Erased type index OOB - " << e.m_origin << " " << e.m_index << " >= " << node_ptr.m_erased_types.size());
+                    // TODO: Emit checks on bounds
+                    rv = node_ptr.m_erased_types[e.m_index].clone();
                     return true;
+                }
                 return false;
-                }) )
-            {
-                // Monomorphise erased type
-                ::HIR::TypeRef  new_ret_type = clone_ty_with(sp, ret_type, [&](const auto& tpl, auto& rv) {
-                    if( tpl.m_data.is_ErasedType() )
-                    {
-                        const auto& e = tpl.m_data.as_ErasedType();
-                        ASSERT_BUG(sp, e.m_index < node_ptr.m_erased_types.size(), "Erased type index OOB - " << e.m_origin << " " << e.m_index << " >= " << node_ptr.m_erased_types.size());
-                        // TODO: Emit checks on bounds
-                        rv = node_ptr.m_erased_types[e.m_index].clone();
-                        return true;
-                    }
-                    return false;
-                    });
+                });
+            m_resolve.expand_associated_types(sp, new_ret_type);
 
-                check_types_equal(sp, new_ret_type, node_ptr->m_res_type);
-            }
-            else
-            {
-                check_types_equal(sp, ret_type, node_ptr->m_res_type);
-            }
+            check_types_equal(sp, new_ret_type, node_ptr->m_res_type);
         }
         
         void visit(::HIR::ExprNode_Block& node) override
diff --git a/src/mir/check.cpp b/src/mir/check.cpp
index 80909902..8c101041 100644
--- a/src/mir/check.cpp
+++ b/src/mir/check.cpp
@@ -607,12 +607,12 @@ namespace {
             auto _ = this->m_resolve.set_item_generics(item.m_params);
             if( item.m_code ) {
                 DEBUG("Function code " << p);
+                // TODO: Get span without needing hir/expr.hpp
                 static Span sp;
                 
                 // Replace ErasedType instances in `ret_type`
-                ::HIR::TypeRef  tmp;
                 const auto& ret_type = item.m_return;
-                const auto& ret_type_v = ( !visit_ty_with(ret_type, [](const auto& t){ return t.m_data.is_ErasedType(); }) ? ret_type : tmp = clone_ty_with(sp, ret_type, [&](const auto& tpl, auto& rv) {
+                auto ret_type_v = clone_ty_with(sp, ret_type, [&](const auto& tpl, auto& rv) {
                     if( tpl.m_data.is_ErasedType() )
                     {
                         const auto& e = tpl.m_data.as_ErasedType();
@@ -621,7 +621,8 @@ namespace {
                         return true;
                     }
                     return false;
-                    }) );
+                    });
+                this->m_resolve.expand_associated_types(sp, ret_type_v);
     
                 MIR_Validate(m_resolve, p, *item.m_code.m_mir, item.m_args, ret_type_v);
             }