{-# LANGUAGE UndecidableInstances #-}

module Cubix.Sin.Compdata.Annotation (
    Annotated(..)
  , getAnn
  , MonadAnnotater(..)
  , AnnotateDefault
  , pattern AnnotateDefault
  , runAnnotateDefault
  , annotateM
  , propAnnSigFun
  , annotateOuter
  ) where

import Control.Monad.Identity ( Identity(..) )
import Control.Monad.Trans ( MonadTrans(..) )
import Data.Default ( Default(..) )
import Data.Comp.Multi ( All, caseCxt, Node, Cxt(..), Context, (:=>), CxtFunM, SigFun, appSigFunM, appCxt, HFix, AnnHFix )
import Data.Comp.Multi.HTraversable ( HTraversable )
import Data.Comp.Multi.Ops ((:&:)(..), Sum (..))

import Cubix.Sin.Compdata.Instances ()

--------------------------------------------------------------------------------------


-------------------------------------------------------------------
------------------------- Getting annotations ---------------------
-------------------------------------------------------------------

---- This exists so you can constrain a functor to be annotated without also
---- naming the unannotated functor. This makes it more convenient for inclusion
---- in constraint synonyms
class Annotated a (f :: Node) | f -> a where
  getAnn' :: f e l -> a

instance Annotated a (f :&: a) where
  getAnn' :: forall (e :: Family) (l :: Sort). (:&:) f a e l -> a
getAnn' (f e l
_ :&: a
x) = a
x

