{-# LANGUAGE
      DerivingVia
  #-}

module Data.Mapping.MapWithDefault where

import Prelude hiding (Applicative(..), Foldable(..))
import Control.Applicative (Applicative(..))
import Data.Algebra.Boolean
import Data.Foldable (Foldable(..))
import Data.List (groupBy)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as M
import qualified Data.Map.Merge.Strict as M
import Data.Mapping
import Data.Maybe (fromMaybe, mapMaybe)
import qualified Data.Set as S
import Data.Mapping.Util


-- | Mappings constant except on an enumerated set of values
data MapWithDefault k v = MapWithDefault {
  forall k v. MapWithDefault k v -> v
common :: v,
  forall k v. MapWithDefault k v -> Map k v
exceptions :: Map k v
} deriving (MapWithDefault k v -> MapWithDefault k v -> Bool
(MapWithDefault k v -> MapWithDefault k v -> Bool)
-> (MapWithDefault k v -> MapWithDefault k v -> Bool)
-> Eq (MapWithDefault k v)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k v.
(Eq v, Eq k) =>
MapWithDefault k v -> MapWithDefault k v -> Bool
$c== :: forall k v.
(Eq v, Eq k) =>
MapWithDefault k v -> MapWithDefault k v -> Bool
== :: MapWithDefault k v -> MapWithDefault k v -> Bool
$c/= :: forall k v.
(Eq v, Eq k) =>
MapWithDefault k v -> MapWithDefault k v -> Bool
/= :: MapWithDefault k v -> MapWithDefault k v -> Bool
Eq, Eq (MapWithDefault k v)
Eq (MapWithDefault k v) =>
(MapWithDefault k v -> MapWithDefault k v -> Ordering)
-> (MapWithDefault k v -> MapWithDefault k v -> Bool)
-> (MapWithDefault k v -> MapWithDefault k v -> Bool)
-> (MapWithDefault k v -> MapWithDefault k v -> Bool)
-> (MapWithDefault k v -> MapWithDefault k v -> Bool)
-> (MapWithDefault k v -> MapWithDefault k v -> MapWithDefault k v)
-> (MapWithDefault k v -> MapWithDefault k v -> MapWithDefault k v)
-> Ord (MapWithDefault k v)
MapWithDefault k v -> MapWithDefault k v -> Bool
MapWithDefault k v -> MapWithDefault k v -> Ordering
MapWithDefault k v -> MapWithDefault k v -> MapWithDefault k v
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall k v. (Ord v, Ord k) => Eq (MapWithDefault k v)
forall k v.
(Ord v, Ord k) =>
MapWithDefault k v -> MapWithDefault k v -> Bool
forall k v.
(Ord v, Ord k) =>
MapWithDefault k v -> MapWithDefault k v -> Ordering
forall k v.
(Ord v, Ord k) =>
MapWithDefault k v -> MapWithDefault k v -> MapWithDefault k v
$ccompare :: forall k v.
(Ord v, Ord k) =>
MapWithDefault k v -> MapWithDefault k v -> Ordering
compare :: MapWithDefault k v -> MapWithDefault k v -> Ordering
$c< :: forall k v.
(Ord v, Ord k) =>
MapWithDefault k v -> MapWithDefault k v -> Bool
< :: MapWithDefault k v -> MapWithDefault k v -> Bool
$c<= :: forall k v.
(Ord v, Ord k) =>
MapWithDefault k v -> MapWithDefault k v -> Bool
<= :: MapWithDefault k v -> MapWithDefault k v -> Bool
$c> :: forall k v.
(Ord v, Ord k) =>
MapWithDefault k v -> MapWithDefault k v -> Bool
> :: MapWithDefault k v -> MapWithDefault k v -> Bool
$c>= :: forall k v.
(Ord v, Ord k) =>
MapWithDefault k v -> MapWithDefault k v -> Bool
>= :: MapWithDefault k v -> MapWithDefault k v -> Bool
$cmax :: forall k v.
(Ord v, Ord k) =>
MapWithDefault k v -> MapWithDefault k v -> MapWithDefault k v
max :: MapWithDefault k v -> MapWithDefault k v -> MapWithDefault k v
$cmin :: forall k v.
(Ord v, Ord k) =>
MapWithDefault k v -> MapWithDefault k v -> MapWithDefault k v
min :: MapWithDefault k v -> MapWithDefault k v -> MapWithDefault k v
Ord)

fromList :: (Ord k, Eq v) => v -> [(k,v)] -> MapWithDefault k v
fromList :: forall k v. (Ord k, Eq v) => v -> [(k, v)] -> MapWithDefault k v
fromList v
a = v -> Map k v -> MapWithDefault k v
forall k v. v -> Map k v -> MapWithDefault k v
MapWithDefault v
a (Map k v -> MapWithDefault k v)
-> ([(k, v)] -> Map k v) -> [(k, v)] -> MapWithDefault k v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(k, v)] -> Map k v
forall k a. Ord k => [(k, a)] -> Map k a
M.fromList ([(k, v)] -> Map k v)
-> ([(k, v)] -> [(k, v)]) -> [(k, v)] -> Map k v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k, v) -> Maybe (k, v)) -> [(k, v)] -> [(k, v)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe ((v -> Maybe v) -> (k, v) -> Maybe (k, v)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> (k, a) -> f (k, b)
traverse (v -> v -> Maybe v
forall a. Eq a => a -> a -> Maybe a
nonDefault v
a))

instance (Show k, Show v) => Show (MapWithDefault k v) where
  showsPrec :: Int -> MapWithDefault k v -> ShowS
showsPrec Int
d (MapWithDefault v
x Map k v
l) =
    (String
"fromList " <>) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    Int -> v -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
d v
x ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    [(k, v)] -> ShowS
forall a. Show a => [a] -> ShowS
showList (Map k v -> [(k, v)]
forall k a. Map k a -> [(k, a)]
M.toList Map k v
l)

fromListWithKey :: (Ord k, Eq v) => v -> (k -> u -> v -> v) -> [(k, u)] -> MapWithDefault k v
fromListWithKey :: forall k v u.
(Ord k, Eq v) =>
v -> (k -> u -> v -> v) -> [(k, u)] -> MapWithDefault k v
fromListWithKey v
a k -> u -> v -> v
f = let
  g :: Map k v -> (k, u) -> Map k v
g Map k v
m (k
k, u
x) = (Maybe v -> Maybe v) -> k -> Map k v -> Map k v
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
M.alter (v -> v -> Maybe v
forall a. Eq a => a -> a -> Maybe a
nonDefault v
a (v -> Maybe v) -> (Maybe v -> v) -> Maybe v -> Maybe v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. k -> u -> v -> v
f k
k u
x (v -> v) -> (Maybe v -> v) -> Maybe v -> v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. v -> Maybe v -> v
forall a. a -> Maybe a -> a
fromMaybe v
a) k
k Map k v
m
  in v -> Map k v -> MapWithDefault k v
forall k v. v -> Map k v -> MapWithDefault k v
MapWithDefault v
a (Map k v -> MapWithDefault k v)
-> ([(k, u)] -> Map k v) -> [(k, u)] -> MapWithDefault k v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map k v -> (k, u) -> Map k v) -> Map k v -> [(k, u)] -> Map k v
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Map k v -> (k, u) -> Map k v
g Map k v
forall k a. Map k a
M.empty

