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


{-| 
    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>.

    For an example see examples/CSE2.hs 

    IMPORTANT: We would like to be able to use RebinadbleSyntax 
               to use Arrow's Syntax in this variant of TTTAS,
               but this extension still doesn't work well with Arrows.
-}

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

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

               -- ** Create New References
               newSRef, 
               extEnv, 
               castSRef,
               updateSRef,

               -- ** State-like operations on the Final Environment
               getFinalEnv, putFinalEnv, updateFinalEnv,
 
               -- ** Run a Trafo
               runTrafo,
               
               -- ** Arrow-style Combinators
               Pair(..), Arrow2(..), ArrowLoop2(..),
               (>>>),
               List(..), sequenceA, returnA

             ) where
 
#if __GLASGOW_HASKELL__ >= 710
import Prelude hiding (lookup, sequenceA) 
#else
import Prelude hiding (lookup) 
#endif

import Language.AbstractSyntax.TTTAS.Common




-- | 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  Trafo  m t a b =  
  Trafo (forall env1 . m env1 -> TrafoE m t env1 a b)
data  TrafoE m t env1 a b = 
  forall env2 .  TrafoE
                 (m env2)
                 (forall s .  a s -> T env2 s -> Env t s env1 -> FinalEnv t s
                              -> (b s, T env1 s, Env t s env2,   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 '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'.
runTrafo   ::  Trafo m t a b -> m () -> (forall s . a s) 
           ->  Result m t b
runTrafo trafo m a = 
 case trafo of
  Trafo trf -> case trf m of
                   TrafoE m2 f -> 
                      let  (rb, _, env2, fenv) =  f a (T id) Empty env2
                      in   Result m2  rb fenv


-- |  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.  
newSRef :: Trafo Unit t (t a) (Ref a)
newSRef 
    =  Trafo 
       (\Unit ->  TrafoE  
                  Unit
                  (\ta (T tr) env fenv ->
                      ( tr Zero
                      , T (tr . Suc)
                      , Ext env ta
                      , fenv
       )          )   )


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

-- | 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  e x (Ref a)
castSRef m r = TrafoE m $ (\   _ (T t) decls fe ->  (t r, T t, decls, fe))

data FUpd i t a  = FUpd (forall s. i -> t a s -> t a s)
newtype SI i s = SI i 

-- | 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 -> FUpd i t a -> TrafoE m t e (SI i) (Ref a)
updateSRef m r (FUpd f) = TrafoE m $ \(SI i) (T t) decls fs -> (t r, T t, updateEnv (f i) r decls, fs)


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

newtype FinalEnv2 t s = Final (FinalEnv t s)

-- | The function 'updateFinalEnv' returns a 'Trafo' that introduces a function
--   (@(UpdFinalEnv t)@) to update the final environment.
updateFinalEnv ::  Trafo m t (UpdFinalEnv t) Unit
updateFinalEnv = Trafo $ \m -> (TrafoE m (\(Upd u) t e fe -> (Unit, t, e, u fe)))

-- | Change the final environment by the one passed in the input.
putFinalEnv ::  Trafo m t (FinalEnv2 t) Unit
putFinalEnv =   Trafo $ \m -> (TrafoE m (\(Final fe) t e _ -> (Unit, t, e, fe)))

-- | Return as output the final environment.
getFinalEnv ::  Trafo m t Unit (FinalEnv2 t)
getFinalEnv = Trafo $ \m -> (TrafoE m (\_ t e fe -> ((Final fe), t, e, fe)))

{-
-- | 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 = proc f ->
                  do fe  <- getFinalEnv -< ()
                     putFinalEnv -< f fe
 -- Trafo $ \m -> (TrafoE m (\f t e fe -> ((), t, e, f fe)))
-}

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
  arr    :: (forall s . a s -> b s) -> arr a b
  first  :: arr a b -> arr (Pair a c) (Pair b c)
  second :: 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
  loop   :: arr (Pair a c) (Pair b c) -> arr a b


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

  (.:.) (Trafo sb) (Trafo sa) = 
    Trafo 
    (\m1 ->
       case sa m1 of 
        TrafoE m2 f1 -> case sb m2 of 
         TrafoE m3 f2 -> 
           TrafoE  
           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 (Trafo m t) where
  arr  f 
    =  Trafo  (\m -> TrafoE m (\a t e u -> (f a, t, e, u)) )

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

  second  f = arr swap >>> first f >>> arr swap
              where   swap :: Pair c a s -> Pair a c s
                      swap ~(P (x, y)) = P (y, x)
             
  f ***  g =  first f >>> second g  
 
  f &&&  g = arr (\b -> P (b, b)) >>> (f *** g)      


instance ArrowLoop2 (Trafo m t) where
  loop (Trafo st) =
    Trafo 
    (\m -> case st m of
            TrafoE m1 f1 ->
             TrafoE  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.
sequenceA :: [Trafo m t a b] -> Trafo m t a (List b)
sequenceA [] = arr (const (List []))
sequenceA (x:xs) 
       =  (x &&& sequenceA xs) >>> 
          arr (\(P (a,List as)) -> List (a:as))

returnA :: Arrow2 arr => arr a a
returnA = arr id