    {-# 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
    segment :: forall a. QuerySegment a -> Segment a
    segmentType :: forall a. QuerySegment a -> SegmentType
    ..} qS :: QueryState
    qS@QueryState{Value
    Query -> QueryState -> Vector Value
    rootVal :: Value
    curVal :: Value
    executeQuery :: Query -> QueryState -> Vector Value
    executeQuery :: QueryState -> Query -> QueryState -> Vector Value
    curVal :: QueryState -> Value
    rootVal :: QueryState -> 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 = 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
    segment :: forall a. QuerySegment a -> Segment a
    segmentType :: forall a. QuerySegment a -> SegmentType
    segmentType :: SegmentType
    segment :: Segment Query
    ..} qS :: QueryState
    qS@QueryState{Value
    Query -> QueryState -> Vector Value
    executeQuery :: QueryState -> Query -> QueryState -> Vector Value
    curVal :: QueryState -> Value
    rootVal :: QueryState -> 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 = 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]
    "]"