{-# 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 (a -> OMap k b -> OMap k a
(a -> b) -> OMap k a -> OMap k b
(forall a b. (a -> b) -> OMap k a -> OMap k b)
-> (forall a b. a -> OMap k b -> OMap k a) -> Functor (OMap k)
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
<$ :: a -> OMap k b -> OMap k a
$c<$ :: forall k a b. a -> OMap k b -> OMap k a
fmap :: (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 (a -> Val b -> Val a
(a -> b) -> Val a -> Val b
(forall a b. (a -> b) -> Val a -> Val b)
-> (forall a b. a -> Val b -> Val a) -> Functor Val
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
<$ :: a -> Val b -> Val a
$c<$ :: forall a b. a -> Val b -> Val a
fmap :: (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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10)
        (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$ String -> ShowS
showString String
"fromList "
        ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [(k, v)] -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
d (OMap k v -> [(k, v)]
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 = [(k, v)] -> [(k, v)] -> Bool
forall a a. (Eq a, Eq a) => [(a, a)] -> [(a, a)] -> Bool
go (OMap k v -> [(k, v)]
forall k v. OMap k v -> [(k, v)]
toAscList OMap k v
xs) (OMap k v -> [(k, v)]
forall k v. OMap k v -> [(k, v)]
toAscList OMap k v
ys) where
        go :: [(a, a)] -> [(a, a)] -> Bool
go [] [] = Bool
True
        go [(a, a)]
_  [] = Bool
False
        go [] [(a, a)]
_  = Bool
False
        go ((a
k1, a
v1) : [(a, a)]
kvs1) ((a
k2, a
v2) : [(a, a)]
kvs2) =
            a
k1 a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
k2 Bool -> Bool -> Bool
&& a
v1 a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
v2 Bool -> Bool -> Bool
&& [(a, a)] -> [(a, a)] -> Bool
go [(a, a)]
kvs1 [(a, a)]
kvs2

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

instance NFData v => NFData (Val v) where
    rnf :: Val v -> ()
rnf (Val Int
_ v
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) = Map k (Val v) -> ()
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 = Gen (OMap k v)
forall (f :: * -> *) a. (Arbitrary1 f, Arbitrary a) => Gen (f a)
QC.arbitrary1
    shrink :: OMap k v -> [OMap k v]
shrink    = OMap k v -> [OMap k v]
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 :: Gen a -> Gen (OMap k a)
liftArbitrary Gen a
arb = ([(k, a)] -> OMap k a) -> Gen [(k, a)] -> Gen (OMap k a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(k, a)] -> OMap k a
forall k v. Ord k => [(k, v)] -> OMap k v
fromList (Gen [(k, a)] -> Gen (OMap k a)) -> Gen [(k, a)] -> Gen (OMap k a)
forall a b. (a -> b) -> a -> b
$ Gen (k, a) -> Gen [(k, a)]
forall (f :: * -> *) a. Arbitrary1 f => Gen a -> Gen (f a)
QC.liftArbitrary (Gen a -> Gen (k, a)
forall (f :: * -> *) a. Arbitrary1 f => Gen a -> Gen (f a)
QC.liftArbitrary Gen a
arb)
    liftShrink :: (a -> [a]) -> OMap k a -> [OMap k a]
liftShrink a -> [a]
shr OMap k a
m  = ([(k, a)] -> OMap k a) -> [[(k, a)]] -> [OMap k a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [(k, a)] -> OMap k a
forall k v. Ord k => [(k, v)] -> OMap k v
fromList ([[(k, a)]] -> [OMap k a]) -> [[(k, a)]] -> [OMap k a]
forall a b. (a -> b) -> a -> b
$ ((k, a) -> [(k, a)]) -> [(k, a)] -> [[(k, a)]]
forall (f :: * -> *) a. Arbitrary1 f => (a -> [a]) -> f a -> [f a]
QC.liftShrink ((a -> [a]) -> (k, a) -> [(k, a)]
forall (f :: * -> *) a. Arbitrary1 f => (a -> [a]) -> f a -> [f a]
QC.liftShrink a -> [a]
shr) ([(k, a)] -> [[(k, a)]]) -> [(k, a)] -> [[(k, a)]]
forall a b. (a -> b) -> a -> b
$ OMap k a -> [(k, a)]
forall k v. OMap k v -> [(k, v)]
toList OMap k a
m

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

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

-- | Elements in key ascending order.
elems :: OMap k v -> [v]
elems :: OMap k v -> [v]
elems (OMap Map k (Val v)
m) = ((k, Val v) -> v) -> [(k, Val v)] -> [v]
forall a b. (a -> b) -> [a] -> [b]
map ((k, v) -> v
forall a b. (a, b) -> b
snd ((k, v) -> v) -> ((k, Val v) -> (k, v)) -> (k, Val v) -> v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (k, Val v) -> (k, v)
forall k v. (k, Val v) -> (k, v)
getVal) ([(k, Val v)] -> [v]) -> [(k, Val v)] -> [v]
forall a b. (a -> b) -> a -> b
$ Map k (Val v) -> [(k, Val v)]
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 :: OMap k v -> [(k, v)]
toAscList (OMap Map k (Val v)
m) = ((k, Val v) -> (k, v)) -> [(k, Val v)] -> [(k, v)]
forall a b. (a -> b) -> [a] -> [b]
map (k, Val v) -> (k, v)
forall k v. (k, Val v) -> (k, v)
getVal ([(k, Val v)] -> [(k, v)]) -> [(k, Val v)] -> [(k, v)]
forall a b. (a -> b) -> a -> b
$ Map k (Val v) -> [(k, Val v)]
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 :: OMap k v -> [(k, v)]
toList (OMap Map k (Val v)
m) = ((k, Val v) -> (k, v)) -> [(k, Val v)] -> [(k, v)]
forall a b. (a -> b) -> [a] -> [b]
map (k, Val v) -> (k, v)
forall k v. (k, Val v) -> (k, v)
getVal ([(k, Val v)] -> [(k, v)]) -> [(k, Val v)] -> [(k, v)]
forall a b. (a -> b) -> a -> b
$ ((k, Val v) -> (k, Val v) -> Ordering)
-> [(k, Val v)] -> [(k, Val v)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (((k, Val v) -> Int) -> (k, Val v) -> (k, Val v) -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (k, Val v) -> Int
forall k v. (k, Val v) -> Int
getIdx) ([(k, Val v)] -> [(k, Val v)]) -> [(k, Val v)] -> [(k, Val v)]
forall a b. (a -> b) -> a -> b
$ Map k (Val v) -> [(k, Val v)]
forall k a. Map k a -> [(k, a)]
Map.toList Map k (Val v)
m

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

getVal :: (k, Val v) -> (k, v)
getVal :: (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 :: [(k, v)] -> OMap k v
fromList [(k, v)]
kvs = Map k (Val v) -> OMap k v
forall k v. Map k (Val v) -> OMap k v
OMap ([(k, Val v)] -> Map k (Val v)
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ((Int -> (k, v) -> (k, Val v)) -> [Int] -> [(k, v)] -> [(k, Val v)]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> (k, v) -> (k, Val v)
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, Int -> v -> Val v
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 :: (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) = Map k (Val c) -> OMap k c
forall k v. Map k (Val v) -> OMap k v
OMap ((These (Val a) (Val b) -> Val c)
-> Map k (Val a) -> Map k (Val b) -> Map k (Val c)
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))            = Int -> c -> Val c
forall v. Int -> v -> Val v
Val Int
i (These a b -> c
f (a -> These a b
forall a b. a -> These a b
This a
x))
        g (That (Val Int
j b
y))            = Int -> c -> Val c
forall v. Int -> v -> Val v
Val Int
j (These a b -> c
f (b -> These a b
forall a b. b -> These a b
That b
y))
        g (These (Val Int
i a
x) (Val Int
j b
y)) = Int -> c -> Val c
forall v. Int -> v -> Val v
Val (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
i Int
j) (These a b -> c
f (a -> b -> These a b
forall a b. a -> b -> These a b
These a
x b
y))