Session types in Idris

sess.idr

module sess

import System.Concurrency.Channels

infix 3 :!:, :?:
infix 2 :+:, :&:

data Protocol : Type where
    Eps : Protocol
    (:!:) : Type -> Protocol -> Protocol
    (:?:) : Type -> Protocol -> Protocol
    (:+:) : Protocol -> Protocol -> Protocol
    (:&:) : Protocol -> Protocol -> Protocol
    Rec : Protocol -> Protocol
    Var : Nat -> Protocol

%name Protocol p,q

dual : Protocol -> Protocol
dual (t :!: p) = t :?: dual p
dual (t :?: p) = t :!: dual p
dual (p :+: q) = dual p :&: dual q
dual (p :&: q) = dual p :+: dual q
dual (Rec p) = Rec (dual p)
dual p = p

dualInvolution : p = dual (dual p)
dualInvolution {p = Eps} = Refl
dualInvolution {p = (t :!: p)} = rewrite dualInvolution {p=p} in Refl
dualInvolution {p = (t :?: p)} = rewrite dualInvolution {p=p} in Refl
dualInvolution {p = (p :+: q)} = 
    rewrite dualInvolution {p=p} in rewrite dualInvolution {p=q} in Refl
dualInvolution {p = (p :&: q)} = 
    rewrite dualInvolution {p=p} in rewrite dualInvolution {p=q} in Refl
dualInvolution {p = (Rec p)} = rewrite dualInvolution {p=p} in Refl
dualInvolution {p = (Var k)} = Refl

data Capacity : Type where
    Cap : List Protocol -> Protocol -> Capacity
    Done : Capacity

data Session : Capacity -> Capacity -> Type -> Type where
    Send : a -> Session (Cap e (a :!: p)) (Cap e p) ()
    Recv : Session (Cap e (a :?: p)) (Cap e p) a
    Close : Session (Cap e Eps) Done ()
    Inl : Session (Cap e (p :+: q)) (Cap e p) ()
    Inr : Session (Cap e (p :+: q)) (Cap e q) ()
    Offer : Session (Cap e p) c a -> Session (Cap e q) c a -> Session (Cap e (p :&: q)) c a
    Enter : Session (Cap e (Rec p)) (Cap (p::e) p) ()
    Zero : Session (Cap (p::e) (Var Z)) (Cap (p::e) p) ()
    Succ : Session (Cap (p::e) (Var (S k))) (Cap e (Var k)) ()
    Loop : Inf (Session (Cap (p::e) p) c a) -> Session (Cap (p::e) p) c a
    Lift : IO a -> Session c c a
    Pure : a -> Session c c a
    (>>=) : Session c1 c2 a -> (a -> Session c2 c3 b) -> Session c1 c3 b

data Fuel = Dry | More (Lazy Fuel)

partial
forever : Fuel
forever = More forever

run : Fuel -> Channel -> Session c1 c2 t -> IO (Maybe t)
run Dry _ _ = pure Nothing
run fuel chan (Send x) = 
    do True <- unsafeSend chan x | False => pure Nothing
       pure (Just ())
run fuel chan Recv = unsafeRecv _ chan
run fuel chan Close = pure (Just ())
run fuel chan Inl = 
    do True <- unsafeSend chan True | False => pure Nothing
       pure (Just ())
run fuel chan Inr = 
    do True <- unsafeSend chan False | False => pure Nothing
       pure (Just ())
run fuel chan (Offer x y) = 
    do Just choice <- unsafeRecv Bool chan | Nothing => pure Nothing
       if choice then run fuel chan x else run fuel chan y
run fuel chan Enter = pure (Just ())
run fuel chan Zero = pure (Just ())
run fuel chan Succ = pure (Just ())
run (More fuel) chan (Loop x) = run fuel chan x
run fuel chan (Lift m) = do x <- m; pure (Just x)
run fuel chan (Pure x) = pure (Just x)
run fuel chan (x >>= f) = 
    do Just r <- run fuel chan x | Nothing => pure Nothing
       run fuel chan (f r)

AtmInner : Protocol
AtmInner = atmDeposit :&: (atmWithdraw :&: Eps)
where
    atmDeposit : Protocol
    atmDeposit = Int :?: (Int :!: Var Z)
    atmWithdraw : Protocol
    atmWithdraw = Int :?: (Var Z :+: Var Z)

Atm : Protocol
Atm = String :?: (Rec AtmInner :+: Eps)

Server : Protocol -> Type
Server p = Session (Cap [] p) Done ()

Client : Protocol -> Type
Client p = Session (Cap [] (dual p)) Done ()

approved : String -> Bool
approved s = s == "myId"

atm : Int -> Server Atm
atm balance = do id <- Recv
                 if approved id then 
                    do Inl; Enter; atmInner balance
                 else
                    do Inr; Close
where
    atmInner balance = 
        Offer 
          (do amt <- Recv
              let newBalance = amt + balance
              Send newBalance
              Zero; Loop (atmInner newBalance)) 
          (Offer 
            (do amt <- Recv
                if balance >= amt then
                    do Inl; Zero; Loop (atmInner (balance - amt))
                else
                    do Inr; Zero; Loop (atmInner balance))
            Close)

atmClient : Client Atm
atmClient = do Lift (putStr "ID: ")
               id <- Lift getLine
               Send id
               Offer 
                 (do Enter; atmClientLoop) 
                 (do Lift (putStrLn "Unrecognized ID."); Close)
where
    atmClientLoop = 
      do Lift (putStrLn "Deposit/Withdraw/Quit")
         cmd <- Lift getLine
         case toLower cmd of
           "deposit" => 
             do Inl; Lift (putStr "Amount: ")
                input <- Lift getLine
                let amt = cast input
                Send amt
                balance <- Recv
                Lift (putStr "Current balance: ")
                Lift (printLn balance)
                Zero; Loop atmClientLoop
           "withdraw" => 
             do Inr; Inl; Lift (putStr "Amount: ")
                input <- Lift getLine
                let amt = cast input
                Send amt
                Offer 
                  (do Lift (putStrLn "Success.")
                      Zero; Loop atmClientLoop)
                  (do Lift (putStrLn "Insufficient funds.")
                      Zero; Loop atmClientLoop)
           "quit" => do Inr; Inr; Close
           _ => do Lift (putStrLn "Unrecognized input.")
                   Loop atmClientLoop

accept : Server p -> IO ()
accept s = do Just chan <- listen 1 | Nothing => accept s
              run forever chan s; pure ()

request : PID -> Client p -> IO ()
request pid s = do Just chan <- connect pid | Nothing => request pid s
                   run forever chan s; pure ()

main : IO ()
main = do Just srv <- spawn (accept (atm 0)) | Nothing => main
          request srv atmClient