{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}

module Annotations (
    -- * Fixed points of functors
    Fix(..), compos,
    Algebra, cata,
    Coalgebra, ana,
    ErrorAlgebra, cascade,
    -- * Annotations
    Ann(..), AnnFix, AnnFix1, mkAnnFix, unannotate, errorCata
  ) where


import Except

import Data.Traversable
import Data.Monoid


-- | Fixpoint of functors.
newtype Fix fT  = In    { out      :: fT (Fix fT)  }

deriving instance Show (f (Fix f)) => Show (Fix f)
deriving instance Eq (f (Fix f)) => Eq (Fix f)

mapFix :: (f (Fix f) -> g (Fix g)) -> Fix f -> Fix g
mapFix f = In . f . out

-- | Lifted annotation of functors.
data Ann x f a = Ann x (f a)
  deriving (Eq, Show)

instance Functor f => Functor (Ann x f) where
  fmap f (Ann x t) = Ann x (fmap f t)

-- | A fully annotated tree.
type AnnFix   xT fT  = Fix (Ann xT fT)

-- | A functor with fully annotated children.
type AnnFix1  xT fT  = fT (AnnFix xT fT)

-- | Supply a tree with an annotation at the top level.
mkAnnFix :: x -> AnnFix1 x f -> AnnFix x f
mkAnnFix x = In . Ann x

-- | Recursively discard annotations.
unannotate :: Functor f => AnnFix x f -> Fix f
unannotate (In (Ann _ tree)) = In (fmap unannotate tree)

-- | Algebras for catamorphisms.
type Algebra fT aT = fT aT -> aT

-- | Reduces a tree to a value according to the algebra.
cata :: Functor fT => Algebra fT aT -> Fix fT -> aT
cata f = f . fmap (cata f) . out

-- | Coalgebras for anamorphisms.
type Coalgebra fT aT = aT -> fT aT

ana :: Functor fT => Coalgebra fT aT -> aT -> Fix fT
ana f = In . fmap (ana f) . f

-- | Apply a transformation to a tree's direct children.
compos :: Functor f => (Fix f -> Fix f) -> Fix f -> Fix f
compos = mapFix . fmap

-- | Algebras for error catamorphisms.
type ErrorAlgebra fT eT aT = fT aT -> Either eT aT

-- | Reduces a tree to a value according to the algebra, propagating potential errors.
cascade :: (Traversable fT, Monoid eT) =>
  ErrorAlgebra fT eT aT -> Algebra fT (Except eT aT)
cascade alg expr = case sequenceA expr of
  Failed xs  -> Failed xs
  OK tree'   -> case alg tree' of
    Left xs    -> Failed xs
    Right res  -> OK res

-- | Reduces a tree to a value according to the algebra, collecting potential
--   errors. The errors are combined with the annotations in the tree at the
--   positions at which the errors occurred.
errorCata :: Traversable fT => ErrorAlgebra fT eT aT -> AnnFix xT fT -> Except [(eT, xT)] aT
errorCata alg (In (Ann x expr)) =
  case traverse (errorCata alg) expr of
    Failed xs  -> Failed xs
    OK expr'   -> case alg expr' of
      Left x'   -> Failed [(x', x)]
      Right v   -> OK v