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

        Copyright:              Bernie Pope 2004

        Module:                 Primitives.

        Description:            Implementation of Baskell's primitive 
                                functions --- those useful functions that
                                cannot be implemented in Baskell syntax
                                (either conveniently or at all).

                                Examples are integer addition, and if-then-else.

                                Also defines a type for each primitive 
                                function.

        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 Primitives
   ( primDecls 
   , primTypes
   )
   where

import AST
   ( Decl (..)
   , Exp (..)
   , Lit (..)
   , (=:)
   , var
   , lam
   , (@@)
   , app
   , litI
   , litC
   , litB
   , list
   , tup
   , prim
   )  

import TypeCheck
   ( Binding (LetBound) 
   , Constraint 
   , SolverType (..)
   )
   
--------------------------------------------------------------------------------

-- all the primitive function declarations
primDecls :: [Decl]
primDecls
   = [primHead, primTail, primNull, primPlus, primSub, primMult, primITE,
      primFst, primSnd, primLT, primGT, primEQInt, primDiv, primMod,
      primIsTuple
     ]

-- types of all the primitive functions
primTypes :: [Constraint]
primTypes
   = [primHeadTy, primTailTy, primNullTy, primPlusTy, primSubTy, primMultTy, primITETy,
      primFstTy, primSndTy, primLTTy, primGTTy, primEQIntTy, primDivTy, primModTy,
      primIsTupleTy
     ]

-- helper for making unary prims
unaryPrim :: String -> (Exp -> Maybe Exp) -> Exp 
unaryPrim name impl 
   = lam ["x"] (prim name impl @@ var "x")

-- helper for making binary prims
binaryPrim :: String -> (Exp -> Maybe Exp) -> Exp 
binaryPrim name impl
   = lam ["x","y"] (prim name impl @@ var "x" @@ var "y")

-- helper for making ternary prims
ternaryPrim :: String -> (Exp -> Maybe Exp) -> Exp 
ternaryPrim name impl
   = lam ["x","y","z"] (prim name impl @@ var "x" @@ var "y" @@ var "z")

-- helper for making types of prims
primType :: String -> SolverType -> Constraint 
primType name t 
   = (TypeOf LetBound name, t)

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

-- head of a list

headName = "head"

primHead :: Decl
primHead 
   = headName =: unaryPrim headName thisPrim
   where 
   thisPrim :: Exp -> Maybe Exp
   thisPrim (App (App (Literal LitCons) hd) _tail) = Just hd
   thisPrim other = Nothing

primHeadTy :: Constraint 
primHeadTy = primType headName (TFun (TList (TVar 1)) (TVar 1))

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

-- tail of a list

tailName = "tail"

primTail :: Decl
primTail
   = tailName =: unaryPrim tailName thisPrim 
   where
   thisPrim :: Exp -> Maybe Exp
   thisPrim (App (App (Literal LitCons) _hd) tail) = Just tail 
   thisPrim other = Nothing

primTailTy :: Constraint
primTailTy = primType tailName (TFun (TList (TVar 1)) (TList (TVar 1)))

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

-- test if a list is empty

nullName = "null"

primNull :: Decl
primNull
   = nullName =: unaryPrim nullName thisPrim 
   where
   thisPrim :: Exp -> Maybe Exp
   thisPrim (Literal LitNil) = Just $ litB True
   thisPrim (App (App (Literal LitCons) _hd) _tail) = Just $ litB False 
   thisPrim other = Nothing

primNullTy :: Constraint
primNullTy = primType nullName (TFun (TList (TVar 1)) TBool)

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

-- first item in a tuple

fstName = "fst"

primFst :: Decl
primFst
   = fstName =: unaryPrim fstName thisPrim
   where
   thisPrim :: Exp -> Maybe Exp
   thisPrim (Tuple [x,_y])  = Just x 
   thisPrim other = Nothing

primFstTy :: Constraint
primFstTy = primType fstName (TFun (TTuple [TVar 1, TVar 2]) (TVar 1))

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

-- second item in a tuple

sndName = "snd"

primSnd :: Decl
primSnd
   = sndName =: unaryPrim sndName thisPrim
   where
   thisPrim :: Exp -> Maybe Exp
   thisPrim (Tuple [_x,y])  = Just y 
   thisPrim other = Nothing

primSndTy :: Constraint
primSndTy = primType sndName (TFun (TTuple [TVar 1, TVar 2]) (TVar 2))

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

-- test for a tuple 
-- note this function has type: a -> Bool
-- this is not available in Haskell, it would normally need some kind
-- of meta-programming facility

isTupleName = "isTuple"

primIsTuple :: Decl
primIsTuple
   = isTupleName =: unaryPrim isTupleName thisPrim
   where
   thisPrim :: Exp -> Maybe Exp
   thisPrim (Tuple _) = Just $ Literal $ LitBool True
   thisPrim other = Just $ Literal $ LitBool False 

primIsTupleTy :: Constraint
primIsTupleTy = primType isTupleName (TFun (TVar 1) TBool)

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

-- integer addition

plusName = "plus"

primPlus :: Decl
primPlus
   = plusName =: binaryPrim plusName thisPrim
   where
   thisPrim :: Exp -> Maybe Exp
   thisPrim (Literal (LitInt i))
      = Just $ unaryPrim (plusName ++ "_partial(1)") (addition i) 
      where
      addition :: Int -> Exp -> Maybe Exp
      addition i (Literal (LitInt j))
         = Just $ litI $ i + j 
      addition i other = Nothing
   thisPrim other = Nothing

