{-# LANGUAGE TypeFamilies #-}
{- | The type 'RMap' @a@ @b@ is like 'M.Map' @a@ @b@, but allows recursive definitions:

>>> :{
  let m1 = RM.insert 23 "Hello" m2
      m2 = RM.insert 42 "World" m1
  in RM.get m1
 :}
fromList [(23,"Hello"),(42,"World")]

All functions in this API are monotone with regard to the ordering of maps that
uses the /discrete/ order on its elements. Furthermore, we only include
functions where the key set does not depend on the actual values of the maps.

This means that maps defined recursively using functions like 'RM.insertWith'
can be used to construct cyclic data structures:

>>> :{
  let m = RM.insertWith (++) 23 "Hi" m
  in take 20 $ RM.get m M.! 23
 :}
"HiHiHiHiHiHiHiHiHiHi"

And because the APIs provided by this package work similar to cyclic data
structures, we can use them inside these maps:

>>> :{
  let m = RM.insertWith RS.union 23 (RS.singleton "Hi") m
  in RM.get m
 :}
fromList [(23,fromList ["Hi"])]

I am looking for a concice but useful example for this feature to be put here!

An alternative would be to order these maps using a pointwise order on the maps
of elements (and do a simple fixed-point iteration underneath). But then we
could not provide a general 'RM.unionWith' function, because not every function
passed to it would be monotone.

-}
module Data.Recursive.Map
  ( RMap
  , get
  , mk
  , empty
  , singleton
  , insert
  , insertWith
  , insertWithKey
  , delete
  , adjust
  , adjustWithKey
  , union
  , unionWith
  , unionWithKey
  , intersection
  , intersectionWith
  , intersectionWithKey
  , member
  , notMember
  , disjoint
  , Data.Recursive.Map.null
  , map
  , mapWithKey
  , fromSet
  , keysSet
  , restrictKeys
  ) where

import Prelude hiding (map)
import qualified Data.Map as M

import Data.Recursive.Internal
import qualified Data.Recursive.Set as RS

-- $setup
-- >>> :load Data.Recursive.Set Data.Recursive.Map Data.Recursive.Bool Data.Recursive.DualBool
-- >>> :module - Data.Recursive.Set Data.Recursive.Map Data.Recursive.Bool Data.Recursive.DualBool
-- >>> import qualified Data.Recursive.Set as RS
-- >>> import qualified Data.Recursive.Map as RM
-- >>> import qualified Data.Recursive.Bool as RB
-- >>> import qualified Data.Recursive.DualBool as RDB
-- >>> import qualified Data.Set as S
-- >>> import qualified Data.Map as M
-- >>> :set -XFlexibleInstances
-- >>> :set -XScopedTypeVariables
-- >>> import Test.QuickCheck
-- >>> instance (Ord a, Arbitrary a) => Arbitrary (RS.RSet a) where arbitrary = RS.mk <$> arbitrary
-- >>> instance (Ord a, Show a) => Show (RS.RSet a) where show = show . RS.get
-- >>> instance (Ord a, Arbitrary a, Arbitrary b) => Arbitrary (RM.RMap a b) where arbitrary = RM.mk <$> arbitrary
-- >>> instance (Ord a, Show a, Show b) => Show (RM.RMap a b) where show = show . RM.get

-- | Extracts the value of a 'MSet'
get :: RMap a b -> M.Map a b
get :: forall a b. RMap a b -> Map a b
get (RMap RSet a
_s Map a b
m) = Map a b
m

-- | prop> RM.get (RM.mk m) === m
mk :: M.Map a b -> RMap a b
mk :: forall a b. Map a b -> RMap a b
mk Map a b
m = RSet a -> Map a b -> RMap a b
forall a b. RSet a -> Map a b -> RMap a b
RMap (Set a -> RSet a
forall a. Set a -> RSet a
RS.mk (Map a b -> Set a
forall k a. Map k a -> Set k
M.keysSet Map a b
m)) Map a b
m

-- | prop> RM.get RM.empty === M.empty
empty :: RMap a b
empty :: forall a b. RMap a b
empty = RSet a -> Map a b -> RMap a b
forall a b. RSet a -> Map a b -> RMap a b
RMap RSet a
forall a. RSet a
RS.empty Map a b
forall k a. Map k a
M.empty

-- | prop> RM.get (RM.singleton k v) === M.singleton k v
singleton :: a -> b -> RMap a b
singleton :: forall a b. a -> b -> RMap a b
singleton a
k b
v = RSet a -> Map a b -> RMap a b
forall a b. RSet a -> Map a b -> RMap a b
RMap (a -> RSet a
forall a. a -> RSet a
RS.singleton a
k) (a -> b -> Map a b
forall k a. k -> a -> Map k a
M.singleton a
k b
v)

build :: Ord a => RS.RSet a -> M.Map a b -> RMap a b
build :: forall a b. Ord a => RSet a -> Map a b -> RMap a b
build RSet a
s Map a b
m = RSet a -> Map a b -> RMap a b
forall a b. RSet a -> Map a b -> RMap a b
RMap RSet a
s ((a -> b) -> Set a -> Map a b
forall k a. (k -> a) -> Set k -> Map k a
M.fromSet (Map a b
m Map a b -> a -> b
forall k a. Ord k => Map k a -> k -> a
M.!) (RSet a -> Set a
forall a. RSet a -> Set a
RS.get RSet a
s))

-- | prop> RM.get (RM.insert k v m) === M.insert k v (RM.get m)
insert :: Ord a => a -> b -> RMap a b -> RMap a b
insert :: forall a b. Ord a => a -> b -> RMap a b -> RMap a b
insert a
k b
v ~(RMap RSet a
rs Map a b
m) = RSet a -> Map a b -> RMap a b
forall a b. Ord a => RSet a -> Map a b -> RMap a b
build (a -> RSet a -> RSet a
forall a. Ord a => a -> RSet a -> RSet a
RS.insert a
k RSet a
rs) (a -> b -> Map a b -> Map a b
forall k a. Ord k => k -> a -> Map k a -> Map k a
M.insert a
k b
v Map a b
m)

-- | prop> RM.get (RM.insertWith (applyFun2 f) k v m) === M.insertWith (applyFun2 f) k v (RM.get m)
insertWith :: Ord a => (b -> b -> b) -> a -> b -> RMap a b -> RMap a b
insertWith :: forall a b.
Ord a =>
(b -> b -> b) -> a -> b -> RMap a b -> RMap a b
insertWith b -> b -> b
f a
k b
v ~(RMap RSet a
rs Map a b
m) = RSet a -> Map a b -> RMap a b
forall a b. Ord a => RSet a -> Map a b -> RMap a b
build (a -> RSet a -> RSet a
forall a. Ord a => a -> RSet a -> RSet a
RS.insert a
k RSet a
rs) ((b -> b -> b) -> a -> b -> Map a b -> Map a b
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWith b -> b -> b
f a
k b
v Map a b
m)

-- | prop> RM.get (RM.insertWithKey (applyFun3 f) k v m) === M.insertWithKey (applyFun3 f) k v (RM.get m)
insertWithKey :: Ord a => (a -> b -> b -> b) -> a -> b -> RMap a b -> RMap a b
insertWithKey :: forall a b.
Ord a =>
(a -> b -> b -> b) -> a -> b -> RMap a b -> RMap a b
insertWithKey a -> b -> b -> b
f a
k b
v ~(RMap RSet a
rs Map a b
m) = RSet a -> Map a b -> RMap a b
forall a b. Ord a => RSet a -> Map a b -> RMap a b
build (a -> RSet a -> RSet a
forall a. Ord a => a -> RSet a -> RSet a
RS.insert a
k RSet a
rs) ((a -> b -> b -> b) -> a -> b -> Map a b -> Map a b
forall k a.
Ord k =>
(k -> a -> a -> a) -> k -> a -> Map k a -> Map k a
M.insertWithKey a -> b -> b -> b
f a
k b
v Map a b
m)

-- | prop> RM.get (RM.delete k m) === M.delete k (RM.get m)
delete :: Ord a => a -> RMap a b -> RMap a b
delete :: forall a b. Ord a => a -> RMap a b -> RMap a b
delete a
k ~(RMap RSet a
rs Map a b
m) = RSet a -> Map a b -> RMap a b
forall a b. Ord a => RSet a -> Map a b -> RMap a b
build (a -> RSet a -> RSet a
forall a. Ord a => a -> RSet a -> RSet a
RS.delete a
k RSet a
rs) (a -> Map a b -> Map a b
forall k a. Ord k => k -> Map k a -> Map k a
M.delete a
k Map a b
m)

-- | prop> RM.get (RM.adjust (applyFun f) k m) === M.adjust (applyFun f) k (RM.get m)
adjust :: Ord a => (b -> b) -> a -> RMap a b -> RMap a b
adjust :: forall a b. Ord a => (b -> b) -> a -> RMap a b -> RMap a b
adjust b -> b
f a
k ~(RMap RSet a
rs Map a b
m) = RSet a -> Map a b -> RMap a b
forall a b. Ord a => RSet a -> Map a b -> RMap a b
build (RSet a -> RSet a
forall a. RSet a -> RSet a
RS.id RSet a
rs) ((b -> b) -> a -> Map a b -> Map a b
forall k a. Ord k => (a -> a) -> k -> Map k a -> Map k a
M.adjust b -> b
f a
k Map a b
m)

-- | prop> RM.get (RM.adjustWithKey (applyFun2 f) k m) === M.adjustWithKey (applyFun2 f) k (RM.get m)
adjustWithKey :: Ord a => (a -> b -> b) -> a -> RMap a b -> RMap a b
adjustWithKey :: forall a b. Ord a => (a -> b -> b) -> a -> RMap a b -> RMap a b
adjustWithKey a -> b -> b
f a
k ~(RMap RSet a
rs Map a b
m) = RSet a -> Map a b -> RMap a b
forall a b. Ord a => RSet a -> Map a b -> RMap a b
build (RSet a -> RSet a
forall a. RSet a -> RSet a
RS.id RSet a
rs) ((a -> b -> b) -> a -> Map a b -> Map a b
forall k a. Ord k => (k -> a -> a) -> k -> Map k a -> Map k a
M.adjustWithKey a -> b -> b
f a
k Map a b
m)

-- | prop> RM.get (RM.union m1 m2) === M.union (RM.get m1) (RM.get m2)
union :: Ord a => RMap a b -> RMap a b -> RMap a b
union :: forall a b. Ord a => RMap a b -> RMap a b -> RMap a b
union ~(RMap RSet a
rs1 Map a b
m1) ~(RMap RSet a
rs2 Map a b
m2) = RSet a -> Map a b -> RMap a b
forall a b. Ord a => RSet a -> Map a b -> RMap a b
build (RSet a -> RSet a -> RSet a
forall a. Ord a => RSet a -> RSet a -> RSet a
RS.union RSet a
rs1 RSet a
rs2) (Map a b -> Map a b -> Map a b
forall k a. Ord k => Map k a -> Map k a -> Map k a
M.union Map a b
m1 Map a b
m2)

-- | prop> RM.get (RM.unionWith (applyFun2 f) m1 m2) === M.unionWith (applyFun2 f) (RM.get m1) (RM.get m2)
unionWith :: Ord a => (b -> b -> b) -> RMap a b -> RMap a b -> RMap a b
unionWith :: forall a b.
Ord a =>
(b -> b -> b) -> RMap a b -> RMap a b -> RMap a b
unionWith b -> b -> b
f ~(RMap RSet a
rs1 Map a b
m1) ~(RMap RSet a
rs2 Map a b
m2) = RSet a -> Map a b -> RMap a b
forall a b. Ord a => RSet a -> Map a b -> RMap a b
build (RSet a -> RSet a -> RSet a
forall a. Ord a => RSet a -> RSet a -> RSet a
RS.union RSet a
rs1 RSet a
rs2) ((b -> b -> b) -> Map a b -> Map a b -> Map a b
forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWith b -> b -> b
f Map a b
m1 Map a b
m2)

-- | prop> RM.get (RM.unionWithKey (applyFun3 f) m1 m2) === M.unionWithKey (applyFun3 f) (RM.get m1) (RM.get m2)
unionWithKey :: Ord a => (a -> b -> b -> b) -> RMap a b -> RMap a b -> RMap a b
unionWithKey :: forall a b.
Ord a =>
(a -> b -> b -> b) -> RMap a b -> RMap a b -> RMap a b
unionWithKey a -> b -> b -> b
f ~(RMap RSet a
rs1 Map a b
m1) ~(RMap RSet a
rs2 Map a b
m2) = RSet a -> Map a b -> RMap a b
forall a b. Ord a => RSet a -> Map a b -> RMap a b
build (RSet a -> RSet a -> RSet a
forall a. Ord a => RSet a -> RSet a -> RSet a
RS.union RSet a
rs1 RSet a
rs2) ((a -> b -> b -> b) -> Map a b -> Map a b -> Map a b
forall k a.
Ord k =>
(k -> a -> a -> a) -> Map k a -> Map k a -> Map k a
M.unionWithKey a -> b -> b -> b
f Map a b
m1 Map a b
m2)

-- | prop> RM.get (RM.intersection m1 m2) === M.intersection (RM.get m1) (RM.get m2)
intersection :: Ord a => RMap a b -> RMap a b -> RMap a b
intersection :: forall a b. Ord a => RMap a b -> RMap a b -> RMap a b
intersection ~(RMap RSet a
rs1 Map a b
m1) ~(RMap RSet a
rs2 Map a b
m2) = RSet a -> Map a b -> RMap a b
forall a b. Ord a => RSet a -> Map a b -> RMap a b
build (RSet a -> RSet a -> RSet a
forall a. Ord a => RSet a -> RSet a -> RSet a
RS.intersection RSet a
rs1 RSet a
rs2) (Map a b -> Map a b -> Map a b
forall k a b. Ord k => Map k a -> Map k b -> Map k a
M.intersection Map a b
m1 Map a b
m2)

-- | prop> RM.get (RM.intersectionWith (applyFun2 f) m1 m2) === M.intersectionWith (applyFun2 f) (RM.get m1) (RM.get m2)
intersectionWith :: Ord a => (b -> b -> b) -> RMap a b -> RMap a b -> RMap a b
intersectionWith :: forall a b.
Ord a =>
(b -> b -> b) -> RMap a b -> RMap a b -> RMap a b
intersectionWith b -> b -> b
f ~(RMap RSet a
rs1 Map a b
m1) ~(RMap RSet a
rs2 Map a b
m2) = RSet a -> Map a b -> RMap a b
forall a b. Ord a => RSet a -> Map a b -> RMap a b
build (RSet a -> RSet a -> RSet a
forall a. Ord a => RSet a -> RSet a -> RSet a
RS.intersection RSet a
rs1 RSet a
rs2) ((b -> b -> b) -> Map a b -> Map a b -> Map a b
forall k a b c.
Ord k =>
(a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWith b -> b -> b
f Map a b
m1 Map a b
m2)

-- | prop> RM.get (RM.intersectionWithKey (applyFun3 f) m1 m2) === M.intersectionWithKey (applyFun3 f) (RM.get m1) (RM.get m2)
intersectionWithKey :: Ord a => (a -> b -> b -> b) -> RMap a b -> RMap a b -> RMap a b
intersectionWithKey :: forall a b.
Ord a =>
(a -> b -> b -> b) -> RMap a b -> RMap a b -> RMap a b
intersectionWithKey a -> b -> b -> b
f ~(RMap RSet a
rs1 Map a b
m1) ~(RMap RSet a
rs2 Map a b
m2) = RSet a -> Map a b -> RMap a b
forall a b. Ord a => RSet a -> Map a b -> RMap a b
build (RSet a -> RSet a -> RSet a
forall a. Ord a => RSet a -> RSet a -> RSet a
RS.intersection RSet a
rs1 RSet a
rs2) ((a -> b -> b -> b) -> Map a b -> Map a b -> Map a b
forall k a b c.
Ord k =>
(k -> a -> b -> c) -> Map k a -> Map k b -> Map k c
M.intersectionWithKey a -> b -> b -> b
f Map a b
m1 Map a b
m2)

-- | prop> RM.get (RM.map (applyFun f) m) === M.map (applyFun f) (RM.get m)
map :: Ord k => (a -> b) -> RMap k a -> RMap k b
map :: forall k a b. Ord k => (a -> b) -> RMap k a -> RMap k b
map a -> b
f ~(RMap RSet k
rs Map k a
m) = RSet k -> Map k b -> RMap k b
forall a b. Ord a => RSet a -> Map a b -> RMap a b
build (RSet k -> RSet k
forall a. RSet a -> RSet a
RS.id RSet k
rs) ((a -> b) -> Map k a -> Map k b
forall a b k. (a -> b) -> Map k a -> Map k b
M.map a -> b
f Map k a
m)

-- | prop> RM.get (RM.mapWithKey (applyFun2 f) m) === M.mapWithKey (applyFun2 f) (RM.get m)
mapWithKey :: Ord k => (k -> a -> b) -> RMap k a -> RMap k b
mapWithKey :: forall k a b. Ord k => (k -> a -> b) -> RMap k a -> RMap k b
mapWithKey k -> a -> b
f ~(RMap RSet k
rs Map k a
m) = RSet k -> Map k b -> RMap k b
forall a b. Ord a => RSet a -> Map a b -> RMap a b
build (RSet k -> RSet k
forall a. RSet a -> RSet a
RS.id RSet k
rs) ((k -> a -> b) -> Map k a -> Map k b
forall k a b. (k -> a -> b) -> Map k a -> Map k b
M.mapWithKey k -> a -> b
f Map k a
m)

-- | prop> RM.get (RM.singleton k v) === M.singleton k v
fromSet :: (a -> b) -> RS.RSet a -> RMap a b
fromSet :: forall a b. (a -> b) -> RSet a -> RMap a b
fromSet a -> b
f RSet a
s = RSet a -> Map a b -> RMap a b
forall a b. RSet a -> Map a b -> RMap a b
RMap RSet a
s ((a -> b) -> Set a -> Map a b
forall k a. (k -> a) -> Set k -> Map k a
M.fromSet a -> b
f (RSet a -> Set a
forall a. RSet a -> Set a
RS.get RSet a
s))

-- | prop> RS.get (RM.keysSet m) === M.keysSet (RM.get m)
keysSet :: RMap a b -> RS.RSet a
keysSet :: forall a b. RMap a b -> RSet a
keysSet ~(RMap RSet a
rs Map a b
_) = RSet a -> RSet a
forall a. RSet a -> RSet a
RS.id RSet a
rs
  -- better use RS.id either here or in fromSet, to avoid unproductive loops

-- | prop> RM.get (RM.restrictKeys m s) === M.restrictKeys (RM.get m) (RS.get s)
restrictKeys :: Ord a => RMap a b -> RS.RSet a -> RMap a b
restrictKeys :: forall a b. Ord a => RMap a b -> RSet a -> RMap a b
restrictKeys ~(RMap RSet a
rs Map a b
m) RSet a
s2 =
    RSet a -> Map a b -> RMap a b
forall a b. Ord a => RSet a -> Map a b -> RMap a b
build (RSet a
rs RSet a -> RSet a -> RSet a
forall a. Ord a => RSet a -> RSet a -> RSet a
`RS.intersection` RSet a
s2) (Map a b -> Set a -> Map a b
forall k a. Ord k => Map k a -> Set k -> Map k a
M.restrictKeys Map a b
m (RSet a -> Set a
forall a. RSet a -> Set a
RS.get RSet a
s2))

-- | prop> RB.get (RM.member k m) === M.member k (RM.get m)
member :: Ord a => a -> RMap a b -> RBool
member :: forall a b. Ord a => a -> RMap a b -> RBool
member a
x ~(RMap RSet a
rs Map a b
_) = a -> RSet a -> RBool
forall a. Ord a => a -> RSet a -> RBool
RS.member a
x RSet a
rs

-- | prop> RDB.get (RM.notMember n r1) === M.notMember n (RM.get r1)
notMember :: Ord a => a -> RMap a b -> RDualBool
notMember :: forall a b. Ord a => a -> RMap a b -> RDualBool
notMember a
x ~(RMap RSet a
rs Map a b
_) = a -> RSet a -> RDualBool
forall a. Ord a => a -> RSet a -> RDualBool
RS.notMember a
x RSet a
rs

-- | prop> RDB.get (RM.disjoint m1 m2) === M.disjoint (RM.get m1) (RM.get m2)
disjoint :: Ord a => RMap a b -> RMap a b -> RDualBool
disjoint :: forall a b. Ord a => RMap a b -> RMap a b -> RDualBool
disjoint ~(RMap RSet a
rs1 Map a b
_ ) ~(RMap RSet a
rs2 Map a b
_) = RSet a -> RSet a -> RDualBool
forall a. Ord a => RSet a -> RSet a -> RDualBool
RS.disjoint RSet a
rs1 RSet a
rs2

-- | prop> RDB.get (RM.null m) === M.null (RM.get m)
null :: Ord a =>  RMap a b -> RDualBool
null :: forall a b. Ord a => RMap a b -> RDualBool
null ~(RMap RSet a
rs Map a b
_) = RSet a -> RDualBool
forall a. RSet a -> RDualBool
RS.null RSet a
rs