{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Generics.Annotations.Yield (MonadYield(..), YieldT, Yield, runYield, runYieldG, runYieldT, runYieldTG) where

import Control.Monad.State
import Control.Monad.Identity

import Generics.MultiRec hiding (show)
import Generics.MultiRec.HFix

import Generics.Annotations.Annotations
import Generics.Annotations.Any

-- | Monads that allow yielding recursively annotated values.
class Monad m => MonadYield m where
  -- | Yielded values have types in this datatype family.
  type YieldFam m :: * -> *
  
  -- | The type of the annotation.
  type AnnType  m :: *
  
  -- | Yields a value with its annotation.
  yield :: (YieldFam m) ix -> AnnType m -> ix -> m ix


-- | The Yield transformer. Allows yielding generic values in family @fam@ with annotations of type @x@.
newtype YieldT x fam m a = YieldT (StateT [AnyAnn x fam] m a)
  deriving (Functor, Monad)

runYieldTG :: Monad m => YieldT x fam m a -> m (a, Maybe (AnyAnn x fam))
runYieldTG (YieldT y) = do
  (x, zs) <- runStateT y []
  return $ (,) x $ case zs of
    [z] -> Just z
    _   -> Nothing

runYieldT :: (Monad m, EqS fam) => fam a -> YieldT x fam m a -> m (Maybe (Ann x fam a))
runYieldT p y = do
  (_, x) <- runYieldTG y
  return (x >>= matchAnyF p)

runYieldG :: YieldT x fam Identity a -> (a, Maybe (AnyAnn x fam))
runYieldG = runIdentity . runYieldTG

runYield :: (EqS fam) => fam a -> YieldT x fam Identity a -> Maybe (Ann x fam a)
runYield p = runIdentity . runYieldT p

-- | Yield over the identity monad.
type Yield x fam = YieldT x fam Identity

instance MonadTrans (YieldT x fam) where
  lift = YieldT . lift


instance (Monad m, HFunctor fam (PF fam), EqS fam, Fam fam) => MonadYield (YieldT x fam m) where
  type YieldFam (YieldT x fam m) = fam
  type AnnType (YieldT x fam m) = x
  yield = doYield

doYield :: forall m fam x ix. (Monad m, HFunctor fam (PF fam), EqS fam, Fam fam) =>
  fam ix -> x -> ix -> YieldT x fam m ix
doYield p bounds x = YieldT $ do
  let pfx = from p x
  let n = length (children p pfx)
  stack <- get
  if (length stack < n)
    then do
      fail ("yield: required " ++ show n ++
        " accumulated children but found only " ++
        show (length stack))
    else do
      let (cs, cs') = splitAt n stack
      let newChild = evalState (hmapM distribute p pfx) (reverse cs)
      put (AnyF p (HIn (K bounds :*: newChild)) : cs')
      return x


distribute :: EqS s => s ix -> a -> State [AnyAnn x s] (Ann x s ix)
distribute w _ = do
  xs <- get
  case xs of
    [] -> error "structure mismatch: too few children"
    AnyF w' x : xs' -> case eqS w w' of
      Nothing   -> error "structure mismatch: incompatible child type"
      Just Refl -> do put xs'; return x