{-# LANGUAGE CPP #-}

-- | The core types of Fixplate.
module Data.Generics.Fixplate.Base where

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

import Control.Applicative
import Control.Monad ( liftM , ap )
import Data.Foldable
import Data.Traversable
import Prelude hiding ( foldl , foldr , mapM , mapM_ , concat , concatMap )

import Text.Show ()
import Text.Read

import Data.Generics.Fixplate.Misc

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

-- | The fixed-point type.
newtype Mu f = Fix { unFix :: f (Mu f) }

-- | We call a tree \"atomic\" if it has no subtrees.
isAtom :: Foldable f => Mu f -> Bool
isAtom = null . toList . unFix

--------------------------------------------------------------------------------
-- * Annotations

-- | Type of annotations
data Ann f a b = Ann
  { attr  :: a           -- ^ the annotation
  , unAnn :: f b         -- ^ the original functor
  }
  deriving (Eq,Ord,Show)

-- | Annotated fixed-point type. Equivalent to @CoFree f a@
type Attr f a = Mu (Ann f a)

-- | Lifting natural transformations to annotations.
liftAnn :: (f e -> g e) -> Ann f a e -> Ann g a e
liftAnn trafo (Ann a x) = Ann a (trafo x)

--------------------------------------------------------------------------------
-- * Co-annotations

-- | Categorical dual of 'Ann'.
data CoAnn f a b
  = Pure a
  | CoAnn (f b)
  deriving (Eq,Ord,Show)

-- | Categorical dual of 'Attr'. Equivalent to @Free f a@
type CoAttr f a = Mu (CoAnn f a)

-- | Lifting natural transformations to annotations.
liftCoAnn :: (f e -> g e) -> CoAnn f a e -> CoAnn g a e
liftCoAnn trafo x = case x of
  Pure  x -> Pure  x
  CoAnn t -> CoAnn (trafo t)

--------------------------------------------------------------------------------
-- * Annotated trees

-- | The attribute of the root node.
attribute :: Attr f a -> a
attribute = attr . unFix

-- | A function forgetting all the attributes from an annotated tree.
forget :: Functor f => Attr f a -> Mu f
forget = Fix . fmap forget . unAnn . unFix

--------------------------------------------------------------------------------
-- * Holes

-- | This a data type defined to be a place-holder for childs.
-- It is used in tree drawing, hashing, and 'Shape'.
-- 
-- It is deliberately not made an instance of 'Show', so that 
-- you can choose your preferred style. For example, an acceptable choice is
--
-- > instance Show Hole where show _ = "_"
--
data Hole = Hole deriving (Eq,Ord)

--------------------------------------------------------------------------------
-- * Higher-order type classes

-- | \"Functorised\" versions of standard type classes. 
-- If you have your a structure functor, for example
--
-- > Expr e 
-- >   = Kst Int 
-- >   | Var String 
-- >   | Add e e 
-- >   deriving (Eq,Ord,Read,Show,Functor,Foldable,Traversable)
--
-- you should make it an instance of these, so that the 
-- fixed-point type @Mu Expr@ can be an instance of
-- @Eq@, @Ord@ and @Show@. Doing so is very easy:
--
-- > instance EqF   Expr where equalF     = (==)
-- > instance OrdF  Expr where compareF   = compare
-- > instance ShowF Expr where showsPrecF = showsPrec
--
-- The @Read@ instance depends on whether we are using GHC or not.
-- The Haskell98 version is
--
-- > instance ReadF Expr where readsPrecF = readsPrec
--
-- while the GHC version is
--
-- > instance ReadF Expr where readPrecF  = readPrec
--
class          EqF   f  where equalF     ::  Eq   a  =>  f a -> f a -> Bool
class EqF f => OrdF  f  where compareF   ::  Ord  a  =>  f a -> f a -> Ordering
class          ShowF f  where showsPrecF ::  Show a  =>  Int -> f a -> ShowS
class          ReadF f  where
#ifdef __GLASGOW_HASKELL__
                              readPrecF  ::  Read a  =>  ReadPrec (f a)
#else
                              readsPrecF ::  Read a  =>  Int -> ReadS (f a)
#endif                        

showF :: (ShowF f, Show a) => f a -> String
showF x = showsF x ""

showsF :: (ShowF f, Show a) => f a -> ShowS
showsF = showsPrecF 0

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

instance EqF f => Eq (Mu f) where Fix x == Fix y = equalF x y
instance OrdF   f  => Ord   (Mu f) where compare (Fix x) (Fix y)  = compareF x y
instance ShowF  f  => Show  (Mu f) where
  showsPrec d (Fix x) = showParen (d>app_prec)
    $ showString "Fix "
    . showsPrecF (app_prec+1) x

instance ReadF  f  => Read  (Mu f) where
#ifdef __GLASGOW_HASKELL__
  readPrec = parens $
    (prec app_prec $ do
      { Ident "Fix" <- lexP
      ; m <- step readPrecF
      ; return (Fix m)
      })
#else                                  
  readsPrec d r = readParen (d > app_prec)
     (\r -> [ (Fix m, t)
            | ("Fix", s) <- lex r
            , (m,t) <- readsPrecF (app_prec+1) s]) r
#endif

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

-- | NOTE: The 'EqF' instance for annotations compares both the annotations and the original part. 
instance (Eq a, EqF f) => EqF (Ann f a) where
  equalF (Ann a x) (Ann b y) = a == b && equalF x y

-- | NOTE: The 'OrdF' instance for annotations first compares the /annotations/, and then 
-- the functor part. If this is not the desired behaviour (it's not clear to me at the moment
-- what is the good default here), you can use the standard newtype trick to define a new
-- behaviour.
instance (Ord a, OrdF f) => OrdF (Ann f a) where
  compareF (Ann a x) (Ann b y) = case compare a b of
    LT -> LT
    GT -> GT
    EQ -> compareF x y

instance (Show a, ShowF f) => ShowF (Ann f a) where
  showsPrecF d (Ann a t)
    = showParen (d>app_prec)
      $ showString "Ann "
      . (showsPrec (app_prec+1) a)
      . showChar ' '
      . (showsPrecF (app_prec+1) t)

instance (Read a, ReadF f) => ReadF (Ann f a) where
#ifdef __GLASGOW_HASKELL__
  readPrecF = parens $
    (prec app_prec $ do
      { Ident "Ann" <- lexP
      ; x <- step readPrec
      ; m <- step readPrecF
      ; return (Ann x m)
      })
#else                                  
  readsPrecF d r = readParen (d > app_prec)
     (\r -> [ (Ann x m, u)
            | ("Ann", s) <- lex r
            , (x,t) <- readsPrec  (app_prec+1) s]) r
            , (m,u) <- readsPrecF (app_prec+1) t]) r