instance Foldable (MapWithDefault k) where
  foldMap :: forall m a. Monoid m => (a -> m) -> MapWithDefault k a -> m
foldMap a -> m
p (MapWithDefault a
a Map k a
f) = a -> m
p a
a m -> m -> m
forall a. Semigroup a => a -> a -> a
<> (a -> m) -> Map k a -> m
forall m a. Monoid m => (a -> m) -> Map k a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap a -> m
p Map k a
f

instance Ord k => Mapping k (MapWithDefault k) where
  cst :: forall v. v -> MapWithDefault k v
cst v
x = v -> Map k v -> MapWithDefault k v
forall k v. v -> Map k v -> MapWithDefault k v
MapWithDefault v
x Map k v
forall k a. Map k a
M.empty
  mmap :: forall v u.
Ord v =>
(u -> v) -> MapWithDefault k u -> MapWithDefault k v
mmap u -> v
p (MapWithDefault u
a Map k u
f) = let
    b :: v
b = u -> v
p u
a
    q :: u -> Maybe v
q u
x = let
      y :: v
y = u -> v
p u
x
      in if v
b v -> v -> Bool
forall a. Eq a => a -> a -> Bool
== v
y then Maybe v
forall a. Maybe a
Nothing else v -> Maybe v
forall a. a -> Maybe a
Just v
y
    in v -> Map k v -> MapWithDefault k v