-- The funny constraint is to get the functional dependency working
instance ( Annotated a f
         , All (Annotated a) fs
         ) => Annotated a (Sum (f ': fs)) where
  getAnn' :: forall (e :: Family) (l :: Sort). Sum (f : fs) e l -> a
getAnn' = forall (cxt :: Node -> Constraint) (fs :: [Node]) (a :: Family)
       (e :: Sort) (b :: Sort).
All cxt fs =>
(forall (f :: Node). cxt f => f a e -> b) -> Sum fs a e -> b
caseCxt @(Annotated a) f e l -> a
forall (a :: Sort) (f :: Node) (e :: Family) (l :: Sort).
Annotated a f =>
f e l -> a
forall (e :: Family) (l :: Sort). f e l -> a
forall (f :: Node). Annotated a f => f e l -> a
getAnn'


getAnn :: (Annotated a f) => HFix f :=> a
getAnn :: forall (a :: Sort) (f :: Node). Annotated a f => HFix f :=> a
getAnn (Term f (Cxt NoHole f (K ())) i
x) = f (Cxt NoHole f (K ())) i -> a
forall (a :: Sort) (f :: Node) (e :: Family) (l :: Sort).
Annotated a f =>
f e l -> a
forall (e :: Family) (l :: Sort). f e l -> a
getAnn' f (Cxt NoHole f (K ())) i
x

-------------------------------------------------------------------
------------------------- Adding annotations ---------------------
-------------------------------------------------------------------

-----------------------------------
------------------ MonadAnnotater
------------------------------------

-------------
--- Class
-------------

class (Monad m) => MonadAnnotater a m where
  annM :: forall (f :: Node) e l. f e l -> m ((f :&: a) e l)

instance (MonadAnnotater a m, MonadTrans t, Monad (t m)) => MonadAnnotater a (t m) where
  annM :: forall (f :: Node) (e :: Family) (l :: Sort).
f e l -> t m ((:&:) f a e l)
annM = m ((:&:) f a e l) -> t m ((:&:) f a e l)
forall (m :: Family) (a :: Sort). Monad m => m a -> t m a
forall (t :: Node) (m :: Family) (a :: Sort).
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m ((:&:) f a e l) -> t m ((:&:) f a e l))
-> (f e l -> m ((:&:) f a e l)) -> f e l -> t m ((:&:) f a e l)
forall (b :: Sort) (c :: Sort) (a :: Sort).
(b -> c) -> (a -> b) -> a -> c
. f e l -> m ((:&:) f a e l)
forall (a :: Sort) (m :: Family) (f :: Node) (e :: Family)
       (l :: Sort).
MonadAnnotater a m =>
f e l -> m ((:&:) f a e l)
forall (f :: Node) (e :: Family) (l :: Sort).
f e l -> m ((:&:) f a e l)
annM

-------------
--- Operations
-------------

annotateM :: (HTraversable f, MonadAnnotater a m) => CxtFunM m f (f :&: a)
annotateM :: forall (f :: Node) (a :: Sort) (m :: Family).
(HTraversable f, MonadAnnotater a m) =>
CxtFunM m f (f :&: a)
annotateM = SigFunM m f (f :&: a) -> CxtFunM m f (f :&: a)
forall (f :: Node) (g :: Node) (m :: Family).
(HTraversable f, Monad m) =>
SigFunM m f g -> CxtFunM m f g
appSigFunM f a i -> m ((:&:) f a a i)
forall (a :: Sort) (m :: Family) (f :: Node) (e :: Family)
       (l :: Sort).
MonadAnnotater a m =>
f e l -> m ((:&:) f a e l)
SigFunM m f (f :&: a)
forall (f :: Node) (e :: Family) (l :: Sort).
f e l -> m ((:&:) f a e l)
annM

----------------------------------
------------------- AnnotateDefault
----------------------------------

newtype AnnotateDefault x = AnnotateDefault' { forall (x :: Sort). AnnotateDefault x -> Identity x
runAnnotateDefault' :: Identity x}
  deriving ( (forall (a :: Sort) (b :: Sort).
 (a -> b) -> AnnotateDefault a -> AnnotateDefault b)
-> (forall (a :: Sort) (b :: Sort).
    a -> AnnotateDefault b -> AnnotateDefault a)
-> Functor AnnotateDefault
forall (a :: Sort) (b :: Sort).
a -> AnnotateDefault b -> AnnotateDefault a
forall (a :: Sort) (b :: Sort).
(a -> b) -> AnnotateDefault a -> AnnotateDefault b
forall (f :: Family).
(forall (a :: Sort) (b :: Sort). (a -> b) -> f a -> f b)
-> (forall (a :: Sort) (b :: Sort). a -> f b -> f a) -> Functor f
$cfmap :: forall (a :: Sort) (b :: Sort).
(a -> b) -> AnnotateDefault a -> AnnotateDefault b
fmap :: forall (a :: Sort) (b :: Sort).
(a -> b) -> AnnotateDefault a -> AnnotateDefault b
$c<$ :: forall (a :: Sort) (b :: Sort).
a -> AnnotateDefault b -> AnnotateDefault a
<$ :: forall (a :: Sort) (b :: Sort).
a -> AnnotateDefault b -> AnnotateDefault a
Functor, Functor AnnotateDefault
Functor AnnotateDefault =>
(forall (a :: Sort). a -> AnnotateDefault a)
-> (forall (a :: Sort) (b :: Sort).
    AnnotateDefault (a -> b) -> AnnotateDefault a -> AnnotateDefault b)
-> (forall (a :: Sort) (b :: Sort) (c :: Sort).
    (a -> b -> c)
    -> AnnotateDefault a -> AnnotateDefault b -> AnnotateDefault c)
-> (forall (a :: Sort) (b :: Sort).
    AnnotateDefault a -> AnnotateDefault b -> AnnotateDefault b)
-> (forall (a :: Sort) (b :: Sort).
    AnnotateDefault a -> AnnotateDefault b -> AnnotateDefault a)
-> Applicative AnnotateDefault
forall (a :: Sort). a -> AnnotateDefault a
forall (a :: Sort) (b :: Sort).
AnnotateDefault a -> AnnotateDefault b -> AnnotateDefault a
forall (a :: Sort) (b :: Sort).
AnnotateDefault a -> AnnotateDefault b -> AnnotateDefault b
forall (a :: Sort) (b :: Sort).
AnnotateDefault (a -> b) -> AnnotateDefault a -> AnnotateDefault b
forall (a :: Sort) (b :: Sort) (c :: Sort).
(a -> b -> c)
-> AnnotateDefault a -> AnnotateDefault b -> AnnotateDefault c
forall (f :: Family).
Functor f =>
(forall (a :: Sort). a -> f a)
-> (forall (a :: Sort) (b :: Sort). f (a -> b) -> f a -> f b)
-> (forall (a :: Sort) (b :: Sort) (c :: Sort).
    (a -> b -> c) -> f a -> f b -> f c)
-> (forall (a :: Sort) (b :: Sort). f a -> f b -> f b)
-> (forall (a :: Sort) (b :: Sort). f a -> f b -> f a)
-> Applicative f
$cpure :: forall (a :: Sort). a -> AnnotateDefault a
pure :: forall (a :: Sort). a -> AnnotateDefault a
$c<*> :: forall (a :: Sort) (b :: Sort).
AnnotateDefault (a -> b) -> AnnotateDefault a -> AnnotateDefault b
<*> :: forall (a :: Sort) (b :: Sort).
AnnotateDefault (a -> b) -> AnnotateDefault a -> AnnotateDefault b
$cliftA2 :: forall (a :: Sort) (b :: Sort) (c :: Sort).
(a -> b -> c)
-> AnnotateDefault a -> AnnotateDefault b -> AnnotateDefault c
liftA2 :: forall (a :: Sort) (b :: Sort) (c :: Sort).
(a -> b -> c)
-> AnnotateDefault a -> AnnotateDefault b -> AnnotateDefault c
$c*> :: forall (a :: Sort) (b :: Sort).
AnnotateDefault a -> AnnotateDefault b -> AnnotateDefault b
*> :: forall (a :: Sort) (b :: Sort).
AnnotateDefault a -> AnnotateDefault b -> AnnotateDefault b
$c<* :: forall (a :: Sort) (b :: Sort).
AnnotateDefault a -> AnnotateDefault b -> AnnotateDefault a
<* :: forall (a :: Sort) (b :: Sort).
AnnotateDefault a -> AnnotateDefault b -> AnnotateDefault a
Applicative, Applicative AnnotateDefault
Applicative AnnotateDefault =>
(forall (a :: Sort) (b :: Sort).
 AnnotateDefault a -> (a -> AnnotateDefault b) -> AnnotateDefault b)
-> (forall (a :: Sort) (b :: Sort).
    AnnotateDefault a -> AnnotateDefault b -> AnnotateDefault b)
-> (forall (a :: Sort). a -> AnnotateDefault a)
-> Monad AnnotateDefault
forall (a :: Sort). a -> AnnotateDefault a
forall (a :: Sort) (b :: Sort).
AnnotateDefault a -> AnnotateDefault b -> AnnotateDefault b
forall (a :: Sort) (b :: Sort).
AnnotateDefault a -> (a -> AnnotateDefault b) -> AnnotateDefault b
forall (m :: Family).
Applicative m =>
(forall (a :: Sort) (b :: Sort). m a -> (a -> m b) -> m b)
-> (forall (a :: Sort) (b :: Sort). m a -> m b -> m b)
-> (forall (a :: Sort). a -> m a)
-> Monad m
$c>>= :: forall (a :: Sort) (b :: Sort).
AnnotateDefault a -> (a -> AnnotateDefault b) -> AnnotateDefault b
>>= :: forall (a :: Sort) (b :: Sort).
AnnotateDefault a -> (a -> AnnotateDefault b) -> AnnotateDefault b
$c>> :: forall (a :: Sort) (b :: Sort).
AnnotateDefault a -> AnnotateDefault b -> AnnotateDefault b
>> :: forall (a :: Sort) (b :: Sort).
AnnotateDefault a -> AnnotateDefault b -> AnnotateDefault b
$creturn :: forall (a :: Sort). a -> AnnotateDefault a
return :: forall (a :: Sort). a -> AnnotateDefault a
Monad )

pattern AnnotateDefault :: x -> AnnotateDefault x
pattern $mAnnotateDefault :: forall {r} {x :: Sort}.
AnnotateDefault x -> (x -> r) -> ((# #) -> r) -> r
$bAnnotateDefault :: forall (a :: Sort). a -> AnnotateDefault a
AnnotateDefault x = AnnotateDefault' (Identity x)

runAnnotateDefault :: AnnotateDefault (AnnHFix a f l) -> AnnHFix a f l
runAnnotateDefault :: forall (a :: Sort) (f :: Node) (l :: Sort).
AnnotateDefault (AnnHFix a f l) -> AnnHFix a f l
runAnnotateDefault = Identity (AnnHFix a f l) -> AnnHFix a f l
forall (a :: Sort). Identity a -> a
runIdentity (Identity (AnnHFix a f l) -> AnnHFix a f l)
-> (AnnotateDefault (AnnHFix a f l) -> Identity (AnnHFix a f l))
-> AnnotateDefault (AnnHFix a f l)
-> AnnHFix a f l
forall (b :: Sort) (c :: Sort) (a :: Sort).
(b -> c) -> (a -> b) -> a -> c
. AnnotateDefault (AnnHFix a f l) -> Identity (AnnHFix a f l)
forall (x :: Sort). AnnotateDefault x -> Identity x
runAnnotateDefault'

-- | Specializing annotation to Maybe a to aid instance selection
instance (Default a) => MonadAnnotater a AnnotateDefault where
  annM :: forall (f :: Node) (e :: Family) (l :: Sort).
f e l -> AnnotateDefault ((:&:) f a e l)
annM f e l
x = (:&:) f a e l -> AnnotateDefault ((:&:) f a e l)
forall (a :: Sort). a -> AnnotateDefault a
forall (m :: Family) (a :: Sort). Monad m => a -> m a
return (f e l
x f e l -> a -> (:&:) f a e l
forall (f :: Node) (a :: Sort) (g :: Family) (e :: Sort).
f g e -> a -> (:&:) f a g e
:&: a
forall (a :: Sort). Default a => a
def)



-----------------------------------
------------------ Lifting other functions to annotated versions
------------------------------------

propAnnSigFun :: SigFun f g -> SigFun (f :&: a) (g :&: a)
propAnnSigFun :: forall (f :: Node) (g :: Node) (a :: Sort).
SigFun f g -> SigFun (f :&: a) (g :&: a)
propAnnSigFun SigFun f g
f (f a i
t :&: a
a) = (f a i -> g a i
SigFun f g
f f a i
t) g a i -> a -> (:&:) g a a i
forall (f :: Node) (a :: Sort) (g :: Family) (e :: Sort).
f g e -> a -> (:&:) f a g e
:&: a
a

annotateOuter :: (HTraversable f, MonadAnnotater a m) => Context f (AnnHFix a f) l -> m (AnnHFix a f l)
annotateOuter :: forall (f :: Node) (a :: Sort) (m :: Family) (l :: Sort).
(HTraversable f, MonadAnnotater a m) =>
Context f (AnnHFix a f) l -> m (AnnHFix a f l)
annotateOuter = (Context (f :&: a) (Cxt NoHole (f :&: a) (K ())) l
 -> AnnHFix a f l)
-> m (Context (f :&: a) (Cxt NoHole (f :&: a) (K ())) l)
-> m (AnnHFix a f l)
forall (a :: Sort) (b :: Sort). (a -> b) -> m a -> m b
forall (f :: Family) (a :: Sort) (b :: Sort).
Functor f =>
(a -> b) -> f a -> f b
fmap Context (f :&: a) (Cxt NoHole (f :&: a) (K ())) l -> AnnHFix a f l
Context (f :&: a) (Cxt NoHole (f :&: a) (K ()))
:-> Cxt NoHole (f :&: a) (K ())
forall (f :: Node) (h :: Sort) (a :: Family).
HFunctor f =>
Context f (Cxt h f a) :-> Cxt h f a
appCxt (m (Context (f :&: a) (Cxt NoHole (f :&: a) (K ())) l)
 -> m (AnnHFix a f l))
-> (Context f (Cxt NoHole (f :&: a) (K ())) l
    -> m (Context (f :&: a) (Cxt NoHole (f :&: a) (K ())) l))
-> Context f (Cxt NoHole (f :&: a) (K ())) l
-> m (AnnHFix a f l)
forall (b :: Sort) (c :: Sort) (a :: Sort).
(b -> c) -> (a -> b) -> a -> c
. Context f (Cxt NoHole (f :&: a) (K ())) l
-> m (Context (f :&: a) (Cxt NoHole (f :&: a) (K ())) l)
CxtFunM m f (f :&: a)
forall (f :: Node) (a :: Sort) (m :: Family).
(HTraversable f, MonadAnnotater a m) =>
CxtFunM m f (f :&: a)
annotateM