{-# LANGUAGE CPP           #-}
{-# LANGUAGE DeriveFunctor #-}
-- | Map which remembers the 'fromList' order.
-- This module is minimal on purpose.
module Data.TreeDiff.OMap (
    -- * Ordered map
    OMap,
    -- * Conversions
    toAscList,
    toList,
    fromList,
    -- * Construction
    empty,
    -- * Query
    elems,
) where

import Data.List      (sortBy)
import Data.Ord       (comparing)
import Data.Semialign (Semialign (..))
import Data.These     (These (..))
import Control.DeepSeq  (NFData (..))

#if MIN_VERSION_containers(0,5,0)
import qualified Data.Map.Strict as Map
#else
import qualified Data.Map as Map
#endif

import qualified Test.QuickCheck as QC

-- $setup
-- >>> import Data.Semialign (alignWith)

-------------------------------------------------------------------------------
-- Types
-------------------------------------------------------------------------------

newtype OMap k v = OMap (Map.Map k (Val v))
  deriving (forall a b. a -> OMap k b -> OMap k a
forall a b. (a -> b) -> OMap k a -> OMap k b
forall k a b. a -> OMap k b -> OMap k a
forall k a b. (a -> b) -> OMap k a -> OMap k b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> OMap k b -> OMap k a
$c<$ :: forall k a b. a -> OMap k b -> OMap k a
fmap :: forall a b. (a -> b) -> OMap k a -> OMap k b
$cfmap :: forall k a b. (a -> b) -> OMap k a -> OMap k b
Functor)

-- Value with its index
data Val v = Val !Int v
  deriving (forall a b. a -> Val b -> Val a
forall a b. (a -> b) -> Val a -> Val b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> Val b -> Val a
$c<$ :: forall a b. a -> Val b -> Val a
fmap :: forall a b. (a -> b) -> Val a -> Val b
$cfmap :: forall a b. (a -> b) -> Val a -> Val b
Functor)

-- | Note: The instance uses 'toList', so 'Eq'ual 'OMap's can be shown differently.
instance (Show k, Show v) => Show (OMap k v) where
    showsPrec :: Int -> OMap k v -> ShowS
showsPrec Int
d OMap k v
m = Bool -> ShowS -> ShowS
showParen (Int
d forall a. Ord a => a -> a -> Bool
> Int
10)
        forall a b. (a -> b) -> a -> b
$ String -> ShowS
showString String
"fromList "
        forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Show a => Int -> a -> ShowS
showsPrec Int
d (forall k v. OMap k v -> [(k, v)]
toList OMap k v
m)

-- |
--
-- >>> xs = toAscList $ fromList [('a', "alpha"), ('b', "beta"), ('g', "gamma")]
-- >>> ys = toAscList $ fromList [('g', "gamma"), ('b', "beta"), ('a', "alpha")]
-- >>> xs == ys
-- True
--
-- >>> zs = toAscList $ fromList [('d', "delta"), ('b', "beta"), ('a', "alpha")]
-- >>> xs == zs
-- False
--
instance (Eq k, Eq v) => Eq (OMap k v) where
    OMap k v
xs == :: OMap k v -> OMap k v -> Bool
== OMap k v
ys = forall k v. OMap k v -> [(k, v)]
toAscList OMap k v
xs forall a. Eq a => a -> a -> Bool
== forall k v. OMap k v -> [(k, v)]
toAscList OMap k v
ys

instance (Ord k, Ord v) => Ord (OMap k v) where
    compare :: OMap k v -> OMap k v -> Ordering
compare OMap k v
xs OMap k v
ys = forall a. Ord a => a -> a -> Ordering
compare (forall k v. OMap k v -> [(k, v)]
toAscList OMap k v
xs) (forall k v. OMap k v -> [(k, v)]
toAscList OMap k v
ys)

-------------------------------------------------------------------------------
-- deepseq
-------------------------------------------------------------------------------

instance NFData v => NFData (Val v) where
    rnf :: Val v -> ()
rnf (Val Int
_ v
v) = forall a. NFData a => a -> ()
rnf v
v

instance (NFData k, NFData v) => NFData (OMap k v) where
    rnf :: OMap k v -> ()
rnf (OMap Map k (Val v)
m) = forall a. NFData a => a -> ()
rnf Map k (Val v)
m

-------------------------------------------------------------------------------
-- QuickCheck
-------------------------------------------------------------------------------

instance (Ord k, QC.Arbitrary k, QC.Arbitrary v) => QC.Arbitrary (OMap k v) where
    arbitrary :: Gen (OMap k v)
arbitrary = forall (f :: * -> *) a. (Arbitrary1 f, Arbitrary a) => Gen (f a)
QC.arbitrary1
    shrink :: OMap k v -> [OMap k v]
shrink    = forall (f :: * -> *) a. (Arbitrary1 f, Arbitrary a) => f a -> [f a]
QC.shrink1

instance (Ord k, QC.Arbitrary k) => QC.Arbitrary1 (OMap k) where
    liftArbitrary :: forall a. Gen a -> Gen (OMap k a)
liftArbitrary Gen a
arb = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall k v. Ord k => [(k, v)] -> OMap k v
fromList forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Arbitrary1 f => Gen a -> Gen (f a)
QC.liftArbitrary (forall (f :: * -> *) a. Arbitrary1 f => Gen a -> Gen (f a)
QC.liftArbitrary Gen a
arb)
    liftShrink :: forall a. (a -> [a]) -> OMap k a -> [OMap k a]
liftShrink a -> [a]
shr OMap k a
m  = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall k v. Ord k => [(k, v)] -> OMap k v
fromList forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Arbitrary1 f => (a -> [a]) -> f a -> [f a]
QC.liftShrink (forall (f :: * -> *) a. Arbitrary1 f => (a -> [a]) -> f a -> [f a]
QC.liftShrink a -> [a]
shr) forall a b. (a -> b) -> a -> b
$ forall k v. OMap k v -> [(k, v)]
toList OMap k a
m

-------------------------------------------------------------------------------
-- Combinators
-------------------------------------------------------------------------------

-- |
--
-- >>> empty :: OMap String Integer
-- fromList []
--
empty :: OMap k v
empty :: forall k v. OMap k v
empty = forall k v. Map k (Val v) -> OMap k v
OMap forall k a. Map k a
Map.empty

-- | Elements in key ascending order.
elems :: OMap k v -> [v]
elems :: forall k v. OMap k v -> [v]
elems (OMap Map k (Val v)
m) = forall a b. (a -> b) -> [a] -> [b]
map (forall a b. (a, b) -> b
snd forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall k v. (k, Val v) -> (k, v)
getVal) forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
Map.toAscList Map k (Val v)
m

-- |
--
-- >>> toAscList $ fromList [('a', "alpha"), ('b', "beta"), ('g', "gamma")]
-- [('a',"alpha"),('b',"beta"),('g',"gamma")]
--
-- >>> toAscList $ fromList [('g', "gamma"), ('b', "beta"), ('a', "alpha")]
-- [('a',"alpha"),('b',"beta"),('g',"gamma")]
--
toAscList :: OMap k v -> [(k, v)]
toAscList :: forall k v. OMap k v -> [(k, v)]
toAscList (OMap Map k (Val v)
m) = forall a b. (a -> b) -> [a] -> [b]
map forall k v. (k, Val v) -> (k, v)
getVal forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
Map.toAscList Map k (Val v)
m

-- | /O(n log n)/. List in creation order.
-- Doesn't respect 'Eq' instance.
--
-- >>> toList $ fromList [('a', "alpha"), ('b', "beta"), ('g', "gamma")]
-- [('a',"alpha"),('b',"beta"),('g',"gamma")]
--
-- >>> toList $ fromList [('g', "gamma"), ('b', "beta"), ('a', "alpha")]
-- [('g',"gamma"),('b',"beta"),('a',"alpha")]
--
toList :: OMap k v -> [(k, v)]
toList :: forall k v. OMap k v -> [(k, v)]
toList (OMap Map k (Val v)
m) = forall a b. (a -> b) -> [a] -> [b]
map forall k v. (k, Val v) -> (k, v)
getVal forall a b. (a -> b) -> a -> b
$ forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing forall k v. (k, Val v) -> Int
getIdx) forall a b. (a -> b) -> a -> b
$ forall k a. Map k a -> [(k, a)]
Map.toList Map k (Val v)
m

getIdx :: (k, Val v) -> Int
getIdx :: forall k v. (k, Val v) -> Int
getIdx (k
_, Val Int
i v
_) = Int
i

getVal :: (k, Val v) -> (k, v)
getVal :: forall k v. (k, Val v) -> (k, v)
getVal (k
k, Val Int
_ v
v) = (k
k, v
v)

-- |
--
-- >>> fromList [('g', "gamma"), ('b', "beta"), ('a', "alpha")]
-- fromList [('g',"gamma"),('b',"beta"),('a',"alpha")]
--
fromList :: Ord k => [(k, v)] -> OMap k v
fromList :: forall k v. Ord k => [(k, v)] -> OMap k v
fromList [(k, v)]
kvs = forall k v. Map k (Val v) -> OMap k v
OMap (forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList (forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith forall {a} {v}. Int -> (a, v) -> (a, Val v)
p [Int
0..] [(k, v)]
kvs)) where
    p :: Int -> (a, v) -> (a, Val v)
p Int
i (a
k, v
v) = (a
k, forall v. Int -> v -> Val v
Val Int
i v
v)

-- |
--
-- >>> let xs = fromList [('a', "alpha"), ('b', "beta")]
-- >>> let ys = fromList [('c', 3 :: Int), ('b', 2)]
-- >>> alignWith id xs ys
-- fromList [('a',This "alpha"),('c',That 3),('b',These "beta" 2)]
--
instance Ord k => Semialign (OMap k) where
    alignWith :: forall a b c. (These a b -> c) -> OMap k a -> OMap k b -> OMap k c
alignWith These a b -> c
f (OMap Map k (Val a)
xs) (OMap Map k (Val b)
ys) = forall k v. Map k (Val v) -> OMap k v
OMap (forall (f :: * -> *) a b c.
Semialign f =>
(These a b -> c) -> f a -> f b -> f c
alignWith These (Val a) (Val b) -> Val c
g Map k (Val a)
xs Map k (Val b)
ys) where
        g :: These (Val a) (Val b) -> Val c
g (This (Val Int
i a
x))            = forall v. Int -> v -> Val v
Val Int
i (These a b -> c
f (forall a b. a -> These a b
This a
x))
        g (That (Val Int
j b
y))            = forall v. Int -> v -> Val v
Val Int
j (These a b -> c
f (forall a b. b -> These a b
That b
y))
        g (These (Val Int
i a
x) (Val Int
j b
y)) = forall v. Int -> v -> Val v
Val (forall a. Ord a => a -> a -> a
min Int
i Int
j) (These a b -> c
f (forall a b. a -> b -> These a b
These a
x b
y))