{-# LANGUAGE RecordWildCards #-}
module Data.Aeson.JSONPath.Query.Selector
  ( qSelector
  , qSelectorLocated
  )
  where

import Data.Aeson     (Value)
import Data.Vector    (Vector)
import Data.Text      (Text)

import qualified Data.Aeson            as JSON
import qualified Data.Aeson.KeyMap     as KM
import qualified Data.Aeson.Key        as K
import qualified Data.Text             as T
import qualified Data.Vector           as V

import Data.Aeson.JSONPath.Types
import Data.Aeson.JSONPath.Query.Filter

import Prelude

qSelector :: Selector Query -> QueryState -> Vector Value
qSelector :: Selector Query -> QueryState -> Vector Value
qSelector (Name Text
key) QueryState{curVal :: QueryState -> Value
curVal=(JSON.Object Object
obj)} = Vector Value
-> (Value -> Vector Value) -> Maybe Value -> Vector Value
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Vector Value
forall a. Vector a
V.empty Value -> Vector Value
forall a. a -> Vector a
V.singleton (Maybe Value -> Vector Value) -> Maybe Value -> Vector Value
forall a b. (a -> b) -> a -> b
$ Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
KM.lookup (Text -> Key
K.fromText Text
key) Object
obj
qSelector (Name Text
_) QueryState
_ = Vector Value
forall a. Vector a
V.empty
qSelector (Index Int
idx) QueryState{curVal :: QueryState -> Value
curVal=(JSON.Array Vector Value
arr)} = Vector Value
-> (Value -> Vector Value) -> Maybe Value -> Vector Value
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Vector Value
forall a. Vector a
V.empty Value -> Vector Value
forall a. a -> Vector a
V.singleton (Maybe Value -> Vector Value) -> Maybe Value -> Vector Value
forall a b. (a -> b) -> a -> b
$ if Int
idx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 then Vector Value -> Int -> Maybe Value
forall a. Vector a -> Int -> Maybe a
(V.!?) Vector Value
arr Int
idx else Vector Value -> Int -> Maybe Value
forall a. Vector a -> Int -> Maybe a
(V.!?) Vector Value
arr (Int
idx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Vector Value -> Int
forall a. Vector a -> Int
V.length Vector Value
arr)
qSelector (Index Int
_) QueryState
_ = Vector Value
forall a. Vector a
V.empty
qSelector (ArraySlice (Maybe Int, Maybe Int, Int)
startEndStep) QueryState{curVal :: QueryState -> Value
curVal=(JSON.Array Vector Value
arr) } = (Maybe Int, Maybe Int, Int) -> Vector Value -> Vector Value
sliceArray (Maybe Int, Maybe Int, Int)
startEndStep Vector Value
arr
qSelector (ArraySlice (Maybe Int, Maybe Int, Int)
_) QueryState
_ = Vector Value
forall a. Vector a
V.empty
qSelector (Filter LogicalOrExpr Query
orExpr) QueryState
qS = LogicalOrExpr Query -> QueryState -> Vector Value
filterOrExpr LogicalOrExpr Query
orExpr QueryState
qS
qSelector Selector Query
WildcardSelector QueryState{Value
Query -> QueryState -> Vector Value
curVal :: QueryState -> Value
rootVal :: Value
curVal :: Value
executeQuery :: Query -> QueryState -> Vector Value
rootVal :: QueryState -> Value
executeQuery :: QueryState -> Query -> QueryState -> Vector Value
..} = case Value
curVal of
    (JSON.Object Object
obj) -> [Value] -> Vector Value
forall a. [a] -> Vector a
V.fromList ([Value] -> Vector Value) -> [Value] -> Vector Value
forall a b. (a -> b) -> a -> b
$ Object -> [Value]
forall v. KeyMap v -> [v]
KM.elems Object
obj
    (JSON.Array Vector Value
arr)  -> Vector Value
arr
    Value
_                 -> Vector Value
forall a. Vector a
V.empty


qSelectorLocated :: Selector Query -> QueryState -> String -> Vector (String,Value)
qSelectorLocated :: Selector Query -> QueryState -> [Char] -> Vector ([Char], Value)
qSelectorLocated (Name Text
key) QueryState{curVal :: QueryState -> Value
curVal=(JSON.Object Object
obj)} [Char]
loc = Vector ([Char], Value)
-> (Value -> Vector ([Char], Value))
-> Maybe Value
-> Vector ([Char], Value)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Vector ([Char], Value)
forall a. Vector a
V.empty (\Value
x-> ([Char], Value) -> Vector ([Char], Value)
forall a. a -> Vector a
V.singleton ([Char] -> Text -> [Char]
toPathKey [Char]
loc Text
key, Value
x)) (Maybe Value -> Vector ([Char], Value))
-> Maybe Value -> Vector ([Char], Value)
forall a b. (a -> b) -> a -> b
$ Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
KM.lookup (Text -> Key
K.fromText Text
key) Object
obj
qSelectorLocated (Name Text
_) QueryState
_ [Char]
_ = Vector ([Char], Value)
forall a. Vector a
V.empty
qSelectorLocated (Index Int
idx) QueryState{curVal :: QueryState -> Value
curVal=(JSON.Array Vector Value
arr)} [Char]
loc = Vector ([Char], Value)
-> (Value -> Vector ([Char], Value))
-> Maybe Value
-> Vector ([Char], Value)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Vector ([Char], Value)
forall a. Vector a
V.empty (\Value
x-> ([Char], Value) -> Vector ([Char], Value)
forall a. a -> Vector a
V.singleton ([Char]
newLocation, Value
x)) (Maybe Value -> Vector ([Char], Value))
-> Maybe Value -> Vector ([Char], Value)
forall a b. (a -> b) -> a -> b
$ Vector Value -> Int -> Maybe Value
forall a. Vector a -> Int -> Maybe a
(V.!?) Vector Value
arr (Int -> Int
getIndex Int
idx)
  where
    newLocation :: [Char]
newLocation = [Char]
loc [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"[" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show (Int -> Int
getIndex Int
idx) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"]"
    getIndex :: Int -> Int
getIndex Int
i = if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 then Int
i else Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Vector Value -> Int
forall a. Vector a -> Int
V.length Vector Value
arr
qSelectorLocated (Index Int
_) QueryState
_ [Char]
_ = Vector ([Char], Value)
forall a. Vector a
V.empty
qSelectorLocated (ArraySlice (Maybe Int
start,Maybe Int
end,Int
step)) QueryState{curVal :: QueryState -> Value
curVal=(JSON.Array Vector Value
arr)} [Char]
loc = (Maybe Int, Maybe Int, Int)
-> Vector ([Char], Value) -> Vector ([Char], Value)
sliceArrayLocated (Maybe Int
start,Maybe Int
end,Int
step) (Vector ([Char], Value) -> Vector ([Char], Value))
-> Vector ([Char], Value) -> Vector ([Char], Value)
forall a b. (a -> b) -> a -> b
$ Vector [Char] -> Vector Value -> Vector ([Char], Value)
forall a b. Vector a -> Vector b -> Vector (a, b)
V.zip ([[Char]] -> Vector [Char]
forall a. [a] -> Vector a
V.fromList [[Char]]
locs) Vector Value
arr
  where
    locs :: [[Char]]
locs = [ [Char]
loc [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"[" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
i [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"]" | Int
i <- [Int]
indices ]
    indices :: [Int]
indices = [Int
0..(Vector Value -> Int
forall a. Vector a -> Int
V.length Vector Value
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)]
qSelectorLocated (ArraySlice (Maybe Int, Maybe Int, Int)
_) QueryState
_ [Char]
_ = Vector ([Char], Value)
forall a. Vector a
V.empty
qSelectorLocated (Filter LogicalOrExpr Query
orExpr) QueryState
qS [Char]
loc = LogicalOrExpr Query
-> QueryState -> [Char] -> Vector ([Char], Value)
filterOrExprLocated LogicalOrExpr Query
orExpr QueryState
qS [Char]
loc
qSelectorLocated Selector Query
WildcardSelector QueryState{Value
Query -> QueryState -> Vector Value
curVal :: QueryState -> Value
rootVal :: QueryState -> Value
executeQuery :: QueryState -> Query -> QueryState -> Vector Value
rootVal :: Value
curVal :: Value
executeQuery :: Query -> QueryState -> Vector Value
..} [Char]
loc = case Value
curVal of
    (JSON.Object Object
obj) -> [([Char], Value)] -> Vector ([Char], Value)
forall a. [a] -> Vector a
V.fromList ([([Char], Value)] -> Vector ([Char], Value))
-> [([Char], Value)] -> Vector ([Char], Value)
forall a b. (a -> b) -> a -> b
$ [[Char]] -> [Value] -> [([Char], Value)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Object -> [[Char]]
forall {v}. KeyMap v -> [[Char]]
locsWithKeys Object
obj) (Object -> [Value]
forall v. KeyMap v -> [v]
KM.elems Object
obj)
    (JSON.Array Vector Value
arr)  -> Vector [Char] -> Vector Value -> Vector ([Char], Value)
forall a b. Vector a -> Vector b -> Vector (a, b)
V.zip ([[Char]] -> Vector [Char]
forall a. [a] -> Vector a
V.fromList (Vector Value -> [[Char]]
forall {a}. Vector a -> [[Char]]
locsWithIdxs Vector Value
arr)) Vector Value
arr
    Value
_                 -> Vector ([Char], Value)
forall a. Vector a
V.empty
    where
      locsWithKeys :: KeyMap v -> [[Char]]
locsWithKeys KeyMap v
obj = (Key -> [Char]) -> [Key] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map ([Char] -> Text -> [Char]
toPathKey [Char]
loc (Text -> [Char]) -> (Key -> Text) -> Key -> [Char]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Key -> Text
K.toText) (KeyMap v -> [Key]
forall v. KeyMap v -> [Key]
KM.keys KeyMap v
obj)
      locsWithIdxs :: Vector a -> [[Char]]
locsWithIdxs Vector a
arr = (Int -> [Char]) -> [Int] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map (\Int
x -> [Char]
loc [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"[" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ Int -> [Char]
forall a. Show a => a -> [Char]
show Int
x [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"]") [Int
0..(Vector a -> Int
forall a. Vector a -> Int
V.length Vector a
arr)]


sliceArray :: (Maybe Int, Maybe Int, Int) -> Vector Value -> Vector Value
sliceArray :: (Maybe Int, Maybe Int, Int) -> Vector Value -> Vector Value
sliceArray (Maybe Int
start,Maybe Int
end,Int
step) Vector Value
vec =
  case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
step Int
0 of
    Ordering
GT -> Int -> Int -> Int -> Vector Value -> Vector Value
forall {a}. Int -> Int -> Int -> Vector a -> Vector a
getSliceForward (Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
0 Int -> Int
normalize Maybe Int
start) (Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
len Int -> Int
normalize Maybe Int
end) Int
step Vector Value
vec
    Ordering
LT -> Int -> Int -> Int -> Vector Value -> Vector Value
forall {a}. Int -> Int -> Int -> Vector a -> Vector a
getSliceReverse (Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Int
lenInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int -> Int
normalize Maybe Int
start) (Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (-Int
1) Int -> Int
normalize Maybe Int
end) Int
step Vector Value
vec
    Ordering
EQ -> Vector Value
forall a. Vector a
V.empty
    where
      -- TODO: Looks kinda ugly, make it pretty <3
      len :: Int
len = Vector Value -> Int
forall a. Vector a -> Int
V.length Vector Value
vec
      normalize :: Int -> Int
normalize Int
i = if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 then Int
i else Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i

      getSliceForward :: Int -> Int -> Int -> Vector a -> Vector a
getSliceForward Int
st Int
en Int
stp Vector a
arr = Int -> Vector a -> Vector a
loop Int
lower Vector a
forall a. Vector a
V.empty
        where
          (Int
lower,Int
upper) = (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
st Int
0) Int
len, Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
en Int
0) Int
len)

          loop :: Int -> Vector a -> Vector a
loop Int
i Vector a
acc =
            if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
upper
              then Int -> Vector a -> Vector a
loop (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
stp) (Vector a -> Vector a) -> Vector a -> Vector a
forall a b. (a -> b) -> a -> b
$ Vector a -> a -> Vector a
forall a. Vector a -> a -> Vector a
V.snoc Vector a
acc (a -> Vector a) -> a -> Vector a
forall a b. (a -> b) -> a -> b
$ Vector a -> Int -> a
forall a. Vector a -> Int -> a
(V.!) Vector a
arr (Int -> Int
normalize Int
i)
            else
              Vector a
acc

      getSliceReverse :: Int -> Int -> Int -> Vector a -> Vector a
getSliceReverse Int
st Int
en Int
stp Vector a
arr = Int -> Vector a -> Vector a
loop Int
upper Vector a
forall a. Vector a
V.empty
        where
          (Int
lower,Int
upper) = (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
en (-Int
1)) (Int
lenInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1), Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
st (-Int
1)) (Int
lenInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1))

          loop :: Int -> Vector a -> Vector a
loop Int
i Vector a
acc =
            if Int
lower Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
i
              then Int -> Vector a -> Vector a
loop (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
stp) (Vector a -> Vector a) -> Vector a -> Vector a
forall a b. (a -> b) -> a -> b
$ Vector a -> a -> Vector a
forall a. Vector a -> a -> Vector a
V.snoc Vector a
acc (a -> Vector a) -> a -> Vector a
forall a b. (a -> b) -> a -> b
$ Vector a -> Int -> a
forall a. Vector a -> Int -> a
(V.!) Vector a
arr (Int -> Int
normalize Int
i)
            else
              Vector a
acc


sliceArrayLocated :: (Maybe Int, Maybe Int, Int) -> Vector (String,Value) -> Vector (String,Value)
sliceArrayLocated :: (Maybe Int, Maybe Int, Int)
-> Vector ([Char], Value) -> Vector ([Char], Value)
sliceArrayLocated (Maybe Int
start,Maybe Int
end,Int
step) Vector ([Char], Value)
vec =
  case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
step Int
0 of
    Ordering
GT -> Int
-> Int -> Int -> Vector ([Char], Value) -> Vector ([Char], Value)
forall {a}. Int -> Int -> Int -> Vector a -> Vector a
getSliceForward (Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
0 Int -> Int
normalize Maybe Int
start) (Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
len Int -> Int
normalize Maybe Int
end) Int
step Vector ([Char], Value)
vec
    Ordering
LT -> Int
-> Int -> Int -> Vector ([Char], Value) -> Vector ([Char], Value)
forall {a}. Int -> Int -> Int -> Vector a -> Vector a
getSliceReverse (Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Int
lenInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int -> Int
normalize Maybe Int
start) (Int -> (Int -> Int) -> Maybe Int -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (-Int
1) Int -> Int
normalize Maybe Int
end) Int
step Vector ([Char], Value)
vec
    Ordering
EQ -> Vector ([Char], Value)
forall a. Vector a
V.empty
    where
      -- TODO: Looks kinda ugly, make it pretty <3
      len :: Int
len = Vector ([Char], Value) -> Int
forall a. Vector a -> Int
V.length Vector ([Char], Value)
vec
      normalize :: Int -> Int
normalize Int
i = if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 then Int
i else Int
len Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i

      getSliceForward :: Int -> Int -> Int -> Vector a -> Vector a
getSliceForward Int
st Int
en Int
stp Vector a
arr = Int -> Vector a -> Vector a
loop Int
lower Vector a
forall a. Vector a
V.empty
        where
          (Int
lower,Int
upper) = (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
st Int
0) Int
len, Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
en Int
0) Int
len)

          loop :: Int -> Vector a -> Vector a
