From ea8bbbc56a4b151e3f6599bf1969c2b9058907d5 Mon Sep 17 00:00:00 2001
From: knaito <knaito@asial.co.jp>
Date: Sat, 26 Sep 2020 18:31:28 +0900
Subject: [PATCH] feat: support 'spago run' and 'spago test'

---
 app/Main.hs               | 37 +++++++++++++++++++++++++++++--------
 src/CodeGen/IL/Printer.hs |  6 +++---
 2 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/app/Main.hs b/app/Main.hs
index 1b49e033f..8eaa1022b 100644
--- a/app/Main.hs
+++ b/app/Main.hs
@@ -43,8 +43,8 @@ import CodeGen.IL.Printer
 
 import Tests
 
-data Command = Build | Run
-
+data Command = Build | Run (Maybe String)
+ 
 parseJson :: Text -> Value
 parseJson text
   | Just fileJson <- decode . L.encodeUtf8 $ L.fromStrict text = fileJson
@@ -56,10 +56,30 @@ jsonToModule value =
     Success (_, r) -> r
     _ -> error "failed"
 
+splitArgs :: [String] -> ([String], Maybe String, [String])
+splitArgs args =
+  let (opts, files) = partition (isPrefixOf "--") args
+      opts' = (map . map) toLower opts
+      (moduleName, files') =
+        if null files
+          then (Nothing, files)
+          else
+            if "--run" `elem` opts' && length (files) >= 1
+              then do
+                let file0 = head files
+                    files' = tail files
+                case file0 of
+                  "Main.main" -> (Just "Main", files')
+                  "Test.Main.main" -> (Just "Test.Main", files')
+                  _ -> (Nothing, files)
+              else (Nothing, files)
+    in (opts, moduleName, files')
+
 main :: IO ()
 main = do
   args <- getArgs
-  let (opts, files) = partition (isPrefixOf "--") args
+  -- let (opts, files) = partition (isPrefixOf "--") args
+  let (opts, moduleName, files) = splitArgs args
       opts' = (map . map) toLower opts
   if "--tests" `elem` opts'
     then runTests
@@ -81,7 +101,7 @@ main = do
             else
               processFiles opts' files
           if "--run" `elem` opts
-            then runBuildTool Run
+            then runBuildTool $ Run moduleName
             else when ("--no-build" `notElem` opts) $
                  runBuildTool Build
   return ()
@@ -109,11 +129,12 @@ main = do
     return ()
   runBuildTool :: Command -> IO ()
   runBuildTool cmd = do
-    let command = case cmd of
-                    Build -> "build"
-                    Run -> "run"
     project <- projectEnv
-    callProcess "go" [command, T.unpack project </> modPrefix' </> "Main"]
+    case cmd of
+      Build -> do
+        callProcess "go" ["build", T.unpack project </> modPrefix' </> "Main"]
+      Run moduleName -> do
+        callProcess "go" ["run", T.unpack project </> modPrefix' </> fromMaybe "Main" moduleName]
     return ()
 
 generateCode :: [String] -> FilePath -> FilePath -> IO ()
diff --git a/src/CodeGen/IL/Printer.hs b/src/CodeGen/IL/Printer.hs
index 35b1074c6..d6a0fdfa9 100644
--- a/src/CodeGen/IL/Printer.hs
+++ b/src/CodeGen/IL/Printer.hs
@@ -449,9 +449,9 @@ interfaceSource _ _ _ = ""
 implHeaderSource :: Text -> [Text] -> Text -> Text
 implHeaderSource mn imports otherPrefix =
   "// Code generated by psgo. DO NOT EDIT.\n\n" <>
-  "package " <> (if mn == "Main" then "main" else mn) <> "\n\n" <>
+  "package " <> (if mn == "Main" || mn == "Test_Main" then "main" else mn) <> "\n\n" <>
   "import . \"" <> runtime <> "\"\n" <>
-  (if mn == "Main"
+  (if mn == "Main" || mn == "Test_Main"
       then "import _ \"" <> otherPrefix <> "/" <> ffiLoader <> "\"\n"
       else "\n") <>
   "import (\n" <>
@@ -482,7 +482,7 @@ implFooterSource mn foreigns =
                         "    })\n" <>
                         "    return " <> valueName name <> "\n" <> 
                         "}\n\n") <$> foreigns))) <>
-  if mn == "Main" then mainSource else "\n"
+  if mn == "Main" || mn == "Test.Main" then mainSource else "\n"
   where
   mainSource :: Text
   mainSource = "\