{-# LANGUAGE RecordWildCards #-}

module TOML.Utils.Map (
  getPathLens,
  getPath,
) where

import Data.Foldable (foldlM)
import Data.List.NonEmpty (NonEmpty ((:|)))
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map

import TOML.Utils.NonEmpty (zipHistory)

{- |
For a non-empty list of keys, iterate through the given 'Map' and return
the possibly missing value at the path and a function to set the value at
the given path and return the modified input 'Map'.

@
let obj = undefined -- { "a": { "b": { "c": 1 } } }
(mValue, setValue) <- getPathLens doRecurse ["a", "b", "c"] obj

print mValue -- Just 1
print (setValue 2) -- { "a": { "b": { "c": 2 } } }
@
-}
getPathLens ::
  (Monad m, Ord k) =>
  -- | How to get and set the next Map from the possibly missing value.
  -- Passes in the path taken so far.
  (NonEmpty k -> Maybe v -> m (Map k v, Map k v -> v)) ->
  NonEmpty k ->
  Map k v ->
  m (Maybe v, v -> Map k v)
getPathLens :: (NonEmpty k -> Maybe v -> m (Map k v, Map k v -> v))
-> NonEmpty k -> Map k v -> m (Maybe v, v -> Map k v)
getPathLens =
  ((v -> Map k v) -> (Map k v -> v) -> k -> Map k v -> v -> Map k v)
-> (k -> Map k v -> v -> Map k v)
-> (NonEmpty k -> Maybe v -> m (Map k v, Map k v -> v))
-> NonEmpty k
-> Map k v
-> m (Maybe v, v -> Map k v)
forall (m :: * -> *) k b a v.
(Monad m, Ord k) =>
(b -> a -> k -> Map k v -> b)
-> (k -> Map k v -> b)
-> (NonEmpty k -> Maybe v -> m (Map k v, a))
-> NonEmpty k
-> Map k v
-> m (Maybe v, b)
getPathLensWith (\v -> Map k v
setVal Map k v -> v
fromMap -> (Map k v -> Map k v) -> k -> Map k v -> v -> Map k v
forall k a t. Ord k => (Map k a -> t) -> k -> Map k a -> a -> t
mkSetter (v -> Map k v
setVal (v -> Map k v) -> (Map k v -> v) -> Map k v -> Map k v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map k v -> v
fromMap)) ((Map k v -> Map k v) -> k -> Map k v -> v -> Map k v
forall k a t. Ord k => (Map k a -> t) -> k -> Map k a -> a -> t
mkSetter Map k v -> Map k v
forall a. a -> a
id)
  where
    mkSetter :: (Map k a -> t) -> k -> Map k a -> a -> t
mkSetter Map k a -> t
setMap k
k Map k a
kvs = \a
v -> Map k a -> t
setMap (Map k a -> t) -> Map k a -> t
forall a b. (a -> b) -> a -> b
$ k -> a -> Map k a -> Map k a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
k a
v Map k a
kvs

-- | Same as 'getPathLens', except without the setter.
getPath ::
  (Monad m, Ord k) =>
  (NonEmpty k -> Maybe v -> m (Map k v)) ->
  NonEmpty k ->
  Map k v ->
  m (Maybe v)
getPath :: (NonEmpty k -> Maybe v -> m (Map k v))
-> NonEmpty k -> Map k v -> m (Maybe v)
getPath NonEmpty k -> Maybe v -> m (Map k v)
doRecurse NonEmpty k
path Map k v
originalMap =
  (Maybe v, ()) -> Maybe v
forall a b. (a, b) -> a
fst ((Maybe v, ()) -> Maybe v) -> m (Maybe v, ()) -> m (Maybe v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (() -> () -> k -> Map k v -> ())
-> (k -> Map k v -> ())
-> (NonEmpty k -> Maybe v -> m (Map k v, ()))
-> NonEmpty k
-> Map k v
-> m (Maybe v, ())
forall (m :: * -> *) k b a v.
(Monad m, Ord k) =>
(b -> a -> k -> Map k v -> b)
-> (k -> Map k v -> b)
-> (NonEmpty k -> Maybe v -> m (Map k v, a))
-> NonEmpty k
-> Map k v
-> m (Maybe v, b)
getPathLensWith (\()
_ ()
_ k
_ Map k v
_ -> ()) (\k
_ Map k v
_ -> ()) NonEmpty k -> Maybe v -> m (Map k v, ())
doRecurse' NonEmpty k
path Map k v
originalMap
  where
    doRecurse' :: NonEmpty k -> Maybe v -> m (Map k v, ())
doRecurse' NonEmpty k
history Maybe v
mVal = do
      Map k v
x <- NonEmpty k -> Maybe v -> m (Map k v)
doRecurse NonEmpty k
history Maybe v
mVal
      (Map k v, ()) -> m (Map k v, ())
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Map k v
x, ())

getPathLensWith ::
  (Monad m, Ord k) =>
  (b -> a -> (k -> Map k v -> b)) ->
  (k -> Map k v -> b) ->
  (NonEmpty k -> Maybe v -> m (Map k v, a)) ->
  NonEmpty k ->
  Map k v ->
  m (Maybe v, b)
getPathLensWith :: (b -> a -> k -> Map k v -> b)
-> (k -> Map k v -> b)
-> (NonEmpty k -> Maybe v -> m (Map k v, a))
-> NonEmpty k
-> Map k v
-> m (Maybe v, b)
getPathLensWith b -> a -> k -> Map k v -> b
mkAnn k -> Map k v -> b
mkFirstAnn NonEmpty k -> Maybe v -> m (Map k v, a)
doRecurse NonEmpty k
path Map k v
originalMap =
  let (NonEmpty k
_, k
k) :| [(NonEmpty k, k)]
ks = NonEmpty k -> NonEmpty (NonEmpty k, k)
forall a. NonEmpty a -> NonEmpty (NonEmpty a, a)
zipHistory NonEmpty k
path
   in ((Maybe v, b) -> (NonEmpty k, k) -> m (Maybe v, b))
-> (Maybe v, b) -> [(NonEmpty k, k)] -> m (Maybe v, b)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM (Maybe v, b) -> (NonEmpty k, k) -> m (Maybe v, b)
go (k -> (k -> Map k v -> b) -> Map k v -> (Maybe v, b)
forall t a b.
Ord t =>
t -> (t -> Map t a -> b) -> Map t a -> (Maybe a, b)
buildLens k
k k -> Map k v -> b
mkFirstAnn Map k v
originalMap) [(NonEmpty k, k)]
ks
  where
    go :: (Maybe v, b) -> (NonEmpty k, k) -> m (Maybe v, b)
go (Maybe v
mVal, b
b) (NonEmpty k
history, k
k) = do
      (Map k v
nextMap, a
a) <- NonEmpty k -> Maybe v -> m (Map k v, a)
doRecurse NonEmpty k
history Maybe v
mVal
      (Maybe v, b) -> m (Maybe v, b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Maybe v, b) -> m (Maybe v, b)) -> (Maybe v, b) -> m (Maybe v, b)
forall a b. (a -> b) -> a -> b
$ k -> (k -> Map k v -> b) -> Map k v -> (Maybe v, b)
forall t a b.
Ord t =>
t -> (t -> Map t a -> b) -> Map t a -> (Maybe a, b)
buildLens k
k (b -> a -> k -> Map k v -> b
mkAnn b
b a
a) Map k v
nextMap

    buildLens :: t -> (t -> Map t a -> b) -> Map t a -> (Maybe a, b)
buildLens t
k t -> Map t a -> b
mkAnn' Map t a
kvs = (t -> Map t a -> Maybe a
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup t
k Map t a
kvs, t -> Map t a -> b
mkAnn' t
k Map t a
kvs)