Trans - Move monomorphisation from codegen pass to its own pass, and do a second inlining pass after monomorph.

8 files changed, 225 insertions, 22 deletions

diff --git a/src/main.cpp b/src/main.cpp
index 479bbc14..c63ce0d6 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -117,6 +117,8 @@ void init_debug_list()
 
     g_debug_disable_map.insert( "HIR Serialise" );
     g_debug_disable_map.insert( "Trans Enumerate" );
+    g_debug_disable_map.insert( "Trans Monomorph" );
+    g_debug_disable_map.insert( "MIR Optimise Inline" );
     g_debug_disable_map.insert( "Trans Codegen" );
 
     // Mutate this map using an environment variable
@@ -643,6 +645,9 @@ int main(int argc, char *argv[])
             #if 1
             // Generate a .o
             TransList   items = CompilePhase<TransList>("Trans Enumerate", [&]() { return Trans_Enumerate_Public(*hir_crate); });
+            CompilePhaseV("Trans Monomorph", [&]() { Trans_Monomorphise_List(*hir_crate, items); });
+            CompilePhaseV("MIR Optimise Inline", [&]() { MIR_OptimiseCrate_Inlining(*hir_crate, items); });
+            //CompilePhaseV("Trans Enumerate Cleanup", [&]() { Trans_Enumerate_Cleanup(*hir_crate, items); });
             CompilePhaseV("Trans Codegen", [&]() { Trans_Codegen(params.outfile + ".o", trans_opt, *hir_crate, items, false); });
             #endif
 
@@ -658,6 +663,8 @@ int main(int argc, char *argv[])
             #if 1
             // Generate a .o
             TransList   items = CompilePhase<TransList>("Trans Enumerate", [&]() { return Trans_Enumerate_Public(*hir_crate); });
+            CompilePhaseV("Trans Monomorph", [&]() { Trans_Monomorphise_List(*hir_crate, items); });
+            CompilePhaseV("MIR Optimise Inline", [&]() { MIR_OptimiseCrate_Inlining(*hir_crate, items); });
             CompilePhaseV("Trans Codegen", [&]() { Trans_Codegen(params.outfile + ".o", trans_opt, *hir_crate, items, false); });
             #endif
             // Save a loadable HIR dump
@@ -674,9 +681,13 @@ int main(int argc, char *argv[])
             // Can just emit the metadata and do miri?
             // - Requires MIR for EVERYTHING, not feasable.
             TransList items = CompilePhase<TransList>("Trans Enumerate", [&]() { return Trans_Enumerate_Public(*hir_crate); });
+            CompilePhaseV("Trans Monomorph", [&]() { Trans_Monomorphise_List(*hir_crate, items); });
+            CompilePhaseV("MIR Optimise Inline", [&]() { MIR_OptimiseCrate_Inlining(*hir_crate, items); });
             CompilePhaseV("Trans Codegen", [&]() { Trans_Codegen(params.outfile + ".o", trans_opt, *hir_crate, items, false); });
 
             TransList items2 = CompilePhase<TransList>("Trans Enumerate", [&]() { return Trans_Enumerate_Main(*hir_crate); });
+            CompilePhaseV("Trans Monomorph", [&]() { Trans_Monomorphise_List(*hir_crate, items2); });
+            CompilePhaseV("MIR Optimise Inline", [&]() { MIR_OptimiseCrate_Inlining(*hir_crate, items2); });
             CompilePhaseV("Trans Codegen", [&]() { Trans_Codegen(params.outfile + "-plugin", trans_opt, *hir_crate, items2, true); });
 
             hir_crate->m_lang_items.clear();    // Make sure that we're not exporting any lang items
@@ -686,9 +697,11 @@ int main(int argc, char *argv[])
             // Generate a binary
             // - Enumerate items for translation
             TransList items = CompilePhase<TransList>("Trans Enumerate", [&]() { return Trans_Enumerate_Main(*hir_crate); });
+            // - Monomorphise
+            CompilePhaseV("Trans Monomorph", [&]() { Trans_Monomorphise_List(*hir_crate, items); });
+            CompilePhaseV("MIR Optimise Inline", [&]() { MIR_OptimiseCrate_Inlining(*hir_crate, items); });
             // - Perform codegen
             CompilePhaseV("Trans Codegen", [&]() { Trans_Codegen(params.outfile, trans_opt, *hir_crate, items, true); });
-            // - Invoke linker?
             break;
         }
     }
diff --git a/src/mir/main_bindings.hpp b/src/mir/main_bindings.hpp
index 4e7e3b78..9f160b12 100644
--- a/src/mir/main_bindings.hpp
+++ b/src/mir/main_bindings.hpp
@@ -11,6 +11,7 @@
 namespace HIR {
 class Crate;
 }
