{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE Rank2Types            #-}
{-# LANGUAGE FlexibleInstances     #-}

module FoldEAlg where

import Except

import Generics.MultiRec.Base
import Generics.MultiRec.HFunctor
import Generics.MultiRec.HFix

-- * The type family of convenient algebras.

-- | The type family we use to describe the convenient algebras.
type family AlgE (f :: (* -> *) -> * -> *) 
                 (e :: *)           -- error type
                 (r :: * -> *)      -- recursive positions
                 (ix :: *)          -- index
                 :: *

-- | For a constant, we take the constant value to a result.
type instance AlgE (K a) e (r :: * -> *) ix = a -> Either e (r ix)

-- | For a unit, no arguments are available.
type instance AlgE U e (r :: * -> *) ix = Either e (r ix)

-- | For an identity, we turn the recursive result into a final result.
--   Note that the index can change.
type instance AlgE (I xi) e r ix = r xi -> Either e (r ix)

-- | For a sum, the algebra is a pair of two algebras.
type instance AlgE (f :+: g) e r ix = (AlgE f e r ix, AlgE g e r ix)

-- | For a product where the left hand side is a constant, we
--   take the value as an additional argument.
type instance AlgE (K a :*: g) e r ix = a -> AlgE g e r ix

-- | For a product where the left hand side is an identity, we
--   take the recursive result as an additional argument.
type instance AlgE (I xi :*: g) e r ix = r xi -> AlgE g e r ix

-- | A tag changes the index of the final result.
type instance AlgE (f :>: xi) e r ix = AlgE f e r xi

-- | Constructors are ignored.
type instance AlgE (C c f) e r ix = AlgE f e r ix

-- * The class to turn convenient algebras into standard algebras.

-- | The class fold explains how to convert a convenient algebra
--   'Alg' back into a function from functor to result, as required
--   by the standard fold function.
class Fold (f :: (* -> *) -> * -> *) where
  alg :: AlgE f e r ix -> f r ix -> Either e (r ix)

instance Fold (K a) where
  alg f (K x) = f x

instance Fold U where
  alg f U     = f

instance Fold (I xi) where
  alg f (I x) = f x

instance (Fold f, Fold g) => Fold (f :+: g) where
  alg (f, g) (L x) = alg f x
  alg (f, g) (R x) = alg g x

instance (Fold g) => Fold (K a :*: g) where
  alg f (K x :*: y) = alg (f x) y

instance (Fold g) => Fold (I xi :*: g) where
  alg f (I x :*: y) = alg (f x) y

instance (Fold f) => Fold (f :>: xi) where
  alg f (Tag x) = alg f x

instance (Fold f) => Fold (C c f) where
  alg f (C x) = alg f x

-- | The algebras passed to the fold have to work for all index types
--   in the family. The additional witness argument is required only
--   to make GHC's typechecker happy.
type ErrorAlgebra phi f e r = forall ix. phi ix -> AlgE f e r ix

foldE :: (HFunctor phi f, Fold f) => ErrorAlgebra phi f e r ->
  phi ix -> HFix (K x :*: f) ix -> Except [(e, x)] (r ix)
foldE a proof (HIn (K k :*: f)) =
  case hmapA (foldE a) proof f of
    Failed xs -> Failed xs
    OK expr' -> case alg (a proof) expr' of
      Left x'  -> Failed [(x', k)]
      Right v  -> OK v

infixr 5 &

-- | For constructing algebras that are made of nested pairs rather
--   than n-ary tuples, it is helpful to use this pairing combinator.
(&) :: a -> b -> (a, b)
(&) = (,)