forall k v. v -> Map k v -> MapWithDefault k v
MapWithDefault v
b (Map k v -> MapWithDefault k v) -> Map k v -> MapWithDefault k v
forall a b. (a -> b) -> a -> b
$ (u -> Maybe v) -> Map k u -> Map k v
forall a b k. (a -> Maybe b) -> Map k a -> Map k b
M.mapMaybe u -> Maybe v
q Map k u
f
  mtraverse :: forall (f :: * -> *) v u.
(Applicative f, Ord v) =>
(u -> f v) -> MapWithDefault k u -> f (MapWithDefault k v)
mtraverse u -> f v
p (MapWithDefault u
a Map k u
f) = let
    b :: f v
b = u -> f v
p u
a
    e :: a -> a -> Maybe a
e a
x a
y = if a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y then Maybe a
forall a. Maybe a
Nothing else a -> Maybe a
forall a. a -> Maybe a
Just a
y
    g :: p -> u -> f (Maybe v)
g p
_ u
x = (v -> v -> Maybe v) -> f v -> f v -> f (Maybe v)
forall a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 v -> v -> Maybe v
forall a. Eq a => a -> a -> Maybe a
e f v
b (u -> f v
p u
x)
    in (v -> Map k v -> MapWithDefault k v)
-> f v -> f (Map k v) -> f (MapWithDefault k v)
forall a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 v -> Map k v -> MapWithDefault k v
forall k v. v -> Map k v -> MapWithDefault k v
MapWithDefault f v
b (f (Map k v) -> f (MapWithDefault k v))
-> f (Map k v) -> f (MapWithDefault k v)
forall a b. (a -> b) -> a -> b
$ (k -> u -> f (Maybe v)) -> Map k u -> f (Map k v)
forall (f :: * -> *) k a b.
Applicative f =>
(k -> a -> f (Maybe b)) -> Map k a -> f (Map k b)
M.traverseMaybeWithKey k -> u -> f (Maybe v)
forall {p}. p -> u -> f (Maybe v)
g Map k u
f
  act :: forall v. MapWithDefault k v -> k -> v
act (MapWithDefault v
a Map k v
f) k
x = v -> Maybe v -> v
forall a. a -> Maybe a -> a
fromMaybe v
a (k -> Map k v -> Maybe v
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup k
x Map k v
f)
  isConst :: forall v. Ord v => MapWithDefault k v -> Maybe v
isConst (MapWithDefault v
a Map k v
f) = if Map k v -> Bool
forall k a. Map k a -> Bool
M.null Map k v
f then v -> Maybe v
forall a. a -> Maybe a
Just v
a else Maybe v
forall a. Maybe a
Nothing
  mergeA :: forall (f :: * -> *) w u v.
(Applicative f, Ord w) =>
(u -> v -> f w)
-> MapWithDefault k u
-> MapWithDefault k v
-> f (MapWithDefault k w)
mergeA u -> v -> f w
h (MapWithDefault u
a Map k u
f) (MapWithDefault v
b Map k v
g) = let
    e :: a -> a -> Maybe a
e a
x a
y = if a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y then a -> Maybe a
forall a. a -> Maybe a
Just a
x else Maybe a
forall a. Maybe a
Nothing
    c :: f w
c = u -> v -> f w
h u
a v
b
    l :: WhenMissing f k u w
l = (k -> u -> f w) -> WhenMissing f k u w
forall (f :: * -> *) k x y.
Applicative f =>
(k -> x -> f y) -> WhenMissing f k x y
M.traverseMissing (\k
_ u
x -> u -> v -> f w
h u
x v
b)
    r :: WhenMissing f k v w
