-- | Haskell Translation
--   Extracts SSTG from Haskell source.
module SSTG.Core.Translation.Haskell
    ( mkTargetBindings
    , mkIOStr
    ) where

import qualified SSTG.Core.Syntax.Language as SL

import BasicTypes
import Coercion
import CorePrep
import CoreSyn
import CoreToStg
import CostCentre
import DataCon
import FastString
import ForeignCall
import GHC
import GHC.Paths
import GhcMonad
import HscTypes
import Literal
import Module
import Name
import Outputable
import Pair
import PrimOp
import SrcLoc
import StgSyn
import TyCon
import TyCoRep
import Type
import UniqSet
import Unique
import Var as V

import System.IO

import qualified Data.IntMap as IM
import qualified Data.Maybe  as MB

-- | Make IO String from Outputable
mkIOStr :: (Outputable a) => a -> IO String
mkIOStr obj = runGhc (Just libdir) $ do
    dflags <- getSessionDynFlags
    let ppr_str = showPpr dflags obj
    return ppr_str

-- | Make Target Bindings
--   Given project directory and source target, make binds, with dependencies.
mkTargetBindings :: FilePath -> FilePath -> IO [SL.Binding]
mkTargetBindings proj src = do
    (sums_gutss, dflags, env) <- mkCompileClosure proj src
    let (sums, gutss) = (map fst sums_gutss, map snd sums_gutss)
    let mod_lcs = map (\s -> (ms_mod s, ms_location s)) sums
    let m_bndss = map mg_binds gutss
    let m_tcss  = map mg_tcs gutss

    let z1 = zip3 mod_lcs m_bndss m_tcss
    preps <- sequence $ map (\((m, l), b, t) -> corePrepPgm env m l b t) z1

    let z2 = zip (map fst mod_lcs) preps
    stg_bndss <- sequence $ map (\(m, p) -> coreToStg dflags m p) z2

    let sl_bnds = map mkBinding (concat stg_bndss)
    return sl_bnds

-- | Make Compilation Closure
--   Captures a snapshot of the DynFlags and HscEnv in addition to
--   the ModGuts in the ModuleGraph. This allows compilation to be, intheory,
--   more portable across different applications, since ModGuts is a crucial
--   intermediary for compilation in general.
mkCompileClosure :: FilePath -> FilePath ->
                    IO ([(ModSummary, ModGuts)], DynFlags, HscEnv)
mkCompileClosure proj src = runGhc (Just libdir) $ do
    beta_flags <- getSessionDynFlags
    let dflags = beta_flags { importPaths = [proj] }
    setSessionDynFlags dflags
    env    <- getSession
    target <- guessTarget src Nothing
    setTargets [target]
    load LoadAllTargets

    mod_graph <- getModuleGraph
    pmods <- (sequence . map parseModule) mod_graph
    tmods <- (sequence . map typecheckModule) pmods
    dmods <- (sequence . map desugarModule) tmods
    let mod_gutss = map coreModule dmods

    let zipd = (zip mod_graph mod_gutss, dflags, env)
    return zipd

-- | Make SSTG Expr
mkExpr :: StgExpr -> SL.Expr
mkExpr (StgLit lit) = SL.Atom (SL.LitAtom (mkLit lit))
mkExpr (StgApp occ args)    = SL.FunApp (mkVar occ) (map mkArg args)
mkExpr (StgConApp dc args)  = SL.ConApp (mkData dc) (map mkArg args)
mkExpr (StgOpApp op args _) = SL.PrimApp (mkPrimOp op) (map mkArg args)
mkExpr (StgTick _ expr)     = mkExpr expr
mkExpr (StgLam _ _)         = error "mkExpr: StgLam detected"
mkExpr (StgLet bnd expr)    = SL.Let (mkBinding bnd) (mkExpr expr)
mkExpr (StgLetNoEscape _ _ bnd expr)     = mkExpr (StgLet bnd expr)
mkExpr (StgCase mexpr _ _ bndr _ _ alts) =
    SL.Case (mkExpr mexpr) (mkVar bndr) (map mkAlt alts)

