module STG where

import Data.Maybe
import Text.PrettyPrint.HughesPJ

-- * Data type

type Var = String
type ConstrName = String
type PrimName = String
type Program = Binds
type Binds = [Bind]
type Bind = (Var, LambdaForm)
data LambdaForm 
    = LambdaForm [Var] Bool [Var] Expr
      deriving Show

data Expr
    = Let Binds Expr
    | LetRec Binds Expr
    | Case Expr Alts
    | App Var Atoms
    | Constr ConstrName Atoms
    | Prim PrimName -- Atoms
    | Literal Literal
    | Pos String Expr
      deriving Show

data Alts
    = AAlts [AlgebraicAlt]
    | PAlts [PrimitiveAlt]
      deriving Show

data AlgebraicAlt 
    = AAlt ConstrName [Var] Expr
    | AAltVar (Maybe Var) Expr
      deriving Show

data PrimitiveAlt
    = PAlt Literal Expr
    | PAltVar (Maybe Var) Expr
--     | PAltDefault Expr
      deriving Show

data Literal
    = LitChar Char
    | LitFloat Float
    | LitString String
    | LitInteger Integer
    | LitInt Int
      deriving Show

type Atoms = [Atom]
data Atom
    = AtomVar String
    | AtomLit Literal
      deriving Show

-- * Pretty-Printer

ppBinds :: Binds -> Doc
ppBinds = fsep . map ppBind 

ppBind :: (Var, LambdaForm) -> Doc
ppBind (varName, lf) =
    text varName <+> text "=" $+$
         (nest 4  $ ppLF lf)

ppLF :: LambdaForm -> Doc
ppLF (LambdaForm freeVars update args expr) =
    ppVars freeVars <+> ppUpdate update <+> ppVars args <+> text "->"  $+$ (nest 4 $ ppExpr expr)

ppVars :: [Var] -> Doc
ppVars = braces . hsep . punctuate comma . map ppVar

ppVar :: Var -> Doc
ppVar = text

ppAtoms :: Atoms -> Doc
ppAtoms = braces . hsep . punctuate comma . map ppAtom

ppAtom :: Atom -> Doc
ppAtom (AtomVar varName) = text varName
ppAtom (AtomLit lit) = ppLiteral lit

ppLiteral :: Literal -> Doc
ppLiteral (LitFloat f) = text (show f)
ppLiteral (LitChar c) = text (show c)
ppLiteral (LitString s) = text (show s)
ppLiteral (LitInteger i) = text (show i)
ppLiteral (LitInt i) = text (show i)

ppUpdate :: Bool -> Doc
ppUpdate False = text "\\n"
ppUpdate True = text "\\u"

ppExpr :: Expr -> Doc
ppExpr (Let binds expr) =
    text "let" <+> ppBinds binds $+$ text "in" $+$ (nest 2 $ ppExpr expr)
ppExpr (LetRec binds expr) =
    text "letrec" <+> ppBinds binds $+$ text "in" $+$ (nest 2 $ ppExpr expr)
ppExpr (Case expr alts) =
    text "case" <+> ppExpr expr <+> text "of" $+$
         (nest 2 $ ppAlts alts)
    where
      ppAlts (AAlts aalts) = fsep $ map ppAAlt aalts
      ppAlts (PAlts palts) = fsep $ map ppPAlt palts
      ppAAlt (AAlt constrName vars expr) =
          text constrName <+> (ppVars vars) <+> text "->" $+$
               (nest 2 $ ppExpr expr)
      ppPAlt (PAlt literal expr) =
          ppLiteral literal <+> text "->" $+$
               (nest 2 $ ppExpr expr)
      ppPAlt (PAltVar mVar expr) =
          ppVar (fromMaybe "_" mVar) <+> text "->" $+$
               (nest 2 $ ppExpr expr)
{-
      ppPAltDefault (PAltDefault expr) =
          text "_" <+> text "->" $+$
               (nest 2 $ ppExpr expr)
-}
ppExpr (App var []) = text var
ppExpr (App var atoms) = text var <+> ppAtoms atoms
    
ppExpr (Constr constrName atoms) =
    text constrName <+> ppAtoms atoms
ppExpr (Prim primName) =
    text primName -- <+> ppAtoms atoms
ppExpr (Literal lit) = ppLiteral lit
ppExpr (Pos p e) = {- text p $+$ -} ppExpr e
-- ppExpr (Prim str) = text str

                  
{-
test =
    ppBinds [("map", LambdaForm [] False ["f", "xs"] 
                   (Case (App "xs" []) 
                             [ AAlt "Nil" [] (Constr "Nil" []) 
                             , (AAlt "Cons" ["y","ys"] 
                                     (Let [ ("fy", LambdaForm ["f","y"] True []  (App "f" [AtomVar "y"])) 
                                          , ("mfy", LambdaForm ["f","ys"] True []  (App "map" [AtomVar "f", AtomVar "ys"])) 
                                          ] 
                                          (Constr "Cons" [AtomVar "fy", AtomVar "mfy"])))
                             ]))
            ]

-}