{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TypeFamilies #-}

module Control.Monad.Interactive where

import Control.Applicative
import qualified Control.Monad.Fail as Fail
import Control.Monad.Primitive
import Control.Monad.Reader
import qualified Data.List as L
import System.Environment
import System.IO
import System.Process

class (Monad m) => MonadInteractive m where
  {-# MINIMAL sendStr, recvLine #-}
  sendStr :: String -> m ()
  recvLine :: m String

  send :: (Show a) => a -> m ()
  send = String -> m ()
forall (m :: * -> *). MonadInteractive m => String -> m ()
sendStrLn (String -> m ()) -> (a -> String) -> a -> m ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> String
forall a. Show a => a -> String
show

  sendStrLn :: String -> m ()
  sendStrLn String
cs = String -> m ()
forall (m :: * -> *). MonadInteractive m => String -> m ()
sendStr String
cs m () -> m () -> m ()
forall a b. m a -> m b -> m b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> String -> m ()
forall (m :: * -> *). MonadInteractive m => String -> m ()
sendStr String
"\n"

  recvLn :: (Read a) => m a
  recvLn = String -> a
forall a. Read a => String -> a
read (String -> a) -> m String -> m a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m String
forall (m :: * -> *). MonadInteractive m => m String
recvLine

data InteractiveHandle = InteractiveHandle
  { InteractiveHandle -> Handle
hin :: Handle
  , InteractiveHandle -> Handle
hout :: Handle
  , InteractiveHandle -> Maybe Handle
mherr :: Maybe Handle
  }

createInteractiveHandle :: IO InteractiveHandle
createInteractiveHandle :: IO InteractiveHandle
createInteractiveHandle =
  IO [String]
getArgs IO [String]
-> ([String] -> IO InteractiveHandle) -> IO InteractiveHandle
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
    [] -> InteractiveHandle -> IO InteractiveHandle
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (InteractiveHandle -> IO InteractiveHandle)
-> InteractiveHandle -> IO InteractiveHandle
forall a b. (a -> b) -> a -> b
$ Handle -> Handle -> Maybe Handle -> InteractiveHandle
InteractiveHandle Handle
stdin Handle
stdout Maybe Handle
forall a. Maybe a
Nothing
    [String]
args -> do
      let cmds :: [String]
cmds = (String -> Bool) -> [String] -> [String]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (String -> Bool) -> String -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> String -> Bool
forall a. Eq a => [a] -> [a] -> Bool
L.isPrefixOf String
"-") [String]
args
      (Just Handle
hout, Just Handle
hin, Maybe Handle
_, ProcessHandle
_) <-
        CreateProcess
-> IO (Maybe Handle, Maybe Handle, Maybe Handle, ProcessHandle)
createProcess
          (String -> CreateProcess
shell (String -> CreateProcess)
-> ([String] -> String) -> [String] -> CreateProcess
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [String] -> String
unwords ([String] -> CreateProcess) -> [String] -> CreateProcess
forall a b. (a -> b) -> a -> b
$ String
"python3 local_testing_tool.py" String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [String]
cmds)
            { std_in = CreatePipe
            , std_out = CreatePipe
            }
      let mherr :: Maybe Handle
mherr
            | String
"--verbose" String -> [String] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [String]
args = Handle -> Maybe Handle
forall a. a -> Maybe a
Just Handle
stderr
            | Bool
otherwise = Maybe Handle
forall a. Maybe a
Nothing
      InteractiveHandle -> IO InteractiveHandle
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return (InteractiveHandle -> IO InteractiveHandle)
-> InteractiveHandle -> IO InteractiveHandle
forall a b. (a -> b) -> a -> b
$ InteractiveHandle{Maybe Handle
Handle
hin :: Handle
hout :: Handle
mherr :: Maybe Handle
hout :: Handle
hin :: Handle
mherr :: Maybe Handle
..}

withInteractiveHandle :: ReaderT InteractiveHandle IO a -> IO a
withInteractiveHandle :: forall a. ReaderT InteractiveHandle IO a -> IO a
withInteractiveHandle ReaderT InteractiveHandle IO a
f = IO InteractiveHandle
createInteractiveHandle IO InteractiveHandle -> (InteractiveHandle -> IO a) -> IO a
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= ReaderT InteractiveHandle IO a -> InteractiveHandle -> IO a
forall r (m :: * -> *) a. ReaderT r m a -> r -> m a
runReaderT ReaderT InteractiveHandle IO a
f

sendDebugFormat :: String -> String
sendDebugFormat :: String -> String
sendDebugFormat = (String
"> " ++)

recvDebugFormat :: String -> String
recvDebugFormat :: String -> String
recvDebugFormat = (String
"< " ++)

instance (MonadIO m) => MonadInteractive (ReaderT InteractiveHandle m) where
  sendStr :: String -> ReaderT InteractiveHandle m ()
sendStr String
cs = (InteractiveHandle -> m ()) -> ReaderT InteractiveHandle m ()
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((InteractiveHandle -> m ()) -> ReaderT InteractiveHandle m ())
-> (InteractiveHandle -> m ()) -> ReaderT InteractiveHandle m ()
forall a b. (a -> b) -> a -> b
$ \InteractiveHandle{Maybe Handle
Handle
hin :: InteractiveHandle -> Handle
hout :: InteractiveHandle -> Handle
mherr :: InteractiveHandle -> Maybe Handle
hin :: Handle
hout :: Handle
mherr :: Maybe Handle
..} ->
    IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
      (Handle -> IO ()) -> Maybe Handle -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\Handle
herr -> Handle -> String -> IO ()
hPutStr Handle
herr (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String -> String
sendDebugFormat String
cs) Maybe Handle
mherr
      Handle -> String -> IO ()
hPutStr Handle
hout String
cs
      Handle -> IO ()
hFlush Handle
hout

  sendStrLn :: String -> ReaderT InteractiveHandle m ()
sendStrLn String
cs = (InteractiveHandle -> m ()) -> ReaderT InteractiveHandle m ()
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((InteractiveHandle -> m ()) -> ReaderT InteractiveHandle m ())
-> (InteractiveHandle -> m ()) -> ReaderT InteractiveHandle m ()
forall a b. (a -> b) -> a -> b
$ \InteractiveHandle{Maybe Handle
Handle
hin :: InteractiveHandle -> Handle
hout :: InteractiveHandle -> Handle
mherr :: InteractiveHandle -> Maybe Handle
hin :: Handle
hout :: Handle
mherr :: Maybe Handle
..} ->
    IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
      (Handle -> IO ()) -> Maybe Handle -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\Handle
herr -> Handle -> String -> IO ()
hPutStrLn Handle
herr (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String -> String
sendDebugFormat String
cs) Maybe Handle
mherr
      Handle -> String -> IO ()
hPutStrLn Handle
hout String
cs
      Handle -> IO ()
hFlush Handle
hout
  recvLine :: ReaderT InteractiveHandle m String
recvLine = (InteractiveHandle -> m String)
-> ReaderT InteractiveHandle m String
forall r (m :: * -> *) a. (r -> m a) -> ReaderT r m a
ReaderT ((InteractiveHandle -> m String)
 -> ReaderT InteractiveHandle m String)
-> (InteractiveHandle -> m String)
-> ReaderT InteractiveHandle m String
forall a b. (a -> b) -> a -> b
$ \InteractiveHandle{Maybe Handle
Handle
hin :: InteractiveHandle -> Handle
hout :: InteractiveHandle -> Handle
mherr :: InteractiveHandle -> Maybe Handle
hin :: Handle
hout :: Handle
mherr :: Maybe Handle
..} ->
    IO String -> m String
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO String -> m String) -> IO String -> m String
forall a b. (a -> b) -> a -> b
$ do
      String
