{-# LANGUAGE PatternGuards, OverloadedStrings #-}
import Control.Applicative as A
import Prelude hiding (filter)
import Data.Maybe
import Data.Char
import Data.Monoid
import Data.Aeson
import Data.Aeson.Parser (jstring, value)
import qualified Data.HashMap.Strict as H
import Data.HashMap.Strict (HashMap)
import qualified Data.Text as T
import Data.Text (Text)
import Data.Text.Encoding (encodeUtf8)
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as L
import qualified Data.ByteString.Lazy.Char8 as L8
import qualified Data.Vector as V
import Data.Vector (Vector)
import qualified Data.HashSet as S
import Data.Attoparsec.Char8 hiding (Result, parse)
import qualified Data.Attoparsec.Lazy as L
import System.Environment
type ValueOp = Value -> Value
type ValueBinOp = Value -> Value -> Value
type Filter = [Value] -> [Value]
type Obj = HashMap Text
err2 :: Value -> Value -> String -> a
err2 x y msg = error . unwords $ [show x, "and", show y, msg]
vecDiff :: Vector Value -> Vector Value -> Vector Value
x `vecDiff` y = V.filter p x
where p = not . (`S.member`s)
s = S.fromList (V.toList y)
(+|), (-|), (/|), ( *| ) :: ValueBinOp
Number x +| Number y = Number (x + y)
String x +| String y = String (x <> y)
Array x +| Array y = Array (x <> y)
Object x +| Object y = Object (y <> x) -- Right biased
x +| y = err2 x y "cannot be added"
Number x -| Number y = Number (x - y)
Array x -| Array y = Array (x `vecDiff` y)
x -| y = err2 x y "cannot be subtracted"
Number x *| Number y = Number (x * y)
x *| y = err2 x y "cannot be multiplied"
Number x /| Number y = Number (x / y)
x /| y = err2 x y "cannot be divided"
lengthFi :: Value -> Int
lengthFi Null = 0
lengthFi (Array v) = V.length v
lengthFi (Object o) = length . H.toList $ o
lengthFi (String s) = B.length . encodeUtf8 $ s
lengthFi (Number _) = error "number has no length"
lengthFi (Bool _) = error "boolean has no length"
lengthOp :: ValueOp
lengthOp = toJSON . lengthFi
keysOp :: ValueOp
keysOp (Array v) = toJSON [0.. V.length v - 1]
keysOp (Object o) = toJSON $ H.keys o
keysOp Number{} = error "number has no keys"
keysOp Null = error "null has no keys"
keysOp Bool{} = error "boolean has no keys"
keysOp String{} = error "string has no keys"
at :: Value -> Value -> Maybe Value
Object o `at` String s = H.lookup s o
Array a `at` Number mn
| Success n <- fromJSON (Number mn) = a V.!? n
_ `at` _ = Nothing
atF :: Value -> Filter
atF key = fmap (fromMaybe Null . (`at` key))
allF :: Filter
allF xs = [ y | Array ys <- xs, y <- V.toList ys ]
bothF :: Filter -> Filter -> Filter
bothF f g xs = f xs ++ g xs
arrayF :: Filter -> Filter
arrayF f xs = [Array (V.fromList . f $ xs)]
objectF :: Obj Filter -> Filter
objectF o xs = [Object . H.fromList $ [ (k,y) | (k,f) <- H.toList o, y <- f xs ]]
{- Confusing:
$ jq -n -c '[(1,2,3,4) | .+.]'
[2,4,6,8]
$ jq -n -c '[(1,2,3,4) | .+1]'
[2,3,4,5]
But:
$ jq -n -c '[(1,2)+(2,1)]'
[3,4,2,3]
instead of [3,3]
-}
ap2F :: ValueBinOp -> Filter -> Filter -> Filter
-- ap2F op f g xs = [ op x y | x <- f xs, y <- g xs ]
ap2F op f g xs = zipWith op (f xs) (g xs)
emptyF :: Filter
emptyF _ = []
constF :: Value -> Filter
constF v xs = [ v | _ <- xs ]
-- Filter
data F = IdF -- .
| CompF F F -- f | g
| AtF Value -- .[key]
| AllF -- .[]
| BothF F F -- f, g
| ArrayF F -- [f]
| ObjectF (Obj F) -- {a: f, b: g}
| EmptyF -- empty
| OpF Op -- length, keys
| Ap2F BinOp F F -- f + g, f - g, f / g, f * g
| ConstF Value -- 1, "foo", null
| ErrorF String
deriving (Show)
data Op = Length | Keys
deriving (Show)
data BinOp = Plus | Minus | Times | Div
deriving (Show)
-- .key
keyF :: Text -> F
keyF = AtF . String
-- def map(f): [.[] | f];
mapF :: F -> F
mapF f = ArrayF (AllF `CompF` f)
-- TODO: deal properly with "f + g - h" and "f - g + h"
binOpF :: BinOp -> [F] -> F
binOpF _ [] = IdF
binOpF _ [x] = x
binOpF op [x,y] = Ap2F op x y
binOpF _ _ = error "binOpF: not supported yet"
concatF, composeF, sumF, productF, minusF, divF :: [F] -> F
concatF [] = IdF
concatF xs = foldr1 BothF xs
composeF [] = IdF
composeF xs = foldr1 CompF xs
sumF = binOpF Plus
productF = binOpF Times
minusF = binOpF Minus
divF = binOpF Div
valueBinOp :: BinOp -> ValueBinOp
valueBinOp Plus = (+|)
valueBinOp Times = (*|)
valueBinOp Minus = (-|)
valueBinOp Div = (/|)
valueOp :: Op -> ValueOp
valueOp Keys = keysOp
valueOp Length = lengthOp
filter :: F -> Filter
filter IdF = id
filter (CompF f g) = filter g . filter f
filter (AtF key) = atF key
filter AllF = allF
filter (BothF f g) = bothF (filter f) (filter g)
filter (ArrayF f) = arrayF (filter f)
filter (ObjectF o) = objectF (fmap filter o)
filter (Ap2F op f g) = ap2F (valueBinOp op) (filter f) (filter g)
filter EmptyF = emptyF
filter (OpF op) = fmap (valueOp op)
filter (ConstF v) = constF v
filter (ErrorF err) = error err
parseSimpleFilter, parseTimesFilter, parseDivFilter,
parseMinusFilter, parsePlusFilter, parseCommaFilter,
parseNoCommaFilter, parseFilter, parseDotFilter :: Parser F
parseDotFilter
= pure AllF <* string "[]"
<|> AtF <$> (char '[' *> value <* tok ']')
<|> keyF <$> bareWord
<|> pure IdF
bareWord :: Parser Text
bareWord = T.pack <$> some (satisfy (\c -> isAscii c && isAlpha c))
parseConstFilter :: Parser Value
parseConstFilter
= String <$> jstring
<|> Number <$> number
<|> pure (Bool True) <* string "true"
<|> pure (Bool False) <* string "false"
<|> pure Null <* string "null"
<|> pure (Array mempty) <* string "[]"
<|> pure (Object mempty) <* string "{}"
parseOp :: Parser Op
parseOp = pure Length <* string "length"
<|> pure Keys <* string "keys"
tok :: Char -> Parser Char
tok c = skipSpace *> char c
parseSimpleFilter
= skipSpace *>
( char '.' *> skipSpace *> parseDotFilter
<|> string "empty" *> pure EmptyF
<|> OpF <$> parseOp
<|> mapF <$> (string "map" *> tok '(' *> parseFilter <* tok ')')
<|> ConstF <$> parseConstFilter
<|> ArrayF <$> (char '[' *> parseFilter <* tok ']')
<|> ObjectF <$> obj parseNoCommaFilter
<|> char '(' *> parseFilter <* tok ')'
)
parseTimesFilter = productF <$> parseSimpleFilter `sepBy` (tok '*')
parseDivFilter = divF <$> parseTimesFilter `sepBy` (tok '/')
parsePlusFilter = sumF <$> parseDivFilter `sepBy` (tok '+')
parseMinusFilter = minusF <$> parsePlusFilter `sepBy` (tok '-')
parseCommaFilter = concatF <$> parseMinusFilter `sepBy` (tok ',')
parseNoCommaFilter = composeF <$> parseMinusFilter `sepBy` (tok '|')
parseFilter = composeF <$> parseCommaFilter `sepBy1` (tok '|')
obj :: Parser a -> Parser (Obj a)
obj p = char '{' *> objectValues (skipSpace *> (bareWord <|> jstring)) p
-- From: https://github.com/bos/aeson/blob/master/Data/Aeson/Parser/Internal.hs
objectValues :: Parser Text -> Parser a -> Parser (Obj a)
objectValues str val = do
skipSpace
let pair = do
a <- str <* skipSpace
b <- char ':' *> skipSpace *> val
return (a,b)
vals <- ((pair <* skipSpace) `sepBy` (char ',' *> skipSpace)) <* char '}'
return (H.fromList vals)
{-# INLINE objectValues #-}
parse :: Parser a -> L.ByteString -> Maybe a
parse p s =
case L.parse p s of
L.Done rest v | L.null rest -> Just v
_ -> Nothing
{-# INLINE parse #-}
stream :: Parser [Value]
stream = value `sepBy` skipSpace
main :: IO ()
main = do [arg] <- getArgs
let Just f = parse (parseFilter <* skipSpace) (L8.pack arg)
-- print f
input <- maybe (fail "JSON decoding") return . parse (stream <* skipSpace) =<< L.getContents
mapM_ (L8.putStrLn . encode) $ filter f input