{-# LANGUAGE
      CPP,
      DerivingVia
  #-}

module Data.Mapping.MapWithDefault where

#if MIN_VERSION_GLASGOW_HASKELL(9,6,0,0)
#else
import Control.Applicative (liftA2)
#endif
import Data.Algebra.Boolean
import Data.List (foldl', 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
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
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
$c== :: forall k v.
(Eq v, Eq k) =>
MapWithDefault k v -> MapWithDefault k v -> Bool
Eq, MapWithDefault k v -> MapWithDefault k v -> Bool
MapWithDefault k v -> MapWithDefault k v -> Ordering
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
min :: 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
max :: MapWithDefault k v -> MapWithDefault k v -> MapWithDefault k v
$cmax :: forall k v.
(Ord v, Ord k) =>
MapWithDefault k v -> MapWithDefault k v -> MapWithDefault k v
>= :: 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
$c< :: forall k v.
(Ord v, Ord k) =>
MapWithDefault k v -> MapWithDefault k v -> Bool
compare :: MapWithDefault k v -> MapWithDefault k v -> Ordering
$ccompare :: forall k v.
(Ord v, Ord k) =>
MapWithDefault k v -> MapWithDefault k v -> Ordering
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 = forall k v. v -> Map k v -> MapWithDefault k v
MapWithDefault v
a forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k a. Ord k => [(k, a)] -> Map k a
M.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (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 " <>) forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    forall a. Show a => Int -> a -> ShowS
showsPrec Int
d v
x forall b c a. (b -> c) -> (a -> b) -> a -> c
.
    forall a. Show a => [a] -> ShowS
showList (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) = forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
M.alter (forall a. Eq a => a -> a -> Maybe a
nonDefault v
a forall b c a. (b -> c) -> (a -> b) -> a -> c
. k -> u -> v -> v
f k
k u
x forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. a -> Maybe a -> a
fromMaybe v
a) k
k Map k v
m
  in forall k v. v -> Map k v -> MapWithDefault k v
MapWithDefault v
a forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Map k v -> (k, u) -> Map k v
g 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 forall a. Semigroup a => a -> a -> a
<> 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 = forall k v. v -> Map k v -> MapWithDefault k v
MapWithDefault v
x 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 forall a. Eq a => a -> a -> Bool
== v
y then forall a. Maybe a
Nothing else forall a. a -> Maybe a
Just v
y
    in forall k v. v -> Map k v -> MapWithDefault k v
MapWithDefault v
b forall a b. (a -> b) -> a -> b
$ 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 forall a. Eq a => a -> a -> Bool
== a
y then forall a. Maybe a
Nothing else forall a. a -> Maybe a
Just a
y
    g :: p -> u -> f (Maybe v)
g p
_ u
x = forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 forall a. Eq a => a -> a -> Maybe a
e f v
b (u -> f v
p u
x)
    in forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 forall k v. v -> Map k v -> MapWithDefault k v
MapWithDefault f v
b forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) k a b.
Applicative f =>
(k -> a -> f (Maybe b)) -> Map k a -> f (Map k b)
M.traverseMaybeWithKey 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 = forall a. a -> Maybe a -> a
fromMaybe v
a (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 forall k a. Map k a -> Bool
M.null Map k v
f then forall a. a -> Maybe a
Just v
a else 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 forall a. Eq a => a -> a -> Bool
== a
y then forall a. a -> Maybe a
Just a
x else 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 = 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 = 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 = forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 forall a. Eq a => a -> a -> Maybe a
e f w
c forall a b. (a -> b) -> a -> b
$ u -> v -> f w
h u
x v
y
    t :: WhenMatched f k u v w
t = forall (f :: * -> *) k x y z.
Applicative f =>
(k -> x -> y -> f (Maybe z)) -> WhenMatched f k x y z
M.zipWithMaybeAMatched forall {p}. p -> u -> v -> f (Maybe w)
h'
    combine :: Map k u -> Map k v -> f (Map k w)
combine = 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 forall {k}. WhenMissing f k u w
l forall {k}. WhenMissing f k v w
r forall {k}. WhenMatched f k u v w
t
    in forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 forall k v. v -> Map k v -> MapWithDefault k v
MapWithDefault f w
c 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 = 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 = 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 forall a. Eq a => a -> a -> Bool
== w
c then forall a. Maybe a
Nothing else forall a. a -> Maybe a
Just w
z
    t :: WhenMatched Identity k u v w
t = forall (f :: * -> *) k x y z.
Applicative f =>
(k -> x -> y -> Maybe z) -> WhenMatched f k x y z
M.zipWithMaybeMatched forall {p}. p -> u -> v -> Maybe w
h'
    combine :: Map k u -> Map k v -> Map k w
combine = 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 forall {k}. WhenMissing Identity k u w
l forall {k}. WhenMissing Identity k v w
r forall {k}. WhenMatched Identity k u v w
t
    in forall k v. v -> Map k v -> MapWithDefault k v
MapWithDefault w
c 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
_) = forall a. Enum a => a -> a
succ a
x forall a. Eq a => a -> a -> Bool
== a
y
    d :: [v] -> [(v, v)]
d [v]
l = forall a b. [a] -> [b] -> [(a, b)]
zip ([v
a] forall a. Semigroup a => a -> a -> a
<> [v]
l) ([v]
l forall a. Semigroup a => a -> a -> a
<> [v
a])
    in forall a. Ord a => [a] -> Set a
S.fromList forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ([v] -> [(v, v)]
d forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a b. (a, b) -> b
snd) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy forall {a} {b} {b}. (Eq a, Enum a) => (a, b) -> (a, b) -> Bool
c forall a b. (a -> b) -> a -> b
$ 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)