-- |
--  Module      : KeyTree
--  Copyright   : © Jonathan Lorimer, 2023
--  License     : MIT
--  Maintainer  : jonathanlorimer@pm.me
--  Stability   : stable
--
-- @since 0.0.2.0
--
-- This module contains types for our internal tree representation of types and
-- configurations. It also contains some helper functions for working with
-- these trees. This should make it easier to implement different source
-- providers.
module KeyTree
  ( -- * Types
    KeyTree
  , KeyForest

    -- * Helper Functions
  , foldKeyTree
  , appendFold
  , mayAppendFold
  , appendTraverse
  , mayAppendTraverse

    -- * Type Re-exports
  , Map
  , Free (..)
  )
where

import Control.Monad.Free
import Data.Functor ((<&>))
import Data.Map.Strict

-- | Type alias for our internal tree structure. If this was written directly
-- as a sum type it would look like this:
--
-- @Pure value | Free (Map key (KeyTree key value))@
--
-- @since 0.0.2.0
type KeyTree key value = Free (Map key) value

-- | Type alias for a subtree
--
-- @since 0.0.2.0
type KeyForest key value = Map key (Free (Map key) value)

-- | Right fold on a 'KeyTree'. Uses 'Data.Map.foldrWithKey' under the hood.
--
-- @since 0.0.2.0
foldKeyTree
  :: (Eq k, Eq v)
  => (v -> a)
  -- ^ Function to run on 'Pure' values
  -> (k -> Free (Map k) v -> a -> a)
  -- ^ Step function for fold
  -> a
  -- ^ Initial accumulator
  -> KeyTree k v
  -- ^ KeyTree to fold
  -> a
foldKeyTree :: forall k v a.
(Eq k, Eq v) =>
(v -> a)
-> (k -> Free (Map k) v -> a -> a) -> a -> Free (Map k) v -> a
foldKeyTree v -> a
_ k -> Free (Map k) v -> a -> a
stepF a
acc (Free Map k (Free (Map k) v)
m) = (k -> Free (Map k) v -> a -> a) -> a -> Map k (Free (Map k) v) -> a
forall k a b. (k -> a -> b -> b) -> b -> Map k a -> b
foldrWithKey k -> Free (Map k) v -> a -> a
stepF a
acc Map k (Free (Map k) v)
m
foldKeyTree v -> a
valF k -> Free (Map k) v -> a -> a
_ a
_ (Pure v
val) = v -> a
valF v
val

