{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE NamedFieldPuns #-}

{- | Flatten nested JSON objects into a single JSON object in which the keys
have been joined by the separator.
-}
module Json.Flatten
  ( flatten
  ) where

import Control.Monad.ST (ST)
import Control.Monad.ST.Run (runByteArrayST)
import Data.Builder.Catenable (Builder)
import qualified Data.Builder.Catenable as Builder
import Data.ByteString.Short.Internal (ShortByteString (SBS))
import qualified Data.Bytes as Bytes
import qualified Data.Bytes.Text.Utf8 as Utf8
import qualified Data.Chunks as Chunks
import Data.Primitive (ByteArray (ByteArray), MutableByteArray, SmallArray)
import qualified Data.Primitive as PM
import qualified Data.Primitive.Contiguous as C
import Data.Text.Short (ShortText)
import qualified Data.Text.Short as TS
import qualified Data.Text.Short.Unsafe as TS
import Data.Word (Word8)
import Json (Member (Member))
import qualified Json

{- | Flatten a json value, recursively descending into objects and joining
keys with the separator. For example:

> { "name": "bilbo"
> , "occupation":
>   { "name": "burglar"
>   , "start": "2022-05-30"
>   }
> , "height": 124
> , "favorites": ["adventures","lunch"]
> }

Becomes:

> { "name": "bilbo"
> , "occupation.name": "burglar"
> , "occupation.start": "2022-05-30"
> , "height": 124
> , "favorites": ["adventures","lunch"]
> }

Currently, the implementation of this function throws an exception if
any separator other than period is used. This may be corrected in a future
release.
-}
flatten :: Char -> Json.Value -> Json.Value
flatten :: Char -> Value -> Value
flatten Char
c Value
v = case Char
c of
  Char
'.' -> Value -> Value
flattenPeriod Value
v
  Char
_ -> [Char] -> Value
forall a. [Char] -> a
errorWithoutStackTrace [Char]
"Json.Flatten.flatten: only period is supported"

-- built backwards
data ShortTexts
  = ShortTextsCons !ShortText !ShortTexts
  | ShortTextsBase !ShortText

flattenPeriod :: Json.Value -> Json.Value
flattenPeriod :: Value -> Value
flattenPeriod Value
x = case Value
x of
  Json.Object SmallArray Member
mbrs ->
    let bldr :: Builder Member
bldr = (Member -> Builder Member) -> SmallArray Member -> Builder Member
forall m a. Monoid m => (a -> m) -> SmallArray a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\Member {ShortText
key :: ShortText
key :: Member -> ShortText
key, Value
value :: Value
value :: Member -> Value
value} -> ShortTexts -> Value -> Builder Member
flattenPrefix (ShortText -> ShortTexts
ShortTextsBase ShortText
key) Value
value) SmallArray Member
mbrs
        chunks :: Chunks Member
chunks = Builder Member -> Chunks Member
forall a. Builder a -> Chunks a
Builder.run Builder Member
bldr
        result :: SmallArray Member
result = Chunks Member -> SmallArray Member
forall a. Chunks a -> SmallArray a
Chunks.concat Chunks Member
chunks
     in SmallArray Member -> Value
Json.Object SmallArray Member
result
  Json.Array SmallArray Value
ys -> SmallArray Value -> Value
Json.Array (SmallArray Value -> Value) -> SmallArray Value -> Value
forall a b. (a -> b) -> a -> b
$! (Value -> Value) -> SmallArray Value -> SmallArray Value
forall (arr1 :: * -> *) b (arr2 :: * -> *) c.
(Contiguous arr1, Element arr1 b, Contiguous arr2,
 Element arr2 c) =>
(b -> c) -> arr1 b -> arr2 c
C.map' Value -> Value
flattenPeriod SmallArray Value
ys
  Value
_ -> Value
x

flattenPrefix ::
  ShortTexts -> -- context accumulator
  Json.Value ->
  Builder Json.Member
flattenPrefix :: ShortTexts -> Value -> Builder Member
flattenPrefix !ShortTexts
pre Value
x = case Value
x of
  Json.Object SmallArray Member
mbrs -> ShortTexts -> SmallArray Member -> Builder Member
flattenObject ShortTexts
pre SmallArray Member
mbrs
  Value
_ ->
    let !a :: Value
a = Value -> Value
flattenPeriod Value
x
        !k :: ShortText
k = ShortTexts -> ShortText
runShortTexts ShortTexts
pre
        !mbr :: Member
mbr = Json.Member {key :: ShortText
key = ShortText
k, value :: Value
value = Value
a}
     in Member -> Builder Member -> Builder Member
forall a. a -> Builder a -> Builder a
Builder.Cons Member
mbr Builder Member
forall a. Builder a
Builder.Empty

flattenObject :: ShortTexts -> SmallArray Json.Member -> Builder Json.Member
flattenObject :: ShortTexts -> SmallArray Member -> Builder Member
flattenObject !ShortTexts
pre !SmallArray Member
mbrs =
  (Member -> Builder Member) -> SmallArray Member -> Builder Member