#endif

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

instance (Eq a, EqF f) => EqF (CoAnn f a) where
  equalF (Pure  a) (Pure  b) = a == b
  equalF (CoAnn x) (CoAnn y) = equalF x y
  equalF _         _         = False

instance (Ord a, OrdF f) => OrdF (CoAnn f a) where
  compareF (Pure  a) (Pure  b) = compare a b
  compareF (CoAnn x) (CoAnn y) = compareF x y
  compareF (Pure  _) (CoAnn _) = LT
  compareF (CoAnn _) (Pure  _) = GT

instance (Show a, ShowF f) => ShowF (CoAnn f a) where
  showsPrecF d (CoAnn t)
    = showParen (d>app_prec)
      $ showString "CoAnn "
      . (showsPrecF (app_prec+1) t)
  showsPrecF d (Pure x)
    = showParen (d>app_prec)
      $ showString "Pure "
      . (showsPrec (app_prec+1) x)

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

instance Functor f => Functor (Ann f a) where
  fmap f (Ann attr t) = Ann attr (fmap f t)

instance Foldable f => Foldable (Ann f a) where
  foldl f x (Ann _ t) = foldl f x t
  foldr f x (Ann _ t) = foldr f x t

instance Traversable f => Traversable (Ann f a) where
  traverse f (Ann x t) = Ann x <$> traverse f t
  mapM f (Ann x t) = liftM (Ann x) (mapM f t)

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