-- | A fold that appends a value at the leaf of the 'KeyTree', identifies what to
-- insert at the leaf by running a function on an accumulated value.
--
-- @since 0.0.2.0
appendFold
  :: (Eq k, Eq v)
  => (a -> v -> v')
  -- ^ Function to run on existing 'Pure' leaves @a@ is the accumulator @v@ is the value in @Pure@.
  -> (a -> v')
  -- ^ Function from accumulator @a@ to a value @v@. This happens when an empty subtree is found.
  -> (k -> a -> Map k (Free (Map k) v) -> a)
  -- ^ Step function for fold with key @k@ as an argument.
  -> a
  -- ^ Accumulator.
  -> KeyTree k v
  -- ^ KeyTree to append values to.
  -> KeyTree k v'
appendFold :: forall k v a v'.
(Eq k, Eq v) =>
(a -> v -> v')
-> (a -> v')
-> (k -> a -> Map k (Free (Map k) v) -> a)
-> a
-> Free (Map k) v
-> KeyTree k v'
appendFold a -> v -> v'
valF a -> v'
accF k -> a -> Map k (Free (Map k) v) -> a
stepF a
acc (Free Map k (Free (Map k) v)
m) =
  if Map k (Free (Map k) v)
m Map k (Free (Map k) v) -> Map k (Free (Map k) v) -> Bool
forall a. Eq a => a -> a -> Bool
== Map k (Free (Map k) v)
forall k a. Map k a
empty
    then v' -> Free (Map k) v'
forall (f :: * -> *) a. a -> Free f a
Pure (v' -> Free (Map k) v') -> v' -> Free (Map k) v'
forall a b. (a -> b) -> a -> b
$ a -> v'
accF a
acc
    else Map k (Free (Map k) v') -> Free (Map k) v'
forall (f :: * -> *) a. f (Free f a) -> Free f a
Free (Map k (Free (Map k) v') -> Free (Map k) v')
-> Map k (Free (Map k) v') -> Free (Map k) v'
forall a b. (a -> b) -> a -> b
$ (k -> Free (Map k) v -> Free (Map k) v')
-> Map k (Free (Map k) v) -> Map k (Free (Map k) v')
forall k a b. (k -> a -> b) -> Map k a -> Map k b
mapWithKey (\k
k Free (Map k) v
a -> (a -> v -> v')
-> (a -> v')
-> (k -> a -> Map k (Free (Map k) v) -> a)
-> a
-> Free (Map k) v
-> Free (Map k) v'
forall k v a v'.
(Eq k, Eq v) =>
(a -> v -> v')
-> (a -> v')
-> (k -> a -> Map k (Free (Map k) v) -> a)
-> a
-> Free (Map k) v
-> KeyTree k v'
appendFold a -> v -> v'
valF a -> v'
accF k -> a -> Map k (Free (Map k) v) -> a
stepF (k -> a -> Map k (Free (Map k) v) -> a
stepF k
k a
acc Map k (Free (Map k) v)
m) Free (Map k) v
a) Map k (Free (Map k) v)
m
appendFold a -> v -> v'
valF a -> v'
_ k -> a -> Map k (Free (Map k) v) -> a
_ a
acc (Pure v
v) = v' -> Free (Map k) v'
forall (f :: * -> *) a. a -> Free f a
Pure (v' -> Free (Map k) v') -> v' -> Free (Map k) v'
forall a b. (a -> b) -> a -> b
$ a -> v -> v'
valF a
acc v
v

-- | A fold that appends a value at the leaf of the 'KeyTree' (like
-- 'appendFold'), but the function on the accumulator can return a 'Maybe'. In
-- the @Nothing@ case we just append a 'Free' full of an empty 'Data.Map.Map'.
--
-- @since 0.0.2.0
mayAppendFold
  :: (Eq k, Eq v)
  => (a -> v -> Maybe v')
  -- ^ Function to run on existing 'Pure' leaves.
  -> (a -> Maybe v')
  -- ^ Function to run when an empty node is found.
  -> (k -> a -> Map k (Free (Map k) v) -> a)
  -- ^ Step function for fold with key @k@ as an argument.
  -> a
  -- ^ Accumulator.
  -> KeyTree k v
  -- ^ Tree to be folded.
  -> KeyTree k v'
mayAppendFold :: forall k v a v'.
(Eq k, Eq v) =>
(a -> v -> Maybe v')
-> (a -> Maybe v')
-> (k -> a -> Map k (Free (Map k) v) -> a)
-> a
-> Free (Map k) v
-> KeyTree k v'
mayAppendFold a -> v -> Maybe v'
valF a -> Maybe v'
accF k -> a -> Map k (Free (Map k) v) -> a
stepF a
acc (Free Map k (Free (Map k) v)
m) =
  if Map k (Free (Map k) v)
m Map k (Free (Map k) v) -> Map k (Free (Map k) v) -> Bool
forall a. Eq a => a -> a -> Bool
== Map k (Free (Map k) v)
forall k a. Map k a
empty
    then case a -> Maybe v'
accF a
acc of
      Maybe v'
Nothing -> Map k (KeyTree k v') -> KeyTree k v'
forall (f :: * -> *) a. f (Free f a) -> Free f a
Free Map k (KeyTree k v')
forall k a. Map k a
empty
      Just v'
v -> v' -> KeyTree k v'
forall (f :: * -> *) a. a -> Free f a
Pure v'
v
    else Map k (KeyTree k v') -> KeyTree k v'
forall (f :: * -> *) a. f (Free f a) -> Free f a
Free (Map k (KeyTree k v') -> KeyTree k v')
-> Map k (KeyTree k v') -> KeyTree k v'
forall a b. (a -> b) -> a -> b
$ (k -> Free (Map k) v -> KeyTree k v')
-> Map k (Free (Map k) v) -> Map k (KeyTree k v')
forall k a b. (k -> a -> b) -> Map k a -> Map k b
mapWithKey (\k
k Free (Map k) v
a -> (a -> v -> Maybe v')
-> (a -> Maybe v')
-> (k -> a -> Map k (Free (Map k) v) -> a)
-> a
-> Free (Map k) v
-> KeyTree k v'
forall k v a v'.
(Eq k, Eq v) =>
(a -> v -> Maybe v')
-> (a -> Maybe v')
-> (k -> a -> Map k (Free (Map k) v) -> a)
-> a
-> Free (Map k) v
-> KeyTree k v'
mayAppendFold a -> v -> Maybe v'
valF a -> Maybe v'
accF k -> a -> Map k (Free (Map k) v) -> a
stepF (k -> a -> Map k (Free (Map k) v) -> a
stepF k
k a
acc Map k (Free (Map k) v)
m) Free (Map k) v
a) Map k (Free (Map k) v)
m
mayAppendFold a -> v -> Maybe v'
valF a -> Maybe v'
_ k -> a -> Map k (Free (Map k) v) -> a
_ a
acc (Pure v
v) =
  case a -> v -> Maybe v'
valF a
acc v
v of
    Maybe v'
Nothing -> Map k (KeyTree k v') -> KeyTree k v'
forall (f :: * -> *) a. f (Free f a) -> Free f a
Free Map k (KeyTree k v')
forall k a. Map k a
empty
    Just v'
v' -> v' -> KeyTree k v'
forall (f :: * -> *) a. a -> Free f a
Pure v'
v'

-- | Like 'appendFold' but with functions that return a value wrapped in an
-- 'Applicative' effect @f@. The function is suffixed with \"Traverse\" because
-- we use the 'sequenceA' to wrap the entire tree in a single 'Applicative'
-- effect.
--
-- @since 0.0.2.0
appendTraverse
  :: (Applicative f, Eq k, Eq v)
  => (a -> v -> f v')
  -- ^ Function to run on existing 'Pure' leaves.
  -> (a -> f v')
  -- ^ Function to run when an empty node is found.
  -> (k -> a -> Map k (Free (Map k) v) -> a)
  -- ^ Step function for fold with key @k@ as an argument.
  -> a
  -- ^ Accumulator.
  -> KeyTree k v
  -- ^ Tree to be folded.
  -> f (KeyTree k v')
appendTraverse :: forall (f :: * -> *) k v a v'.
(Applicative f, Eq k, Eq v) =>
(a -> v -> f v')
-> (a -> f v')
-> (k -> a -> Map k (Free (Map k) v) -> a)
-> a
-> Free (Map k) v
-> f (KeyTree k v')
appendTraverse a -> v -> f v'
valF a -> f v'
accF k -> a -> Map k (Free (Map k) v) -> a
stepF a
acc = Free (Map k) (f v') -> f (Free (Map k) v')
forall (t :: * -> *) (f :: * -> *) a.
(Traversable t, Applicative f) =>
t (f a) -> f (t a)
forall (f :: * -> *) a.
Applicative f =>
Free (Map k) (f a) -> f (Free (Map k) a)
sequenceA (Free (Map k) (f v') -> f (Free (Map k) v'))
-> (Free (Map k) v -> Free (Map k) (f v'))
-> Free (Map k) v
-> f (Free (Map k) v')
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> v -> f v')
-> (a -> f v')
-> (k -> a -> Map k (Free (Map k) v) -> a)
-> a
-> Free (Map k) v
-> Free (Map k) (f v')
forall k v a v'.
(Eq k, Eq v) =>
(a -> v -> v')
-> (a -> v')
-> (k -> a -> Map k (Free (Map k) v) -> a)
-> a
-> Free (Map k) v
-> KeyTree k v'
appendFold a -> v -> f v'
valF a -> f v'
accF k -> a -> Map k (Free (Map k) v) -> a
stepF a
acc

-- | Similar to 'appendTraverse' except the 'Applicative' effect can optionally
-- return a result. The function is manually implemented rather than using
-- 'Traversable' methods so that we don't have to require 'Traversable' on @f@.
-- This allows consumers to use effects like 'IO' that don't have a
-- 'Traversable' instance.
--
-- @since 0.0.2.0
mayAppendTraverse
  :: (Applicative f, Eq k, Eq v)
  => (a -> v -> f v')
  -> (a -> f (Maybe v'))
  -> (k -> a -> Map k (Free (Map k) v) -> a)
  -> a
  -> KeyTree k v
  -> f (KeyTree k v')
mayAppendTraverse :: forall (f :: * -> *) k v a v'.
(Applicative f, Eq k, Eq v) =>
(a -> v -> f v')
-> (a -> f (Maybe v'))
-> (k -> a -> Map k (Free (Map k) v) -> a)
-> a
-> Free (Map k) v
-> f (KeyTree k v')
mayAppendTraverse a -> v -> f v'
valF a -> f (Maybe v')
accF k -> a -> Map k (Free (Map k) v) -> a
stepF a
acc (Free Map k (Free (Map k) v)
m) =
  if Map k (Free (Map k) v)
m Map k (Free (Map k) v) -> Map k (Free (Map k) v) -> Bool
forall a. Eq a => a -> a -> Bool
== Map k (Free (Map k) v)
forall k a. Map k a
empty
    then
      a -> f (Maybe v')
accF a
acc f (Maybe v') -> (Maybe v' -> KeyTree k v') -> f (KeyTree k v')
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \case
        Maybe v'
Nothing -> Map k (KeyTree k v') -> KeyTree k v'
forall (f :: * -> *) a. f (Free f a) -> Free f a
Free Map k (KeyTree k v')
forall k a. Map k a
empty
        Just v'
v -> v' -> KeyTree k v'
forall (f :: * -> *) a. a -> Free f a
Pure v'
v
    else Map k (KeyTree k v') -> KeyTree k v'
forall (f :: * -> *) a. f (Free f a) -> Free f a
Free (Map k (KeyTree k v') -> KeyTree k v')
-> f (Map k (KeyTree k v')) -> f (KeyTree k v')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (k -> Free (Map k) v -> f (KeyTree k v'))
-> Map k (Free (Map k) v) -> f (Map k (KeyTree k v'))
forall (t :: * -> *) k a b.
Applicative t =>
(k -> a -> t b) -> Map k a -> t (Map k b)
traverseWithKey (\k
k Free (Map k) v
v -> (a -> v -> f v')
-> (a -> f (Maybe v'))
-> (k -> a -> Map k (Free (Map k) v) -> a)
-> a
-> Free (Map k) v
-> f (KeyTree k v')
forall (f :: * -> *) k v a v'.
(Applicative f, Eq k, Eq v) =>
(a -> v -> f v')
-> (a -> f (Maybe v'))
-> (k -> a -> Map k (Free (Map k) v) -> a)
-> a
-> Free (Map k) v
-> f (KeyTree k v')
mayAppendTraverse a -> v -> f v'
valF a -> f (Maybe v')
accF k -> a -> Map k (Free (Map k) v) -> a
stepF (k -> a -> Map k (Free (Map k) v) -> a
stepF k
k a
acc Map k (Free (Map k) v)
m) Free (Map k) v
v) Map k (Free (Map k) v)
m
mayAppendTraverse a -> v -> f v'
valF a -> f (Maybe v')
_ k -> a -> Map k (Free (Map k) v) -> a
_ a
acc (Pure v
v) = v' -> KeyTree k v'
forall (f :: * -> *) a. a -> Free f a
Pure (v' -> KeyTree k v') -> f v' -> f (KeyTree k v')
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> v -> f v'
valF a
acc v
v