Trans - #[start] attribute

3 files changed, 55 insertions, 16 deletions

diff --git a/src/expand/lang_item.cpp b/src/expand/lang_item.cpp
index a995aec0..0a5c87f5 100644
--- a/src/expand/lang_item.cpp
+++ b/src/expand/lang_item.cpp
@@ -201,8 +201,30 @@ public:
     }
 };
 
+class Decorator_Start:
+    public ExpandDecorator
+{
+public:
+    AttrStage stage() const override { return AttrStage::Post; }
+    void handle(const Span& sp, const AST::MetaItem& attr, AST::Crate& crate, const AST::Path& path, AST::Module& mod, AST::Item& i) const override
+    {
+        TU_IFLET(::AST::Item, i, Function, e,
+            auto rv = crate.m_lang_items.insert(::std::make_pair( ::std::string("mrustc-start"), ::AST::Path(path) ));
+            if( !rv.second )
+            {
+                const auto& other_path = rv.first->second;
+                ERROR(sp, E0000, "Duplicate definition of #[start] - " << other_path << " and " << path);
+            }
+        )
+        else {
+            ERROR(sp, E0000, "#[start] on non-function " << path);
+        }
+    }
+};
+
 
 STATIC_DECORATOR("lang", Decorator_LangItem)
 STATIC_DECORATOR("main", Decorator_Main);
+STATIC_DECORATOR("start", Decorator_Start);
 
 
diff --git a/src/trans/codegen_c.cpp b/src/trans/codegen_c.cpp
index 7308a9f3..28dffa8f 100644
--- a/src/trans/codegen_c.cpp
+++ b/src/trans/codegen_c.cpp
@@ -101,12 +101,19 @@ namespace {
 
             if( is_executable )
             {
-                m_of
-                    << "int main(int argc, const char* argv[]) {\n"
-                    << "\t" << Trans_Mangle( ::HIR::GenericPath(m_resolve.m_crate.get_lang_item_path(Span(), "start")) ) << "("
-                        << "(uint8_t*)" << Trans_Mangle( ::HIR::GenericPath(m_resolve.m_crate.get_lang_item_path(Span(), "mrustc-main")) ) << ", argc, (uint8_t**)argv"
-                        << ");\n"
-                    << "}\n";
+                m_of << "int main(int argc, const char* argv[]) {\n";
+                auto c_start_path = m_resolve.m_crate.get_lang_item_path_opt("mrustc-start");
+                if( c_start_path == ::HIR::SimplePath() )
+                {
+                    m_of << "\t" << Trans_Mangle( ::HIR::GenericPath(m_resolve.m_crate.get_lang_item_path(Span(), "start")) ) << "("
+                            << "(uint8_t*)" << Trans_Mangle( ::HIR::GenericPath(m_resolve.m_crate.get_lang_item_path(Span(), "mrustc-main")) ) << ", argc, (uint8_t**)argv"
+                            << ");\n";
+                }
+                else
+                {
+                    m_of << "\t" << Trans_Mangle(::HIR::GenericPath(c_start_path)) << "(argc, argv);\n";
+                }
+                m_of << "}\n";
             }
 
             m_of.flush();
diff --git a/src/trans/enumerate.cpp b/src/trans/enumerate.cpp
index f3306c65..a6d3edab 100644
--- a/src/trans/enumerate.cpp
+++ b/src/trans/enumerate.cpp
@@ -59,21 +59,31 @@ TransList Trans_Enumerate_Main(const ::HIR::Crate& crate)
 
     EnumState   state { crate };
 
-    // "start" language item
-    // - Takes main, and argc/argv as arguments
+    auto c_start_path = crate.get_lang_item_path_opt("mrustc-start");
+    if( c_start_path == ::HIR::SimplePath() )
     {
-        auto start_path = crate.get_lang_item_path(sp, "start");
-        const auto& fcn = crate.get_function_by_path(sp, start_path);
+        // "start" language item
+        // - Takes main, and argc/argv as arguments
+        {
+            auto start_path = crate.get_lang_item_path(sp, "start");
+            const auto& fcn = crate.get_function_by_path(sp, start_path);
 
-        state.enum_fcn( start_path, fcn, {} );
-    }
+            state.enum_fcn( start_path, fcn, {} );
+        }
+
+        // user entrypoint
+        {
+            auto main_path = crate.get_lang_item_path(Span(), "mrustc-main");
+            const auto& fcn = crate.get_function_by_path(sp, main_path);
 
-    // user entrypoint
+            state.enum_fcn( main_path, fcn, {} );
+        }
+    }
+    else
     {
-        auto main_path = crate.get_lang_item_path(Span(), "mrustc-main");
-        const auto& fcn = crate.get_function_by_path(sp, main_path);
+        const auto& fcn = crate.get_function_by_path(sp, c_start_path);
 
-        state.enum_fcn( main_path, fcn, {} );
+        state.enum_fcn( c_start_path, fcn, {} );
     }
 
     return Trans_Enumerate_CommonPost(state);