content: {-# LANGUAGE GeneralizedNewtypeDeriving, MultiParamTypeClasses,
               FlexibleContexts, OverloadedStrings, CPP #-}
  module Haste.Linker (link) where
  import Haste.Config
  import Haste.Module
  import qualified Data.Map as M
  import qualified Data.Set as S
  import Control.Monad.State.Strict
  import Control.Monad.Trans.Either
  import Haste.AST
  import qualified Data.ByteString.Lazy as B
  import qualified Data.ByteString as BS
  import Data.ByteString.UTF8 (toString, fromString)
  import Data.ByteString.Builder
  import Data.Monoid
  import Data.List (sort, group)
  import System.IO (hPutStrLn, stderr)
  import Crypto.Hash
  
  -- | The program entry point.
  --   This will need to change when we start supporting building "binaries"
  --   using cabal, since we'll have all sorts of funny package names then.
  mainSym :: Name
  mainSym = name "main" (Just ("main", "Main"))
  
  -- | Link a program using the given config and input file name.
  link :: Config -> BS.ByteString -> FilePath -> IO ()
  link cfg pkgid target = do
    let mainmod =
          case mainMod cfg of
           Just (m, p) -> (fromString m, fromString p)
           _           -> error "Haste.Linker.link called without main sym!"
    (spt, ds) <- getAllDefs cfg (targetLibPath cfg : libPaths cfg) mainmod pkgid mainSym
    let myDefs = if wholeProgramOpts cfg
                   then topLevelInline cfg spt ds
                   else ds
        (progText, (spt', myMain')) = prettyProg cfg mainSym spt myDefs
        callMain = stringUtf8 "B(A(" <> myMain' <> stringUtf8 ", [0]));"
        launchApp = appStart cfg (stringUtf8 "hasteMain")
    
    rtslibs <- mapM readFile $ rtsLibs cfg
    extlibs <- mapM readFile $ jsExternals cfg
    let finalProgram = toLazyByteString $ assembleProg (wrapProg cfg)
                                                       extlibs
                                                       rtslibs
                                                       progText
                                                       callMain
                                                       launchApp
                                                       spt'
    B.writeFile (outFile cfg cfg target) finalProgram
    where
      addHashLine p = concat
        [ "var __haste_prog_id = '"
        , show (hash (B.toStrict p) :: Digest SHA3_256)
        , "';\n"
        ]
      assembleProg True extlibs rtslibs progText callMain launchApp spt =
        stringUtf8 (unlines extlibs)
        <> stringUtf8 "function hasteMain() {\n"
        <> (if useStrict cfg then stringUtf8 "\"use strict\";\n" else mempty)
        <> stringUtf8 (addHashLine (toLazyByteString progText))
        <> stringUtf8 "var __haste_script_elem = hasteMain.scriptElem;\n"
        <> stringUtf8 (unlines rtslibs)
        <> progText
        <> mconcat (map addSPT spt)
        <> callMain
        <> stringUtf8 "};\n"
        <> stringUtf8 "hasteMain.scriptElem = typeof document == 'object' ? document.currentScript : null;\n"
        <> launchApp
      assembleProg _ extlibs rtslibs progText callMain launchApp spt =
        (if useStrict cfg then stringUtf8 "\"use strict\";\n" else mempty)
        <> stringUtf8 (addHashLine (toLazyByteString progText))
        <> stringUtf8 "var __haste_script_elem = typeof document == 'object' ? document.currentScript : null;\n"
        <> stringUtf8 (unlines extlibs)
        <> stringUtf8 (unlines rtslibs)
        <> progText
        <> mconcat (map addSPT spt)
        <> stringUtf8 "\nvar hasteMain = function() {" <> callMain
                                                       <> stringUtf8 "};"
        <> launchApp
  
      addSPT ptr = mconcat
        [ stringUtf8 "__spt_insert("
        , ptr
        , stringUtf8 ");\n"
        ]
  
  -- | Produce an info message if verbose reporting is enabled.
  info' :: Config -> String -> IO ()
  info' cfg = when (verbose cfg) . hPutStrLn stderr
  
  -- | Generate a sequence of all assignments needed to run Main.main.
  getAllDefs :: Config
             -> [FilePath]
             -> (BS.ByteString, BS.ByteString)
             -> BS.ByteString
             -> Name
             -> IO ([Name], Stm)
  getAllDefs cfg libpaths mainmod pkgid mainsym =
    runDep cfg mainmod $ addDef libpaths pkgid mainsym
  
  data DepState = DepState {
      mainModule  :: !(BS.ByteString, BS.ByteString),
      defs        :: !(Stm -> Stm),
      alreadySeen :: !(S.Set Name),
      modules     :: !(M.Map BS.ByteString Module),
      infoLogger  :: String -> IO (),
      staticPtrs  :: ![[Name]]
    }
  
  type DepM a = EitherT Name (StateT DepState IO) a
  
  initState :: Config -> (BS.ByteString, BS.ByteString) -> DepState
  initState cfg m = DepState {
      mainModule  = m,
      defs        = id,
      alreadySeen = S.empty,
      modules     = M.empty,
      infoLogger  = info' cfg,
      staticPtrs  = []
    }
  
  -- | Log a message to stdout if verbose reporting is on.
  info :: String -> DepM ()
  info s = do
    st <- get
    liftIO $ infoLogger st s
  
  -- | Run a dependency resolution computation.
  runDep :: Show a => Config -> (BS.ByteString,BS.ByteString) -> DepM a -> IO ([Name], Stm)
  runDep cfg mainmod m = do
      res <- runStateT (runEitherT m) (initState cfg mainmod)
      case res of
        (Right _, st) ->
          return (snubcat $ staticPtrs st, defs st stop)
        (Left (Name f (Just (_, modul))), _) -> do
          error $ msg (toString modul) (toString f)
        (r, _) -> do
          error $ "Impossible result in runDep: " ++ show r
    where
      msg "Main" "main" =
        "Unable to locate a main function.\n" ++
        "If your main function is not `Main.main' you must specify it using " ++
        "`-main-is',\n" ++
        "for instance, `-main-is MyModule.myMain'.\n" ++
        "If your progam intentionally has no main function," ++
        " please use `--dont-link' to avoid this error."
      msg s f =
        "Unable to locate function `" ++ f ++ "' in module `" ++ s ++ "'!"
  
      -- snubcat: sort + nub + concat in O(n log n) instead of O(n²)
      snubcat = map head . group . sort . concat
  
  -- | Return the module the given variable resides in.
  getModuleOf :: [FilePath] -> Name -> DepM Module
  getModuleOf libpaths v@(Name n _) =
    case moduleOf v of
      Just "GHC.Prim" -> return foreignModule
      Just ""         -> return foreignModule
      Nothing         -> return foreignModule
      Just ":Main"    -> do
        (p, m) <- mainModule `fmap` get
        getModuleOf libpaths (Name n (Just (p, m)))
      Just m          -> do
        mm <- getModule libpaths (maybe "main" id $ pkgOf v) m
        case mm of
          Just m' -> return m'
          _       -> left v
  
  -- | Return the module at the given path, loading it into cache if it's not
  --   already there. Modules are preferentially loaded from jslib file.
  getModule :: [FilePath] -> BS.ByteString -> BS.ByteString -> DepM (Maybe Module)
  getModule libpaths pkgid modname = do
      st <- get
      case M.lookup modname (modules st) of
        Just m -> do
          return $ Just m
        _ -> do
          info $ "Linking " ++ toString modname
          go libpaths
    where
      go (libpath:lps) = do
        mm <- liftIO $ readModule libpath (toString pkgid) (toString modname)
        case mm of
          Just m -> do
            st <- get
            put st
              { modules = M.insert modname m (modules st)
              , staticPtrs = modSPT m : staticPtrs st
              }
            mapM_ (addDef libpaths pkgid) (modSPT m)
            return (Just m)
          _ -> do
            go lps
      go [] = do
        return Nothing
  
  -- | Add a new definition and its dependencies. If the given identifier has
  --   already been added, it's just ignored.
  addDef :: [FilePath] -> BS.ByteString -> Name -> DepM ()
  addDef libpaths pkgid v = do
    st <- get
    when (not $ v `S.member` alreadySeen st) $ do
      m <- getModuleOf libpaths v
      -- getModuleOf may update the state, so we need to refresh it
      st' <- get
      let dependencies = maybe S.empty id (M.lookup v (modDeps m))
      put st' {alreadySeen = S.insert v (alreadySeen st')}
      S.foldl' (\a x -> a >> addDef libpaths pkgid x) (return ()) dependencies
  
      -- addDef _definitely_ updates the state, so refresh once again
      st'' <- get
      let Name cmnt _ = v
          defs' =
            maybe (defs st'')
                  (\body -> defs st'' . newVar True (internalVar v cmnt) body)
                  (M.lookup v (modDefs m))
      put st'' {defs = defs'}
  