-- | Make SSTG Arg
mkArg :: StgArg -> SL.Atom
mkArg (StgVarArg occ) = SL.VarAtom (mkVar occ)
mkArg (StgLitArg lit) = SL.LitAtom (mkLit lit)

-- | Make SSTG Name
mkName :: Name -> SL.Name
mkName name = SL.Name occ mod ns unq
  where occ = (occNameString . nameOccName) name
        ns  = (mkNameSpace . occNameSpace . nameOccName) name
        unq = (getKey . nameUnique) name
        mod = case nameModule_maybe name of
            Nothing -> Nothing
            Just md -> Just ((moduleNameString . moduleName) md)

-- | Make SSTG NameSpace
mkNameSpace :: NameSpace -> SL.NameSpace
mkNameSpace ns | isVarNameSpace ns     = SL.VarNSpace
               | isTvNameSpace  ns     = SL.TvNSpace
               | isDataConNameSpace ns = SL.DataNSpace
               | isTcClsNameSpace ns   = SL.TcClsNSpace
               | otherwise = error "mkNameSpace: unrecognized namespace"

-- | Make SSTG Var
mkVar :: Var -> SL.Var
mkVar var = SL.Var vname vtype
  where vname = (mkName . V.varName) var
        vtype = (mkType . varType) var

-- | Make SSTG Binding
mkBinding :: StgBinding -> SL.Binding
mkBinding (StgNonRec bnd rhs) = SL.Binding SL.NonRec [(mkVar bnd, mkRhs rhs)]
mkBinding (StgRec bnds) = SL.Binding SL.Rec
                                     (map (\(b, r) -> (mkVar b, mkRhs r)) bnds)

-- | Make SSTG Rhs
mkRhs :: StgRhs -> SL.BindRhs
mkRhs (StgRhsCon _ dc args) = SL.ConForm (mkData dc) (map mkArg args)
mkRhs (StgRhsClosure _ _ _ _ _ params expr) =
    SL.FunForm (map mkVar params) (mkExpr expr)

-- | Make SSTG Lit
mkLit :: Literal -> SL.Lit
mkLit lit = case lit of
  (MachChar char)   -> SL.MachChar char ((mkType . literalType) lit)
  (MachStr bstr)    -> SL.MachStr (show bstr)   ((mkType . literalType) lit)
  (MachInt i)       -> SL.MachInt (fromInteger i) ((mkType . literalType) lit)
  (MachInt64 i)     -> SL.MachInt (fromInteger i) ((mkType . literalType) lit)
  (MachWord i)      -> SL.MachWord (fromInteger i) ((mkType . literalType) lit)
  (MachWord64 i)    -> SL.MachWord (fromInteger i) ((mkType . literalType) lit)
  (MachFloat rat)   -> SL.MachFloat rat ((mkType . literalType) lit)
  (MachDouble rat)  -> SL.MachDouble rat ((mkType . literalType) lit)
  (LitInteger i _)  -> SL.MachInt (fromInteger i) ((mkType . literalType) lit)
  MachNullAddr      -> SL.MachNullAddr ((mkType . literalType) lit)
  (MachLabel f m _) -> SL.MachLabel (unpackFS f) m ((mkType . literalType) lit)

-- | Make SSTG Data Constructor ID
mkDataId :: DataCon -> SL.ConTag
mkDataId datacon = SL.ConTag name tag
  where name = (mkName . dataConName) datacon
        tag  = dataConTag datacon

-- | Make SSTG Data Constructor
mkData :: DataCon -> SL.DataCon
mkData datacon = SL.DataCon dcid ty args
  where dcid = mkDataId datacon
        ty   = (mkType . dataConRepType) datacon
        args = map mkType (dataConOrigArgTys datacon)

-- | Make SSTG Primitive Operation
mkPrimOp :: StgOp -> SL.PrimFun
mkPrimOp (StgPrimOp op) = SL.PrimFun (SL.Name occ Nothing ns unq) ty
  where occname = primOpOcc op
        occ = occNameString occname
        ns  = (mkNameSpace . occNameSpace) occname
        unq = primOpTag op
        ty  = (mkType . primOpType) op