+struct TransList;
 
 extern void HIR_GenerateMIR(::HIR::Crate& crate);
 extern void MIR_Dump(::std::ostream& sink, const ::HIR::Crate& crate);
@@ -19,3 +20,4 @@ extern void MIR_CheckCrate_Full(/*const*/ ::HIR::Crate& crate);
 
 extern void MIR_CleanupCrate(::HIR::Crate& crate);
 extern void MIR_OptimiseCrate(::HIR::Crate& crate, bool minimal_optimisations);
+extern void MIR_OptimiseCrate_Inlining(const ::HIR::Crate& crate, TransList& list);
diff --git a/src/mir/optimise.cpp b/src/mir/optimise.cpp
index dbc7e9c2..0ddd8d3d 100644
--- a/src/mir/optimise.cpp
+++ b/src/mir/optimise.cpp
@@ -16,6 +16,7 @@
 #include <algorithm>
 #include <iomanip>
 #include <trans/target.hpp>
+#include <trans/trans_list.hpp> // Note: This is included for inlining after enumeration and monomorph
 
 #include <hir/expr.hpp> // HACK
 
@@ -281,8 +282,32 @@ namespace {
             return monomorphise_type_get_cb(sp, self_ty, &impl_params, fcn_params, nullptr);
         }
     };
