unfoldr_v5.hs

gistfile1.txt

{-# LANGUAGE BangPatterns #-}

import Control.Arrow (first, second)
import Data.Char
import Control.Monad.ST
import Data.Vector.Unboxed.Mutable (STVector)
import qualified Data.Vector.Unboxed.Mutable as UMVec
import qualified Data.Vector.Unboxed as UVec
import Data.Vector.Generic (freeze)
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, STVector s Int)

growIfFilled :: Int -> STVector s Int -> ST s (STVector s Int)
growIfFilled i arr =
  if i >= capacity
  then UMVec.grow arr capacity
  else return arr
  where capacity = UMVec.length arr

readIntsST :: State s -> T.Text -> ST s (STVector s Int)
readIntsST (i, arr) t =
  case readInt t of
    (Just x, t') -> do
      arr' <- growIfFilled i arr
      UMVec.write arr' i x
      readIntsST (i+1, arr') t'
    _ ->
      UMVec.clone $ UMVec.slice 0 i arr

readInts :: T.Text -> UVec.Vector Int
readInts t = runST $ do
  empty <- UMVec.new 1024
  arr <- readIntsST (0, empty) t
  freeze arr

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