mkPrimOp otherwise = error "mkPrimOp: got StgPrimCallOp or StgFCallOp"

-- | Make SSTG Alt
mkAlt :: StgAlt -> SL.Alt
mkAlt (a, b, _, e) = SL.Alt (mkAltCon a) (map mkVar b) (mkExpr e)

-- | Make SSTG Alt Constructor
mkAltCon :: AltCon -> SL.AltCon
mkAltCon (DataAlt dc) = SL.DataAlt (mkData dc)
mkAltCon (LitAlt lit) = SL.LitAlt  (mkLit lit)
mkAltCon DEFAULT      = SL.Default

-- | Make SSTG Type
mkType :: Type -> SL.Type
mkType (TyVarTy v) = SL.TyVarTy (mkName (V.varName v)) (mkType (varType v))
mkType (AppTy t1 t2)    = SL.AppTy (mkType t1) (mkType t2)
mkType (TyConApp tc ts) = SL.TyConApp (mkTyCon tc) (map mkType ts)
mkType (ForAllTy b ty)  = SL.ForAllTy (mkTyBndr b) (mkType ty)
mkType (LitTy tlit)     = SL.LitTy (mkTyLit tlit)
mkType (CastTy ty cor)  = SL.CastTy (mkType ty) (mkCoercion cor)
mkType (CoercionTy cor) = SL.CoercionTy (mkCoercion cor)

-- | Make SSTG Kind
mkKind :: Kind -> SL.Type
mkKind = mkType

-- | Make SSTG Type Constructor
mkTyCon :: TyCon -> SL.TyCon
mkTyCon tc | isFunTyCon     tc = SL.FunTyCon     name
           | isAlgTyCon     tc = SL.AlgTyCon     name algrhs
           | isFamilyTyCon  tc = SL.FamilyTyCon  name
           | isPrimTyCon    tc = SL.PrimTyCon    name
           | isTcTyCon      tc = SL.TcTyCon      name
           | isTypeSynonymTyCon tc = SL.SynonymTyCon name
           | isPromotedDataCon tc  = SL.PromotedDataCon name dcon
           | otherwise = error "mkTyCon: unrecognized TyCon"
  where name   = (mkName . tyConName) tc
        kind   = (mkKind . tyConKind) tc
        algrhs = (mkAlgTyConRhs . algTyConRhs) tc
        dcon   = (mkData . MB.fromJust . isPromotedDataCon_maybe) tc

-- | Make SSTG Algebraic Type Constructor RHS
mkAlgTyConRhs :: AlgTyConRhs -> SL.AlgTyRhs
mkAlgTyConRhs (AbstractTyCon b) = SL.AbstractTyCon b
mkAlgTyConRhs (DataTyCon {data_cons = dcs}) = SL.DataTyCon  (map mkDataId dcs)
mkAlgTyConRhs (TupleTyCon {data_con = dc})  = SL.TupleTyCon (mkDataId dc)
mkAlgTyConRhs (NewTyCon {data_con = dc})    = SL.NewTyCon   (mkDataId dc)

-- | make SSTG Type Binder
mkTyBndr :: TyBinder -> SL.TyBinder
mkTyBndr (Anon ty)   = SL.AnonTyBndr  (mkType ty)
mkTyBndr (Named v f) = SL.NamedTyBndr (mkName (V.varName v))
                                      (mkType (varType v))

-- | Make SSTG Type Literal
mkTyLit :: TyLit -> SL.TyLit
mkTyLit (NumTyLit int) = SL.NumTyLit (fromInteger int)
mkTyLit (StrTyLit fs)  = SL.StrTyLit (unpackFS fs)

-- | Make SSTG Coercion
mkCoercion :: Coercion -> SL.Coercion
mkCoercion coer = SL.Coercion (mkType a) (mkType b)
  where (a, b) = (unPair . coercionKind) coer