{-# LANGUAGE AllowAmbiguousTypes   #-} -- For isNode'
{-# LANGUAGE ConstraintKinds       #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms       #-}
{-# LANGUAGE Rank2Types            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeApplications      #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE ViewPatterns          #-}
--------------------------------------------------------------------------------
-- |
-- Module      :  Data.Comp.Multi.Annotation
-- Copyright   :  (c) 2011 Patrick Bahr
-- License     :  BSD3
-- Maintainer  :  Patrick Bahr <paba@diku.dk>
-- Stability   :  experimental
-- Portability :  non-portable (GHC Extensions)
--
-- This module defines annotations on signatures. All definitions are
-- generalised versions of those in "Data.Comp.Annotation".
--
--------------------------------------------------------------------------------

module Data.Comp.Multi.Annotation
    (
     AnnTerm,
     AnnHFix,
     (:&:) (..),
     RemA (..),
     liftA,
     ann,
     liftA',
     stripA,
     propAnn,
     project',
     pattern (::&::),
     isNode',
     inj',
     inject',
     injectOpt,
     caseH',
     caseCxt',
     caseCxt'',
     DistAnn,
     AnnCxt,
     AnnCxtS
    ) where

import Data.Comp.Dict
import Data.Comp.Elem
import Data.Comp.Multi.Algebra
import Data.Comp.Multi.HFunctor
import Data.Comp.Multi.Kinds
import Data.Comp.Multi.Ops
import Data.Comp.Multi.Sum
import Data.Comp.Multi.Term

type AnnHFix a f = HFix (f :&: a)
type AnnTerm a fs = HFix (Sum fs :&: a)

type AnnCxt p h f a = Cxt h (f :&: p) a
-- type AnnContext p f a = AnnCxt p Hole f a

type AnnCxtS p h fs a = AnnCxt p h (Sum fs) a
-- type AnnContextS p fs a = AnnContext p (Sum fs) a

-- | This function transforms a function with a domain constructed
-- from a functor to a function with a domain constructed with the
-- same functor but with an additional annotation.
liftA :: (RemA s s') => (s' a :-> t) -> s a :-> t
liftA :: forall (s :: Node) (s' :: Node) (a :: Family) (t :: Family).
RemA s s' =>
(s' a :-> t) -> s a :-> t
liftA s' a :-> t
f s a i
v = s' a i -> t i
s' a :-> t
f (s a i -> s' a i
forall (a :: Family). s a :-> s' a
forall (s :: Node) (s' :: Node) (a :: Family).
RemA s s' =>
s a :-> s' a
remA s a i
v)

-- | This function annotates each sub term of the given term with the
-- given value (of type a).

ann :: (HFunctor f) => p -> CxtFun f (f :&: p)
ann :: forall (f :: Node) p. HFunctor f => p -> CxtFun f (f :&: p)
ann p
c = SigFun f (f :&: p) -> CxtFun f (f :&: p)
forall (f :: Node) (g :: Node).
HFunctor f =>
SigFun f g -> CxtFun f g
appSigFun (f a i -> p -> (:&:) f p a i
forall (f :: Node) a (g :: Family) e. f g e -> a -> (:&:) f a g e
:&: p
c)

-- | This function transforms a function with a domain constructed
-- from a functor to a function with a domain constructed with the
-- same functor but with an additional annotation.
liftA' :: (HFunctor s)
       => (s a :-> Cxt h s a) -> (s :&: p) a :-> Cxt h (s :&: p) a
liftA' :: forall (s :: Node) (a :: Family) h p.
HFunctor s =>
(s a :-> Cxt h s a) -> (:&:) s p a :-> Cxt h (s :&: p) a
liftA' s a :-> Cxt h s a
f (s a i
v' :&: p
p) = p -> CxtFun s (s :&: p)
forall (f :: Node) p. HFunctor f => p -> CxtFun f (f :&: p)
ann p
p (s a i -> Cxt h s a i
s a :-> Cxt h s a
f s a i
v')

{-| This function strips the annotations from a term over a
functor with annotations. -}

stripA :: (RemA g f, HFunctor g) => CxtFun g f
stripA :: forall (g :: Node) (f :: Node).
(RemA g f, HFunctor g) =>
CxtFun g f
stripA = SigFun g f -> CxtFun g f
forall (f :: Node) (g :: Node).
HFunctor f =>
SigFun f g -> CxtFun f g
appSigFun g a i -> f a i
SigFun g f
forall (s :: Node) (s' :: Node) (a :: Family).
RemA s s' =>
s a :-> s' a
remA


propAnn :: (HFunctor g) => Hom f g -> Hom (f :&: p) (g :&: p)
propAnn :: forall (g :: Node) (f :: Node) p.
HFunctor g =>
Hom f g -> Hom (f :&: p) (g :&: p)
propAnn Hom f g
alg (f a i
f :&: p
p) = p -> CxtFun g (g :&: p)
forall (f :: Node) p. HFunctor f => p -> CxtFun f (f :&: p)
ann p
p (f a i -> Cxt Hole g a i
Hom f g
alg f a i
f)

-- | This function is similar to 'project' but applies to signatures
-- with an annotation which is then ignored.
project' :: (RemA f f', s :<: f') => Cxt h f a i -> Maybe (s (Cxt h f a) i)
project' :: forall (f :: Node) (f' :: Node) (s :: Node) h (a :: Family) i.
(RemA f f', s :<: f') =>
Cxt h f a i -> Maybe (s (Cxt h f a) i)
project' (Term f (Cxt h f a) i
x) = f' (Cxt h f a) i -> Maybe (s (Cxt h f a) i)
forall (a :: Family). NatM Maybe (f' a) (s a)
forall (f :: Node) (g :: Node) (a :: Family).
(f :<: g) =>
NatM Maybe (g a) (f a)
proj (f' (Cxt h f a) i -> Maybe (s (Cxt h f a) i))
-> f' (Cxt h f a) i -> Maybe (s (Cxt h f a) i)
forall a b. (a -> b) -> a -> b
$ f (Cxt h f a) i -> f' (Cxt h f a) i
forall (a :: Family). f a :-> f' a
forall (s :: Node) (s' :: Node) (a :: Family).
RemA s s' =>
s a :-> s' a
remA f (Cxt h f a) i
x
project' Cxt h f a i
_ = Maybe (s (Cxt h f a) i)
forall a. Maybe a
Nothing

pattern (::&::) :: (g :-<: fs) => g (Cxt h (Sum fs :&: x) a) l -> x -> Cxt h (Sum fs :&: x) a l
pattern $m::&:: :: forall {r} {g :: Node} {fs :: Signature} {h} {x} {a :: Family} {l}.
(g :-<: fs) =>
Cxt h (Sum fs :&: x) a l
-> (g (Cxt h (Sum fs :&: x) a) l -> x -> r) -> ((# #) -> r) -> r
$b::&:: :: forall (g :: Node) (fs :: Signature) h x (a :: Family) l.
(g :-<: fs) =>
g (Cxt h (Sum fs :&: x) a) l -> x -> Cxt h (Sum fs :&: x) a l
(::&::) t x <- Term ((proj -> Just t) :&: x) where
  (::&::) g (Cxt h (Sum fs :&: x) a) l
t x
x = (:&:) (Sum fs) x (Cxt h (Sum fs :&: x) a) l
-> Cxt h (Sum fs :&: x) a l
forall (f :: Node) h (a :: Family) i.
f (Cxt h f a) i -> Cxt h f a i
Term (g (Cxt h (Sum fs :&: x) a) l -> Sum fs (Cxt h (Sum fs :&: x) a) l
forall (a :: Family). g a :-> Sum fs a
forall (f :: Node) (g :: Node) (a :: Family).
(f :<: g) =>
f a :-> g a
inj g (Cxt h (Sum fs :&: x) a) l
t Sum fs (Cxt h (Sum fs :&: x) a) l
-> x -> (:&:) (Sum fs) x (Cxt h (Sum fs :&: x) a) l
forall (f :: Node) a (g :: Family) e. f g e -> a -> (:&:) f a g e
:&: x
x)

isNode' :: forall f g g' h a l. (HFunctor g, RemA g g', f :<: g') => Cxt h g a l -> Bool
isNode' :: forall (f :: Node) (g :: Node) (g' :: Node) h (a :: Family) l.
(HFunctor g, RemA g g', f :<: g') =>
Cxt h g a l -> Bool
isNode' Cxt h g a l
t = forall (f :: Node) (g :: Node) h (a :: Family) l.
(f :<: g) =>
Cxt h g a l -> Bool
isNode @f (Cxt h g' a l -> Bool) -> Cxt h g' a l -> Bool
forall a b. (a -> b) -> a -> b
$ Cxt h g a l -> Cxt h g' a l
CxtFun g g'
forall (g :: Node) (f :: Node).
(RemA g f, HFunctor g) =>
CxtFun g f
stripA Cxt h g a l
t

inj' :: (f :<: g) => (f :&: p) e l -> (g :&: p) e l
inj' :: forall (f :: Node) (g :: Node) p (e :: Family) l.
(f :<: g) =>
(:&:) f p e l -> (:&:) g p e l
inj' (f e l
x :&: p
p) = (f e l -> g e l
forall (a :: Family). f a :-> g a
forall (f :: Node) (g :: Node) (a :: Family).
(f :<: g) =>
f a :-> g a
inj f e l
x) g e l -> p -> (:&:) g p e l
forall (f :: Node) a (g :: Family) e. f g e -> a -> (:&:) f a g e
:&: p
p

inject' :: (f :<: g) => (f :&: p) (Cxt h (g :&: p) a) :-> Cxt h (g :&: p) a
inject' :: forall (f :: Node) (g :: Node) p h (a :: Family).
(f :<: g) =>
(:&:) f p (Cxt h (g :&: p) a) :-> Cxt h (g :&: p) a
inject' = (:&:) g p (Cxt h (g :&: p) a) i -> Cxt h (g :&: p) a i
forall (f :: Node) h (a :: Family) i.
f (Cxt h f a) i -> Cxt h f a i
Term ((:&:) g p (Cxt h (g :&: p) a) i -> Cxt h (g :&: p) a i)
-> ((:&:) f p (Cxt h (g :&: p) a) i
    -> (:&:) g p (Cxt h (g :&: p) a) i)
-> (:&:) f p (Cxt h (g :&: p) a) i
-> Cxt h (g :&: p) a i
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (:&:) f p (Cxt h (g :&: p) a) i -> (:&:) g p (Cxt h (g :&: p) a) i
forall (f :: Node) (g :: Node) p (e :: Family) l.
(f :<: g) =>
(:&:) f p e l -> (:&:) g p e l
inj'

injectOpt :: (f :<: g) => f (AnnHFix (Maybe p) g) l -> AnnHFix (Maybe p) g l
injectOpt :: forall (f :: Node) (g :: Node) p l.
(f :<: g) =>
f (AnnHFix (Maybe p) g) l -> AnnHFix (Maybe p) g l
injectOpt f (AnnHFix (Maybe p) g) l
t = (:&:) f (Maybe p) (AnnHFix (Maybe p) g) l
-> Cxt NoHole (g :&: Maybe p) (K ()) l
(:&:) f (Maybe p) (AnnHFix (Maybe p) g) :-> AnnHFix (Maybe p) g
forall (f :: Node) (g :: Node) p h (a :: Family).
(f :<: g) =>
(:&:) f p (Cxt h (g :&: p) a) :-> Cxt h (g :&: p) a
inject' (f (AnnHFix (Maybe p) g) l
t f (AnnHFix (Maybe p) g) l
-> Maybe p -> (:&:) f (Maybe p) (AnnHFix (Maybe p) g) l
forall (f :: Node) a (g :: Family) e. f g e -> a -> (:&:) f a g e
:&: Maybe p
forall a. Maybe a
Nothing)

caseH' :: forall fs a e l t. Alts (DistAnn fs a) e l t -> (Sum fs :&: a) e l -> t
caseH' :: forall (fs :: Signature) a (e :: Family) l t.
Alts (DistAnn fs a) e l t -> (:&:) (Sum fs) a e l -> t
caseH' Alts (DistAnn fs a) e l t
alts = Alts (DistAnn fs a) e l t -> Sum (DistAnn fs a) e l -> t
forall (fs :: Signature) (a :: Family) e b.
Alts fs a e b -> Sum fs a e -> b
caseH Alts (DistAnn fs a) e l t
alts (Sum (DistAnn fs a) e l -> t)
-> ((:&:) (Sum fs) a e l -> Sum (DistAnn fs a) e l)
-> (:&:) (Sum fs) a e l
-> t
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (:&:) (Sum fs) a e l -> Sum (DistAnn fs a) e l
forall (fs :: Signature) a (e :: Family) i.
(:&:) (Sum fs) a e i -> Sum (DistAnn fs a) e i
distAnn

caseCxt' :: forall cxt fs a e l t. (All cxt fs) => (forall f. (cxt f) => (f :&: a) e l -> t) -> (Sum fs :&: a) e l -> t
caseCxt' :: forall (cxt :: Node -> Constraint) (fs :: Signature) a
       (e :: Family) l t.
All cxt fs =>
(forall (f :: Node). cxt f => (:&:) f a e l -> t)
-> (:&:) (Sum fs) a e l -> t
caseCxt' forall (f :: Node). cxt f => (:&:) f a e l -> t
f (Sum Elem f fs
wit f e l
v :&: a
a) =
  (:&:) f a e l -> t
forall (f :: Node). cxt f => (:&:) f a e l -> t
f (f e l
v f e l -> a -> (:&:) f a e l
forall (f :: Node) a (g :: Family) e. f g e -> a -> (:&:) f a g e
:&: a
a) (cxt f => t) -> Dict cxt f -> t
forall {k} (c :: k -> Constraint) (a :: k) r.
(c a => r) -> Dict c a -> r
\\ forall {k} (c :: k -> Constraint) (f :: k) (fs :: [k]).
All c fs =>
Elem f fs -> Dict c f
forall (c :: Node -> Constraint) (f :: Node) (fs :: Signature).
All c fs =>
Elem f fs -> Dict c f
dictFor @cxt Elem f fs
wit

caseCxt'' :: forall cxt fs a e l t. (All cxt (DistAnn fs a)) => (forall f. (cxt (f :&: a)) => (f :&: a) e l -> t) -> (Sum fs :&: a) e l -> t
caseCxt'' :: forall (cxt :: Node -> Constraint) (fs :: Signature) a
       (e :: Family) l t.
All cxt (DistAnn fs a) =>
(forall (f :: Node). cxt (f :&: a) => (:&:) f a e l -> t)
-> (:&:) (Sum fs) a e l -> t
caseCxt'' forall (f :: Node). cxt (f :&: a) => (:&:) f a e l -> t
f (Sum Elem f fs
wit f e l
v :&: a
a) =
  (:&:) f a e l -> t
forall (f :: Node). cxt (f :&: a) => (:&:) f a e l -> t
f (f e l
v f e l -> a -> (:&:) f a e l
forall (f :: Node) a (g :: Family) e. f g e -> a -> (:&:) f a g e
:&: a
a) (cxt (f :&: a) => t) -> Dict cxt (f :&: a) -> t
forall {k} (c :: k -> Constraint) (a :: k) r.
(c a => r) -> Dict c a -> r
\\ forall {k} (c :: k -> Constraint) (f :: k) (fs :: [k]).
All c fs =>
Elem f fs -> Dict c f
forall (c :: Node -> Constraint) (f :: Node) (fs :: Signature).
All c fs =>
Elem f fs -> Dict c f
dictFor @cxt (Elem f fs -> Elem (f :&: a) (DistAnn fs a)
forall (f :: Node). Elem f fs -> Elem (f :&: a) (DistAnn fs a)
annWit Elem f fs
wit)

  where annWit :: Elem f fs -> Elem (f :&: a) (DistAnn fs a)
        annWit :: forall (f :: Node). Elem f fs -> Elem (f :&: a) (DistAnn fs a)
annWit = Elem f fs -> Elem (f :&: a) (DistAnn fs a)
forall {k1} {k2} (f :: k1) (fs :: [k1]) (g :: k2) (gs :: [k2]).
Elem f fs -> Elem g gs
unsafeElem


type family DistAnn (fs :: Signature) (a :: *) :: Signature where
  DistAnn (f ': fs) a = f :&: a ': DistAnn fs a
  DistAnn '[]       _ = '[]

distAnn :: (Sum fs :&: a) e :-> Sum (DistAnn fs a) e
distAnn :: forall (fs :: Signature) a (e :: Family) i.
(:&:) (Sum fs) a e i -> Sum (DistAnn fs a) e i
distAnn (Sum Elem f fs
wit f e i
v :&: a
a) =
  Elem f fs
-> f e i -> (f e :-> (:&:) f a e) -> Sum (DistAnn fs a) e i
forall (f :: Node) (fs :: Signature) (a :: Family) e (g :: Node)
       (gs :: Signature).
Elem f fs -> f a e -> (f a :-> g a) -> Sum gs a e
unsafeMapSum Elem f fs
wit f e i
v (f e i -> a -> (:&:) f a e i
forall (f :: Node) a (g :: Family) e. f g e -> a -> (:&:) f a g e
:&: a
a)