HIR Expand UFCS - Pass that replaces all call types with path-based calls

4 files changed, 251 insertions, 1 deletions

diff --git a/src/hir/expr.cpp b/src/hir/expr.cpp
index 31493a3d..a1bd5e9f 100644
--- a/src/hir/expr.cpp
+++ b/src/hir/expr.cpp
@@ -118,7 +118,10 @@ DEF_VISIT(ExprNode_ArraySized, node,
 )
 
 DEF_VISIT(ExprNode_Closure, node,
-    visit_node_ptr(node.m_code);
+    if(node.m_code)
+    {
+        visit_node_ptr(node.m_code);
+    }
 )
 
 #undef DEF_VISIT
diff --git a/src/hir_expand/main_bindings.hpp b/src/hir_expand/main_bindings.hpp
index 2ead572d..525396a2 100644
--- a/src/hir_expand/main_bindings.hpp
+++ b/src/hir_expand/main_bindings.hpp
@@ -7,3 +7,4 @@ namespace HIR {
 };
 
 extern void HIR_Expand_Closures(::HIR::Crate& crate);
+extern void HIR_Expand_UfcsEverything(::HIR::Crate& crate);
diff --git a/src/hir_expand/ufcs_everything.cpp b/src/hir_expand/ufcs_everything.cpp
new file mode 100644
index 00000000..30e1d531
--- /dev/null
+++ b/src/hir_expand/ufcs_everything.cpp
@@ -0,0 +1,242 @@
+/*
+ * MRustC - Rust Compiler
+ * - By John Hodge (Mutabah/thePowersGang)
+ *
+ * hir_expand/ufcs_everything.cpp
+ * - HIR Expansion - All calls (methods, values, ...) as UFCS path calls
+ */
+#include <hir/visitor.hpp>
+#include <hir/expr.hpp>
+#include <hir_typeck/static.hpp>
+#include <algorithm>
+#include "main_bindings.hpp"
+
+namespace {
+    inline HIR::ExprNodeP mk_exprnodep(HIR::ExprNode* en, ::HIR::TypeRef ty){ en->m_res_type = mv$(ty); return HIR::ExprNodeP(en); }
+}
+#define NEWNODE(TY, CLASS, ...)  mk_exprnodep(new HIR::ExprNode_##CLASS(__VA_ARGS__), TY)
+
+namespace {
+    
+    class ExprVisitor_Mutate:
+        public ::HIR::ExprVisitorDef
+    {
+        const ::HIR::Crate& m_crate;
+        ::HIR::ExprNodeP    m_replacement;
+        
+    public:
+        ExprVisitor_Mutate(const ::HIR::Crate& crate):
+            m_crate(crate)
+        {
+        }
+        void visit_node_ptr(::HIR::ExprPtr& root) {
+            const char* node_ty = typeid(*root).name();
+            TRACE_FUNCTION_FR(&*root << " " << node_ty << " : " << root->m_res_type, node_ty);
+            root->visit(*this);
+            if( m_replacement ) {
+                root.reset( m_replacement.release() );
+            }
+        }
+        
+        void visit_node_ptr(::HIR::ExprNodeP& node) override {
+            const char* node_ty = typeid(*node).name();
+            TRACE_FUNCTION_FR(&*node << " " << node_ty << " : " << node->m_res_type, node_ty);
+            assert( node );
+            node->visit(*this);
+            if( m_replacement ) {
+                node = mv$(m_replacement);
+            }
+        }
+        
+        
+        void visit(::HIR::ExprNode_CallValue& node) override
+        {
+            const auto& sp = node.span();
+            
+            ::HIR::ExprVisitorDef::visit(node);
+            
+            ::HIR::PathParams   trait_args;
+            {
+                ::std::vector< ::HIR::TypeRef>  arg_types;
+                // NOTE: In this case, m_arg_types is just the argument types
+                for(const auto& arg_ty : node.m_arg_types)
+                    arg_types.push_back(arg_ty.clone());
+                trait_args.m_types.push_back( ::HIR::TypeRef(mv$(arg_types)) );
+            }
+            
+            // TODO: You can call via &-ptrs, but that currently isn't handled in typeck
+            TU_IFLET(::HIR::TypeRef::Data, node.m_value->m_res_type.m_data, Closure, e,
+                if( node.m_trait_used == ::HIR::ExprNode_CallValue::TraitUsed::Unknown )
+                {
+                    // NOTE: Closure node still exists, and will do until MIR construction deletes the HIR
+                    switch(e.node->m_class)
+                    {
+                    case ::HIR::ExprNode_Closure::Class::Unknown:
+                        BUG(sp, "References an ::Unknown closure");
+                    case ::HIR::ExprNode_Closure::Class::NoCapture:
+                    case ::HIR::ExprNode_Closure::Class::Shared:
+                        node.m_trait_used = ::HIR::ExprNode_CallValue::TraitUsed::Fn;
+                        // TODO: Add borrow.
+                        break;
+                    case ::HIR::ExprNode_Closure::Class::Mut:
+                        node.m_trait_used = ::HIR::ExprNode_CallValue::TraitUsed::FnMut;
+                        // TODO: Add borrow.
+                        break;
+                    case ::HIR::ExprNode_Closure::Class::Once:
+                        node.m_trait_used = ::HIR::ExprNode_CallValue::TraitUsed::FnOnce;
+                        // TODO: Add borrow.
+                        break;
+                    }
+                }
+            )
+            
+            // Use marking in node to determine trait to use
+            ::HIR::Path   method_path(::HIR::SimplePath{});
+            switch(node.m_trait_used)
+            {
+            case ::HIR::ExprNode_CallValue::TraitUsed::Fn:
+                method_path = ::HIR::Path(
+                    node.m_value->m_res_type.clone(),
+                    ::HIR::GenericPath( m_crate.get_lang_item_path(sp, "fn"), mv$(trait_args) ),
+                    "call"
+                    );
+                break;
+            case ::HIR::ExprNode_CallValue::TraitUsed::FnMut:
+                method_path = ::HIR::Path(
+                    node.m_value->m_res_type.clone(),
+                    ::HIR::GenericPath( m_crate.get_lang_item_path(sp, "fn_mut"), mv$(trait_args) ),
+                    "call_mut"
+                    );
+                break;
+            case ::HIR::ExprNode_CallValue::TraitUsed::FnOnce:
+                method_path = ::HIR::Path(
+                    node.m_value->m_res_type.clone(),
+                    ::HIR::GenericPath( m_crate.get_lang_item_path(sp, "fn_once"), mv$(trait_args) ),
+                    "call_once"
+                    );
+                break;
+            
+            //case ::HIR::ExprNode_CallValue::TraitUsed::Unknown:
+            default:
+                BUG(node.span(), "Encountered CallValue with TraitUsed::Unknown, ty=" << node.m_value->m_res_type);
+            }
+            
+            auto self_arg_type = node.m_value->m_res_type.clone();
+            //auto ret_type = node.m_res_type.clone();
+            // Construct argument list for the output
+            ::std::vector< ::HIR::ExprNodeP>    args;
+            args.reserve( 1 + node.m_args.size() );
+            args.push_back( mv$(node.m_value) );
+            for(auto& arg : node.m_args)
+                args.push_back( mv$(arg) );
+            
+            m_replacement = NEWNODE(mv$(node.m_res_type), CallPath, sp,
+                mv$(method_path),
+                mv$(args)
+                );
+            
+            auto& arg_types = dynamic_cast< ::HIR::ExprNode_CallPath&>(*m_replacement).m_cache.m_arg_types;
+            arg_types.push_back( mv$(self_arg_type) );
+            for(auto& ty : node.m_arg_types)
+                arg_types.push_back( mv$(ty) );
+            //arg_types.push_back( mv$(ret_type) );
+        }
+        void visit(::HIR::ExprNode_CallMethod& node) override
+        {
+            const auto& sp = node.span();
+            
+            ::HIR::ExprVisitorDef::visit(node);
+            
+            ::std::vector< ::HIR::ExprNodeP>    args;
+            args.reserve( 1 + node.m_args.size() );
+            args.push_back( mv$(node.m_value) );
+            for(auto& arg : node.m_args)
+                args.push_back( mv$(arg) );
+            
+            // Replace using known function path
+            m_replacement = NEWNODE(mv$(node.m_res_type), CallPath, sp,
+                mv$(node.m_method_path),
+                mv$(args)
+                );
+            dynamic_cast< ::HIR::ExprNode_CallPath&>(*m_replacement).m_cache = mv$(node.m_cache);
+        }
+    };
+    class OuterVisitor:
+        public ::HIR::Visitor
+    {
+        const ::HIR::Crate& m_crate;
+    public:
+        OuterVisitor(const ::HIR::Crate& crate):
+            m_crate(crate)
+        {
+        }
+        
+        // NOTE: This is left here to ensure that any expressions that aren't handled by higher code cause a failure
+        void visit_expr(::HIR::ExprPtr& exp) {
+            BUG(Span(), "visit_expr hit in OuterVisitor");
+        }
+        
+        void visit_type(::HIR::TypeRef& ty) override
+        {
+            TU_IFLET(::HIR::TypeRef::Data, ty.m_data, Array, e,
+                this->visit_type( *e.inner );
+                DEBUG("Array size " << ty);
+                if( e.size ) {
+                    ExprVisitor_Mutate  ev(m_crate);
+                    ev.visit_node_ptr( e.size );
+                }
+            )
+            else {
+                ::HIR::Visitor::visit_type(ty);
+            }
+        }
+        // ------
+        // Code-containing items
+        // ------
+        void visit_function(::HIR::ItemPath p, ::HIR::Function& item) override {
+            //auto _ = this->m_ms.set_item_generics(item.m_params);
+            if( item.m_code )
+            {
+                DEBUG("Function code " << p);
+                ExprVisitor_Mutate  ev(m_crate);
+                ev.visit_node_ptr( item.m_code );
+            }
+            else
+            {
+                DEBUG("Function code " << p << " (none)");
+            }
+        }
+        void visit_static(::HIR::ItemPath p, ::HIR::Static& item) override {
+            if( item.m_value )
+            {
+                ExprVisitor_Mutate  ev(m_crate);
+                ev.visit_node_ptr(item.m_value);
+            }
+        }
+        void visit_constant(::HIR::ItemPath p, ::HIR::Constant& item) override {
+            if( item.m_value )
+            {
+                ExprVisitor_Mutate  ev(m_crate);
+                ev.visit_node_ptr(item.m_value);
+            }
+        }
+        void visit_enum(::HIR::ItemPath p, ::HIR::Enum& item) override {
+            for(auto& var : item.m_variants)
+            {
+                TU_IFLET(::HIR::Enum::Variant, var.second, Value, e,
+                    DEBUG("Enum value " << p << " - " << var.first);
+                    
+                    ExprVisitor_Mutate  ev(m_crate);
+                    ev.visit_node_ptr(e);
+                )
+            }
+        }
+    };
+}   // namespace
+
+void HIR_Expand_UfcsEverything(::HIR::Crate& crate)
+{
+    OuterVisitor    ov(crate);
+    ov.visit_crate( crate );
+}
+
diff --git a/src/main.cpp b/src/main.cpp
index 24690779..2316358d 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -185,6 +185,10 @@ int main(int argc, char *argv[])
         CompilePhaseV("Expand HIR Closures", [&]() {

             HIR_Expand_Closures(*hir_crate);

             });

+        // - And calls can be turned into UFCS

+        CompilePhaseV("Expand HIR Calls", [&]() {

+            HIR_Expand_UfcsEverything(*hir_crate);

+            });

         // - Ensure that typeck worked (including Fn trait call insertion etc)

         CompilePhaseV("Typecheck Expressions (validate)", [&]() {

             Typecheck_Expressions_Validate(*hir_crate);