r = (k -> v -> f w) -> WhenMissing f k v w
forall (f :: * -> *) k x y.
Applicative f =>
(k -> x -> f y) -> WhenMissing f k x y
M.traverseMissing (\k
_ v
y -> u -> v -> f w
h u
a v
y)
    h' :: p -> u -> v -> f (Maybe w)
h' p
_ u
x v
y = (w -> w -> Maybe w) -> f w -> f w -> f (Maybe w)
forall a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 w -> w -> Maybe w
forall a. Eq a => a -> a -> Maybe a
e f w
c (f w -> f (Maybe w)) -> f w -> f (Maybe w)
forall a b. (a -> b) -> a -> b
$ u -> v -> f w
h u
x v
y
    t :: WhenMatched f k u v w
t = (k -> u -> v -> f (Maybe w)) -> WhenMatched f k u v w
forall (f :: * -> *) k x y z.
Applicative f =>
(k -> x -> y -> f (Maybe z)) -> WhenMatched f k x y z
M.zipWithMaybeAMatched k -> u -> v -> f (Maybe w)
forall {p}. p -> u -> v -> f (Maybe w)
h'
    combine :: Map k u -> Map k v -> f (Map k w)
combine = WhenMissing f k u w
-> WhenMissing f k v w
-> WhenMatched f k u v w
-> Map k u
-> Map k v
-> f (Map k w)
forall (f :: * -> *) k a c b.
(Applicative f, Ord k) =>
WhenMissing f k a c
-> WhenMissing f k b c
-> WhenMatched f k a b c
-> Map k a
-> Map k b
-> f (Map k c)
M.mergeA WhenMissing f k u w
forall {k}. WhenMissing f k u w
l WhenMissing f k v w
forall {k}. WhenMissing f k v w
r WhenMatched f k u v w
forall {k}. WhenMatched f k u v w
t
    in (w -> Map k w -> MapWithDefault k w)
-> f w -> f (Map k w) -> f (MapWithDefault k w)
forall a b c. (a -> b -> c) -> f a -> f b -> f c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 w -> Map k w -> MapWithDefault k w
forall k v. v -> Map k v -> MapWithDefault k v
MapWithDefault f w
c (f (Map k w) -> f (MapWithDefault k w))
-> f (Map k w) -> f (MapWithDefault k w)
forall a b. (a -> b) -> a -> b
$ Map k u -> Map k v -> f (Map k w)
combine Map k u
f Map k v
g
  merge :: forall w u v.
Ord w =>
(u -> v -> w)
-> MapWithDefault k u -> MapWithDefault k v -> MapWithDefault k w
merge u -> v -> w
h (MapWithDefault u
a Map k u
f) (MapWithDefault v
b Map k v
g) = let
    c :: w
c = u -> v -> w
h u
a v
b
    l :: WhenMissing Identity k u w
l = (k -> u -> w) -> WhenMissing Identity k u w
forall (f :: * -> *) k x y.
Applicative f =>
(k -> x -> y) -> WhenMissing f k x y
M.mapMissing (\k
_ u
x -> u -> v -> w
h u
x v
b)
    r :: WhenMissing Identity k v w
r = (k -> v -> w) -> WhenMissing Identity k v w
forall (f :: * -> *) k x y.
Applicative f =>
(k -> x -> y) -> WhenMissing f k x y
M.mapMissing (\k
_ v
y -> u -> v -> w
h u
a v
y)
    h' :: p -> u -> v -> Maybe w
h' p
_ u
x v
y = let
      z :: w
z = u -> v -> w
h u
x v
y
      in if w
z w -> w -> Bool
forall a. Eq a => a -> a -> Bool
== w
c then Maybe w
forall a. Maybe a
Nothing else w -> Maybe w
forall a. a -> Maybe a
Just w
z
    t :: WhenMatched Identity k u v w
