unfoldr_v4.hs

{-# LANGUAGE BangPatterns #-}

import Control.Arrow (first, second)
import Data.Char
import Control.Monad.ST
import Data.Array.ST
import Data.Array.Unboxed
import qualified Data.Text as T
import qualified Data.Text.Read as TR
import qualified Data.Text.IO as TIO
import System.IO
import System.TimeIt

toInt :: T.Text -> Maybe Int
toInt x =
  case TR.decimal x of
    Right (v,_) -> Just v
    Left _ -> Nothing

readInt :: T.Text -> (Maybe Int, T.Text)
readInt = second (T.dropWhile isSpace) . first toInt . T.span isDigit

-- The number of filled elements in the array, and the array.
type State s = (Int, STUArray s Int Int)

resizeArrayTo :: Int -> STUArray s Int Int -> ST s (STUArray s Int Int)
resizeArrayTo newsize arr = do
  xs <- getElems arr
  newListArray (0, newsize-1) xs

growIfFilled :: Int -> STUArray s Int Int -> ST s (STUArray s Int Int)
growIfFilled i arr = do
  (_, upper) <- getBounds arr
  if i >= upper
    then resizeArrayTo (upper*2) arr
    else return arr

readIntsST :: State s -> T.Text -> ST s (STUArray s Int Int)
readIntsST (i, arr) t =
  case readInt t of
    (Just x, t') -> do
      arr' <- growIfFilled i arr
      writeArray arr' i x
      readIntsST (i+1, arr') t'
    _ ->
      resizeArrayTo i arr

readInts :: T.Text -> Array Int Int
readInts t = runST $ do
  empty <- newArray_ (0,1024)
  arr <- readIntsST (0, empty) t
  freeze arr

main :: IO ()
main = do
  !text <- TIO.hGetContents stdin
  timeIt $ print $ bounds $ readInts text