-    const ::MIR::Function* get_called_mir(const ::MIR::TypeResolve& state, const ::HIR::Path& path, ParamsSet& params)
+    const ::MIR::Function* get_called_mir(const ::MIR::TypeResolve& state, const TransList* list, const ::HIR::Path& path, ParamsSet& params)
     {
+        // If a TransList is avaliable, then all referenced functions must be in it.
+        if( list )
+        {
+            auto it = list->m_functions.find(path);
+            if( it == list->m_functions.end() )
+            {
+                MIR_BUG(state, "Enumeration failure - Function " << path << " not in TransList");
+            }
+            const auto& hir_fcn = *it->second->ptr;
+            if( it->second->monomorphised.code ) {
+                return &*it->second->monomorphised.code;
+            }
+            else if( hir_fcn.m_code.m_mir ) {
+                MIR_ASSERT(state, hir_fcn.m_params.m_types.empty(), "Enumeration failure - Function had params, but wasn't monomorphised - " << path);
+                // TODO: Check for trait methods too?
+                return &*hir_fcn.m_code.m_mir;
+            }
+            else {
+                MIR_ASSERT(state, !hir_fcn.m_code, "LowerMIR failure - No MIR but HIR is present?! - " << path);
+                // External function (no MIR present)
+                return nullptr;
+            }
+        }
+
         TU_MATCHA( (path.m_data), (pe),
         (Generic,
             const auto& fcn = state.m_crate.get_function_by_path(state.sp, pe.m_path);
@@ -492,8 +517,10 @@ namespace {
     }
 }
 
+// TODO: Move this block of definitions+code above the namespace above.
+
 bool MIR_Optimise_BlockSimplify(::MIR::TypeResolve& state, ::MIR::Function& fcn);
-bool MIR_Optimise_Inlining(::MIR::TypeResolve& state, ::MIR::Function& fcn, bool minimal);
+bool MIR_Optimise_Inlining(::MIR::TypeResolve& state, ::MIR::Function& fcn, bool minimal, const TransList* list=nullptr);
 bool MIR_Optimise_SplitAggregates(::MIR::TypeResolve& state, ::MIR::Function& fcn);
 bool MIR_Optimise_PropagateSingleAssignments(::MIR::TypeResolve& state, ::MIR::Function& fcn);
 bool MIR_Optimise_PropagateKnownValues(::MIR::TypeResolve& state, ::MIR::Function& fcn);
@@ -541,6 +568,40 @@ void MIR_OptimiseMin(const StaticTraitResolve& resolve, const ::HIR::ItemPath& p
 #endif
     return ;
 }
+/// Optimise doing inlining then cleaning up the mess
+///
+/// Returns true if any optimisation was performed
+///
+/// NOTE: This function can only be called after enumeration and monomorphisation, so it takes the TransList by reference not nullable pointer
+bool MIR_OptimiseInline(const StaticTraitResolve& resolve, const ::HIR::ItemPath& path, ::MIR::Function& fcn, const ::HIR::Function::args_t& args, const ::HIR::TypeRef& ret_type, const TransList& list)
+{
+    static Span sp;
+    bool rv = false;
+    TRACE_FUNCTION_FR(path, rv);
+    ::MIR::TypeResolve   state { sp, resolve, FMT_CB(ss, ss << path;), ret_type, args, fcn };
+
+    while( MIR_Optimise_Inlining(state, fcn, false, &list) )
+    {
+        MIR_Cleanup(resolve, path, fcn, args, ret_type);
+#if CHECK_AFTER_ALL
+        MIR_Validate(resolve, path, fcn, args, ret_type);
+#endif
+        rv = true;
+    }
+
+    MIR_Optimise_BlockSimplify(state, fcn);
+    MIR_Optimise_UnifyBlocks(state, fcn);
+
+    MIR_Optimise_GarbageCollect(state, fcn);
+    //MIR_Validate_Full(resolve, path, fcn, args, ret_type);
+    MIR_SortBlocks(resolve, path, fcn);
+
+#if CHECK_AFTER_DONE > 1
+    MIR_Validate(resolve, path, fcn, args, ret_type);
+#endif
+
+    return rv;
+}
 void MIR_Optimise(const StaticTraitResolve& resolve, const ::HIR::ItemPath& path, ::MIR::Function& fcn, const ::HIR::Function::args_t& args, const ::HIR::TypeRef& ret_type)
 {
     static Span sp;
@@ -772,7 +833,7 @@ bool MIR_Optimise_BlockSimplify(::MIR::TypeResolve& state, ::MIR::Function& fcn)
 // --------------------------------------------------------------------
 // If two temporaries don't overlap in lifetime (blocks in which they're valid), unify the two
 // --------------------------------------------------------------------
-bool MIR_Optimise_Inlining(::MIR::TypeResolve& state, ::MIR::Function& fcn, bool minimal)
+bool MIR_Optimise_Inlining(::MIR::TypeResolve& state, ::MIR::Function& fcn, bool minimal, const TransList* list/*=nullptr*/)
 {
     bool inline_happened = false;
     TRACE_FUNCTION_FR("", inline_happened);
@@ -1191,7 +1252,7 @@ bool MIR_Optimise_Inlining(::MIR::TypeResolve& state, ::MIR::Function& fcn, bool
             const auto& path = te->fcn.as_Path();
 
             Cloner  cloner { state.sp, state.m_resolve, *te };
-            const auto* called_mir = get_called_mir(state, path,  cloner.params);
+            const auto* called_mir = get_called_mir(state, list, path,  cloner.params);
             if( !called_mir )
                 continue ;
             if( called_mir == &fcn )
@@ -3655,3 +3716,45 @@ void MIR_OptimiseCrate(::HIR::Crate& crate, bool do_minimal_optimisation)
     ov.visit_crate(crate);
 }
 
+void MIR_OptimiseCrate_Inlining(const ::HIR::Crate& crate, TransList& list)
+{
+    ::StaticTraitResolve    resolve { crate };
+
+    bool did_inline_on_pass;
+
+    size_t  MAX_ITERATIONS = 5; // TODO: Tune this.
+    size_t  num_iterations = 0;
+    do
+    {
+        did_inline_on_pass = false;
+
+        for(auto& fcn_ent : list.m_functions)
+        {
+            const auto& path = fcn_ent.first;
+            const auto& pp = fcn_ent.second->pp;
+            auto& hir_fcn = *const_cast<::HIR::Function*>(fcn_ent.second->ptr);
+            auto& mono_fcn = fcn_ent.second->monomorphised;
+
+            ::std::string s = FMT(path);
+            ::HIR::ItemPath ip(s);
+
+            if( mono_fcn.code )
+            {
+                did_inline_on_pass |= MIR_OptimiseInline(resolve, ip, *mono_fcn.code, mono_fcn.arg_tys, mono_fcn.ret_ty, list);
+            }
+            else if( hir_fcn.m_code.m_mir)
+            {
+                did_inline_on_pass |= MIR_OptimiseInline(resolve, ip, *hir_fcn.m_code.m_mir, hir_fcn.m_args, hir_fcn.m_return, list);
+            }
+            else
+            {
+                // Extern, no optimisations
+            }
+        }
+    } while( did_inline_on_pass && num_iterations < MAX_ITERATIONS );
+
+    if( did_inline_on_pass )
+    {
+        DEBUG("Ran inlining optimise pass to exhaustion (maximum of " << MAX_ITERATIONS << " hit");
+    }
+}
diff --git a/src/trans/codegen.cpp b/src/trans/codegen.cpp
index 12252465..9e93caba 100644
--- a/src/trans/codegen.cpp
+++ b/src/trans/codegen.cpp
@@ -168,29 +168,18 @@ void Trans_Codegen(const ::std::string& outfile, const TransOptions& opt, const
             const auto& pp = ent.second->pp;
             TRACE_FUNCTION_F(path);
             DEBUG("FUNCTION CODE " << path);
+            // `is_extern` is set if there's no HIR (i.e. this function is from an external crate)
             bool is_extern = ! static_cast<bool>(fcn.m_code);
             // If this is a provided trait method, it needs to be monomorphised too.
             bool is_method = ( fcn.m_args.size() > 0 && visit_ty_with(fcn.m_args[0].second, [&](const auto& x){return x == ::HIR::TypeRef("Self",0xFFFF);}) );
             if( pp.has_types() || is_method )
             {
-                ::StaticTraitResolve    resolve { crate };
-                auto ret_type = pp.monomorph(resolve, fcn.m_return);
-                ::HIR::Function::args_t args;
-                for(const auto& a : fcn.m_args)
-                    args.push_back(::std::make_pair( ::HIR::Pattern{}, pp.monomorph(resolve, a.second) ));
-                auto mir = Trans_Monomorphise(resolve, pp, fcn.m_code.m_mir);
-                ::std::string s = FMT(path);
-                ::HIR::ItemPath ip(s);
-                MIR_Validate(resolve, ip, *mir, args, ret_type);
-                MIR_Cleanup(resolve, ip, *mir, args, ret_type);
-                MIR_Optimise(resolve, ip, *mir, args, ret_type);
-                MIR_Validate(resolve, ip, *mir, args, ret_type);
+                ASSERT_BUG(sp, ent.second->monomorphised.code, "Function that required monomorphisation wasn't monomorphised");
+
                 // TODO: Flag that this should be a weak (or weak-er) symbol?
-                // - If it's from an external crate, it should be weak
-                codegen->emit_function_code(path, fcn, ent.second->pp, is_extern,  mir);
+                // - If it's from an external crate, it should be weak, but what about local ones?
+                codegen->emit_function_code(path, fcn, ent.second->pp, is_extern,  ent.second->monomorphised.code);
             }
-            // TODO: Detect if the function was a #[inline] function from another crate, and don't emit if that is the case?
-            // - Emiting is nice, but it should be emitted as a weak symbol
             else {
                 codegen->emit_function_code(path, fcn, pp, is_extern,  fcn.m_code.m_mir);
             }
diff --git a/src/trans/enumerate.cpp b/src/trans/enumerate.cpp
index 4c91f5f3..4a6e5767 100644
--- a/src/trans/enumerate.cpp
+++ b/src/trans/enumerate.cpp
@@ -332,6 +332,51 @@ TransList Trans_Enumerate_Public(::HIR::Crate& crate)
     return rv;
 }
 
+void Trans_Enumerate_Cleanup(const ::HIR::Crate& crate, TransList& list)
+{
+    EnumState   state { crate };
+
+    // TODO: Get a list of "root" functions (e.g. main, public functions, things used by public generics) and re-enumerate based on that.
+
+    // Visit every function used
+    for(const auto& ent : list.m_functions)
+    {
+        if( ent.second->monomorphised.code )
+        {
+            Trans_Enumerate_FillFrom_MIR(state, *ent.second->monomorphised.code, {});
+        }
+        else if( ent.second->ptr->m_code.m_mir )
+        {
+            Trans_Enumerate_FillFrom_MIR(state, *ent.second->ptr->m_code.m_mir, {});
+        }
+        else
+        {
+        }
+    }
+
+    // Remove any item in `list.m_functions` that doesn't appear in `state.rv.m_functions`
+    for(auto it = list.m_functions.begin(); it != list.m_functions.end();)
+    {
+        auto it2 = state.rv.m_functions.find(it->first);
+        if( it2 == state.rv.m_functions.end() )
+        {
+            DEBUG("Remove " << it->first);
+            it = list.m_functions.erase(it);
+        }
+        else
+        {
+            DEBUG("Keep " << it->first);
+            ++ it;
+        }
+    }
+
+    // Sanity check: all items in `state.rv.m_functions` must exist in `list.m_functions`
+    for(const auto& e : state.rv.m_functions)
+    {
+        auto it = list.m_functions.find(e.first);
+        ASSERT_BUG(Span(), it != list.m_functions.end(), "Enumerate Error - New function appeared after monomorphisation - " << e.first);
+    }
+}
 
 /// Common post-processing
 void Trans_Enumerate_CommonPost_Run(EnumState& state)
diff --git a/src/trans/main_bindings.hpp b/src/trans/main_bindings.hpp
index d78991d2..46d2cdb1 100644
--- a/src/trans/main_bindings.hpp
+++ b/src/trans/main_bindings.hpp
@@ -25,8 +25,12 @@ struct TransOptions
 };
 
 extern TransList Trans_Enumerate_Main(const ::HIR::Crate& crate);
-extern TransList Trans_Enumerate_Test(const ::HIR::Crate& crate);
 // NOTE: This also sets the saveout flags
 extern TransList Trans_Enumerate_Public(::HIR::Crate& crate);
 
+/// Re-run enumeration on monomorphised functions, removing now-unused items
+extern void Trans_Enumerate_Cleanup(const ::HIR::Crate& crate, TransList& list);
+
+extern void Trans_Monomorphise_List(const ::HIR::Crate& crate, TransList& list);
+
 extern void Trans_Codegen(const ::std::string& outfile, const TransOptions& opt, const ::HIR::Crate& crate, const TransList& list, bool is_executable);
diff --git a/src/trans/monomorphise.cpp b/src/trans/monomorphise.cpp
index fd05c85b..892cb730 100644
--- a/src/trans/monomorphise.cpp
+++ b/src/trans/monomorphise.cpp
@@ -6,8 +6,10 @@
  * - MIR monomorphisation
  */
 #include "monomorphise.hpp"
+#include "hir_typeck/static.hpp"
 #include <mir/mir.hpp>
 #include <hir/hir.hpp>
+#include <mir/operations.hpp>   // Needed for post-monomorph checks and optimisations
 
 namespace {
     ::MIR::LValue monomorph_LValue(const ::StaticTraitResolve& resolve, const Trans_Params& params, const ::MIR::LValue& tpl)
@@ -333,3 +335,41 @@ namespace {
 
     return ::MIR::FunctionPointer( box$(output).release() );
 }
+
+/// Monomorphise all functions in a TransList
+void Trans_Monomorphise_List(const ::HIR::Crate& crate, TransList& list)
+{
+    ::StaticTraitResolve    resolve { crate };
+    for(auto& fcn_ent : list.m_functions)
+    {
+        const auto& fcn = *fcn_ent.second->ptr;
+        // Trait methods (which are the only case where `Self` can exist in the argument list at this stage) always need to be monomorphised.
+        bool is_method = ( fcn.m_args.size() > 0 && visit_ty_with(fcn.m_args[0].second, [&](const auto& x){return x == ::HIR::TypeRef("Self",0xFFFF);}) );
+        if(fcn_ent.second->pp.has_types() || is_method)
+        {
+            const auto& path = fcn_ent.first;
+            const auto& pp = fcn_ent.second->pp;
+            TRACE_FUNCTION_FR(path, path);
+
+            auto mir = Trans_Monomorphise(resolve, fcn_ent.second->pp, fcn.m_code.m_mir);
+
+            // TODO: Should these be moved to their own pass? Potentially not, the extra pass should just be an inlining optimise pass
+            auto ret_type = pp.monomorph(resolve, fcn.m_return);
+            ::HIR::Function::args_t args;
+            for(const auto& a : fcn.m_args)
+                args.push_back(::std::make_pair( ::HIR::Pattern{}, pp.monomorph(resolve, a.second) ));
+
+            ::std::string s = FMT(path);
+            ::HIR::ItemPath ip(s);
+            MIR_Validate(resolve, ip, *mir, args, ret_type);
+            MIR_Cleanup(resolve, ip, *mir, args, ret_type);
+            MIR_Optimise(resolve, ip, *mir, args, ret_type);
+            MIR_Validate(resolve, ip, *mir, args, ret_type);
+
+            fcn_ent.second->monomorphised.ret_ty = ::std::move(ret_type);
+            fcn_ent.second->monomorphised.arg_tys = ::std::move(args);
+            fcn_ent.second->monomorphised.code = ::std::move(mir);
+        }
+    }
+}
+
diff --git a/src/trans/trans_list.hpp b/src/trans/trans_list.hpp
index 2ef876d6..48274f87 100644
--- a/src/trans/trans_list.hpp
+++ b/src/trans/trans_list.hpp
@@ -41,10 +41,17 @@ struct Trans_Params
     }
 };
 
+struct CachedFunction {
+    ::HIR::TypeRef  ret_ty;
+    ::HIR::Function::args_t arg_tys;
+    ::MIR::FunctionPointer  code;
+};
 struct TransList_Function
 {
     const ::HIR::Function*  ptr;
     Trans_Params    pp;
+    // If `pp.has_types` is true, the below is valid
+    CachedFunction  monomorphised;
 };
 struct TransList_Static
 {