t = (k -> u -> v -> Maybe w) -> WhenMatched Identity k u v w
forall (f :: * -> *) k x y z.
Applicative f =>
(k -> x -> y -> Maybe z) -> WhenMatched f k x y z
M.zipWithMaybeMatched k -> u -> v -> Maybe w
forall {p}. p -> u -> v -> Maybe w
h'
    combine :: Map k u -> Map k v -> Map k w
combine = SimpleWhenMissing k u w
-> SimpleWhenMissing k v w
-> SimpleWhenMatched k u v w
-> Map k u
-> Map k v
-> Map k w
forall k a c b.
Ord k =>
SimpleWhenMissing k a c
-> SimpleWhenMissing k b c
-> SimpleWhenMatched k a b c
-> Map k a
-> Map k b
-> Map k c
M.merge SimpleWhenMissing k u w
forall {k}. WhenMissing Identity k u w
l SimpleWhenMissing k v w
forall {k}. WhenMissing Identity k v w
r SimpleWhenMatched k u v w
forall {k}. WhenMatched Identity k u v w
t
    in w -> Map k w -> MapWithDefault k w
forall k v. v -> Map k v -> MapWithDefault k v
MapWithDefault w
c (Map k w -> MapWithDefault k w) -> Map k w -> MapWithDefault k w
forall a b. (a -> b) -> a -> b
$ Map k u -> Map k v -> Map k w
combine Map k u
f Map k v
g

-- | This instance assumes that k is unbounded
--
-- It would be possible to do something valid in greater generality (for
-- example, a MaybeBounded class), which might be a good idea.
instance (Enum k, Eq k) => Neighbourly (MapWithDefault k) where
  neighbours :: forall v. Ord v => MapWithDefault k v -> Set (v, v)
neighbours (MapWithDefault v
a Map k v
f) = let
    c :: (a, b) -> (a, b) -> Bool
c (a
x,b
_) (a
y,b
_) = a -> a
forall a. Enum a => a -> a
succ a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y
    d :: [v] -> [(v, v)]
d [v]
l = [v] -> [v] -> [(v, v)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([v
a] [v] -> [v] -> [v]
forall a. Semigroup a => a -> a -> a
<> [v]
l) ([v]
l [v] -> [v] -> [v]
forall a. Semigroup a => a -> a -> a
<> [v
a])
    in [(v, v)] -> Set (v, v)
forall a. Ord a => [a] -> Set a
S.fromList ([(v, v)] -> Set (v, v))
-> ([(k, v)] -> [(v, v)]) -> [(k, v)] -> Set (v, v)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([(k, v)] -> [(v, v)]) -> [[(k, v)]] -> [(v, v)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([v] -> [(v, v)]
d ([v] -> [(v, v)]) -> ([(k, v)] -> [v]) -> [(k, v)] -> [(v, v)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k, v) -> v) -> [(k, v)] -> [v]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (k, v) -> v
forall a b. (a, b) -> b
snd) ([[(k, v)]] -> [(v, v)])
-> ([(k, v)] -> [[(k, v)]]) -> [(k, v)] -> [(v, v)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((k, v) -> (k, v) -> Bool) -> [(k, v)] -> [[(k, v)]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (k, v) -> (k, v) -> Bool
forall {a} {b} {b}. (Eq a, Enum a) => (a, b) -> (a, b) -> Bool
c ([(k, v)] -> Set (v, v)) -> [(k, v)] -> Set (v, v)
forall a b. (a -> b) -> a -> b
$ Map k v -> [(k, v)]
forall k a. Map k a -> [(k, a)]
M.toAscList Map k v
f

deriving via (AlgebraWrapper k (MapWithDefault k) b)
  instance (Ord k, Ord b, Semigroup b) => Semigroup (MapWithDefault k b)

deriving via (AlgebraWrapper k (MapWithDefault k) b)
  instance (Ord k, Ord b, Monoid b) => Monoid (MapWithDefault k b)

deriving via (AlgebraWrapper k (MapWithDefault k) b)
  instance (Ord k, Ord b, Num b) => Num (MapWithDefault k b)

deriving via (AlgebraWrapper k (MapWithDefault k) b)
  instance (Ord k, Ord b, Boolean b) => Boolean (MapWithDefault k b)