Fast JSON

fast-json.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PartialTypeSignatures #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ViewPatterns     #-}

module Lib.JSObject where

import qualified Data.Text.Encoding              as TE
import Data.Text (Text)
import Data.Aeson
import qualified Data.Attoparsec.ByteString as P
import Data.Attoparsec.ByteString.Char8
import Control.Monad.Loops
import Data.Functor
import Data.Scientific
import Control.Lens.Internal.ByteString
import qualified Data.ByteString.Char8 as BSC8
import Data.ByteString (ByteString)
import qualified Data.Vector as V
import Data.Vector (Vector)
import Data.Bool
import qualified Data.HashMap.Strict as HM
import Control.Error.Util (hush)
import Control.Monad
import Debug.Trace
import Data.List

collectFromJSObject :: ByteString -> [ Text ] -> Maybe [ (Text, Value) ]
collectFromJSObject bs keys =
  hush $ P.parseOnly objParser bs
  where
  objParser = do
    char '{' *> skipWhile isSpace
    loop keys []
  {-# INLINE loop #-}
  loop :: [ Text ] -> [ (Text, Value) ] -> Parser [ (Text, Value) ]
  loop searchKeys foundKeys =
    if null searchKeys
    then pure foundKeys
    else do
      kv@(k, _) <- parseKVPairs <* skipWhile isSpace
      c <- satisfy (\c -> c == ',' || c == '}') <* skipWhile isSpace
      if k `elem` searchKeys
      then loop (delete k searchKeys) (kv : foundKeys)
      else case c of
        ',' -> loop searchKeys foundKeys
        '}' -> pure foundKeys

lookupInJSObject :: ByteString -> Text -> Maybe Value
lookupInJSObject bs key = join $ hush $ P.parseOnly objParser bs
  where
  objParser = do
    char '{' *> skipWhile isSpace
    loop
  loop = do
    (k, v) <- parseKVPairs
    if k == key
    then pure (Just v)
    else do
      skipWhile isSpace
      c <- satisfy (\c -> c == ',' || c == '}') <* skipWhile isSpace
      case c of
        ',' -> loop
        '}' -> pure Nothing

{-# INLINE parseKVPairs #-}
parseKVPairs :: Parser (Text, Value)
parseKVPairs = do
  key <-
    char '"'
    *> (takeTill (== '"'))
    <* char '"'
    <* skipWhile isSpace
    <* char ':'
    <* skipWhile isSpace
  value <-
    parseJSValue
  pure (TE.decodeUtf8 key, value)

{-# INLINE parseJSValue #-}
parseJSValue :: Parser Value
parseJSValue = do
  c <- peekChar'
  if (isDigit c || c == '-')
  then Number <$> parseNumberValue
  else case c of
    '"' -> String <$> parseStringValue
    '[' -> Array <$> parseVector
    '{' -> (Object . HM.fromList . V.toList) <$> parseObject
    'n' -> parseNull
    't' -> parseTrue
    'f' -> parseFalse
    x -> fail ("Unexpected char " <> [x])

data StringParserState
  = NonSpecial
  | BackSlash
  | Finished deriving Eq

{-# INLINE parseStringValue #-}
parseStringValue :: Parser Text
parseStringValue = fmap TE.decodeUtf8 $
  char '"' -- The initial '"'
  *> (fst <$> iterateUntilM ((== Finished) . snd) parseChunk (mempty, NonSpecial))
  where
  parseChunk (bs, !sps) =
    case sps of
      NonSpecial -> do
        s <- (bs <>) <$> P.takeWhile (\(toEnum . fromEnum -> c) -> c /= '\\' && c /= '"')
        anyChar <&> \case
          '"' -> (s, Finished)
          '\\' -> (s <> BSC8.singleton '\\', BackSlash)
      BackSlash -> do
        s <- (bs <>) <$> P.takeWhile ((== '\\') . toEnum . fromEnum)
        (\c -> (s <> BSC8.singleton c, NonSpecial)) <$> anyChar

{-# INLINE parseNumberValue #-}
parseNumberValue :: Parser Scientific
parseNumberValue =
  read . unpackStrict8 <$>
    P.takeWhile (\(toEnum . fromEnum -> c) -> isDigit c || c == '.' || c == 'e' || c == '-')

parseNull :: Parser Value
parseNull = Null <$ string "null"

parseTrue :: Parser Value
parseTrue = Bool True <$ string "true"

parseFalse :: Parser Value
parseFalse = Bool False <$ string "false"

{-# INLINE parseVector #-}
parseVector :: Parser (Vector Value)
parseVector =
  char '[' *> go <* char ']'
  where
    go =
      V.unfoldrM
        (bool
          (do
            v <- skipWhile isSpace *> parseJSValue <* skipWhile isSpace
            skipWhile isSpace
            peekChar' >>= \case
              ',' -> char ',' *> skipWhile isSpace $> (Just (v, False))
              ']' -> pure (Just (v, True)))
          (pure Nothing))
          False

parseObject :: Parser (Vector (Text, Value))
parseObject =
  char '{' *> go <* char '}'
  where
    go =
      V.unfoldrM
        (bool
          (do
            kv <- skipWhile isSpace *> parseKVPairs <* skipWhile isSpace
            peekChar' >>= \case
              ',' -> char ',' *> skipWhile isSpace $> (Just (kv, False))
              '}' -> pure (Just (kv, True)))
          (pure Nothing))
          False