{-# LANGUAGE CPP #-}
module Data.Aeson.Flatten
    ( flattenObject
    , unflattenObject
    , mergeObject
    , mergeValue
    , objToHm
    , hmToObj
    ) where

import Data.Aeson
#if MIN_VERSION_aeson(2,0,0)
import Data.Aeson.KeyMap
#endif
import Data.HashMap.Strict qualified as HM
import Data.Text qualified as Text

-- | Convert an object to a hashmap. Compatibility shim for pre-/post-aeson-2.
objToHm :: Object -> HM.HashMap Text.Text Value
#if MIN_VERSION_aeson(2,0,0)
objToHm :: Object -> HashMap Text Value
objToHm = Object -> HashMap Text Value
forall v. KeyMap v -> HashMap Text v
toHashMapText
#else
objToHm = id
#endif

-- | Convert a hashmap to an object. Compatibility shim for pre-/post-aeson-2.
hmToObj :: HM.HashMap Text.Text Value -> Object
#if MIN_VERSION_aeson(2,0,0)
hmToObj :: HashMap Text Value -> Object
hmToObj = HashMap Text Value -> Object
forall v. HashMap Text v -> KeyMap v
fromHashMapText
#else
hmToObj = id
#endif

-- | Turn a nested object into a "flat" object where the keys represent paths into the original
-- object. The keys in the result will be the keys in the original path, separated by `sep`.
-- The inverse of 'unflattenObject'.
flattenObject :: Text.Text -> Object -> Object
flattenObject :: Text -> Object -> Object
flattenObject Text
sep Object
o = HashMap Text Value -> Object
hmToObj (HashMap Text Value -> Object) -> HashMap Text Value -> Object
forall a b. (a -> b) -> a -> b
$ Maybe Text -> HashMap Text Value -> HashMap Text Value
go Maybe Text
forall a. Maybe a
Nothing (Object -> HashMap Text Value
objToHm Object
o)
    where
        go :: Maybe Text.Text -> HM.HashMap Text.Text Value -> HM.HashMap Text.Text Value
        go :: Maybe Text -> HashMap Text Value -> HashMap Text Value
go Maybe Text
mprefix = (Text -> Value -> HashMap Text Value)
-> HashMap Text Value -> HashMap Text Value
forall m k v. Monoid m => (k -> v -> m) -> HashMap k v -> m
HM.foldMapWithKey ((Text -> Value -> HashMap Text Value)
 -> HashMap Text Value -> HashMap Text Value)
-> (Text -> Value -> HashMap Text Value)
-> HashMap Text Value
-> HashMap Text Value
forall a b. (a -> b) -> a -> b
$ \Text
k Value
v ->
            let newName :: Text
newName = case Maybe Text
mprefix of
                    Just Text
prefix -> Text
prefix Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
sep Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
k
                    Maybe Text
Nothing     -> Text
k
            in case Value
v of
                Object Object
o' -> Maybe Text -> HashMap Text Value -> HashMap Text Value
go (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
newName) (HashMap Text Value -> HashMap Text Value)
-> HashMap Text Value -> HashMap Text Value
forall a b. (a -> b) -> a -> b
$ Object -> HashMap Text Value
objToHm Object
o'
                Value
leaf      -> Text -> Value -> HashMap Text Value
forall k v. Hashable k => k -> v -> HashMap k v
HM.singleton Text
newName Value
leaf

-- | Turn a "flat" object whose keys represent paths into an unflattened object.
-- The keys in the result will be the resulting path, separated by `sep`.
-- The inverse of 'flattenObject'.
unflattenObject :: Text.Text -> Object -> Object
unflattenObject :: Text -> Object -> Object
unflattenObject Text
sep Object
o =
    (Object -> Text -> Value -> Object)
-> Object -> HashMap Text Value -> Object
forall a k v. (a -> k -> v -> a) -> a -> HashMap k v -> a
HM.foldlWithKey (\Object
acc Text
k Value
v -> Object -> Object -> Object
mergeObject Object
acc (Text -> Value -> Object
mkPathObject Text
k Value
v)) Object
forall a. Monoid a => a
mempty (Object -> HashMap Text Value
objToHm Object
o)
    where
        mkPathObject :: Text.Text -> Value -> Object
        mkPathObject :: Text -> Value -> Object
mkPathObject Text
k Value
value =
            let path :: [Text]
path = HasCallStack => Text -> Text -> [Text]
Text -> Text -> [Text]
Text.splitOn Text
sep Text
k
            in HashMap Text Value -> Object
hmToObj (HashMap Text Value -> Object) -> HashMap Text Value -> Object
forall a b. (a -> b) -> a -> b
$ [Text] -> Value -> HashMap Text Value
go [Text]
path Value
value
            where
                go :: [Text.Text] -> Value -> HM.HashMap Text.Text Value
                go :: [Text] -> Value -> HashMap Text Value
go [] Value
_        = [Char] -> HashMap Text Value
forall a. HasCallStack => [Char] -> a
error [Char]
"empty path"
                go [Text
n] Value
v       = Text -> Value -> HashMap Text Value
forall k v. Hashable k => k -> v -> HashMap k v
HM.singleton Text
n Value
v
                go (Text
n:Text
n':[Text]
xs) Value
v = Text -> Value -> HashMap Text Value
forall k v. Hashable k => k -> v -> HashMap k v
HM.singleton Text
n (Value -> HashMap Text Value) -> Value -> HashMap Text Value
forall a b. (a -> b) -> a -> b
$ Object -> Value
Object (Object -> Value) -> Object -> Value
forall a b. (a -> b) -> a -> b
$ HashMap Text Value -> Object
hmToObj (HashMap Text Value -> Object) -> HashMap Text Value -> Object
forall a b. (a -> b) -> a -> b
$ [Text] -> Value -> HashMap Text Value
go (Text
n'Text -> [Text] -> [Text]
forall a. a -> [a] -> [a]
:[Text]
xs) Value
v

-- | Merge two objects, merging the values where both sides have an entry for a key rather than
-- taking the first.
mergeObject :: Object -> Object -> Object
mergeObject :: Object -> Object -> Object
mergeObject Object
o1 Object
o2 = HashMap Text Value -> Object
hmToObj (HashMap Text Value -> Object) -> HashMap Text Value -> Object
forall a b. (a -> b) -> a -> b
$ (Value -> Value -> Value)
-> HashMap Text Value -> HashMap Text Value -> HashMap Text Value
forall k v.
Eq k =>
(v -> v -> v) -> HashMap k v -> HashMap k v -> HashMap k v
HM.unionWith Value -> Value -> Value
mergeValue (Object -> HashMap Text Value
objToHm Object
o1) (Object -> HashMap Text Value
objToHm Object
o2)

-- | Merge two values, merging the objects using 'mergeObject'. Can't merge anything else.
mergeValue :: Value -> Value -> Value
mergeValue :: Value -> Value -> Value
mergeValue (Object Object
o1) (Object Object
o2) = Object -> Value
Object (Object -> Value) -> Object -> Value
forall a b. (a -> b) -> a -> b
$ Object -> Object -> Object
mergeObject Object
o1 Object
o2
mergeValue Value
_ Value
_                     = [Char] -> Value
forall a. HasCallStack => [Char] -> a
error [Char]
"can't merge"