res <- Handle -> IO String
hGetLine Handle
hin
      (Handle -> IO ()) -> Maybe Handle -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\Handle
herr -> Handle -> String -> IO ()
hPutStrLn Handle
herr (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String -> String
recvDebugFormat String
res) Maybe Handle
mherr
      String -> IO String
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return String
res

type Result m r = m r
type Accepted m r = Result m r
type Failed m r = Result m r
type Running a m r = a -> Result m r
type JudgeInternal m a r =
  Accepted m r -> Failed m r -> Running a m r -> Result m r

newtype Judge m a = Judge {forall {k} (m :: k -> *) a.
Judge m a -> forall (r :: k). JudgeInternal m a r
unJudge :: forall r. JudgeInternal m a r}

instance Functor (Judge m) where
  fmap :: forall a b. (a -> b) -> Judge m a -> Judge m b
fmap a -> b
f Judge m a
m = (forall (r :: k). JudgeInternal m b r) -> Judge m b
forall {k} (m :: k -> *) a.
(forall (r :: k). JudgeInternal m a r) -> Judge m a
Judge ((forall (r :: k). JudgeInternal m b r) -> Judge m b)
-> (forall (r :: k). JudgeInternal m b r) -> Judge m b
forall a b. (a -> b) -> a -> b
$ \Accepted m r
ac Accepted m r
wa Running b m r
k ->
    Judge m a -> forall (r :: k). JudgeInternal m a r