loop Int
i Vector a
acc =
            if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
upper
              then Int -> Vector a -> Vector a
loop (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
stp) (Vector a -> Vector a) -> Vector a -> Vector a
forall a b. (a -> b) -> a -> b
$ Vector a -> a -> Vector a
forall a. Vector a -> a -> Vector a
V.snoc Vector a
acc (a -> Vector a) -> a -> Vector a
forall a b. (a -> b) -> a -> b
$ Vector a -> Int -> a
forall a. Vector a -> Int -> a
(V.!) Vector a
arr (Int -> Int
normalize Int
i)
            else
              Vector a
acc

      getSliceReverse :: Int -> Int -> Int -> Vector a -> Vector a
getSliceReverse Int
st Int
en Int
stp Vector a
arr = Int -> Vector a -> Vector a
loop Int
upper Vector a
forall a. Vector a
V.empty
        where
          (Int
lower,Int
upper) = (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
en (-Int
1)) (Int
lenInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1), Int -> Int -> Int
forall a. Ord a => a -> a -> a
min (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
st (-Int
1)) (Int
lenInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1))

          loop :: Int -> Vector a -> Vector a
loop Int
i Vector a
acc =
            if Int
lower Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
i
              then Int -> Vector a -> Vector a
loop (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
stp) (Vector a -> Vector a) -> Vector a -> Vector a
forall a b. (a -> b) -> a -> b
$ Vector a -> a -> Vector a
forall a. Vector a -> a -> Vector a
V.snoc Vector a
acc (a -> Vector a) -> a -> Vector a
forall a b. (a -> b) -> a -> b
$ Vector a -> Int -> a
forall a. Vector a -> Int -> a
(V.!) Vector a
arr (Int -> Int
normalize Int
i)
            else
              Vector a
acc


toPathKey :: String -> Text -> String
toPathKey :: [Char] -> Text -> [Char]
toPathKey [Char]
loc Text
key = [Char]
loc [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"['" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char] -> [Char]
escapeEscapees (Text -> [Char]
T.unpack Text
key) [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"']"
  where
    escapeEscapees :: String -> String
    escapeEscapees :: [Char] -> [Char]
escapeEscapees [] = []
    escapeEscapees (Char
x:[Char]
xs) = Char -> [Char]
checkChar Char
x [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char] -> [Char]
escapeEscapees [Char]
xs
      where
        -- TODO: Do we need to escape unicode chars?
        checkChar :: Char -> [Char]
checkChar Char
'\\' = [Char
'\\', Char
'\\']
        checkChar Char
'\'' = [Char
'\\', Char
'\'']
        checkChar Char
'\b' = [Char
'\\', Char
'b']
        checkChar Char
'\r' = [Char
'\\', Char
'r']
        checkChar Char
'\t' = [Char
'\\', Char
't']
        checkChar Char
'\f' = [Char
'\\', Char
'f']
        checkChar Char
'\n' = [Char
'\\', Char
'n']
        checkChar Char
c = [Char
c]