instance Functor f => Functor (CoAnn f a) where
  fmap f (CoAnn t) = CoAnn (fmap f t)
  fmap f (Pure  x) = Pure  x

instance Foldable f => Foldable (CoAnn f a) where
  foldl f a (CoAnn t) = foldl f a t
  foldl f a (Pure  x) = a

  foldr f a (CoAnn t) = foldr f a t
  foldr f a (Pure  x) = a

instance Traversable f => Traversable (CoAnn f a) where
  traverse f (CoAnn t) = CoAnn <$> traverse f t
  traverse f (Pure  x) = pure (Pure x)

  mapM f (CoAnn t) = liftM CoAnn (mapM f t)
  mapM f (Pure  x) = return (Pure x)

--------------------------------------------------------------------------------
-- * Attrib (cofree comonad)

-- | A newtype wrapper around @Attr f a@ so that we can make @Attr f@ 
-- an instance of Functor, Foldable and Traversable (and Comonad). This is necessary
-- since Haskell does not allow partial application of type synonyms.
--
-- Equivalent to the co-free comonad.
newtype Attrib f a = Attrib { unAttrib :: Attr f a }

instance (ShowF f, Show a) => Show (Attrib f a) where
  showsPrec d (Attrib x)
    = showParen (d>app_prec)
      $ showString "Attrib "
      . (showsPrec (app_prec+1) x)

instance Functor f => Functor (Attrib f) where
  fmap h (Attrib y) = Attrib (go y) where
    go (Fix (Ann x t)) = Fix $ Ann (h x) (fmap go t)

instance Foldable f => Foldable (Attrib f) where
  foldl h a (Attrib y) = go a y where go b (Fix (Ann x t)) = foldl go (h b x) t
  foldr h a (Attrib y) = go y a where go (Fix (Ann x t)) b = h x (foldr go b t)

instance Traversable f => Traversable (Attrib f) where
  traverse h (Attrib y) = Attrib <$> go y where
    go (Fix (Ann x t)) = Fix <$> (Ann <$> h x <*> traverse go t)

--------------------------------------------------------------------------------
-- * CoAttrib (free monad)

-- | Categorial dual of 'Attrib'. Equivalent to the free monad.
newtype CoAttrib f a = CoAttrib { unCoAttrib :: CoAttr f a }

instance (ShowF f, Show a) => Show (CoAttrib f a) where
  showsPrec d (CoAttrib x)
    = showParen (d>app_prec)
      $ showString "CoAttrib "
      . (showsPrec (app_prec+1) x)

instance Functor f => Functor (CoAttrib f) where
  fmap h (CoAttrib y) = CoAttrib (go y) where
    go (Fix (CoAnn t)) = Fix $ CoAnn (fmap go t)
    go (Fix (Pure  x)) = Fix $ Pure  (h x)

instance Foldable f => Foldable (CoAttrib f) where
  foldl h a (CoAttrib y) = go a y where
     go b (Fix (CoAnn t)) = foldl go b t
     go b (Fix (Pure  x)) = h b x
  foldr h a (CoAttrib y) = go y a where
     go (Fix (CoAnn t)) b = foldr go b t
     go (Fix (Pure  x)) b = h x b

instance Traversable f => Traversable (CoAttrib f) where
  traverse h (CoAttrib y) = CoAttrib <$> go y where
    go (Fix (CoAnn t)) = Fix <$> (CoAnn <$> traverse go t)
    go (Fix (Pure  x)) = Fix <$> (Pure  <$> h x)

instance Functor f => Applicative (CoAttrib f) where
  pure x = CoAttrib (Fix (Pure x))
  (<*>)  = ap

instance Functor f => Monad (CoAttrib f) where
  return x = CoAttrib (Fix (Pure x))
  CoAttrib (Fix (CoAnn t))  >>=  u  =  CoAttrib (Fix (CoAnn (fmap (unCoAttrib . (>>=u) . CoAttrib) t)))
  CoAttrib (Fix (Pure  x))  >>=  u  =  u x


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