{-# LANGUAGE AllowAmbiguousTypes, FlexibleContexts, ScopedTypeVariables, TypeApplications #-}
{-|
== Fast input

Reads space separated values from ASCII input, the typical format in competitive programming.
This is overkill in most situations. For more information see
https://github.com/meooow25/haccepted/wiki/Performance-tips#input-and-output.

-}

{-
Implementation notes:
* Yes, it uses strict Bytestring and not lazy, because it's faster. This means the input stream is
  fully loaded into memory, but that's fine for typical input file sizes (< ~10 MB).
-}

module Scanner
    ( S
    , runS
    , charS
    , byteStrS
    , intS
    , ioArrS
    , arrS
    , uArrayS
    , arrayS
    , graphS
    , graphDirS
    ) where

import Control.Monad.State.Strict
import Data.Array
import Data.Array.Base
import Data.Array.IO
import Data.Graph
import Data.Maybe
import qualified Data.ByteString.Char8 as C

import Misc ( modifyArray' )

type S = StateT C.ByteString IO

-- | Runs the given action with standard input.
runS :: S a -> IO a
runS :: S a -> IO a
runS S a
s = IO ByteString
C.getContents IO ByteString -> (ByteString -> IO a) -> IO a
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= S a -> ByteString -> IO a
forall (m :: * -> *) s a. Monad m => StateT s m a -> s -> m a
evalStateT S a
s

-- | Reads a Char.
charS :: S Char
charS :: S Char
charS = (ByteString -> (Char, ByteString)) -> S Char
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state ((ByteString -> (Char, ByteString)) -> S Char)
-> (ByteString -> (Char, ByteString)) -> S Char
forall a b. (a -> b) -> a -> b
$ (Char, ByteString)
-> Maybe (Char, ByteString) -> (Char, ByteString)
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> (Char, ByteString)
forall a. HasCallStack => [Char] -> a
error [Char]
"charS: no char") (Maybe (Char, ByteString) -> (Char, ByteString))
-> (ByteString -> Maybe (Char, ByteString))
-> ByteString
-> (Char, ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Maybe (Char, ByteString)
C.uncons (ByteString -> Maybe (Char, ByteString))
-> (ByteString -> ByteString)
-> ByteString
-> Maybe (Char, ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> ByteString -> ByteString
C.dropWhile (Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<=Char
' ')

-- | Reads a ByteString.
byteStrS :: S C.ByteString
byteStrS :: S ByteString
byteStrS = (ByteString -> (ByteString, ByteString)) -> S ByteString
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state ((ByteString -> (ByteString, ByteString)) -> S ByteString)
-> (ByteString -> (ByteString, ByteString)) -> S ByteString
forall a b. (a -> b) -> a -> b
$ (Char -> Bool) -> ByteString -> (ByteString, ByteString)
C.break (Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<=Char
' ') (ByteString -> (ByteString, ByteString))
-> (ByteString -> ByteString)
-> ByteString
-> (ByteString, ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> ByteString -> ByteString
C.dropWhile (Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<=Char
' ')

-- | Reads an Int.
intS :: S Int
intS :: S Int
intS = (ByteString -> (Int, ByteString)) -> S Int
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state ((ByteString -> (Int, ByteString)) -> S Int)
-> (ByteString -> (Int, ByteString)) -> S Int
forall a b. (a -> b) -> a -> b
$ (Int, ByteString) -> Maybe (Int, ByteString) -> (Int, ByteString)
forall a. a -> Maybe a -> a
fromMaybe ([Char] -> (Int, ByteString)
forall a. HasCallStack => [Char] -> a
error [Char]
"intS: no int") (Maybe (Int, ByteString) -> (Int, ByteString))
-> (ByteString -> Maybe (Int, ByteString))
-> ByteString
-> (Int, ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Maybe (Int, ByteString)
C.readInt (ByteString -> Maybe (Int, ByteString))
-> (ByteString -> ByteString)
-> ByteString
-> Maybe (Int, ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> ByteString -> ByteString
C.dropWhile (Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
<=Char
' ')

-- | Reads an array mutable in IO.
ioArrS :: (MArray a e IO, Ix i) => (i,i) -> S e -> S (a i e)
ioArrS :: (i, i) -> S e -> S (a i e)
ioArrS (i, i)
bnds S e
se = do
    a i e
a <- IO (a i e) -> S (a i e)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (IO (a i e) -> S (a i e)) -> IO (a i e) -> S (a i e)
forall a b. (a -> b) -> a -> b
$ (i, i) -> IO (a i e)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> m (a i e)
newArray_ (i, i)
bnds
    Int
n <- IO Int -> S Int
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (IO Int -> S Int) -> IO Int -> S Int
forall a b. (a -> b) -> a -> b
$ a i e -> IO Int
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m Int
getNumElements a i e
a
    [Int]
-> (Int -> StateT ByteString IO ()) -> StateT ByteString IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Int
0..Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] ((Int -> StateT ByteString IO ()) -> StateT ByteString IO ())
-> (Int -> StateT ByteString IO ()) -> StateT ByteString IO ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> S e
se S e -> (e -> StateT ByteString IO ()) -> StateT ByteString IO ()
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IO () -> StateT ByteString IO ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (IO () -> StateT ByteString IO ())
-> (e -> IO ()) -> e -> StateT ByteString IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a i e -> Int -> e -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> Int -> e -> m ()
unsafeWrite a i e
a Int
i
    a i e -> S (a i e)
forall (f :: * -> *) a. Applicative f => a -> f a
pure a i e
a
{-# INLINE ioArrS #-}

-- | Reads an array.
arrS :: forall a a' e i. (MArray a e IO, IArray a' e, Ix i) => (i,i) -> S e -> S (a' i e)
arrS :: (i, i) -> S e -> S (a' i e)
arrS (i, i)
bnds S e
se = (i, i) -> S e -> S (a i e)
forall (a :: * -> * -> *) e i.
(MArray a e IO, Ix i) =>
(i, i) -> S e -> S (a i e)
ioArrS @a (i, i)
bnds S e
se S (a i e) -> (a i e -> S (a' i e)) -> S (a' i e)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IO (a' i e) -> S (a' i e)
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (IO (a' i e) -> S (a' i e))
-> (a i e -> IO (a' i e)) -> a i e -> S (a' i e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a i e -> IO (a' i e)
forall i (a :: * -> * -> *) e (m :: * -> *) (b :: * -> * -> *).
(Ix i, MArray a e m, IArray b e) =>
a i e -> m (b i e)
unsafeFreeze
{-# INLINE arrS #-}

-- | Reads a UArray.
uArrayS :: (MArray IOUArray e IO, IArray UArray e, Ix i) => (i,i) -> S e -> S (UArray i e)
uArrayS :: (i, i) -> S e -> S (UArray i e)
uArrayS = forall (a' :: * -> * -> *) e i.
(MArray IOUArray e IO, IArray a' e, Ix i) =>
(i, i) -> S e -> S (a' i e)
forall (a :: * -> * -> *) (a' :: * -> * -> *) e i.
(MArray a e IO, IArray a' e, Ix i) =>
(i, i) -> S e -> S (a' i e)
arrS @IOUArray

-- | Reads an Array.
arrayS :: Ix i => (i,i) -> S e -> S (Array i e)
arrayS :: (i, i) -> S e -> S (Array i e)
arrayS = forall (a' :: * -> * -> *) e i.
(MArray IOArray e IO, IArray a' e, Ix i) =>
(i, i) -> S e -> S (a' i e)
forall (a :: * -> * -> *) (a' :: * -> * -> *) e i.
(MArray a e IO, IArray a' e, Ix i) =>
(i, i) -> S e -> S (a' i e)
arrS @IOArray

-- | Reads an undirected graph from an edge list.
graphS :: Bounds -> Int -> S Graph
graphS :: Bounds -> Int -> S Graph
graphS Bounds
bnds Int
m = do
    IOArray Int [Int]
g <- IO (IOArray Int [Int]) -> StateT ByteString IO (IOArray Int [Int])
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (IO (IOArray Int [Int])
 -> StateT ByteString IO (IOArray Int [Int]))
-> IO (IOArray Int [Int])
-> StateT ByteString IO (IOArray Int [Int])
forall a b. (a -> b) -> a -> b
$ Bounds -> [Int] -> IO (IOArray Int [Int])
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray Bounds
bnds [] :: S (IOArray Vertex [Vertex])
    Int -> StateT ByteString IO () -> StateT ByteString IO ()
forall (m :: * -> *) a. Applicative m => Int -> m a -> m ()
replicateM_ Int
m (StateT ByteString IO () -> StateT ByteString IO ())
-> StateT ByteString IO () -> StateT ByteString IO ()
forall a b. (a -> b) -> a -> b
$ do
        Int
a <- S Int
intS
        Int
b <- S Int
intS
        IO () -> StateT ByteString IO ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (IO () -> StateT ByteString IO ())
-> IO () -> StateT ByteString IO ()
forall a b. (a -> b) -> a -> b
$ IOArray Int [Int] -> Int -> ([Int] -> [Int]) -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> (e -> e) -> m ()
modifyArray' IOArray Int [Int]
g Int
a (Int
bInt -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:) IO () -> IO () -> IO ()
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> IOArray Int [Int] -> Int -> ([Int] -> [Int]) -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> (e -> e) -> m ()
modifyArray' IOArray Int [Int]
g Int
b (Int
aInt -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:)
    IO Graph -> S Graph
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (IO Graph -> S Graph) -> IO Graph -> S Graph
forall a b. (a -> b) -> a -> b
$ IOArray Int [Int] -> IO Graph
forall i (a :: * -> * -> *) e (m :: * -> *) (b :: * -> * -> *).
(Ix i, MArray a e m, IArray b e) =>
a i e -> m (b i e)
unsafeFreeze IOArray Int [Int]
g

-- | Reads a directed graph from an edge list.
graphDirS :: Bounds -> Int -> S Graph
graphDirS :: Bounds -> Int -> S Graph
graphDirS Bounds
bnds Int
m = do
    IOArray Int [Int]
g <- IO (IOArray Int [Int]) -> StateT ByteString IO (IOArray Int [Int])
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (IO (IOArray Int [Int])
 -> StateT ByteString IO (IOArray Int [Int]))
-> IO (IOArray Int [Int])
-> StateT ByteString IO (IOArray Int [Int])
forall a b. (a -> b) -> a -> b
$ Bounds -> [Int] -> IO (IOArray Int [Int])
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> e -> m (a i e)
newArray Bounds
bnds [] :: S (IOArray Vertex [Vertex])
    Int -> StateT ByteString IO () -> StateT ByteString IO ()
forall (m :: * -> *) a. Applicative m => Int -> m a -> m ()
replicateM_ Int
m (StateT ByteString IO () -> StateT ByteString IO ())
-> StateT ByteString IO () -> StateT ByteString IO ()
forall a b. (a -> b) -> a -> b
$ do
        Int
a <- S Int
intS
        Int
b <- S Int
intS
        IO () -> StateT ByteString IO ()
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (IO () -> StateT ByteString IO ())
-> IO () -> StateT ByteString IO ()
forall a b. (a -> b) -> a -> b
$ IOArray Int [Int] -> Int -> ([Int] -> [Int]) -> IO ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> (e -> e) -> m ()
modifyArray' IOArray Int [Int]
g Int
a (Int
bInt -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:)
    IO Graph -> S Graph
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (IO Graph -> S Graph) -> IO Graph -> S Graph
forall a b. (a -> b) -> a -> b
$ IOArray Int [Int] -> IO Graph
forall i (a :: * -> * -> *) e (m :: * -> *) (b :: * -> * -> *).
(Ix i, MArray a e m, IArray b e) =>
a i e -> m (b i e)
unsafeFreeze IOArray Int [Int]
g