{- 
 - 	Monadic Constraint Programming
 - 	http://www.cs.kuleuven.be/~toms/MCP/
 - 	Pieter Wuille
 -}



{-# LANGUAGE MultiParamTypeClasses #-}

{-# LANGUAGE FlexibleInstances #-}

{-# LANGUAGE FlexibleContexts #-}



module Data.Expr.Sugar (

  (@+), (@-), (@*), (@/), (@%), (@?), (@??), (@:),

  (!), (@!!), (@++), (@..), size, slice, xhead, xtail, xmap, xfold, list, channel, xsum,

  (@||), (@&&), inv,

  (@/=), (@>), (@<), (@>=), (@<=), (@=), 

  loopall, loopany, forall, forany,

  Expr(), ColExpr(), BoolExpr(),

  ToExpr(..), ToColExpr(..), ToBoolExpr(..),

  sorted, sSorted, allDiff, allDiffD,

  ExprClass, ExprRange,

) where 



import Data.Expr.Data

import Data.Expr.Util



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

-- | Built-in class instances | --

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



instance (Eq s, Eq c, Eq b, Show s, Show c, Show b) => Num (Expr s c b) where

  a + b = simplify $ a `Plus` b

  a - b = simplify $ a `Minus` b

  a * b = simplify $ a `Mult` b

  abs a = simplify $ Abs a

  negate a = simplify $ (Const 0) `Minus` a

  fromInteger c = Const $ fromInteger c

  signum (Const a) = Const $ signum a

  signum a = error "signum not possible for generic Expr"



instance (Ord s, Ord c, Ord b, Eq s, Eq c, Eq b, Show s, Show c, Show b) => Real (Expr s c b) where

  toRational (Const x) = toRational x

  toRational _ = error "toRational not possible for generic Expr"



instance (Eq s, Eq c, Eq b) => Enum (Expr s c b) where

  succ a = simplify $ a `Plus` (Const 1)

  pred a = simplify $ a `Minus` (Const 1)

  toEnum = Const . toEnum

  fromEnum (Const a) = fromEnum a

  fromEnum _ = error "fromEnum not possible for generic Expr"



instance (Ord s, Ord c, Ord b, Eq s, Eq c, Eq b, Show s, Show c, Show b) => Integral (Expr s c b) where

  toInteger (Const a) = toInteger a

  toInteger _ = error "toInteger not possible for generic Expr"

  divMod a b = (simplify $ a `Div` b, simplify $ a `Mod` b)

  quotRem (Const a) (Const b) = case quotRem a b of (c,d) -> (Const c,Const d)

  quotRem (Const 0) b = (Const 0,Const 0)

  quotRem a (Const 1) = (a,Const 0)

  quotRem a (Const (-1)) = (negate a,Const 0)

  quotRem _ _ = error "quotRem not possible for generic Expr"



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

-- | convertion from/to expression types | --

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



-- convertible to expressions:

class ToExpr tt cc bb t where

  toExpr :: t -> Expr tt cc bb



-- convertible to collection-expressions:

class ToColExpr tt cc bb c where

  toColExpr :: c -> ColExpr tt cc bb



-- convertible to boolean expressions:

class ToBoolExpr tt cc bb b where

  toBoolExpr :: b -> BoolExpr tt cc bb



-- infix 4 @=, @/=



class (Eq tt, Eq cc, Eq bb) => ExprClass tt cc bb a where

  (@=)  :: a -> a -> BoolExpr tt cc bb

  (@/=) :: a -> a -> BoolExpr tt cc bb

  a @/= b = boolSimplify $ BoolNot $ a @= b



class (Eq tt, Eq cc, Eq bb) => ExprRange tt cc bb r where

  (@:)  :: Expr tt cc bb -> r -> BoolExpr tt cc bb



-- integers can be used as constant expressions

instance ToExpr tt cc bb Integer where

  toExpr = Const



-- expressions themselves are trivially convertible to expressions

instance ToExpr t a b (Expr t a b) where

  toExpr = id



-- ints can be used as constant expressions

instance ToExpr tt cc bb Int where

  toExpr = Const . toInteger



-- boolean expressions can be used as integer expressions (being 0 or 1)

instance (Eq t, Eq a, Eq b) => ToExpr t a b (BoolExpr t a b) where

  toExpr = simplify . Channel



-- collection expressions themselves are trivially convertible to collection expressions

instance ToColExpr t a b (ColExpr t a b) where

  toColExpr = id



-- an expression can be used as a collection of one expressions

instance (Eq t, Eq a, Eq b) => ToColExpr t a b (Expr t a b) where

  toColExpr a = colSimplify $ ColList [a]



-- a list of expressions van be used as a collection

instance (Eq b, Eq a, Eq t) => ToColExpr t a b [Expr t a b] where

  toColExpr = colSimplify . ColList



-- a boolean constant can be used as a constant boolean expression

instance ToBoolExpr tt cc bb Bool where

  toBoolExpr = BoolConst



-- boolean expressions are trivially convertible to boolean expressions

instance ToBoolExpr t a b (BoolExpr t a b) where

  toBoolExpr = id



-- the integer terms used by an expression can be used as interger expressions

instance ToExpr t a b t where

  toExpr = Term



-- the collections terms used by an expression can be used as collection expressions

instance ToColExpr t a b a where

  toColExpr = ColTerm



-- the boolean terms used by an expression can be used as boolean expressions

instance ToBoolExpr t a b b where

  toBoolExpr = BoolTerm



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

-- | integer operators/functions | --

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



-- @+ @- @* @/ @% are identical to + - * / % for integer expressions, except

-- that they also accept types convertible to expressions, instead of only

-- expressions themselves



infixl 6 @+, @-

infixl 7 @*

infixl 7 @/

infixl 7 @%



(@+) :: (Eq t, Eq c, Eq b, ToExpr t c b p, ToExpr t c b q) => p -> q -> Expr t c b

(@-) :: (Eq t, Eq c, Eq b, ToExpr t c b p, ToExpr t c b q) => p -> q -> Expr t c b

(@*) :: (Eq t, Eq c, Eq b, ToExpr t c b p, ToExpr t c b q) => p -> q -> Expr t c b

(@/) :: (Eq t, Eq c, Eq b, ToExpr t c b p, ToExpr t c b q) => p -> q -> Expr t c b

(@%) :: (Eq t, Eq c, Eq b, ToExpr t c b p, ToExpr t c b q) => p -> q -> Expr t c b



a @+ b = simplify $ (toExpr a) `Plus` (toExpr b)

a @- b = simplify $ (toExpr a) `Minus` (toExpr b)

a @* b = simplify $ (toExpr a) `Mult` (toExpr b)

a @/ b = simplify $ (toExpr a) `Div` (toExpr b)

a @% b = simplify $ (toExpr a) `Mod` (toExpr b)



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

-- | list operators/functions | --

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



infix 9 !

infix 9 @!!

infix 9 @..

infixr 5 @++

infix 4 @?

infix 4 @??

infix 5 @:



(!) :: (Eq t, Eq c, Eq b) => ColExpr t c b -> Expr t c b -> Expr t c b

(@!!) :: (Eq t, Eq c, Eq b) => ColExpr t c b -> Integer -> Expr t c b

(@..) :: (Eq t, Eq c, Eq b) => Expr t c b -> Expr t c b -> ColExpr t c b

(@++) :: (Eq t, Eq c, Eq b) => ColExpr t c b -> ColExpr t c b -> ColExpr t c b



(@?) :: (Eq t, Eq c, Eq b) => BoolExpr t c b -> (Expr t c b, Expr t c b) -> Expr t c b

c @? (t,f) = simplify $ Cond c t f



(@??) :: (Eq t, Eq c, Eq b) => BoolExpr t c b -> (BoolExpr t c b, BoolExpr t c b) -> BoolExpr t c b

c @?? (t,f) = boolSimplify $ BoolCond c t f



c!p = simplify $ At c p

c @!! p = simplify $ At c (Const p)

a @.. b = colSimplify $ ColRange (toExpr a) (toExpr b)

a @++ b = colSimplify $ ColCat (toColExpr a) (toColExpr b)



size :: (Eq t, Eq c, Eq b) => ColExpr t c b -> Expr t c b

size a = simplify $ ColSize a



xfold :: (Eq t, Eq c, Eq b) => (Expr t c b -> Expr t c b -> Expr t c b) -> Expr t c b -> ColExpr t c b -> Expr t c b

xfold f i c = simplify $ Fold (\a b -> f a b) i c



xsum :: (Num (Expr t c b), Eq t, Eq c, Eq b) => ColExpr t c b -> Expr t c b

xsum c = xfold (+) (Const 0) c



list :: (Eq t, Eq c, Eq b) => [Expr t c b] -> ColExpr t c b

list x = colSimplify $ ColList x



xhead :: (Eq t, Eq c, Eq b, ToColExpr t c b p) => p -> Expr t c b

xhead c = simplify $ At (toColExpr c) (Const 0)



xtail :: (Eq t, Eq c, Eq b, ToColExpr t c b p) => p -> ColExpr t c b

xtail c = let cc = toColExpr c in colSimplify $ ColSlice (\x -> simplify (x `Plus` (Const 1))) (simplify $ (size cc) `Minus` (Const 1)) cc



slice :: (Eq t, Eq c, Eq b) => ColExpr t c b -> ColExpr t c b -> ColExpr t c b

slice c p = case (c,p) of

  (_,ColRange l h) -> colSimplify $ ColSlice (\x -> simplify (l `Plus` x)) (simplify $ Const 1 `Plus` (simplify $ h `Minus` l)) c

  (_,ColMap f (ColRange l h)) -> colSimplify $ ColSlice (\i -> simplify $ f $ simplify (l `Plus` i)) (simplify $ Const 1 `Plus` (simplify $ h `Minus` l)) c

  (_,ColSlice f n c2) -> colSimplify $ ColSlice (\i -> simplify $ c2 `At` (f i)) n c

  _ -> xmap (\i -> simplify $ c `At` i) p



xmap :: (Eq t, Eq c, Eq b) => (Expr t c b -> Expr t c b) -> ColExpr t c b -> ColExpr t c b

xmap f c = colSimplify $ ColMap f c



loopall :: (Eq t, Eq c, Eq b) => (Expr t c b,Expr t c b) -> (Expr t c b -> BoolExpr t c b) -> BoolExpr t c b

loopall (l,h) f = boolSimplify $ BoolAll f $ colSimplify $ ColRange l h



loopany :: (Eq t, Eq c, Eq b) => (Expr t c b,Expr t c b) -> (Expr t c b -> BoolExpr t c b) -> BoolExpr t c b

loopany (l,h) f = boolSimplify $ BoolAny f $ colSimplify $ ColRange l h



forall :: (Eq t, Eq c, Eq b) => (ColExpr t c b) -> (Expr t c b -> BoolExpr t c b) -> BoolExpr t c b

forall c f = boolSimplify $ BoolAll f c



forany :: (Eq t, Eq c, Eq b) => (ColExpr t c b) -> (Expr t c b -> BoolExpr t c b) -> BoolExpr t c b

forany c f = boolSimplify $ BoolAny f c



channel :: (Eq t, Eq c, Eq b) => BoolExpr t c b -> Expr t c b

channel = simplify . Channel 



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

-- | boolean operators/functions | --

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



-- infixr 1 /\

-- infixr 1 \/

infixr 2 @||

infixr 3 @&&



-- (\/) :: (Eq t, Eq c, Eq b, ToBoolExpr t c b p, ToBoolExpr t c b q) => p -> q -> BoolExpr t c b

-- (/\) :: (Eq t, Eq c, Eq b, ToBoolExpr t c b p, ToBoolExpr t c b q) => p -> q -> BoolExpr t c b

inv :: (Eq t, Eq c, Eq b, ToBoolExpr t c b p) => p -> BoolExpr t c b



a @|| b = boolSimplify $ BoolOr (toBoolExpr a) (toBoolExpr b)

a @&& b = boolSimplify $ BoolAnd (toBoolExpr a) (toBoolExpr b)

inv a = boolSimplify $ BoolNot (toBoolExpr a)

-- a \/ b = a @|| b

-- a /\ b = a @&& b



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

-- | relational operators/functions | --

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



instance (Eq t, Eq c, Eq b) => ExprClass t c b (Expr t c b) where

  a @= b = boolSimplify $ Rel a EREqual b



instance (Eq t, Eq c, Eq b) => ExprClass t c b (BoolExpr t c b) where

  a @= b = boolSimplify $ BoolEqual a b



instance (Eq t, Eq c, Eq b) => ExprClass t c b (ColExpr t c b) where

  a @= b = boolSimplify $ ColEqual a b



  

infixr 4 @<,@<=,@>,@>=

(@<) ::  (Eq t, Eq c, Eq b) => Expr t c b -> Expr t c b -> BoolExpr t c b

(@>) ::  (Eq t, Eq c, Eq b) => Expr t c b -> Expr t c b -> BoolExpr t c b

(@<=) :: (Eq t, Eq c, Eq b) => Expr t c b -> Expr t c b -> BoolExpr t c b

(@>=) :: (Eq t, Eq c, Eq b) => Expr t c b -> Expr t c b -> BoolExpr t c b



a @< b = boolSimplify $ Rel a ERLess b

a @> b = boolSimplify $ Rel b ERLess a

a @<= b = boolSimplify $ Rel a ERLess (simplify $ b `Plus` (Const 1))

a @>= b = boolSimplify $ Rel b ERLess (simplify $ a `Plus` (Const 1))



sorted c = boolSimplify $ Sorted False c

sSorted c = boolSimplify $ Sorted True c

allDiff c = boolSimplify $ AllDiff False c

allDiffD c = boolSimplify $ AllDiff True c



instance (Eq t, Eq c, Eq b) => ExprRange t c b (Expr t c b,Expr t c b) where

  a @: (l,h) = (a @>= l) @&& (a @<= h)



instance (Eq t, Eq c, Eq b) => ExprRange t c b (ColExpr t c b) where

  a @: c = boolSimplify $ Dom a c