{-# OPTIONS_GHC -fglasgow-exts -fallow-overlapping-instances #-}

----------------------------------------------------------------
-- Code for the article
-- "Modularity and Implementation of
--  Mathematical Operational Semantics"
--
-- Mauro Jaskelioff   Neil Ghani   Graham Hutton
----------------------------------------------------------------

module MOS where

import Prelude hiding (seq) 
import Control.Monad

infixr 5 :+:

data L1 = Put1 A | Seq1 L1 L1 | Alt1 L1 L1

type A      = Char
----------------------------------------------------------------
-- Terms as Free Monads 
----------------------------------------------------------------

data L2 a = Put2 A | Seq2 a a | Alt2 a a 

instance Functor L2 where
    fmap _ (Put2 c)    = Put2 c 
    fmap f (Seq2 p q)  = Seq2  (f p)  (f q)
    fmap f (Alt2 p q)  = Alt2  (f p)  (f q)

data Term f x  =  Var x |  Con (f (Term f x))            

instance Functor f  => Functor  (Term f) where
    fmap f (Var x)  = Var (f x)
    fmap f (Con t)  = Con (fmap (fmap f) t)

instance Functor f  => Monad    (Term f) where
    return          = Var
    (Var x) >>= f   = f x
    (Con t) >>= f   = Con (fmap (>>= f) t)  

foldTerm :: Functor f => (a -> b) -> (f b -> b) -> Term f a -> b
foldTerm var _ (Var a)      = var a
foldTerm var con (Con fta)  = con (fmap (foldTerm var con) fta)

type Program f = Term f Zero

data Zero

empty :: Zero -> a
empty _ = undefined

----------------------------------------------------------------
-- Coproducts of Free Monads 
----------------------------------------------------------------

data (f :+: g) a     = Inl (f a) | Inr (g a) 

instance (Functor f, Functor g) => Functor (f :+: g) where
    fmap h (Inl fx)  = Inl (fmap h fx)
    fmap h (Inr gx)  = Inr (fmap h gx)

copair               :: (f a -> b) -> (g a -> b) -> (f :+: g) a -> b
copair f _ (Inl fa)  = f fa
copair _ g (Inr ga)  = g ga

type P       = Put :+: Seq :+: Alt

data Put  a  = Put Char
data Seq  a  = Seq  a  a
data Alt  a  = Alt  a  a

instance Functor Put where
    fmap _ (Put c)    = Put c

instance Functor Alt where
    fmap f (Alt p q)  = Alt (f p) (f q)

instance Functor Seq where
    fmap f (Seq p q)  = Seq (f p) (f q)

prog               :: Program P
prog               = alt1 (put1 'a' `seq1` put1 'b')  (put1 'c')

put1               :: Char -> Term P a
put1 c             = Con (Inl (Put c))

seq1, alt1         :: Term P a -> Term P a -> Term P a
seq1 p q           = Con (Inr (Inl (Seq p q)))
alt1 p q           = Con (Inr (Inr (Alt p q)))

data N a    = N Int
data Add a  = Add a a
data Ifz a  = Ifz a a a

data Thr a  = Thr
data Cat a  = Cat a a 


type Z      = N :+: Add :+: Ifz

type E      = Thr :+: Cat

instance Functor N where
    fmap f (N n) = N n

instance Functor Add where
    fmap f (Add p q) = Add (f p) (f q)

instance Functor Ifz where
    fmap f (Ifz c t e) = Ifz (f c) (f t) (f e)

instance Functor Thr where
    fmap f Thr = Thr

instance Functor Cat where
    fmap f (Cat t u) = Cat (f t) (f u)

type EP   = Thr :+: Cat :+: Put :+: Seq :+: Alt
type EZ   = Thr :+: Cat :+: N :+: Add :+: Ifz
type NDZ  = N :+: Add :+: Alt

put'    :: Char -> Term EP a
put' c  = Con (Inr (Inr (Inl (Put c))))
-----------------------------------------------------------------------------
-- SubType relation
-----------------------------------------------------------------------------

infixl 4 >>>=

class (Functor sub, Functor sup) => SubType sub sup where
    inj :: sub a -> sup a
    prj :: sup a -> Maybe (sub a)

instance (Functor f) => SubType f f where
    inj  = id
    prj  = Just

instance (Functor f, Functor g) => SubType f (f :+: g) where
    inj          = Inl
    prj (Inl f)  = Just f
    prj _        = Nothing

instance (Functor h, SubType f g) => SubType f (h :+: g) where
    inj          = Inr . inj
    prj (Inr a)  = prj a
    prj _        = Nothing

(>>>=)         :: (SubType sub sup) => (sup x -> a) -> (sub x -> a) -> sup x -> a
(f >>>= g) x   = case  prj x of 
                       Nothing  -> f x
                       Just y   -> g y 

con          :: (SubType s t) => s (Term t x) -> Term t x
con          = Con . inj

put          :: (SubType Put s) => Char -> Term s x
put c        = con (Put c)

seq          :: (SubType Seq s) => Term s x -> Term s x -> Term s x
seq p q      = con (Seq p q)

alt          :: (SubType Alt s) => Term s x -> Term s x -> Term s x
alt p q      = con (Alt p q)

n            :: (SubType N s) => Int -> Term s x
n m          = con (N m)

add          :: (SubType Add s) => Term s x -> Term s x -> Term s x
add p q      = con (Add p q)

ifz          :: (SubType Ifz s) => Term s x -> (Term s x, Term s x) -> Term s x
ifz c (t,e)  = con (Ifz c t e)

thr          :: (SubType Thr s) => Term s x
thr          = con Thr

cat          :: (SubType Cat s) => Term s x -> Term s x -> Term s x
cat p q      = con (Cat p q)

-------------------------------------------------------------------------
-- Composition of functors
-------------------------------------------------------------------------

infixr 7 :.:

data (h :.: g) x = Comp {deComp :: (h (g x)) }

instance (Functor h, Functor g) => Functor (h :.: g) where
    fmap f (Comp c) = Comp (fmap (fmap f) c)

data Pr a   = Pr A a
data KA a   = KA A

instance Functor Pr  where
    fmap f (Pr c a)   = Pr c (f a)

instance Functor KA  where
    fmap _ (KA c)     = KA c

data KI a    = KI Int

instance Functor KI where
    fmap _ (KI c) = KI c
data KE a    = KE

instance Functor KE where
    fmap _ KE = KE

----------------------------------------------------------------
-- Execution of a Transition System
----------------------------------------------------------------

data Nu f = Nu (f (Nu f)) 


unfold    :: Functor b => (x -> b x) -> x -> Nu b
unfold g  = Nu. fmap (unfold g) . g

----------------------------------------------------------------
-- Mathematical Operational Semantics 
----------------------------------------------------------------

type OR s b =  forall x y.  (x -> y) ->     -- Term environment
                            (x -> b y) ->   -- Behaviour environment
                            s x ->          -- Source of the conclusion
                            b (Term s y)    -- Target of the conclusion

type (:..:) f g x = f (g x)

orP                    :: OR L2 ([] :..: (KA :+: Pr))
orP _ _    (Put2 c)    =   [Inl (KA c)]
orP te be  (Seq2 p q)  =   be' p >>= copair  
                           (\(KA c)     -> [Inr (Pr c (te' q))])
                           (\(Pr c p')  -> [Inr (Pr c (Con (Seq2 p' (te' q))))])
                           where  te'  = Var . te
                                  be'  = map (fmap Var). be
orP te be  (Alt2 p q)  =  map (fmap Var) (be p ++ be q)

----------------------------------------------------------------
-- Obtaining a transition relation 
----------------------------------------------------------------

opMonad       ::  (Functor s, Functor b) =>
                  OR s b -> (x -> b x) -> Term s x -> b (Term s x)
opMonad op k  = snd. foldTerm  (pair Var (fmap Var. k)) 
                               (pair (Con. fmap fst) (fmap join. op fst snd))
                  where  pair f g a = (f a, g a)

run         ::  (Functor s, Functor b) => OR s b -> Program s -> Nu b
run op      = unfold (opMonad op empty)


----------------------------------------------------------------
-- Modular Operational Semantics
----------------------------------------------------------------

type MOR s t m b  =  forall x y.  (x -> y) -> 
                                  (x -> m (b y)) ->
                                  s x -> m (b (Term t y))

ossify            :: MOR s s m b -> OR s (m :.: b)
ossify mor te be  = Comp . mor te (deComp . be)

----------------------------------------------------------------
-- Joining MOR
----------------------------------------------------------------

infixr 5 <+>

(<+>)                ::  MOR s t m b -> MOR s' t m b -> MOR (s :+: s') t m b
(op1 <+> op2) te be  = copair (op1 te be) (op2 te be)

----------------------------------------------------------------
-- Defining Modular Operational Rules 
----------------------------------------------------------------

morPut                      :: (SubType Put s, SubType KA b, Monad m) => MOR Put s m b
morPut _ _ (Put c)          = return (inj (KA c))

morSeq                      ::  (SubType Seq s, SubType KA b, SubType Pr b, Monad m)
                                => MOR Seq s m b
morSeq te be (Seq p q)      = be' p >>=  return . fmap (`seq` te' q) >>>= \(KA c)  -> 
                                        return (inj (Pr c (te' q)))
                              where  te'  = Var . te
                                     be'  = liftM (fmap Var). be

morAlt                      :: (SubType Alt s, Functor b, MonadPlus m) 
                               => MOR Alt s m b
morAlt te be (Alt p q)      = liftM (fmap Var) (be p `mplus` be q)

addl, addr          :: (SubType Add s) => Term s x -> Term s x -> Term s x
addl = add
addr = add

ifzr          :: (SubType Ifz s) => Term s x -> (Term s x, Term s x) -> Term s x
ifzr = ifz

catr          :: (SubType Cat s) => Term s x -> Term s x -> Term s x
catr = cat

morN                        ::(SubType N s, SubType KI b, Monad m) => MOR N s m b
morN _ _ (N n)              = return (inj (KI n))

morAdd                      :: (SubType Add s, SubType KI b, Monad m) => MOR Add s m b
morAdd te be (Add p q)      =  be' p >>=  return . fmap (`addr` te' q)   >>>= \(KI n) -> 
                               be' q >>=  return . fmap (te' p `addl`)  >>>= \(KI m) -> 
                               return (inj (KI (n+m)))
                               where  te'  = Var . te
                                      be'  = liftM (fmap Var). be

morIfz                      :: (SubType Ifz s, SubType KI b, Monad m) => MOR Ifz s m b
morIfz te be (Ifz c t e)    = be' c  >>=  return . fmap (`ifzr` (te' t,te' e)) 
                                          >>>= \(KI n)  -> if n==0 then be' t else be' e
                              where   te'   = Var . te
                                      be'   = liftM (fmap Var). be

morThr                  :: (SubType Thr s, SubType KE b, Monad m) => MOR Thr s m b
morThr _ _ Thr          = return (inj KE)

morCat                  :: (SubType Cat s, SubType KE b, Monad m) => MOR Cat s m b
morCat te be (Cat p q)  = be' p >>=  return . fmap (`catr` te' q) >>>= \KE -> 
                                     be' q
                          where  te'  = return . te
                                 be'  = liftM (fmap return). be

-------------------------------------------------------------------------
-- Putting it all together
-------------------------------------------------------------------------

morP  ::   (SubType Put s, SubType Seq s, SubType Alt s, SubType KA b, SubType Pr b, MonadPlus m) =>
           MOR P s m b
morP  =  morPut <+> morSeq <+> morAlt

orP'  :: OR P ([] :.: (Pr :+: KA))
orP'  = ossify morP

morZ ::  (SubType N s,  SubType Ifz s, SubType Add s,  Monad m, SubType KI b ) => MOR Z s m b
morZ = morN <+> morAdd <+> morIfz

z  :: OR Z (Id :.: KI)
z  = ossify morZ

  ep  :: OR EP ([] :.: (KE :+: Pr :+: KA))
ep  = ossify (morThr <+> morCat <+> morP)

ez  :: OR EZ (Id :.: (KE :+: KI))
ez  = ossify (morThr <+> morCat <+> morZ)

zp  :: OR (Put :+: Seq :+: Z) (Id :.: (KA :+: Pr :+: KI)) 
zp  = ossify (morPut <+> morSeq <+> morZ)

zP  :: OR (Put :+: Seq :+: Alt :+: Z) ([] :.: (KA :+: Pr :+: KI)) 
zP  = ossify (morPut <+> morSeq <+> morAlt <+> morZ)

addND a b = alt (a `add` b)  (b `add` a)

type AOR s b    = forall a. s (a , b a) -> b (Term s a)

fromOR          :: OR s b -> AOR s b
fromOR os       =  os fst snd

toOR            :: (Functor s) => AOR s b -> OR s b
toOR aor te be  = aor. fmap (pair te be)
                where  pair f g a = (f a, g a)

-------------------------------------------------------------------------
-- Identity Monad
-------------------------------------------------------------------------
  
data Id a = Id a

instance Functor Id where
    fmap f (Id a) = Id (f a)

instance Monad Id where
    return        = Id
    (Id a) >>= f  = f a


-------------------------------------------------------------------------
-- Showing fixpoints
-------------------------------------------------------------------------

class Functor f => PreservesString f where
    preservesString :: f String -> String

instance (Show a, PreservesString f) => Show (Term f a) where
    show (Var a)  = show a
    show (Con x)  = preservesString (fmap show x)

instance (PreservesString f, PreservesString g) => PreservesString (f :+: g) where
    preservesString (Inl x)  = preservesString x
    preservesString (Inr x)  = preservesString x

instance Show Zero where

-------------------------------------------------------------------------
-- Showing Behaviour and Syntax
-------------------------------------------------------------------------

instance PreservesString [] where
    preservesString = ("["++). h
        where h []        = "]"
              h [x]       = x++"]"
              h (x:y:xs)  = x++", "++h (y:xs)

instance PreservesString KA where
    preservesString (KA c)  = "-"++show c++"->."

instance PreservesString KI where
    preservesString (KI i)  = show i

instance PreservesString KE where
    preservesString KE  = "Exception!"

instance PreservesString Pr where
    preservesString (Pr c a)   = "-"++[c]++"-> "++a

instance PreservesString N where
    preservesString (N n) = show n

instance PreservesString Thr where
    preservesString Thr = "Throw"

instance PreservesString Cat where
    preservesString (Cat p q)  = "(cat "++p++" : "++q++")"

instance PreservesString Add where
    preservesString (Add p q)  = "("++p++" + "++q++")"

instance PreservesString Ifz where
    preservesString (Ifz c p q)  = "(if "++c++" then " ++p++" else "++q++")"

instance PreservesString Put where
    preservesString (Put c)    = [c]

instance PreservesString Seq where
    preservesString (Seq p q)  = "("++p++" ; "++q++")"

instance PreservesString Alt where
    preservesString (Alt p q)  = "("++p++" or "++q++")"


-----------------------------------------------------------
-- Example Programs
-----------------------------------------------------------

-- P
p0,p1,p2  :: Program P
p0  = seq (put 'b') (put 'c')
p1  = alt (put 'a') (alt p0 (put 'd'))
p2  = seq (alt (put 'a') (put 'd')) p0

-- EP
p3,p4,p5 :: Program EP
p3 = (put 'a' `seq` thr) `alt` (put 'b' `seq` put 'c')
p4 = cat p3 (put 'X')
p5 = (cat (put 'a' `seq` thr) (put 'Y')) `alt` (put 'b' `seq` put 'c')

-- EZ
p10,p11,p12,p13 :: Program EZ
p10 = n 1 `add` (n 2 `add` n 4)
p11 = cat (n 1 `add` thr) (n 4 `add` n 6)
p12 = add p10 p11
p13 = ifz (n 0) (n 1, thr)

-- printZ
p20 :: Program (Put :+: Seq :+: Z)
p20 = (put 'o' `seq` put 'w') `add` (n 2 `add` (put 'f' `seq` n 4))