{-# OPTIONS -XKindSignatures -XRankNTypes -XArrows -XGADTs #-}
{-# LANGUAGE CPP #-}

{-| 
    Library for Typed Transformations of Typed Abstract Syntax.

    The library is documented in the paper: /Typed Transformations of Typed Abstract Syntax/

    Bibtex entry: <http://www.cs.uu.nl/wiki/bin/viewfile/Center/TTTAS?rev=1;filename=TTTAS.bib>

    For more documentation see the TTTAS webpage: 
    <http://www.cs.uu.nl/wiki/bin/view/Center/TTTAS>.
-}

module Language.AbstractSyntax.TTTAS (
               -- * Typed References and Environments

               -- ** Typed References
               Ref(..), Equal(..), 
               match, lookup, update,

               -- ** Declarations
               Env(..), FinalEnv, T(..),
               lookupEnv, updateEnv,

               -- * Transformation Library
               
               -- ** Trafo
               Trafo(..), TrafoE(..),

               -- ** Create New References
               Unit(..),
               newSRef, extEnv, castSRef, updateSRef,

               -- ** Update the Final Environment
               updateFinalEnv,
 
               -- ** Run a Trafo
               Result(..), 
               runTrafo,

               -- ** Other Combinators
               sequenceA,

               -- * Alternative Transformation Library
               
               -- ** Trafo2
               Trafo2(..), TrafoE2(..),

               -- ** Create New References
               newSRef2,

               -- ** Update the Final Environment
               UpdFinalEnv(..), updateFinalEnv2,

               -- ** Run a Trafo2 
               runTrafo2,

               -- ** Arrow-style Combinators
               Pair(..), Arrow2(..), ArrowLoop2(..),
               (>>>>),
               List(..), sequenceA2

             ) where
 
import Unsafe.Coerce ( unsafeCoerce )
import qualified Prelude as P

#if __GLASGOW_HASKELL__ >= 609 
import Control.Category 
import Prelude hiding (lookup,(.), id) 
#endif   

import Control.Arrow 
#if __GLASGOW_HASKELL__ < 609 
  hiding (pure) 
import Prelude hiding (lookup)
#endif


-- | The 'Ref' type for represents typed indices which are
--   labeled with both the type of value to which they
--   refer and the type of the environment (a nested
--   Cartesian product, growing to the right) in which 
--   this value lives.
 --  The constructor 'Zero' expresses that the first
 --   element of the environment has to be of type @a@.
 --  The constructor 'Suc' does not care about the type
 --   of the first element in the environment, 
 --   being polymorphic in the type @b@.
data Ref a env where
 Zero  ::                      Ref a (env',a)
 Suc   ::  Ref a      env' ->  Ref a (env',b)

-- | The 'Equal' type encodes type equality. 
data Equal :: * -> * -> * where 
 Eq :: Equal a a

-- | The function 'match' compares two references for equality.
--   If they refer to the same element in the environment
--   the value @Just Eq@ is returned, expressing the fact that
--   the types of the referred values are the same too.
match :: Ref a env -> Ref b env -> Maybe (Equal a b)
match  Zero     Zero        =  Just Eq
match  (Suc x)  (Suc y)     =  match x y
match  _        _           =  Nothing

-- | The function 'lookup' returns the element indexed in the
--   environment parameter by the 'Ref' parameter. The types
--   guarantee that the lookup succeeds.
lookup :: Ref a env -> env -> a
lookup  Zero     (_,a)  =  a
lookup  (Suc r)  (e,_)  =  lookup r e

-- | The function 'update' takes an additional function as
--   argument, which is used to update the value the 
--   reference addresses.
update :: (a -> a) -> Ref a env -> env -> env  
update f  Zero     (e,a)   =  (e,f a)
update f  (Suc r)  (e,x)   =  (update f r e,x)


-- | The type @Env term use def@ represents a sequence of
--   instantiations of type @forall a. term a use@, where
--   all the instances of @a@ are stored in the type parameter
--   @def@. The type @use@ is a sequence containing the 
--   types to which may be referred from within terms of type
--   @term a use@. 
data Env term use def where
  Empty  ::  Env t use  ()
  Ext    ::  Env t use def' ->  t a use  
         ->  Env t use  (def',a)


lookupEnv  :: Ref a env -> Env t s env -> t a s
lookupEnv  Zero     (Ext _ t)   =  t
lookupEnv  (Suc r)  (Ext ts _)  =  lookupEnv r ts
lookupEnv  _        _           =  error "Error: The impossible happened!"

updateEnv  ::  (t a s -> t a s) -> Ref a env 
           ->  Env t s env ->  Env t s env
updateEnv  f  Zero     (Ext ts t)  
              =        Ext ts (f t) 
updateEnv  f  (Suc r)  (Ext ts t)  
              =        Ext (updateEnv f r ts)  t
updateEnv  _  _        _ 
              =  error "Error: The impossible happened!"


-- | When the types @def@ and @use@ of an 'Env' coincide,
--   we can be sure that the references in the terms do not
--   point to values outside the environment but point
--   to terms representing the right type. This kind of
--   environment is the /final environment/ of a transformation.
type FinalEnv t usedef = Env t usedef usedef
 
-- | The type 'T' encodes a 'Ref'-transformer. It is usually used
--   to transform references from an actual environment to 
--   the final one.
newtype T e s  = T {unT :: forall x . Ref x e -> Ref x s}

-- | The type 'Trafo' is the type of the transformation steps on a heterogeneous collection.
--   The argument @m@ stands for the type of the meta-data.  
--   A |Trafo| takes the meta-data on the current environment |env1| as input and 
--   yields  meta-data for the (possibly extended) environment |env2|.
--   The type @t@ is the type of the terms stored in the environment.
--   The type variable @s@ represents the type of the final result, which we do expose. 
--   Its role is similar to the @s@ in the type @ST s a@.
--   The arguments @a@ and @b@ are the Arrow's input and output, respectively.
data  Trafo  m t s a b =  
  Trafo (forall env1 . m env1 -> TrafoE m t s env1 a b)

-- | The type 'TrafoE' is used to introduce an existential quantifier into
--   the definition of 'Trafo'.  
--   It can be seen that a 'Trafo' is a function taking as arguments: the input (@a@),
--   a 'Ref'-transformer (@T env2 s@) from the environment constructed in this step 
--   to the final environment, the environment (@Env t s env1@) where the current 
--   transformation starts and a function (@FinalEnv t s -> FinalEnv t s@) to update 
--   (modify) the final environment. The function returns: the output (@b@), 
--   a 'Ref'-transformer (@T env1 s@) from the initial environment of this step to the final 
--   environment, the environment (@Env t s env2@) constructed in this step and a function 
--   (@FinalEnv t s -> FinalEnv t s@) to update (modify) the final environment.
--   NOTE: The function (@FinalEnv t s -> FinalEnv t s@) was introduced in version 0.3. 
--   It's carried throw the transformation steps and can be modified (composed to another function)
--   using the function 'updateFinalEnv'.
data  TrafoE m t s env1 a b = 
 forall env2 . TrafoE   (  m env2)
                        (       a  ->  T env2 s ->  Env t s env1 -> (FinalEnv t s -> FinalEnv t s)
                          -> (  b,     T env1 s,    Env t s env2, (FinalEnv t s -> FinalEnv t s))
                        )

data Unit s         = Unit

-- |  The Trafo 'newSRef' takes a typed term as input, adds it to the environment 
--    and yields a reference pointing to this value.
--    No meta-information on the environment is recorded by 'newSRef';
--    therefore we use the type 'Unit' for the meta-data.  
newSRef :: Trafo Unit t s (t a s) (Ref a s)
newSRef 
 =  Trafo  (\ _->  extEnv Unit)


-- | The function 'updateFinalEnv' returns a 'Trafo' that introduces a function
--   (@FinalEnv t s -> FinalEnv t s@) to update the final environment.
updateFinalEnv ::  Trafo m t s (FinalEnv t s -> FinalEnv t s) ()
updateFinalEnv = Trafo $ \m -> (TrafoE m (\f' t e f -> ((), t, e, f' . f)))


-- | The function 'extEnv' returns a 'TrafoE' that extends the current environment.
extEnv :: m (e,a) -> TrafoE m t s e (t a s) (Ref a s)
extEnv m = TrafoE m $ \ta  (T tr) env  f ->  (tr Zero,  T (tr P.. Suc),  Ext env ta, f )

-- | The function 'castSRef'  returns a 'TrafoE' that casts the reference 
--   passed as parameter (in the constructed environment) to one in the final environment.
castSRef :: m e -> Ref a e -> TrafoE m t s e x (Ref a s)
castSRef m r = TrafoE m $ (\   _ (T t) decls f ->  (t r, T t, decls, f))

-- | The function 'updateSRef' returns a 'TrafoE' that updates the value pointed
--   by the reference passed as parameter into the current environment.
updateSRef :: m e -> Ref a e -> (i -> t a s -> t a s) -> TrafoE m t s e i (Ref a s)
updateSRef m r f = TrafoE m $ \i (T t) decls fs -> (t r, T t, updateEnv (f i) r decls, fs)


instance Functor (TrafoE m t s e a) where
 fmap f (TrafoE m step) = TrafoE m $ \i t e fs -> case step i t e fs of
                                                  (i',t',e',fs') -> (f i',t',e',fs')

-- | The type 'Result' is the type of the result of \"running\" a 'Trafo'. 
--   Because @s@ could be anything we have to hide it using existential quantification.
data Result  m t b   
   =  forall s . Result  (m s) (b s) (FinalEnv t s)

-- | The function 'runTrafo' takes as arguments the 'Trafo' we want to run, meta-information 
--   for the empty environment, and an input value. 
--   The result of 'runTrafo' (type 'Result') is the final environment (@Env t s s@) together 
--   with the resulting meta-data (@m s@), and the output value (@b s@). 
--   The rank-2 type for 'runTrafo' ensures that transformation steps cannot make 
--   any assumptions about the type of final environment (@s@).
runTrafo   ::  (forall s . Trafo m t s a (b s)) ->  m () -> a 
           ->  Result m t b
runTrafo trafo m a = case trafo of
  Trafo trf -> case trf m of
                   TrafoE m2 f -> 
                     case  f a (T unsafeCoerce) Empty P.id of  -- the function could also be passed as argument
                       (rb, _, env2, fenvs) -> 
                         Result  (unsafeCoerce m2)  
                                 rb 
                                 (fenvs $ unsafeCoerce env2) 

#if __GLASGOW_HASKELL__ >= 609

instance Category (Trafo m t s) where
 -- |(.) :: Trafo m t s b c  ->  Trafo m t s a b ->  Trafo m t s a c|
 Trafo t2 . Trafo t1 = 
  Trafo 
  (\m1 -> case t1 m1 of
                                      TrafoE m2 f1 -> case t2 m2 of
                                       TrafoE m3 f2 ->
                                        TrafoE 
                                        m3 
                                        (\a  tt env1 fs ->
                                             let  (b,tt1, env2, fs')   = f1 a tt2 env1 fs
                                                  (c,tt2, env3, fs'')  = f2 b tt env2 fs'
                                             in   (c,tt1, env3, fs'')                                                                    
                                        )
  )

 -- |id :: Trafo m t s a a|
 id =  Trafo  (\m -> TrafoE m (\a t e f -> (a, t, e, f)) )

#endif


instance Arrow (Trafo m t s) where

 -- |arr :: (a  -> b) -> Trafo m t s a b|
 arr f =  Trafo  (\m -> TrafoE m (\a t e fs -> (f a, t, e, fs)) )

#if __GLASGOW_HASKELL__ < 609 

 Trafo t1 >>> Trafo t2 = 
  Trafo 
  (\m1 -> case t1 m1 of
                                      TrafoE m2 f1 -> case t2 m2 of
                                       TrafoE m3 f2 ->
                                        TrafoE 
                                        m3 
                                        (\a  tt env1 fs ->
                                             let  (b,tt1, env2, fs')   = f1 a tt2 env1 fs
                                                  (c,tt2, env3, fs'')  = f2 b tt env2 fs'
                                             in   (c,tt1, env3, fs'')                                                                    
                                        )
  )

#endif 

 -- |first :: Trafo m t s a b -> Trafo m t s (a, c) (b, c)|
 first (Trafo tr) 
  = Trafo (\m1 -> case tr m1 of
        TrafoE m2 f -> 
           TrafoE  
           m2 
           (\  ~(a,c) tt env1 fs -> 
               let (b,tt1,env2, fs') = f a tt env1 fs
               in  ((b,c),tt1, env2, fs'))) 


instance ArrowLoop (Trafo m t s) where
 -- |loop :: Trafo m t s (a, x) (b, x) -> Trafo m t s a b|
 loop (Trafo st) =
   Trafo 
   (\m -> case st m of
           TrafoE m1 f1 ->
            TrafoE  m1
            (\a t e f -> 
              let ((b, x),t1,e1,f') = f1 (a, x) t e f                                                                            
              in (b,t1,e1,f')
            ))



-- | The combinator 'sequenceA' sequentially composes a list
--   of 'Trafo's into a 'Trafo' that yields a list of outputs.
--   Its use is analogous to the combinator 'sequence' combinator
--   for 'Monad's.
sequenceA :: [Trafo m t s a b] -> Trafo m t s a [b]
sequenceA [] = arr (const [])
sequenceA (x:xs) 
     = proc a ->
        do b  <- x -< a
           bs <- sequenceA xs -< a
           returnA -< (b:bs)  




-- | Alternative version of 'Trafo' where the universal quantification 
--   over |s| is moved inside the quantification over |env2|.
--   Note that the type variables |a| and |b| are now labelled with |s|, 
--   and hence have kind |(* -> *)|.
data  Trafo2  m t a b =  
  Trafo2 (forall env1 . m env1 -> TrafoE2 m t env1 a b)
data  TrafoE2 m t env1 a b = 
  forall env2 .  TrafoE2
                 (m env2)
                 (forall s .  a s -> T env2 s -> Env t s env1 -> UpdFinalEnv t s
                              -> (b s, T env1 s, Env t s env2, UpdFinalEnv t s)
                 )

newtype UpdFinalEnv t s = Upd (FinalEnv t s -> FinalEnv t s)

-- | The function 'runTrafo2' takes as arguments the 'Trafo2' we want to run, meta-information 
--   for the empty environment, and an input value. 
--   The result of 'runTrafo2' (type 'Result') is the final environment (@Env t s s@) together 
--   with the resulting meta-data (@m s@), and the output value (@b s@). 
--   The rank-2 type for 'runTrafo2' ensures that transformation steps cannot make 
--   any assumptions about the type of final environment (@s@).
--   It is an alternative version of 'runTrafo' which does not use
--   'unsafeCoerce'.
runTrafo2  ::  Trafo2 m t a b -> m () -> (forall s . a s) 
           ->  Result m t b
runTrafo2 trafo m a = 
 case trafo of
  Trafo2 trf -> case trf m of
                   TrafoE2 m2 f -> 
                      let  (rb, _, env2, Upd upds) =  f a (T P.id) Empty (Upd P.id)
                      in   Result m2  rb (upds env2)


-- |  The Trafo2 'newSRef2' takes a typed term as input, adds it to the environment 
--    and yields a reference pointing to this value.
--    No meta-information on the environment is recorded by 'newSRef2';
--    therefore we use the type 'Unit' for the meta-data.  
newSRef2 :: Trafo2 Unit t (t a) (Ref a)
newSRef2 
    =  Trafo2  
       (\Unit ->  TrafoE2  
                  Unit
                  (\ta (T tr) env upds ->
                      ( tr Zero
                      , T (tr P.. Suc)
                      , Ext env ta
                      , upds
       )          )   )


-- | The function 'updateFinalEnv2' returns a 'Trafo2' that introduces a function
--   (@(UpdFinalEnv t)@) to update the final environment.
updateFinalEnv2 ::  Trafo2 m t (UpdFinalEnv t) Unit
updateFinalEnv2 = Trafo2 $ \m -> (TrafoE2 m (\(Upd u') t e (Upd u) -> (Unit, t, e, (Upd $ u' P.. u))))


newtype Pair a b s = P (a s, b s)

class Category2 (cat :: (* -> *) -> (* -> *) -> *) where
  id2     :: cat a a
  (.:.)   :: cat b c -> cat a b -> cat a c

class Category2 arr => Arrow2 (arr :: (* -> *) -> (* -> *) -> *) where
  arr2    :: (forall s . a s -> b s) -> arr a b
  first2  :: arr a b -> arr (Pair a c) (Pair b c)
  second2 :: arr a b -> arr (Pair c a) (Pair c b)
  (****)  :: arr a b -> arr a' b'
          -> arr (Pair a a') (Pair b b')
  (&&&&)  :: arr a b -> arr a b' -> arr a (Pair b b')

class Arrow2 arr => ArrowLoop2 arr where
  loop2   :: arr (Pair a c) (Pair b c) -> arr a b


instance Category2 (Trafo2 m t) where
  id2   =  Trafo2  (\m -> TrafoE2 m (\a t e u -> (a, t, e, u)))

  (.:.) (Trafo2 sb) (Trafo2 sa) = 
    Trafo2 
    (\m1 ->
       case sa m1 of 
        TrafoE2 m2 f1 -> case sb m2 of 
         TrafoE2 m3 f2 -> 
           TrafoE2  
           m3 
           (\a t3s e1 u1 -> let  (b, t1s, e2, u2) = f1 a t2s e1 u1
                                 (c, t2s, e3, u3) = f2 b t3s e2 u2
                            in   (c, t1s, e3, u3)
           ))


(>>>>) :: Category2 cat => cat a b -> cat b c -> cat a c
f >>>> g = g .:. f

 
instance Arrow2 (Trafo2 m t) where
  arr2 f 
    =  Trafo2  (\m -> TrafoE2 m (\a t e u -> (f a, t, e, u)) )

  
  first2 (Trafo2 s) 
    =  Trafo2
       (\m1 -> case s m1 of
          TrafoE2 m2 f -> 
             TrafoE2  m2
                    (\(P (a, c)) t2s e1 u1 -> 
                                      let (b,t12,e2,u2) = f a t2s e1 u1 
                                      in  (P (b, c),t12,e2,u2)
                    )              
       ) 

  second2 f = arr2 swap >>>> first2 f >>>> arr2 swap
              where   swap :: Pair c a s -> Pair a c s
                      swap ~(P (x, y)) = P (y, x)
             
  f **** g =  first2 f >>>> second2 g  
 
  f &&&& g = arr2 (\b -> P (b, b)) >>>> (f **** g)      


instance ArrowLoop2 (Trafo2 m t) where
  loop2 (Trafo2 st) =
    Trafo2 
    (\m -> case st m of
            TrafoE2 m1 f1 ->
             TrafoE2  m1
             (\a t e u -> 
               let (P (b, x),t1,e1,u1) = f1 (P (a, x)) t e u                                                                            
               in (b,t1,e1,u1)
             ) 
    )


newtype List a s = List [a s]

-- | The combinator 'sequenceA2' sequentially composes a list
--   of 'Trafo2's into a 'Trafo2' that yields a 'List' of outputs.
--   Its use is analogous to the combinator 'sequence' combinator
--   for 'Monad's.
sequenceA2 :: [Trafo2 m t a b] -> Trafo2 m t a (List b)
sequenceA2 [] = arr2 (const (List []))
sequenceA2 (x:xs) 
       =  (x &&&& sequenceA2 xs) >>>> 
          arr2 (\(P (a,List as)) -> List (a:as))