{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE EmptyCase #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}

module Data.Vessel.Disperse where

import Control.Arrow ((***))
import Data.Witherable
import Data.Align
import Data.Map.Monoidal (MonoidalMap (..))
import qualified Data.Map.Monoidal as Map
import Data.These
import Data.Foldable hiding (null)

import Data.Vessel.Internal

-- | Disperse is a simplified version of View for ordinary containers. This is used as a stepping stone to obtain the View instance for MapV.
class Disperse row where
  disperse :: (Foldable col, Filterable col, Functor col) => col (row a) -> row (col a)
  condense :: (Align col) => row (col a) -> col (row a)

instance Ord k => Disperse (MonoidalMap k) where
  disperse :: forall col a. (Foldable col, Filterable col, Functor col)
           => col (MonoidalMap k a)
           -> MonoidalMap k (col a)
  disperse :: col (MonoidalMap k a) -> MonoidalMap k (col a)
disperse col (MonoidalMap k a)
col = MonoidalMap k a -> col (MonoidalMap k a) -> MonoidalMap k (col a)
disperse' (Map k a -> MonoidalMap k a
forall k a. Map k a -> MonoidalMap k a
Map.MonoidalMap (col (Map k a) -> Map k a
forall (t :: * -> *) m. (Foldable t, Monoid m) => t m -> m
fold ((MonoidalMap k a -> Map k a)
-> col (MonoidalMap k a) -> col (Map k a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap MonoidalMap k a -> Map k a
forall k a. MonoidalMap k a -> Map k a
Map.getMonoidalMap col (MonoidalMap k a)
col))) col (MonoidalMap k a)
col
    where
      disperse'
        :: MonoidalMap k a
        -> col (MonoidalMap k a)
        -> MonoidalMap k (col a)
      disperse' :: MonoidalMap k a -> col (MonoidalMap k a) -> MonoidalMap k (col a)
disperse' MonoidalMap k a
folded col (MonoidalMap k a)
col' = case MonoidalMap k a -> Pivot k a
forall k a. Ord k => MonoidalMap k a -> Pivot k a
findPivot MonoidalMap k a
folded of
        Pivot k a
None -> Map k (col a) -> MonoidalMap k (col a)
forall k a. Map k a -> MonoidalMap k a
Map.MonoidalMap Map k (col a)
forall a. Monoid a => a
mempty
        One k
k a
_ -> k -> col a -> MonoidalMap k (col a)
forall k a. k -> a -> MonoidalMap k a
Map.singleton k
k ((MonoidalMap k a -> Maybe a) -> col (MonoidalMap k a) -> col a
forall (f :: * -> *) a b.
Filterable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe (k -> MonoidalMap k a -> Maybe a
forall k a. Ord k => k -> MonoidalMap k a -> Maybe a
Map.lookup k
k) col (MonoidalMap k a)
col')
        Split k
pivot MonoidalMap k a
l MonoidalMap k a
r -> (MonoidalMap k (col a)
 -> MonoidalMap k (col a) -> MonoidalMap k (col a))
-> (MonoidalMap k (col a), MonoidalMap k (col a))
-> MonoidalMap k (col a)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry MonoidalMap k (col a)
-> MonoidalMap k (col a) -> MonoidalMap k (col a)
forall k a.
Ord k =>
MonoidalMap k a -> MonoidalMap k a -> MonoidalMap k a
unionDistinctAsc ((MonoidalMap k (col a), MonoidalMap k (col a))
 -> MonoidalMap k (col a))
-> (MonoidalMap k (col a), MonoidalMap k (col a))
-> MonoidalMap k (col a)
forall a b. (a -> b) -> a -> b
$ (MonoidalMap k a -> col (MonoidalMap k a) -> MonoidalMap k (col a)
disperse' MonoidalMap k a
l (col (MonoidalMap k a) -> MonoidalMap k (col a))
-> (col (MonoidalMap k a) -> MonoidalMap k (col a))
-> (col (MonoidalMap k a), col (MonoidalMap k a))
-> (MonoidalMap k (col a), MonoidalMap k (col a))
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** MonoidalMap k a -> col (MonoidalMap k a) -> MonoidalMap k (col a)
disperse' MonoidalMap k a
r) ((col (MonoidalMap k a), col (MonoidalMap k a))
 -> (MonoidalMap k (col a), MonoidalMap k (col a)))
-> (col (MonoidalMap k a), col (MonoidalMap k a))
-> (MonoidalMap k (col a), MonoidalMap k (col a))
forall a b. (a -> b) -> a -> b
$ k
-> col (MonoidalMap k a)
-> (col (MonoidalMap k a), col (MonoidalMap k a))
(Ord k, Filterable col) =>
k
-> col (MonoidalMap k a)
-> (col (MonoidalMap k a), col (MonoidalMap k a))
split k
pivot col (MonoidalMap k a)
col'
      split
        :: (Ord k, Filterable col)
        => k -> col (MonoidalMap k a)
        -> (col (MonoidalMap k a), col (MonoidalMap k a))
      split :: k
-> col (MonoidalMap k a)
-> (col (MonoidalMap k a), col (MonoidalMap k a))
split k
pivot col (MonoidalMap k a)
col' = col (These (MonoidalMap k a) (MonoidalMap k a))
-> (col (MonoidalMap k a), col (MonoidalMap k a))
forall (f :: * -> *) a b.
Filterable f =>
f (These a b) -> (f a, f b)
unalignProperly (col (These (MonoidalMap k a) (MonoidalMap k a))
 -> (col (MonoidalMap k a), col (MonoidalMap k a)))
-> col (These (MonoidalMap k a) (MonoidalMap k a))
-> (col (MonoidalMap k a), col (MonoidalMap k a))
forall a b. (a -> b) -> a -> b
$ (MonoidalMap k a
 -> Maybe (These (MonoidalMap k a) (MonoidalMap k a)))
-> col (MonoidalMap k a)
-> col (These (MonoidalMap k a) (MonoidalMap k a))
forall (f :: * -> *) a b.
Filterable f =>
(a -> Maybe b) -> f a -> f b
mapMaybe (k
-> MonoidalMap k a
-> Maybe (These (MonoidalMap k a) (MonoidalMap k a))
Ord k =>
k
-> MonoidalMap k a
-> Maybe (These (MonoidalMap k a) (MonoidalMap k a))
splitOne k
pivot) col (MonoidalMap k a)
col'
      splitOne
        :: Ord k
        => k
        -> MonoidalMap k a
        -> Maybe (These (MonoidalMap k a) (MonoidalMap k a))
      splitOne :: k
-> MonoidalMap k a
-> Maybe (These (MonoidalMap k a) (MonoidalMap k a))
splitOne k
pivot MonoidalMap k a
m =
        let (MonoidalMap k a
l, MonoidalMap k a
r) = k -> MonoidalMap k a -> (MonoidalMap k a, MonoidalMap k a)
forall k a.
Ord k =>
k -> MonoidalMap k a -> (MonoidalMap k a, MonoidalMap k a)
splitLT k
pivot MonoidalMap k a
m
        in case (MonoidalMap k a -> Bool
forall k a. MonoidalMap k a -> Bool
Map.null MonoidalMap k a
l, MonoidalMap k a -> Bool
forall k a. MonoidalMap k a -> Bool
Map.null MonoidalMap k a
r) of
          (Bool
True, Bool
True) -> Maybe (These (MonoidalMap k a) (MonoidalMap k a))
forall a. Maybe a
Nothing
          (Bool
False, Bool
True) -> These (MonoidalMap k a) (MonoidalMap k a)
-> Maybe (These (MonoidalMap k a) (MonoidalMap k a))
forall a. a -> Maybe a
Just (These (MonoidalMap k a) (MonoidalMap k a)
 -> Maybe (These (MonoidalMap k a) (MonoidalMap k a)))
-> These (MonoidalMap k a) (MonoidalMap k a)
-> Maybe (These (MonoidalMap k a) (MonoidalMap k a))
forall a b. (a -> b) -> a -> b
$ MonoidalMap k a -> These (MonoidalMap k a) (MonoidalMap k a)
forall a b. a -> These a b
This MonoidalMap k a
l
          (Bool
True, Bool
False) -> These (MonoidalMap k a) (MonoidalMap k a)
-> Maybe (These (MonoidalMap k a) (MonoidalMap k a))
forall a. a -> Maybe a
Just (These (MonoidalMap k a) (MonoidalMap k a)
 -> Maybe (These (MonoidalMap k a) (MonoidalMap k a)))
-> These (MonoidalMap k a) (MonoidalMap k a)
-> Maybe (These (MonoidalMap k a) (MonoidalMap k a))
forall a b. (a -> b) -> a -> b
$ MonoidalMap k a -> These (MonoidalMap k a) (MonoidalMap k a)
forall a b. b -> These a b
That MonoidalMap k a
r
          (Bool
False, Bool
False) -> These (MonoidalMap k a) (MonoidalMap k a)
-> Maybe (These (MonoidalMap k a) (MonoidalMap k a))
forall a. a -> Maybe a
Just (These (MonoidalMap k a) (MonoidalMap k a)
 -> Maybe (These (MonoidalMap k a) (MonoidalMap k a)))
-> These (MonoidalMap k a) (MonoidalMap k a)
-> Maybe (These (MonoidalMap k a) (MonoidalMap k a))
forall a b. (a -> b) -> a -> b
$ MonoidalMap k a
-> MonoidalMap k a -> These (MonoidalMap k a) (MonoidalMap k a)
forall a b. a -> b -> These a b
These MonoidalMap k a
l MonoidalMap k a
r

  condense :: forall col a. (Align col)
           => MonoidalMap k (col a)
           -> col (MonoidalMap k a)
  condense :: MonoidalMap k (col a) -> col (MonoidalMap k a)
condense MonoidalMap k (col a)
row = case MonoidalMap k (col a) -> Pivot k (col a)
forall k a. Ord k => MonoidalMap k a -> Pivot k a
findPivot MonoidalMap k (col a)
row of
    Pivot k (col a)
None -> col (MonoidalMap k a)
forall (f :: * -> *) a. Align f => f a
nil
    One k
k col a
v -> (a -> MonoidalMap k a) -> col a -> col (MonoidalMap k a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (k -> a -> MonoidalMap k a
forall k a. k -> a -> MonoidalMap k a
Map.singleton k
k) col a
v
    Split k
pivot MonoidalMap k (col a)
_l MonoidalMap k (col a)
_r -> (col (MonoidalMap k a)
 -> col (MonoidalMap k a) -> col (MonoidalMap k a))
-> (col (MonoidalMap k a), col (MonoidalMap k a))
-> col (MonoidalMap k a)
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry ((These (MonoidalMap k a) (MonoidalMap k a) -> MonoidalMap k a)
-> col (MonoidalMap k a)
-> col (MonoidalMap k a)
-> col (MonoidalMap k a)
forall (f :: * -> *) a b c.
Semialign f =>
(These a b -> c) -> f a -> f b -> f c
alignWith ((MonoidalMap k a -> MonoidalMap k a -> MonoidalMap k a)
-> These (MonoidalMap k a) (MonoidalMap k a) -> MonoidalMap k a
forall a. (a -> a -> a) -> These a a -> a
mergeThese MonoidalMap k a -> MonoidalMap k a -> MonoidalMap k a
forall k a.
Ord k =>
MonoidalMap k a -> MonoidalMap k a -> MonoidalMap k a
unionDistinctAsc)) ((col (MonoidalMap k a), col (MonoidalMap k a))
 -> col (MonoidalMap k a))
-> (col (MonoidalMap k a), col (MonoidalMap k a))
-> col (MonoidalMap k a)
forall a b. (a -> b) -> a -> b
$ MonoidalMap k (col a) -> col (MonoidalMap k a)
forall (row :: * -> *) (col :: * -> *) a.
(Disperse row, Align col) =>
row (col a) -> col (row a)
condense (MonoidalMap k (col a) -> col (MonoidalMap k a))
-> (MonoidalMap k (col a) -> col (MonoidalMap k a))
-> (MonoidalMap k (col a), MonoidalMap k (col a))
-> (col (MonoidalMap k a), col (MonoidalMap k a))
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** MonoidalMap k (col a) -> col (MonoidalMap k a)
forall (row :: * -> *) (col :: * -> *) a.
(Disperse row, Align col) =>
row (col a) -> col (row a)
condense ((MonoidalMap k (col a), MonoidalMap k (col a))
 -> (col (MonoidalMap k a), col (MonoidalMap k a)))
-> (MonoidalMap k (col a), MonoidalMap k (col a))
-> (col (MonoidalMap k a), col (MonoidalMap k a))
forall a b. (a -> b) -> a -> b
$ k
-> MonoidalMap k (col a)
-> (MonoidalMap k (col a), MonoidalMap k (col a))
forall k a.
Ord k =>
k -> MonoidalMap k a -> (MonoidalMap k a, MonoidalMap k a)
splitLT k
pivot MonoidalMap k (col a)
row