{-# LANGUAGE RecordWildCards #-}
module Data.Aeson.JSONPath.Query.Segment
  ( qQuerySegment
  , qQuerySegmentLocated
  )
  where

import Control.Monad (join)
import Data.Aeson    (Value)
import Data.Vector   (Vector)

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

import Data.Aeson.JSONPath.Types
import Data.Aeson.JSONPath.Query.Selector

import Prelude

qQuerySegment :: QuerySegment Query -> QueryState -> Vector Value
qQuerySegment :: QuerySegment Query -> QueryState -> Vector Value
qQuerySegment QuerySegment{Segment Query
SegmentType
segmentType :: SegmentType
segment :: Segment Query
segmentType :: forall a. QuerySegment a -> SegmentType
segment :: forall a. QuerySegment a -> Segment a
..} qS :: QueryState
qS@QueryState{Value
Query -> QueryState -> Vector Value
rootVal :: Value
curVal :: Value
executeQuery :: Query -> QueryState -> Vector Value
rootVal :: QueryState -> Value
curVal :: QueryState -> Value
executeQuery :: QueryState -> Query -> QueryState -> Vector Value
..} = case SegmentType
segmentType of
  SegmentType
Child      -> Vector Value -> Vector Value
joinAfterMap (Vector Value -> Vector Value) -> Vector Value -> Vector Value
forall a b. (a -> b) -> a -> b
$ Value -> Vector Value
forall a. a -> Vector a
V.singleton Value
curVal
  SegmentType
Descendant -> Vector Value -> Vector Value
joinAfterMap (Vector Value -> Vector Value) -> Vector Value -> Vector Value
forall a b. (a -> b) -> a -> b
$ Value -> Vector Value
allElemsRecursive Value
curVal
  where
    joinAfterMap :: Vector Value -> Vector Value
joinAfterMap Vector Value
vec = Vector (Vector Value) -> Vector Value
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (Vector (Vector Value) -> Vector Value)
-> Vector (Vector Value) -> Vector Value
forall a b. (a -> b) -> a -> b
$ (Value -> Vector Value) -> Vector Value -> Vector (Vector Value)
forall a b. (a -> b) -> Vector a -> Vector b
V.map (\Value
cur -> Segment Query -> QueryState -> Vector Value
qSegment Segment Query
segment QueryState
qS{ curVal :: Value
curVal = Value
cur }) Vector Value
vec

qQuerySegmentLocated :: QuerySegment Query -> QueryState-> String -> Vector (String, Value)
qQuerySegmentLocated :: QuerySegment Query
-> QueryState -> [Char] -> Vector ([Char], Value)
qQuerySegmentLocated QuerySegment{Segment Query
SegmentType
segmentType :: forall a. QuerySegment a -> SegmentType
segment :: forall a. QuerySegment a -> Segment a
segmentType :: SegmentType
segment :: Segment Query
..} qS :: QueryState
qS@QueryState{Value
Query -> QueryState -> Vector Value
rootVal :: QueryState -> Value
curVal :: QueryState -> Value
executeQuery :: QueryState -> Query -> QueryState -> Vector Value
rootVal :: Value
curVal :: Value
executeQuery :: Query -> QueryState -> Vector Value
..} [Char]
loc = case SegmentType
segmentType of
  SegmentType
Child      -> Vector ([Char], Value) -> Vector ([Char], Value)
joinAfterMap (Vector ([Char], Value) -> Vector ([Char], Value))
-> Vector ([Char], Value) -> Vector ([Char], Value)
forall a b. (a -> b) -> a -> b
$ ([Char], Value) -> Vector ([Char], Value)
forall a. a -> Vector a
V.singleton ([Char]
loc,Value
curVal)
  SegmentType
Descendant -> Vector ([Char], Value) -> Vector ([Char], Value)
joinAfterMap (Vector ([Char], Value) -> Vector ([Char], Value))
-> Vector ([Char], Value) -> Vector ([Char], Value)
forall a b. (a -> b) -> a -> b
$ ([Char], Value) -> Vector ([Char], Value)
allElemsRecursiveLocated ([Char]
loc,Value
curVal)
  where
    joinAfterMap :: Vector ([Char], Value) -> Vector ([Char], Value)
joinAfterMap Vector ([Char], Value)
vec = Vector (Vector ([Char], Value)) -> Vector ([Char], Value)
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (Vector (Vector ([Char], Value)) -> Vector ([Char], Value))
-> Vector (Vector ([Char], Value)) -> Vector ([Char], Value)
forall a b. (a -> b) -> a -> b
$ (([Char], Value) -> Vector ([Char], Value))
-> Vector ([Char], Value) -> Vector (Vector ([Char], Value))
forall a b. (a -> b) -> Vector a -> Vector b
V.map (\([Char]
location,Value
cur) -> Segment Query -> QueryState -> [Char] -> Vector ([Char], Value)
qSegmentLocated Segment Query
segment QueryState
qS{ curVal :: Value
curVal = Value
cur } [Char]
location) Vector ([Char], Value)
vec

qSegment :: Segment Query -> QueryState -> Vector Value
qSegment :: Segment Query -> QueryState -> Vector Value
qSegment (Bracketed [Selector Query]
sels) QueryState
qS = [Vector Value] -> Vector Value
forall a. [Vector a] -> Vector a
V.concat ([Vector Value] -> Vector Value) -> [Vector Value] -> Vector Value
forall a b. (a -> b) -> a -> b
$ (Selector Query -> Vector Value)
-> [Selector Query] -> [Vector Value]
forall a b. (a -> b) -> [a] -> [b]
map (Selector Query -> QueryState -> Vector Value
`qSelector` QueryState
qS) [Selector Query]
sels
qSegment (Dotted Text
key) QueryState
qS = Selector Query -> QueryState -> Vector Value
qSelector (Text -> Selector Query
forall a. Text -> Selector a
Name Text
key) QueryState
qS
qSegment Segment Query
WildcardSegment QueryState
qS = Selector Query -> QueryState -> Vector Value
qSelector Selector Query
forall a. Selector a
WildcardSelector QueryState
qS

qSegmentLocated :: Segment Query -> QueryState -> String -> Vector (String,Value)
qSegmentLocated :: Segment Query -> QueryState -> [Char] -> Vector ([Char], Value)
qSegmentLocated (Bracketed [Selector Query]
sels) QueryState
qS [Char]
loc = [Vector ([Char], Value)] -> Vector ([Char], Value)
forall a. [Vector a] -> Vector a
V.concat ([Vector ([Char], Value)] -> Vector ([Char], Value))
-> [Vector ([Char], Value)] -> Vector ([Char], Value)
forall a b. (a -> b) -> a -> b
$ (Selector Query -> Vector ([Char], Value))
-> [Selector Query] -> [Vector ([Char], Value)]
forall a b. (a -> b) -> [a] -> [b]
map (\Selector Query
sel -> Selector Query -> QueryState -> [Char] -> Vector ([Char], Value)
qSelectorLocated Selector Query
sel QueryState
qS [Char]
loc) [Selector Query]
sels
qSegmentLocated (Dotted Text
key) QueryState
qS [Char]
loc = Selector Query -> QueryState -> [Char] -> Vector ([Char], Value)
qSelectorLocated (Text -> Selector Query
forall a. Text -> Selector a
Name Text
key) QueryState
qS [Char]
loc
qSegmentLocated Segment Query
WildcardSegment QueryState
qS [Char]
loc = Selector Query -> QueryState -> [Char] -> Vector ([Char], Value)
qSelectorLocated Selector Query
forall a. Selector a
WildcardSelector QueryState
qS [Char]
loc

-- TODO: Looks kinda ugly, make it pretty <3
allElemsRecursive :: Value -> Vector Value
allElemsRecursive :: Value -> Vector Value
allElemsRecursive o :: Value
o@(JSON.Object Object
obj) = [Vector Value] -> Vector Value
forall a. [Vector a] -> Vector a
V.concat [
    Value -> Vector Value
forall a. a -> Vector a
V.singleton Value
o,
    [Vector Value] -> Vector Value
forall a. [Vector a] -> Vector a
V.concat ([Vector Value] -> Vector Value) -> [Vector Value] -> Vector Value
forall a b. (a -> b) -> a -> b
$ (Value -> Vector Value) -> [Value] -> [Vector Value]
forall a b. (a -> b) -> [a] -> [b]
map Value -> Vector Value
allElemsRecursive (Object -> [Value]
forall v. KeyMap v -> [v]
KM.elems Object
obj)
  ]
allElemsRecursive a :: Value
a@(JSON.Array Vector Value
arr) = [Vector Value] -> Vector Value
forall a. [Vector a] -> Vector a
V.concat [
    Value -> Vector Value
forall a. a -> Vector a
V.singleton Value
a,
    [Vector Value] -> Vector Value
forall a. [Vector a] -> Vector a
V.concat ([Vector Value] -> Vector Value) -> [Vector Value] -> Vector Value
forall a b. (a -> b) -> a -> b
$ (Value -> Vector Value) -> [Value] -> [Vector Value]
forall a b. (a -> b) -> [a] -> [b]
map Value -> Vector Value
allElemsRecursive (Vector Value -> [Value]
forall a. Vector a -> [a]
V.toList Vector Value
arr)
  ]
allElemsRecursive Value
_ = Vector Value
forall a. Vector a
V.empty

-- TODO: Looks kinda ugly, make it pretty <3
allElemsRecursiveLocated :: (String,Value) -> Vector (String,Value)
allElemsRecursiveLocated :: ([Char], Value) -> Vector ([Char], Value)
allElemsRecursiveLocated ([Char]
loc, o :: Value
o@(JSON.Object Object
obj)) = [Vector ([Char], Value)] -> Vector ([Char], Value)
forall a. [Vector a] -> Vector a
V.concat [
    ([Char], Value) -> Vector ([Char], Value)
forall a. a -> Vector a
V.singleton ([Char]
loc,Value
o),
    [Vector ([Char], Value)] -> Vector ([Char], Value)
forall a. [Vector a] -> Vector a
V.concat ([Vector ([Char], Value)] -> Vector ([Char], Value))
-> [Vector ([Char], Value)] -> Vector ([Char], Value)
forall a b. (a -> b) -> a -> b
$ ([Char] -> Value -> Vector ([Char], Value))
-> [[Char]] -> [Value] -> [Vector ([Char], Value)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith ((([Char], Value) -> Vector ([Char], Value))
-> [Char] -> Value -> Vector ([Char], Value)
forall a b c. ((a, b) -> c) -> a -> b -> c
curry ([Char], Value) -> Vector ([Char], Value)
allElemsRecursiveLocated) [[Char]]
keys (Object -> [Value]
forall v. KeyMap v -> [v]
KM.elems Object
obj)
  ]
  where
    keys :: [[Char]]
keys = (Key -> [Char]) -> [Key] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map ([Char] -> Key -> [Char]
toPathKey [Char]
loc) ([Key] -> [[Char]]) -> [Key] -> [[Char]]
forall a b. (a -> b) -> a -> b
$ Object -> [Key]
forall v. KeyMap v -> [Key]
KM.keys Object
obj
allElemsRecursiveLocated ([Char]
loc, a :: Value
a@(JSON.Array Vector Value
arr)) = [Vector ([Char], Value)] -> Vector ([Char], Value)
forall a. [Vector a] -> Vector a
V.concat [
    ([Char], Value) -> Vector ([Char], Value)
forall a. a -> Vector a
V.singleton ([Char]
loc,Value
a),
    [Vector ([Char], Value)] -> Vector ([Char], Value)
forall a. [Vector a] -> Vector a
V.concat ([Vector ([Char], Value)] -> Vector ([Char], Value))
-> [Vector ([Char], Value)] -> Vector ([Char], Value)
forall a b. (a -> b) -> a -> b
$ ([Char] -> Value -> Vector ([Char], Value))
-> [[Char]] -> [Value] -> [Vector ([Char], Value)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith ((([Char], Value) -> Vector ([Char], Value))
-> [Char] -> Value -> Vector ([Char], Value)
forall a b c. ((a, b) -> c) -> a -> b -> c
curry ([Char], Value) -> Vector ([Char], Value)
allElemsRecursiveLocated) [[Char]]
indices (Vector Value -> [Value]
forall a. Vector a -> [a]
V.toList Vector Value
arr)
  ]
  where
    indices :: [[Char]]
indices = (Int -> [Char]) -> [Int] -> [[Char]]
forall a b. (a -> b) -> [a] -> [b]
map ([Char] -> Int -> [Char]
toPathIdx [Char]
loc) [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]
allElemsRecursiveLocated ([Char], Value)
_ = Vector ([Char], Value)
forall a. Vector a
V.empty

toPathKey :: String -> KM.Key -> String
toPathKey :: [Char] -> Key -> [Char]
toPathKey [Char]
loc Key
key = [Char]
loc [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"['" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char] -> [Char]
escapeEscapees (Key -> [Char]
K.toString Key
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]

toPathIdx :: String -> Int -> String
toPathIdx :: [Char] -> Int -> [Char]
toPathIdx [Char]
loc Int
idx = [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
idx [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"]"