forall {k} (m :: k -> *) a.
Judge m a -> forall (r :: k). JudgeInternal m a r
unJudge Judge m a
m Accepted m r
ac Accepted m r
wa (Running b m r
k Running b m r -> (a -> b) -> a -> Accepted m r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f)
  {-# INLINE fmap #-}

instance Applicative (Judge m) where
  pure :: forall a. a -> Judge m a
pure a
x = (forall (r :: k). JudgeInternal m a r) -> Judge m a
forall {k} (m :: k -> *) a.
(forall (r :: k). JudgeInternal m a r) -> Judge m a
Judge ((forall (r :: k). JudgeInternal m a r) -> Judge m a)
-> (forall (r :: k). JudgeInternal m a r) -> Judge m a
forall a b. (a -> b) -> a -> b
$ \Accepted m r
_ Accepted m r
_ Running a m r
k -> Running a m r
k a
x
  {-# INLINE pure #-}
  Judge m (a -> b)
mf <*> :: forall a b. Judge m (a -> b) -> Judge m a -> Judge m b
<*> Judge m a
mx = (forall (r :: k). JudgeInternal m b r) -> Judge m b
forall {k} (m :: k -> *) a.
(forall (r :: k). JudgeInternal m a r) -> Judge m a
Judge ((forall (r :: k). JudgeInternal m b r) -> Judge m b)
-> (forall (r :: k). JudgeInternal m b r) -> Judge m b
forall a b. (a -> b) -> a -> b
$ \Accepted m r
ac Accepted m r
wa Running b m r
k ->
    Judge m (a -> b) -> forall (r :: k). JudgeInternal m (a -> b) r
forall {k} (m :: k -> *) a.
Judge m a -> forall (r :: k). JudgeInternal m a r
unJudge Judge m (a -> b)
mf Accepted m r
ac Accepted m r
wa (Running (a -> b) m r -> Accepted m r)
-> Running (a -> b) m r -> Accepted m r
forall a b. (a -> b) -> a -> b
$ \a -> b
f ->
      Judge m a -> forall (r :: k). JudgeInternal m a r
forall {k} (m :: k -> *) a.
Judge m a -> forall (r :: k). JudgeInternal m a r
unJudge Judge m a
mx Accepted m r
ac Accepted m r
wa (Running b m r
k Running b m r -> (a -> b) -> a -> Accepted m r
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f)
  {-# INLINE (<*>) #-}

instance Alternative (Judge m) where
  empty :: forall a. Judge m a
empty = (forall (r :: k). JudgeInternal m a r) -> Judge m a
forall {k} (m :: k -> *) a.
(forall (r :: k). JudgeInternal m a r) -> Judge m a
Judge ((forall (r :: k). JudgeInternal m a r) -> Judge m a)
-> (forall (r :: k). JudgeInternal m a r) -> Judge m a
forall a b. (a -> b) -> a -> b
$ \Accepted m r
ac Accepted m r
_ Running a m r
_ -> Accepted m r
ac
  {-# INLINE empty #-}
  Judge m a
mx <|> :: forall a. Judge m a -> Judge m a -> Judge m a
<|> Judge m a
my = (forall (r :: k). JudgeInternal m a r) -> Judge m a
forall {k} (m :: k -> *) a.
(forall (r :: k). JudgeInternal m a r) -> Judge m a
Judge ((forall (r :: k). JudgeInternal m a r) -> Judge m a)
-> (forall (r :: k). JudgeInternal m a r) -> Judge m a
forall a b. (a -> b) -> a -> b
$ \Accepted m r
ac Accepted m r
wa Running a m r
k ->
    Judge m a -> forall (r :: k). JudgeInternal m a r
forall {k} (m :: k -> *) a.
Judge m a -> forall (r :: k). JudgeInternal m a r
unJudge Judge m a
mx (Judge m a -> forall (r :: k). JudgeInternal m a r
forall {k} (m :: k -> *) a.
Judge m a -> forall (r :: k). JudgeInternal m a r
unJudge Judge m a
my Accepted m r
ac Accepted m r
wa Running a m r
k) Accepted m r
wa Running a m r
k
  {-# INLINE (<|>) #-}

instance Monad (Judge m) where
  Judge m a
mx >>= :: forall a b. Judge m a -> (a -> Judge m b) -> Judge m b
>>= a -> Judge m b
f = (forall (r :: k). JudgeInternal m b r) -> Judge m b
forall {k} (m :: k -> *) a.
(forall (r :: k). JudgeInternal m a r) -> Judge m a
Judge ((forall (r :: k). JudgeInternal m b r) -> Judge m b)
-> (forall (r :: k). JudgeInternal m b r) -> Judge m b
forall a b. (a -> b) -> a -> b
$ \Accepted m r
ac Accepted m r
wa Running b m r
k ->
    Judge m a -> forall (r :: k). JudgeInternal m a r
forall {k} (m :: k -> *) a.
Judge m a -> forall (r :: k). JudgeInternal m a r
unJudge Judge m a
mx Accepted m r
ac Accepted m r
wa (Running a m r -> Accepted m r) -> Running a m r -> Accepted m r
forall a b. (a -> b) -> a -> b
$ \a
x ->
      Judge m b -> forall (r :: k). JudgeInternal m b r
forall {k} (m :: k -> *) a.
Judge m a -> forall (r :: k). JudgeInternal m a r
unJudge (a -> Judge m b
f a
x) Accepted m r
ac Accepted m r
wa Running b m r
k
  {-# INLINE (>>=) #-}

instance Fail.MonadFail (Judge m) where
  fail :: forall a. String -> Judge m a
fail String
_ = (forall (r :: k). JudgeInternal m a r) -> Judge m a
forall {k} (m :: k -> *) a.
(forall (r :: k). JudgeInternal m a r) -> Judge m a
Judge ((forall (r :: k). JudgeInternal m a r) -> Judge m a)
-> (forall (r :: k). JudgeInternal m a r) -> Judge m a
forall a b. (a -> b) -> a -> b
$ \Accepted m r
_ Accepted m r
wa Running a m r
_ -> Accepted m r
wa
  {-# INLINE fail #-}

instance MonadTrans Judge where
  lift :: forall (m :: * -> *) a. Monad m => m a -> Judge m a
lift m a
m = (forall r. JudgeInternal m a r) -> Judge m a
forall {k} (m :: k -> *) a.
(forall (r :: k). JudgeInternal m a r) -> Judge m a
Judge ((forall r. JudgeInternal m a r) -> Judge m a)
-> (forall r. JudgeInternal m a r) -> Judge m a
forall a b. (a -> b) -> a -> b
$ \Accepted m r
_ Accepted m r
_ Running a m r
k -> m a
m m a -> Running a m r -> Accepted m r
forall a b. m a -> (a -> m b) -> m b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= Running a m r
k
  {-# INLINE lift #-}

instance (PrimMonad m) => PrimMonad (Judge m) where
  type PrimState (Judge m) = PrimState m
  primitive :: forall a.
(State# (PrimState (Judge m))
 -> (# State# (PrimState (Judge m)), a #))
-> Judge m a
primitive = m a -> Judge m a
forall (m :: * -> *) a. Monad m => m a -> Judge m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (m a -> Judge m a)
-> ((State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a)
-> (State# (PrimState m) -> (# State# (PrimState m), a #))
-> Judge m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
forall a.
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
forall (m :: * -> *) a.
PrimMonad m =>
(State# (PrimState m) -> (# State# (PrimState m), a #)) -> m a
primitive
  {-# INLINE primitive #-}

type Interactive a = Judge (ReaderT InteractiveHandle IO) a

interactive ::
  (forall r. JudgeInternal (ReaderT InteractiveHandle IO) a r) ->
  Interactive a
interactive :: forall a.
(forall r. JudgeInternal (ReaderT InteractiveHandle IO) a r)
-> Interactive a
interactive = (forall r. JudgeInternal (ReaderT InteractiveHandle IO) a r)
-> Judge (ReaderT InteractiveHandle IO) a
forall {k} (m :: k -> *) a.
(forall (r :: k). JudgeInternal m a r) -> Judge m a
Judge

runInteractive_ :: Interactive a -> ReaderT InteractiveHandle IO ()
runInteractive_ :: forall a. Interactive a -> ReaderT InteractiveHandle IO ()
runInteractive_ Interactive a
m = Interactive a
-> forall r. JudgeInternal (ReaderT InteractiveHandle IO) a r
forall {k} (m :: k -> *) a.
Judge m a -> forall (r :: k). JudgeInternal m a r
unJudge Interactive a
m (() -> ReaderT InteractiveHandle IO ()
forall a. a -> ReaderT InteractiveHandle IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()) (() -> ReaderT InteractiveHandle IO ()
forall a. a -> ReaderT InteractiveHandle IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()) (ReaderT InteractiveHandle IO ()
-> a -> ReaderT InteractiveHandle IO ()
forall a b. a -> b -> a
const (ReaderT InteractiveHandle IO ()
 -> a -> ReaderT InteractiveHandle IO ())
-> ReaderT InteractiveHandle IO ()
-> a
-> ReaderT InteractiveHandle IO ()
forall a b. (a -> b) -> a -> b
$ () -> ReaderT InteractiveHandle IO ()
forall a. a -> ReaderT InteractiveHandle IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ())

withInteractive :: Interactive a -> IO ()
withInteractive :: forall a. Interactive a -> IO ()
withInteractive = ReaderT InteractiveHandle IO () -> IO ()
forall a. ReaderT InteractiveHandle IO a -> IO a
withInteractiveHandle (ReaderT InteractiveHandle IO () -> IO ())
-> (Interactive a -> ReaderT InteractiveHandle IO ())
-> Interactive a
-> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Interactive a -> ReaderT InteractiveHandle IO ()
forall a. Interactive a -> ReaderT InteractiveHandle IO ()
runInteractive_