{-# 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) =>
  (NonEmpty k -> Maybe v -> m (Map k v, Map k v -> v))
  -- ^ How to get and set the next Map from the possibly missing value.
  -- Passes in the path taken so far.
  -> NonEmpty k
  -> Map k v
  -> m (Maybe v, v -> Map k v)
getPathLens :: forall (m :: * -> *) k v.
(Monad m, Ord k) =>
(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 =
  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 -> forall {k} {a} {b}.
Ord k =>
(Map k a -> b) -> k -> Map k a -> a -> b
mkSetter (v -> Map k v
setVal forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map k v -> v
fromMap)) (forall {k} {a} {b}.
Ord k =>
(Map k a -> b) -> k -> Map k a -> a -> b
mkSetter forall a. a -> a
id)
  where
    mkSetter :: (Map k a -> b) -> k -> Map k a -> a -> b
mkSetter Map k a -> b
setMap k
k Map k a
kvs = \a
v -> Map k a -> b
setMap forall a b. (a -> b) -> a -> b
$ 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 :: forall (m :: * -> *) k v.
(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)
doRecurse NonEmpty k
path Map k v
originalMap =
  forall a b. (a, b) -> a
fst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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
      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 :: 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 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 = forall a. NonEmpty a -> NonEmpty (NonEmpty a, a)
zipHistory NonEmpty k
path
   in 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 (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
      forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> 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 = (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)