{-# LANGUAGE DataKinds, FlexibleContexts, FlexibleInstances,
             MultiParamTypeClasses, RankNTypes, TypeFamilies, TypeOperators,
             UndecidableInstances #-}
module Control.Arrow.Elision.Flexible
       ( Nav
       , Conv(..)

         -- * Flexible Elision completion functions
       , complete
       , complete'
       , unelide
       , unelide'

         -- * Re-exports from 'Control.Arrow.Elision.Simple'
       , module Control.Arrow.Elision.Simple
       ) where

import Control.Arrow.Elision.Simple hiding (complete, complete', unelide,
                                     unelide')

import qualified Control.Arrow.Elision.Simple as Simple (unelide)

--------------------------------------------------------------------------------
-- | Determines the transformation requires to take some interpreter @f@ and to
-- transform it into another interpreter @g@. This is used primarily as the
-- kind for 'Nav'.
data Direction
  = L Direction           -- * In @Nav f (g // h)@, @g@ has a common sub expression.
  | R Direction           -- * In @Nav f (g // h)@, @h@ has a common sub expression.
  | X Direction Direction -- * Subexpressions exist for both @f0@ and @f1@ in @g // h@ in @Nav (f0 // f1) (g // h)@
  | Equiv                 -- * In @Nav f g@, @f ~ g@.
  | Divergent             -- * There are no common subtypes.

--------------------------------------------------------------------------------
-- | "Navigate" down a tree of types to find common types which are equivalent.
-- See 'Direction' for the result.
type family Nav f g :: Direction where
    Nav  f        f       = 'Equiv
    Nav (f // g) (h // i) = 'X (Nav f (h // i)) (Nav g (h // i))
    Nav  f       (h // i) =  Choose (Converges (Nav f h)) (Nav f h) (Converges (Nav f i)) (Nav f i)
    Nav  f        g       = 'Divergent

--------------------------------------------------------------------------------
-- | A constraint that can only be satisfied if a direction terminates in
-- 'Equiv'.
class 'True ~ Converges a => Convergent (a :: Direction)
instance Convergent 'Equiv
instance Convergent a => Convergent ('L a)
instance Convergent a => Convergent ('R a)
instance (Convergent a, Convergent b) => Convergent ('X a b)

--------------------------------------------------------------------------------
-- A type level function which determines if @a@ terminates in 'Equiv'.
--
-- Because 'True ~ Converges a' can be cumbersome to type, this module also
-- provides an equivalent typeclass, 'Convergent', which can be used as a
-- constraint on a function.
type family Converges a where
    Converges  'Equiv     = 'True
    Converges  'Divergent = 'False
    Converges ('L k)      =  Converges k
    Converges ('R k)      =  Converges k
    Converges ('X k1 k2)  =  And (Converges k1) (Converges k2)

type family And a b where
    And 'True 'True = 'True
    And  a      b   = 'False

--------------------------------------------------------------------------------
-- Pass in whether or not the function converges along with the type, and
-- return either @a@ or @b@ if they converge, favoring @a@.
type family Choose (exA :: Bool) (a :: Direction) (exB :: Bool) (b :: Direction) where
    Choose 'True  a exB    b = 'L a
    Choose 'False a 'True  b = 'R b
    Choose 'False a 'False b = 'Divergent

--------------------------------------------------------------------------------
-- | In @Conv k f g@, the interpreter for @f@ is completely covered by the
-- interpreter for @g@, so @g@ can be substituted for @f@ with 'conv'.
class (k ~ Nav f g, Convergent k) => Conv k f g where
    conv :: (g a -> m a) -> f a -> m a

instance Conv 'Equiv f f where
    conv = id

instance ('L k ~ Nav f (g // h), Conv k f g) => Conv ('L k) f (g // h) where
    conv f = conv (f . Sum . Left)

instance ('R k ~ Nav f (g // h), Conv k f h) => Conv ('R k) f (g // h) where
    conv f = conv (f . Sum . Right)

instance ( 'X k0 k1 ~ Nav (e // f) (g // h)
         , Conv k0 e (g // h)
         , Conv k1 f (g // h)
         ) => Conv ('X k0 k1) (e // f) (g // h) where
    conv f = conv f ||| conv f ^<< runSum

--------------------------------------------------------------------------------
-- | A flexible version of 'Control.Arrow.Elision.Simple.unelide'
unelide :: (Monad m, Conv k f g) => Elision f a b -> (forall c. g c -> m c) -> a -> m b
unelide el fn = Simple.unelide el (conv fn)

--------------------------------------------------------------------------------
-- | A flexible version of 'Control.Arrow.Elision.Simple.unelide''
unelide' :: (Monad m, Conv k f g) => Elision f () b -> (forall c. g c -> m c) -> m b
unelide' el fn = unelide el fn ()

--------------------------------------------------------------------------------
-- | A flexible version of 'Control.Arrow.Elision.Simple.complete'
complete :: (Monad m, Conv k f g) => (forall c. g c -> m c) -> a -> Elision f a b -> m b
complete fn x el = unelide el fn x

--------------------------------------------------------------------------------
-- | A flexible version of 'Control.Arrow.Elision.Simple.complete''
complete' :: (Monad m, Conv k f g) => (forall c. g c -> m c) -> Elision f () b -> m b
complete' fn el = unelide' el fn