forall m a. Monoid m => (a -> m) -> SmallArray a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap
    ( \Member {ShortText
key :: Member -> ShortText
key :: ShortText
key, Value
value :: Member -> Value
value :: Value
value} -> ShortTexts -> Value -> Builder Member
flattenPrefix (ShortText -> ShortTexts -> ShortTexts
ShortTextsCons ShortText
key ShortTexts
pre) Value
value
    )
    SmallArray Member
mbrs

runShortTexts :: ShortTexts -> ShortText
runShortTexts :: ShortTexts -> ShortText
runShortTexts !ShortTexts
ts0 = Int -> ShortTexts -> ShortText
go Int
0 ShortTexts
ts0
 where
  paste :: MutableByteArray s -> Int -> ShortTexts -> ST s ByteArray
  paste :: forall s. MutableByteArray s -> Int -> ShortTexts -> ST s ByteArray
paste !MutableByteArray s
dst !Int
ix (ShortTextsBase ShortText
t) =
    let len :: Int
len = Bytes -> Int
Bytes.length (ShortText -> Bytes
Utf8.fromShortText ShortText
t)
     in case Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
len of
          Int
0 -> do
            MutableByteArray (PrimState (ST s))
-> Int -> ByteArray -> Int -> Int -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> ByteArray -> Int -> Int -> m ()
PM.copyByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
dst Int
0 (ShortText -> ByteArray
st2ba ShortText
t) Int
0 Int
len
            MutableByteArray (PrimState (ST s)) -> ST s ByteArray
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m) -> m ByteArray
PM.unsafeFreezeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
dst
          Int
_ -> [Char] -> ST s ByteArray
forall a. [Char] -> a
errorWithoutStackTrace [Char]
"Json.Flatten.runShortTexts: implementation mistake"
  paste !MutableByteArray s
dst !Int
ix (ShortTextsCons ShortText
t ShortTexts
ts) = do
    let !len :: Int
len = Bytes -> Int
Bytes.length (ShortText -> Bytes
Utf8.fromShortText ShortText
t)
    let !ixNext :: Int
ixNext = Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
len
    MutableByteArray (PrimState (ST s))
-> Int -> ByteArray -> Int -> Int -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
MutableByteArray (PrimState m)
-> Int -> ByteArray -> Int -> Int -> m ()
PM.copyByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
dst Int
ixNext (ShortText -> ByteArray
st2ba ShortText
t) Int
0 Int
len
    let !ixPred :: Int
ixPred = Int
ixNext Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
    MutableByteArray (PrimState (ST s)) -> Int -> Word8 -> ST s ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutableByteArray (PrimState m) -> Int -> a -> m ()
PM.writeByteArray MutableByteArray s
MutableByteArray (PrimState (ST s))
dst Int
ixPred (Word8
0x2E :: Word8)
    MutableByteArray s -> Int -> ShortTexts -> ST s ByteArray
forall s. MutableByteArray s -> Int -> ShortTexts -> ST s ByteArray
paste MutableByteArray s
dst Int
ixPred ShortTexts
ts
  go :: Int -> ShortTexts -> ShortText
  go :: Int -> ShortTexts -> ShortText
go !Int
byteLenAcc (ShortTextsCons ShortText
t ShortTexts
ts) =
    Int -> ShortTexts -> ShortText
go (Bytes -> Int
Bytes.length (ShortText -> Bytes
Utf8.fromShortText ShortText
t) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
byteLenAcc Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ShortTexts
ts
  go !Int
byteLenAcc (ShortTextsBase ShortText
t) =
    let !(ByteArray ByteArray#
r) = (forall s. ST s ByteArray) -> ByteArray
runByteArrayST ((forall s. ST s ByteArray) -> ByteArray)
-> (forall s. ST s ByteArray) -> ByteArray
forall a b. (a -> b) -> a -> b
$ do
          let totalLen :: Int
totalLen = Bytes -> Int
Bytes.length (ShortText -> Bytes
Utf8.fromShortText ShortText
t) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
byteLenAcc
          MutableByteArray s
dst <- Int -> ST s (MutableByteArray (PrimState (ST s)))
forall (m :: * -> *).
PrimMonad m =>
Int -> m (MutableByteArray (PrimState m))
PM.newByteArray Int
totalLen
          MutableByteArray s -> Int -> ShortTexts -> ST s ByteArray
forall s. MutableByteArray s -> Int -> ShortTexts -> ST s ByteArray
paste MutableByteArray s
dst Int
totalLen ShortTexts
ts0
     in ShortByteString -> ShortText
TS.fromShortByteStringUnsafe (ByteArray# -> ShortByteString
SBS ByteArray#
r)

st2ba :: ShortText -> ByteArray
{-# INLINE st2ba #-}
st2ba :: ShortText -> ByteArray
st2ba ShortText
t = case ShortText -> ShortByteString
TS.toShortByteString ShortText
t of
  SBS ByteArray#
x -> ByteArray# -> ByteArray
ByteArray ByteArray#
x