primPlusTy :: Constraint
primPlusTy = primType plusName (TFun TInt (TFun TInt TInt)) 

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

-- integer subtraction

subName = "sub"

primSub :: Decl
primSub
   = subName =: binaryPrim subName thisPrim
   where
   thisPrim :: Exp -> Maybe Exp
   thisPrim (Literal (LitInt i))
      = Just $ unaryPrim (subName ++ "_partial(1)") (subtraction i) 
      where
      subtraction :: Int -> Exp -> Maybe Exp
      subtraction i (Literal (LitInt j))
         = Just $ litI $ i - j 
      subtraction i other = Nothing
   thisPrim other = Nothing

primSubTy :: Constraint
primSubTy = primType subName (TFun TInt (TFun TInt TInt)) 

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

-- integer multiplication

multName = "mult"

primMult :: Decl
primMult
   = multName =: binaryPrim multName thisPrim
   where
   thisPrim :: Exp -> Maybe Exp
   thisPrim (Literal (LitInt i))
      = Just $ unaryPrim (multName ++ "_partial(1)") (mult i)
      where
      mult :: Int -> Exp -> Maybe Exp
      mult i (Literal (LitInt j))
         = Just $ litI $ i * j  
      mult i other = Nothing
   thisPrim other = Nothing

primMultTy :: Constraint
primMultTy = primType multName (TFun TInt (TFun TInt TInt)) 

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

-- integer division

divName = "div"

primDiv :: Decl
primDiv
   = divName =: binaryPrim divName thisPrim
   where
   thisPrim :: Exp -> Maybe Exp
   thisPrim (Literal (LitInt i))
      = Just $ unaryPrim (divName ++ "_partial(1)") (divide i)
      where
      divide :: Int -> Exp -> Maybe Exp
      divide i (Literal (LitInt j))
         = Just $ litI $ i `div` j
      divide i other = Nothing
   thisPrim other = Nothing

primDivTy :: Constraint
primDivTy = primType divName (TFun TInt (TFun TInt TInt)) 

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

-- integer modulus

modName = "mod"

primMod :: Decl
primMod
   = modName =: binaryPrim modName thisPrim
   where
   thisPrim :: Exp -> Maybe Exp
   thisPrim (Literal (LitInt i))
      = Just $ unaryPrim (modName ++ "_partial(1)") (modulus i)
      where 
      modulus :: Int -> Exp -> Maybe Exp
      modulus i (Literal (LitInt j))
         = Just $ litI $ i `mod` j
      modulus i other = Nothing
   thisPrim other = Nothing

primModTy :: Constraint
primModTy = primType modName (TFun TInt (TFun TInt TInt)) 

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

-- integer less than

ltName = "lt"

primLT :: Decl
primLT
   = ltName =: binaryPrim ltName thisPrim
   where
   thisPrim :: Exp -> Maybe Exp
   thisPrim (Literal (LitInt i))
      = Just $ unaryPrim (ltName ++ "_partial(1)") (lt i) 
      where
      lt :: Int -> Exp -> Maybe Exp
      lt i (Literal (LitInt j))
         = Just $ litB $ i < j 
      lt i other = Nothing
   thisPrim other = Nothing

primLTTy :: Constraint
primLTTy = primType ltName (TFun TInt (TFun TInt TBool)) 

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

-- integer greater than

gtName = "gt"

primGT :: Decl
primGT
   = gtName =: binaryPrim gtName thisPrim
   where
   thisPrim :: Exp -> Maybe Exp
   thisPrim (Literal (LitInt i))
      = Just $ unaryPrim (gtName ++ "_partial(1)") (gt i)
      where
      gt :: Int -> Exp -> Maybe Exp
      gt i (Literal (LitInt j))
         = Just $ litB $ i > j
      gt i other = Nothing
   thisPrim other = Nothing


primGTTy :: Constraint
primGTTy = primType gtName (TFun TInt (TFun TInt TBool)) 

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

-- integer equality

eqIntName = "eqI"

primEQInt :: Decl
primEQInt
   = eqIntName =: binaryPrim eqIntName thisPrim
   where
   thisPrim :: Exp -> Maybe Exp
   thisPrim (Literal (LitInt i))
      = Just $ unaryPrim (eqIntName ++ "_partial(1)") (eq i)
      where
      eq :: Int -> Exp -> Maybe Exp
      eq i (Literal (LitInt j))
         = Just $ litB $ i == j
      eq i other = Nothing
   thisPrim other = Nothing

primEQIntTy :: Constraint
primEQIntTy = primType eqIntName (TFun TInt (TFun TInt TBool)) 

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

-- if-then-else

iteName = "ite"

primITE :: Decl
primITE
   = iteName =: ternaryPrim iteName thisPrim
   where
   thisPrim :: Exp -> Maybe Exp
   thisPrim (Literal (LitBool b))
      = Just $ binaryPrim (iteName ++ "_partial(2)") (ite b) 
      where
      ite :: Bool -> Exp -> Maybe Exp
      ite True exp
         = Just $ unaryPrim (iteName ++ "_partial(1)") (trueBranch exp)
      ite False exp
         = Just $ unaryPrim (iteName ++ "_partial(1)") falseBranch 
      trueBranch :: Exp -> Exp -> Maybe Exp
      trueBranch keep ignore = Just keep
      falseBranch :: Exp -> Maybe Exp
      falseBranch keep = Just keep 
   thisPrim other = Nothing

primITETy :: Constraint
primITETy = primType iteName (TFun TBool (TFun (TVar 1) (TFun (TVar 1) (TVar 1))))