{-------------------------------------------------------------------------------

        Copyright:              Bernie Pope 2004

        Module:                 AST

        Description:            Definition of the Abstract Syntax Tree for
                                Baskell programs.

        Primary Authors:        Bernie Pope

-------------------------------------------------------------------------------}

{-
    This file is part of baskell.

    baskell is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    baskell is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with baskell; if not, write to the Free Software
    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
-}

module AST
   ( Ident
   , Exp (..)
   , Lit (..)
   , Decl (..)
   , Program (..)
   , emptyProgram
   , isEmptyProgram
   , (=:)
   , var
   , lam
   , (@@)
   , app
   , litI
   , litC
   , litB
   , list
   , tup
   , prim
   ) where

import Pretty
   ( parensIf
   , char
   , text
   , (<>)
   , comma
   , Doc
   , (<+>)
   , vcat
   , parens
   , cat
   , render
   , punctuate
   , int
   , Pretty (..)
   )

import Data.List (sort)

import Type (Type)

--------------------------------------------------------------------------------

-- identifiers
type Ident = String

-- expressions
data Exp
   = Var Ident
   | Lam Ident Exp
   | LamStrict Ident Exp
   | App Exp Exp
   | Literal Lit
   | Tuple [Exp]
   | Prim String (Exp -> Maybe Exp)
   deriving Show

instance Show (a -> b) where
   show f = "<function>"

-- literals
data Lit
   = LitInt  Int
   | LitChar Char
   | LitBool Bool
   | LitCons
   | LitNil
   deriving Show

-- declarations
data Decl
   = Decl Ident Exp
   | Sig Ident Type
   deriving Show

instance Eq Decl where
   (Decl ident1 _body1) == (Decl ident2 _body2) = ident1 == ident2
   (Sig ident1 _ty1) == (Sig ident2 _ty2) = ident1 == ident2

instance Ord Decl where
   compare (Decl ident1 _body1) (Decl ident2 _body2)
      = compare ident1 ident2
   compare (Sig ident1 _ty1) (Sig ident2 _ty2)
      = compare ident1 ident2

-- programs
newtype Program
   = Program [Decl]
   deriving Show

emptyProgram :: Program
emptyProgram = Program []

isEmptyProgram :: Program -> Bool
isEmptyProgram (Program decls)
   = null decls

--------------------------------------------------------------------------------

-- pretty printing of the AST

instance Pretty Exp where
   pretty exp = prettyExp False exp

-- this should be a reader monad
prettyExp :: Bool -> Exp -> Doc
prettyExp _parens (Var ident) = text ident
prettyExp parens (Lam ident exp)
   = parensIf parens
        $ char '\\' <>
          text ident <+>
          text "->" <+>
          pretty exp
prettyExp parens (LamStrict ident exp)
   = parensIf parens
        $ char '!' <>
          text ident <+>
          text "->" <+>
          pretty exp
prettyExp _parens exp@(App (App (Literal LitCons) _hd) _tl)
   = prettyList exp
prettyExp parens exp@(App e1 e2)
   = parensIf parens
        $ prettyExp (not $ isApp e1) e1 <+> prettyExp True e2
   where
   isApp :: Exp -> Bool
   isApp (App _ _) = True
   isApp other     = False
prettyExp _parens (Literal lit)
   = pretty lit
prettyExp _parens (Tuple exps)
   = parens $ cat (punctuate comma $ map pretty exps)
prettyExp _parens (Prim name _implementation)
   = text "prim_" <> text name

instance Pretty Lit where
   pretty (LitInt i)  = pretty i
   pretty (LitChar c) = text $ show c
   pretty (LitBool b) = text $ show b
   pretty LitCons     = text $ "Cons"
   pretty LitNil      = text $ "[]"

instance Pretty Decl where
   pretty (Decl name body) = text name <+> text "=" <+> pretty body
   pretty (Sig name ty)    = text name <+> text "::" <+> pretty ty

instance Pretty Program where
   pretty (Program decls) = vcat $ map pretty (sort decls)

prettyList :: Exp -> Doc
prettyList (Literal LitNil) = text "[]"
prettyList (App (App (Literal LitCons) x) tail)
   = char '[' <> pretty x <> prettyList' tail
   where
   prettyList' (Literal LitNil) = char ']'
   prettyList' (App (App (Literal LitCons) h) t)
      = comma <> pretty h <> prettyList' t
   prettyList' other = comma <> pretty other


--------------------------------------------------------------------------------

-- below are some convenience functions that
-- make writing AST values slighlty nicer.

infix 0 =:
(=:) :: Ident -> Exp -> Decl
(=:) = Decl

var :: Ident -> Exp
var = Var

lam :: [Ident] -> Exp -> Exp
lam xs e = foldr Lam e xs

infixl 9 @@
(@@) :: Exp -> Exp -> Exp
e1 @@ e2 = App e1 e2

app :: [Exp] -> Exp
app [] = error "app applied to empty list"
app xs = foldl1 App xs

litI :: Int -> Exp
litI = Literal . LitInt

litC :: Char -> Exp
litC = Literal . LitChar

litB :: Bool -> Exp
litB = Literal . LitBool

litLC :: Exp
litLC = Literal LitCons

litLN :: Exp
litLN = Literal LitNil

list :: [Exp] -> Exp
list = foldr cons litLN

cons :: Exp -> Exp -> Exp
cons hd tl = litLC @@ hd @@ tl

tup :: [Exp] -> Exp
tup = Tuple

prim :: String -> (Exp -> Maybe Exp) -> Exp
prim = Prim