valueparser.hs

{-# LANGUAGE TupleSections #-}
{-# LANGUAGE OverloadedStrings #-}
module Main(main) where

import Control.Monad
import Control.Arrow ((***))
import Data.Int (Int8, Int16, Int32, Int64)
import Data.Char (isSpace, isDigit, isAlpha)
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.IO as T

import Control.Monad.ST

type STVector s = UMVec.STVector s
type Vector = UVec.Vector

dropRestOfLine, dropSpaces :: T.Text -> T.Text
dropRestOfLine = T.drop 1 . T.dropWhile (not . (=='\n'))
dropSpaces t = case T.dropWhile isSpace t of
  t' | "--" `T.isPrefixOf` t' -> dropSpaces $ dropRestOfLine t'
     | otherwise -> t'

type ReadValue v = T.Text -> (Maybe v, T.Text)

symbol :: Char -> T.Text -> Maybe T.Text
symbol c t
  | Just (c', t') <- T.uncons t, c' == c = Just $ dropSpaces t'
  | otherwise = Nothing

-- (Used elements, shape, elements, remaining input)
type State s v = (Int, Vector Int, STVector s v, T.Text)

readArrayElemsST :: UMVec.Unbox v =>
                    Int -> Int -> ReadValue v -> State s v
                 -> ST s (Maybe (Int, State s v))
readArrayElemsST j r rv s = do
  ms <- readRankedArrayOfST r rv s
  case ms of
    Just (i, shape, arr, t)
      | Just t' <- symbol ',' t ->
          readArrayElemsST (j+1) r rv (i, shape, arr, t')
      | otherwise -> return $ Just (j, (i, shape, arr, t))
    _ ->
      return $ Just (0, s)

updateShape :: Int -> Int -> Vector Int -> Maybe (Vector Int)
updateShape d n shape =
  if old_n < 0
  then Just $ shape UVec.// [(r-d, n)]
  else if old_n == n then Just shape
       else Nothing
  where r = UVec.length shape
        old_n = shape UVec.! (r-d)

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

readRankedArrayOfST :: UMVec.Unbox v =>
                 Int -> ReadValue v -> State s v
              -> ST s (Maybe (State s v))
readRankedArrayOfST 0 rv (i, shape, arr, t)
  | (Just v, t') <- rv t = do
      arr' <- growIfFilled i arr
      UMVec.write arr' i v
      return $ Just (i+1, shape, arr', t')
readRankedArrayOfST r rv (i, shape, arr, t)
  | Just t' <- symbol '[' t = do
      ms <- readArrayElemsST 1 (r-1) rv (i, shape, arr, t')
      return $ do
        (j, s) <- ms
        closeArray r j s
readRankedArrayOfST _ _ _ =
  return Nothing

closeArray :: Int -> Int -> State s v -> Maybe (State s v)
closeArray r j (i, shape, arr, t) = do
  t' <- symbol ']' t
  shape' <- updateShape r j shape
  return (i, shape', arr, t')

readRankedArrayOf :: UMVec.Unbox v =>
                     Int -> ReadValue v -> T.Text -> Maybe (Vector Int, Vector v, T.Text)
readRankedArrayOf r rv t = runST $ do
  empty <- UMVec.new 1024
  ms <- readRankedArrayOfST r rv (0, UVec.replicate r (-1), empty, t)
  case ms of
    Just (i, shape, arr, t') -> do
      arr' <- freeze (UMVec.slice 0 i arr)
      return $ Just (shape, arr', t')
    Nothing ->
      return Nothing

data Array = Int32Array (Vector Int) (Vector Int32)
           | BoolArray (Vector Int) (Vector Bool)
           deriving Show

tryRead :: Read a => T.Text -> Maybe a
tryRead s = case reads s' of
  [(x, "")] -> Just x
  _ -> Nothing
  where s' = T.unpack s

readInt32 :: ReadValue Int32
readInt32 = (tryRead *** dropSpaces) . T.span constituent
  where constituent c = c == '-' || isDigit c

readBool :: ReadValue Bool
readBool = (tryRead *** dropSpaces) . T.span isAlpha

readArray :: T.Text -> Maybe (Array, T.Text)
readArray full_t = insideBrackets 0 full_t
  where insideBrackets r t = maybe (tryValueAndReadArray r t) (insideBrackets (r+1)) $ symbol '[' t
        tryWith f mk r t
          | (Just _, _) <- f t = do
              (shape, arr, rest_t) <- readRankedArrayOf r f full_t
              return (mk shape arr, rest_t)
          | otherwise = Nothing
        tryValueAndReadArray r t=
          tryWith readInt32 Int32Array r t `mplus`
          tryWith readBool BoolArray r t

readArrays :: T.Text -> Maybe [Array]
readArrays t
  | T.null t = Just []
  | otherwise = do (a, t') <- readArray t
                   (a:) <$> readArrays t'

main :: IO ()
main = do
  res <- readArrays <$> T.getContents
  case res of
    Nothing -> error "invalid input"
    Just arrs -> mapM_ (putStrLn . describe) arrs
  where describe (Int32Array shape _) = "Int32 array of shape " ++ show shape
        describe (BoolArray shape _) = "Bool array of shape " ++ show shape