From 2345986848a5c64bdb268520f31e113a817d3981 Mon Sep 17 00:00:00 2001
From: Alexander Kuzmenko <alex@stablex.ru>
Date: Wed, 1 Mar 2017 15:08:36 +0300
Subject: [PATCH] added -D JSTACK_MAIN for custom entry points (fixes #4)

---
 README.md           |  5 +++++
 extraParams.hxml    |  2 +-
 haxelib.json        |  4 ++--
 src/jstack/Tools.hx | 49 ++++++++++++++++++++++++++++++++-------------
 4 files changed, 43 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index 7002db3..3e47066 100644
--- a/README.md
+++ b/README.md
@@ -13,3 +13,8 @@ haxelib install jstack
 
 ## Usage
 Just add JStack to compilation with `-lib jstack` compiler flag.
+
+If you don't have `-main` in your build config, then you need to specify entry point like this:
+```
+-D JSTACK_MAIN=my.SomeClass.entryPoint
+```
diff --git a/extraParams.hxml b/extraParams.hxml
index 8177fe9..7dc4aab 100644
--- a/extraParams.hxml
+++ b/extraParams.hxml
@@ -1,3 +1,3 @@
 --macro include('jstack.JStack')
 --macro keep('jstack.JStack')
---macro jstack.Tools.addInjectMetaToMain()
\ No newline at end of file
+--macro jstack.Tools.addInjectMetaToEntryPoint()
\ No newline at end of file
diff --git a/haxelib.json b/haxelib.json
index e7b0ee3..ccd9d52 100644
--- a/haxelib.json
+++ b/haxelib.json
@@ -4,8 +4,8 @@
   "license"      : "MIT",
   "tags"         : ["js", "php7", "stack", "callstack", "stacktrace"],
   "description"  : "Friendly stack traces for JS and PHP7 targets. Makes them point to haxe sources.",
-  "version"      : "2.1.1",
-  "releasenote"  : "Handle uncaught exceptions in nodejs",
+  "version"      : "2.2.0",
+  "releasenote"  : "Custom entry point with -D JSTACK_MAIN=my.SomeClass.entryPoint",
   "classPath"    : "src",
   "contributors" : ["RealyUniqueName"],
   "dependencies" : {
diff --git a/src/jstack/Tools.hx b/src/jstack/Tools.hx
index 5c3322e..53c3aea 100644
--- a/src/jstack/Tools.hx
+++ b/src/jstack/Tools.hx
@@ -23,7 +23,7 @@ class Tools {
     /**
         Inject `json.JStack.onReady()` into app entry point, so that app will not start untill source map is ready.
     **/
-    static public function addInjectMetaToMain() : Void
+    static public function addInjectMetaToEntryPoint() : Void
     {
         #if (display || !(debug || JSTACK_FORCE))
             return;
@@ -32,42 +32,63 @@ class Tools {
         Compiler.define('js_source_map');
         Compiler.define('source_map');
 
-        var main : String = null;
-        var args = Sys.args();
-        for (i in 0...args.length) {
-            if(args[i] == '-main') {
-                main = args[i + 1];
-                break;
+        var entryClass : String = null;
+        var entryMethod : String = null;
+
+        if (Context.defined('JSTACK_MAIN')) {
+            var jstackMain = Context.definedValue('JSTACK_MAIN');
+
+            if (jstackMain == null || jstackMain.length == 0 || jstackMain == '1') {
+                Context.error('JSTACK_MAIN should have a value. E.g.: -D JSTACK_MAIN=my.SomeClass.entryPoint', (macro {}).pos);
+            }
+
+            var parts = jstackMain.split('.');
+            if (parts.length < 2) {
+                Context.error('JSTACK_MAIN value should have a class name and a function name. E.g.: -D JSTACK_MAIN=my.SomeClass.entryPoint', (macro {}).pos);
             }
+            entryMethod = parts.pop();
+            entryClass = parts.join('.');
         }
-        if (main == null) {
-            Context.warning('JStack: Failed to find entry point. Did you specify `-main` directive?', (macro {}).pos);
+
+        if (entryClass == null) {
+            var args = Sys.args();
+            for (i in 0...args.length) {
+                if(args[i] == '-main') {
+                    entryClass = args[i + 1];
+                    entryMethod = 'main';
+                    break;
+                }
+            }
+        }
+
+        if (entryClass == null) {
+            Context.warning('JStack: Failed to find entry point. Did you specify `-main` or `-D JSTACK_MAIN`?', (macro {}).pos);
             return;
         }
 
-        Compiler.addMetadata('@:build(jstack.Tools.injectInMain())', main);
+        Compiler.addMetadata('@:build(jstack.Tools.injectInEntryPoint("$entryMethod"))', entryClass);
     }
 #end
 
-    macro static public function injectInMain() : Array<Field>
+    macro static public function injectInEntryPoint(method:String) : Array<Field>
     {
         var fields = Context.getBuildFields();
         var injected = false;
 
         for (field in fields) {
-            if (field.name != 'main') continue;
+            if (field.name != method) continue;
 
             switch (field.kind) {
                 case FFun(fn):
                     fn.expr = macro jstack.JStack.onReady(function() ${fn.expr});
                     injected = true;
                 case _:
-                    Context.error('JStack: Failed to inject JStack in `main` function.', field.pos);
+                    Context.error('JStack: Failed to inject JStack in `$method` function.', field.pos);
             }
         }
 
         if (!injected) {
-            Context.error('JStack: Failed to find static function main.', (macro {}).pos);
+            Context.error('JStack: Failed to find entry point method "$method".', (macro {}).pos);
         }
 
         return fields;