module Stg.Language (
Program (..),
Binds (..),
LambdaForm (..),
prettyLambda,
UpdateFlag (..),
Rec (..),
Expr (..),
Alts (..),
NonDefaultAlts (..),
AlgebraicAlt (..),
PrimitiveAlt (..),
DefaultAlt (..),
Literal (..),
PrimOp (..),
Var (..),
Atom (..),
Constr (..),
Pretty (..),
classify,
LambdaType(..),
) where
import Control.DeepSeq
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NonEmpty
import Data.Map (Map)
import qualified Data.Map as M
import Data.Monoid hiding (Alt)
import Data.Text (Text)
import qualified Data.Text as T
import GHC.Exts
import GHC.Generics
import Language.Haskell.TH.Lift
import Text.PrettyPrint.ANSI.Leijen hiding ((<>))
data StgAstStyle = StgAstStyle
{ keyword :: Doc -> Doc
, prim :: Doc -> Doc
, variable :: Doc -> Doc
, constructor :: Doc -> Doc
, semicolon :: Doc -> Doc
}
style :: StgAstStyle
style = StgAstStyle
{ keyword = id
, prim = dullgreen
, variable = dullyellow
, constructor = dullmagenta
, semicolon = dullwhite
}
newtype Program = Program Binds
deriving (Eq, Ord, Show, Generic)
instance Monoid Program where
mempty = Program mempty
Program x `mappend` Program y = Program (x <> y)
newtype Binds = Binds (Map Var LambdaForm)
deriving (Eq, Ord, Generic)
instance Monoid Binds where
mempty = Binds mempty
Binds x `mappend` Binds y = Binds (y <> x)
instance Show Binds where
show (Binds binds) = "(Binds " <> show (M.assocs binds) <> ")"
data LambdaForm = LambdaForm ![Var] !UpdateFlag ![Var] !Expr
deriving (Eq, Ord, Show, Generic)
data LambdaType =
LambdaCon
| LambdaFun
| LambdaThunk
deriving (Eq, Ord, Show)
instance Pretty LambdaType where
pretty = \case
LambdaCon -> "Con"
LambdaFun -> "Fun"
LambdaThunk -> "Thunk"
classify :: LambdaForm -> LambdaType
classify = \case
LambdaForm _ _ [] AppC{} -> LambdaCon
LambdaForm _ _ (_:_) _ -> LambdaFun
LambdaForm _ _ [] _ -> LambdaThunk
data UpdateFlag = Update | NoUpdate
deriving (Eq, Ord, Show, Generic, Enum, Bounded)
data Rec = NonRecursive | Recursive
deriving (Eq, Ord, Show, Generic, Enum, Bounded)
data Expr =
Let !Rec !Binds !Expr
| Case !Expr !Alts
| AppF !Var ![Atom]
| AppC !Constr ![Atom]
| AppP !PrimOp !Atom !Atom
| Lit !Literal
deriving (Eq, Ord, Show, Generic)
data Alts = Alts !NonDefaultAlts !DefaultAlt
deriving (Eq, Ord, Show, Generic)
data NonDefaultAlts =
NoNonDefaultAlts
| AlgebraicAlts !(NonEmpty AlgebraicAlt)
| PrimitiveAlts !(NonEmpty PrimitiveAlt)
deriving (Eq, Ord, Show, Generic)
data AlgebraicAlt = AlgebraicAlt !Constr ![Var] !Expr
deriving (Eq, Ord, Show, Generic)
data PrimitiveAlt = PrimitiveAlt !Literal !Expr
deriving (Eq, Ord, Show, Generic)
data DefaultAlt =
DefaultNotBound !Expr
| DefaultBound !Var !Expr
deriving (Eq, Ord, Show, Generic)
newtype Literal = Literal Integer
deriving (Eq, Ord, Show, Generic)
data PrimOp =
Add
| Sub
| Mul
| Div
| Mod
| Eq
| Lt
| Leq
| Gt
| Geq
| Neq
deriving (Eq, Ord, Show, Generic, Bounded, Enum)
newtype Var = Var Text
deriving (Eq, Ord, Show, Generic)
instance IsString Var where fromString = coerce . T.pack
data Atom =
AtomVar !Var
| AtomLit !Literal
deriving (Eq, Ord, Show, Generic)
newtype Constr = Constr Text
deriving (Eq, Ord, Show, Generic)
instance IsString Constr where fromString = coerce . T.pack
deriveLiftMany [ ''Program, ''Literal, ''LambdaForm, ''UpdateFlag, ''Rec
, ''Expr, ''Alts, ''AlgebraicAlt, ''PrimitiveAlt, ''DefaultAlt
, ''PrimOp, ''Atom ]
instance Lift NonDefaultAlts where
lift NoNonDefaultAlts = [| NoNonDefaultAlts |]
lift (AlgebraicAlts alts) =
[| AlgebraicAlts (NonEmpty.fromList $(lift (toList alts))) |]
lift (PrimitiveAlts alts) =
[| PrimitiveAlts (NonEmpty.fromList $(lift (toList alts))) |]
instance Lift Binds where
lift (Binds binds) = [| Binds (M.fromList $(lift (M.assocs binds))) |]
instance Lift Constr where
lift (Constr con) = [| Constr (T.pack $(lift (T.unpack con))) |]
instance Lift Var where
lift (Var var) = [| Var (T.pack $(lift (T.unpack var))) |]
semicolonTerminated :: [Doc] -> Doc
semicolonTerminated = align . vsep . punctuate (semicolon style ";")
instance Pretty Program where
pretty (Program binds) = pretty binds
instance Pretty Binds where
pretty (Binds bs) =
(semicolonTerminated . map prettyBinding . M.assocs) bs
where
prettyBinding (var, lambda) =
pretty var <+> "=" <+> pretty lambda
prettyLambda
:: ([Var] -> Doc)
-> LambdaForm
-> Doc
prettyLambda pprFree (LambdaForm free upd bound expr) =
(prettyExp . prettyUpd . prettyBound . prettyFree) "\\"
where
prettyFree | null free = id
| otherwise = (<> lparen <> pprFree free <> rparen)
prettyUpd = (<+> case upd of Update -> "=>"
NoUpdate -> "->" )
prettyBound | null bound = id
| null free = (<> prettyList bound)
| otherwise = (<+> prettyList bound)
prettyExp = (<+> pretty expr)
instance Pretty LambdaForm where
pretty = prettyLambda prettyList
instance Pretty Rec where
pretty = \case
NonRecursive -> ""
Recursive -> "rec"
instance Pretty Expr where
pretty = \case
Let rec binds expr ->
let inBlock = indent 4 (keyword style "in" <+> pretty expr)
bindingBlock = line <> indent 4 (
keyword style ("let" <> pretty rec) <+> pretty binds )
in vsep [bindingBlock, inBlock]
Case expr alts -> vsep [ hsep [ keyword style "case"
, pretty expr
, keyword style "of" ]
, indent 4 (align (pretty alts)) ]
AppF var [] -> pretty var
AppF var args -> pretty var <+> prettyList args
AppC con [] -> pretty con
AppC con args -> pretty con <+> prettyList args
AppP op arg1 arg2 -> pretty op <+> pretty arg1 <+> pretty arg2
Lit lit -> pretty lit
instance Pretty Alts where
pretty (Alts NoNonDefaultAlts def) = pretty def
pretty (Alts (AlgebraicAlts alts) def) =
semicolonTerminated (map pretty (toList alts) <> [pretty def])
pretty (Alts (PrimitiveAlts alts) def) =
semicolonTerminated (map pretty (toList alts) <> [pretty def])
instance Pretty AlgebraicAlt where
pretty (AlgebraicAlt con [] expr)
= pretty con <+> "->" <+> pretty expr
pretty (AlgebraicAlt con args expr)
= pretty con <+> prettyList args <+> "->" <+> pretty expr
instance Pretty PrimitiveAlt where
pretty (PrimitiveAlt lit expr) =
pretty lit <+> "->" <+> pretty expr
instance Pretty DefaultAlt where
pretty = \case
DefaultNotBound expr -> "default" <+> "->" <+> pretty expr
DefaultBound var expr -> pretty var <+> "->" <+> pretty expr
instance Pretty Literal where
pretty (Literal i) = prim style (integer i <> "#")
instance Pretty PrimOp where
pretty op = prim style (case op of
Add -> "+#"
Sub -> "-#"
Mul -> "*#"
Div -> "/#"
Mod -> "%#"
Eq -> "==#"
Lt -> "<#"
Leq -> "<=#"
Gt -> ">#"
Geq -> ">=#"
Neq -> "/=#" )
instance Pretty Var where
pretty (Var name) = variable style (string (T.unpack name))
prettyList = hsep . map pretty
instance Pretty Atom where
pretty = \case
AtomVar var -> pretty var
AtomLit lit -> pretty lit
prettyList = hsep . map pretty
instance Pretty Constr where
pretty (Constr name) = constructor style (string (T.unpack name))
instance NFData Program
instance NFData Binds
instance NFData LambdaForm
instance NFData UpdateFlag
instance NFData Rec
instance NFData Expr
instance NFData Alts
instance NFData NonDefaultAlts
instance NFData AlgebraicAlt
instance NFData PrimitiveAlt
instance NFData DefaultAlt
instance NFData Literal
instance NFData PrimOp
instance NFData Var
instance NFData Atom
instance NFData Constr