{-# LANGUAGE DeriveFoldable, DeriveFunctor, DeriveGeneric, DeriveTraversable,
             FlexibleInstances, TypeFamilies #-}

-- | This modules exposes a type representing the fixpoint of a functor type,
-- paired with an annotation. It exports the recursion-schemes instances for
-- this type and a few specific utility functions.
--
-- = Motivation
--
-- Suppose one has an Abstract Syntax Tree (AST) for the lambda calculus:
--
-- > data Term = Var String | App Term Term | Lam String Term
--
-- Such a type can easily be used with the recursion-schemes library, but it is
-- not always convenient to use this type. It is often the case that one wants
-- to add extra informations to every node of an AST, such as location or type
-- information. In this case, instead of adding those informations as an
-- extra-field to all data constructors, one could prefer to represent terms as
-- a record of a descriptor and the information present at every node, like so:
--
-- > data TermDesc = Var String | App Term Term | Lam String Term
-- > data Term = Term { termDesc :: TermDesc
-- >                  , termTyp  :: Typ
-- >                  , termLoc  :: Loc
-- >                  , ...
-- >                  }
--
-- This library implements this general pattern through the 'Annot' type,
-- representing the fixpoint of a functor type, paired with some annotation. In
-- this setting, the above example would be represented like so:
--
-- > data TermDesc r = Var String | App r r | Lam String r
-- >
-- > data TermAnn = TermAnn { termTyp :: Typ
-- >                        , termLoc :: Loc
-- >                        , ...
-- >                        }
-- >
-- > type Term = Annot TermDesc TermAnn

module Data.Functor.Annotated where

import Data.Functor.Foldable
import GHC.Generics

-- | The fixpoint type of functor @f@, with some annotation @a@.
data Annot f a = Annot a (f (Annot f a))
               deriving( Generic
                       , Functor
                       , Foldable
                       , Traversable
                       )

type instance Base (Annot f a) = f

instance Functor f => Recursive (Annot f a) where
  project = strip

-- ** Annotating and deannotating

-- | Strips one level of annotation
strip :: Annot f a -> f (Annot f a)
strip (Annot _ d) = d

-- | Extracts the annotation
annotation :: Annot f a -> a
annotation (Annot a _) = a

-- | Annotates all the node of a functor's fixpoint value
annotate :: Functor f => (f (Annot f a) -> a) -> Fix f -> Annot f a
annotate = annotateRec

-- | Generalized version of 'annotate' for instances of 'Recursive'
annotateRec :: Recursive t
            => (Base t (Annot (Base t) a) -> a)
            -> t
            -> Annot (Base t) a
annotateRec ann = cata aux
  where aux d = Annot (ann d) d

-- | Strips all annotations
deannotate :: Functor f => Annot f a -> Fix f
deannotate = deannotateCorec

-- | Generalized version of 'deannotate' for instances of 'Corecursive'
deannotateCorec :: Corecursive t => Annot (Base t) a -> t
deannotateCorec = embed . fmap deannotateCorec . strip

-- ** Specific schemes

-- | 'cata'morphism with access to the current annotation
cataAnn :: Functor f => (a -> f b -> b) -> Annot f a -> b
cataAnn fun = paraAnn (\a ffaab -> fun a (fmap snd ffaab))

-- | 'para'morphism with access to the current annotation
paraAnn :: Functor f
        => (a -> f (Annot f a, b) -> b)
        -> Annot f a
        -> b
paraAnn fun ann = para aux ann (annotation ann)
  where aux fafaab a =
          fun a (fmap (\(afa, ab) -> (afa, ab (annotation afa))) fafaab)