{-# language GADTs #-}
{-# language MagicHash #-}
{-# language LambdaCase #-}
{-# language BangPatterns #-}
{-# language TypeFamilies #-}
{-# language RoleAnnotations #-}
{-# language PatternSynonyms #-}
{-# language FlexibleInstances #-}
{-# language MultiParamTypeClasses #-}
{-# options_ghc -Wno-incomplete-patterns -O0 #-}

--------------------------------------------------------------------------------
-- |
-- Copyright :  (c) Edward Kmett 2017-2018
-- License   :  BSD-2-Clause OR Apache-2.0
-- Maintainer:  Edward Kmett <ekmett@gmail.com>
-- Stability :  experimental
-- Portability: non-portable
--
--------------------------------------------------------------------------------

module Map
  ( Map
  , singleton
  , insert
  , lookup
  , toAscList
  , union
  , split
  , irfoldr
  ) where

import Util.BitQueue
import Control.Lens
import Data.Default
import Data.Hashable
import Data.Function (on)
import Data.Semigroup
import GHC.Exts
import Prelude hiding (lookup)

import Relative
import Delta

type Size = Int

data Map k a
  = Bin {-# unpack #-} !Size {-# unpack #-} !Delta !k !a !(Map k a) !(Map k a)
  | Tip
  deriving Show

type role Map nominal nominal

instance Relative (Map k a) where
  rel _ Tip = Tip
  rel 0 m   = m -- improve sharing
  rel d (Bin s d' k a l r) = Bin s (d+d') k a l r

size :: Map k a -> Int
size (Bin s _ _ _ _ _) = s
size Tip = 0

irfoldr :: (Delta -> k -> a -> r -> r) -> r -> Delta -> Map k a -> r
irfoldr _ z !_ Tip = z
irfoldr f z d (Bin _ d' k a l r) | !d'' <- d <> d' = irfoldr f (f d'' k a (irfoldr f z d'' r)) d'' l

toAscList :: (StrictRelativeOrder k, Relative a) => Map k a -> [(k,a)]
toAscList = irfoldr (\d k x xs -> (rel d k,rel d x):xs) [] 0

instance (StrictRelativeOrder k, Relative a, Eq a) => Eq (Map k a) where (==) = on (==) toAscList
instance (StrictRelativeOrder k, Relative a, Ord a) => Ord (Map k a) where compare = on compare toAscList
instance (StrictRelativeOrder k, RelativeOrder a) => RelativeOrder (Map k a)
instance (StrictRelativeOrder k, StrictRelativeOrder a) => StrictRelativeOrder (Map k a)

--instance (StrictRelativeOrder k, Show k, Relative a, Show a) => Show (Map k a) where
--  showsPrec d m = showParen (d > 10) $ showString "fromList " . showsPrec 11 (toAscList m)

instance (StrictRelativeOrder k, Hashable k, Relative a, Hashable a) => Hashable (Map k a) where
  hashWithSalt d = hashWithSalt d . toAscList

type instance IxValue (Map k a) = a
type instance Index (Map k a) = k

instance (StrictRelativeOrder k, Relative a) => Ixed (Map k a)
instance (StrictRelativeOrder k, Relative a) => At (Map k a) where
  at !k f m = case lookupTrace k m of
    TraceResult mv d q -> f mv <&> \case
      Nothing -> case mv of
        Nothing -> m
        Just old -> deleteAlong old q m
      Just !new | !nd <- negate d -> case mv of
        Nothing -> insertAlong q (rel nd k) (rel nd new) m
        Just _  -> replaceAlong q (rel nd new) m
  {-# inlinable at #-}

singleton :: k -> a -> Map k a
singleton k a = Bin 1 mempty k a Tip Tip
{-# inline singleton #-}

instance Default (Map k a) where
  def = Tip

instance (StrictRelativeOrder k, Relative a) => Semigroup (Map k a) where
  (<>) = union
  stimes = stimesIdempotentMonoid

-- | /O(m*log(n\/m + 1)), m <= n/
instance (StrictRelativeOrder k, Relative a) => Monoid (Map k a) where
  mempty = Tip
  mappend = union

instance (StrictRelativeOrder k, Relative a) => RelativeSemigroup (Map k a)
instance (StrictRelativeOrder k, Relative a) => RelativeMonoid (Map k a)

lookup :: (StrictRelativeOrder k, Relative a) => k -> Map k a -> Maybe a
lookup = go 0 where
  go !_ !_ Tip = Nothing
  go d k (Bin _ d' kx x l r) | !d'' <- d+d' = case compare k (rel d'' kx) of
    LT -> go d'' k l
    GT -> go d'' k r
    EQ -> Just (rel d'' x)
{-# inlinable lookup #-}

union :: (StrictRelativeOrder k, Relative a) => Map k a -> Map k a -> Map k a
union t1 Tip  = t1
union t1 (Bin _ d k x Tip Tip) = insertRD 0 (rel d k) (rel d x) t1
union (Bin _ d k x Tip Tip) t2 = insertD 0 (rel d k) (rel d x) t2
union Tip t2 = t2
union t1@(Bin _ d1 k1 x1 l1 r1) t2 = case split d1 k1 t2 of
  (l2, r2)
    | ptrEq l1l2 l1 && ptrEq r1r2 r1 -> t1
    | otherwise -> link d1 k1 x1 l1l2 r1r2
    where
      !l1l2 = union l1 l2
      !r1r2 = union r1 r2
{-# inlinable union #-}

insert :: StrictRelativeOrder k => k -> a -> Map k a -> Map k a
insert = insertD 0 

insertD :: (Ord k, Relative k) => Delta -> k -> a -> Map k a -> Map k a
insertD !d !kx !x Tip = Bin 1 (negate d) kx x Tip Tip
insertD d kx x t@(Bin sz dy ky y l r) | !d'' <- d+dy = case compare kx (rel d'' ky) of
  LT | ptrEq l' l -> t
     | otherwise -> balanceL dy ky y l' r
     where !l' = insertD d'' kx x l
  GT | ptrEq r' r -> t
     | otherwise -> balanceR dy ky y l r'
     where !r' = insertD d'' kx x r
  EQ -> Bin sz (negate d) kx x (rel d'' l) (rel d'' r)
{-# inlinable insert #-}

insertRD :: (Ord k, Relative k) => Delta -> k -> a -> Map k a -> Map k a
insertRD !d !kx !x Tip = Bin 1 (-d) kx x Tip Tip
insertRD d kx x t@(Bin _ dy ky y l r) | !d'' <- d+dy = case compare kx (rel d'' ky) of
  LT | ptrEq l' l -> t
     | otherwise -> balanceL dy ky y l' r
     where !l' = insertRD d'' kx x l
  GT | ptrEq r' r -> t
     | otherwise -> balanceR dy ky y l r'
     where !r' = insertRD d'' kx x r
  EQ -> t

split :: StrictRelativeOrder k => Delta -> k -> Map k a -> (Map k a, Map k a)
split !d0 !k0 t0 = toPair $ go d0 k0 t0 where
  go :: (Relative k, Ord k) => Delta -> k -> Map k a -> StrictPair (Map k a) (Map k a)
  go d k t = case t of
    Tip -> Tip :*: Tip
    Bin _ dx kx x l r -> case compare (rel (d-dx) k) kx of
      LT -> let lt :*: gt = go (d-dx) k l in lt :*: link dx kx x gt r
      GT -> let lt :*: gt = go (d-dx) k r in link dx kx x l lt :*: gt
      EQ -> rel dx l :*: rel dx r
{-# INLINABLE split #-}

--------------------------------------------------------------------------------
-- Implementation Details
--------------------------------------------------------------------------------

ptrEq :: a -> a -> Bool
ptrEq x y = isTrue# (reallyUnsafePtrEquality# x y)
{-# inline ptrEq #-}


toPair :: StrictPair a b -> (a, b)
toPair (a :*: b) = (a, b)
{-# inline toPair #-}

link :: Delta -> k -> a -> Map k a -> Map k a -> Map k a
link d kx x Tip r  = insertMin d kx x r
link d kx x l Tip  = insertMax d kx x l
link d kx x l@(Bin sizeL dy ky y ly ry) r@(Bin sizeR dz kz z lz rz)
  | delta*sizeL < sizeR  = balanceL (d+dz) kz z (link (negate dz) kx x l (rel dz lz)) rz
  | delta*sizeR < sizeL  = balanceR (d+dy) ky y ly (link (negate dy) kx x (rel dy ry) r)
  | otherwise            = bin d kx x l r

bin :: Delta -> k -> a -> Map k a -> Map k a -> Map k a
bin d kx x l r = Bin (size l + size r + 1) d kx x l r

insertMax,insertMin :: Delta -> k -> a -> Map k a -> Map k a
insertMax d kx x t = case t of
  Tip -> Bin 1 (negate d) kx x Tip Tip
  Bin _ dy ky y l r -> balanceR dy ky y l (insertMax (d+dy) kx x r)

insertMin d kx x t = case t of
  Tip -> Bin 1 (negate d) kx x Tip Tip
  Bin _ dy ky y l r -> balanceL dy ky y (insertMin (d+dy) kx x l) r


data StrictPair a b = !a :*: !b

data TraceResult a = TraceResult !(Maybe a) {-# unpack #-} !Delta {-# unpack #-} !BitQueue

-- Look up a key and return a result indicating whether it was found and what path was taken.
lookupTrace :: (Ord k, Relative k) => k -> Map k a -> TraceResult a
lookupTrace = go mempty emptyQB where
  go :: (Ord k, Relative k) => Delta -> BitQueueB -> k -> Map k a -> TraceResult a
  go !d !q !_ Tip = TraceResult Nothing d (buildQ q)
  go d q k (Bin _ d' kx x l r) | !d'' <- d+d' = case compare k (rel d'' kx) of
    LT -> go d'' (snocQB q False) k l
    GT -> go d'' (snocQB q True) k r
    EQ -> TraceResult (Just x) d'' (buildQ q)
{-# inlinable lookupTrace #-}

insertAlong :: BitQueue -> k -> a -> Map k a -> Map k a
insertAlong !_ kx x Tip = singleton kx x
insertAlong q kx x (Bin _ d' ky y l r) = case unconsQ q of
  Just (False, tl) -> balanceL d' ky y (insertAlong tl kx x l) r
  Just (True,tl) -> balanceR d' ky y l (insertAlong tl kx x r)
  Nothing -> error "Coda.Relative.Map.insertAlong: failure"

deleteAlong :: any -> BitQueue -> Map k a -> Map k a
deleteAlong old !q0 !m = go (bogus old) q0 m where
  go :: Proxy# () -> BitQueue -> Map k a -> Map k a
  go !_ !_ Tip = Tip
  go foom q (Bin _ d' ky y l r) = case unconsQ q of
    Just (False, tl) -> balanceR d' ky y (go foom tl l) r
    Just (True, tl) -> balanceL d' ky y l (go foom tl r)
    Nothing -> glue l r

bogus :: a -> Proxy# ()
bogus _ = proxy#
{-# noinline bogus #-}

replaceAlong :: BitQueue -> a -> Map k a -> Map k a
replaceAlong !_ _ Tip = error "Coda.Relative.Map.replaceAlong: failure"
replaceAlong q x (Bin sz d' ky y l r) = case unconsQ q of
  Just (False, tl) -> Bin sz d' ky y (replaceAlong tl x l) r
  Just (True,tl) -> Bin sz d' ky y l (replaceAlong tl x r)
  Nothing -> Bin sz d' ky x l r

balanceL :: Delta -> k -> a -> Map k a -> Map k a -> Map k a
balanceL d k x l r = case r of
  Tip -> case l of
    Tip -> Bin 1 d k x Tip Tip
    Bin _ _ _ _ Tip Tip -> Bin 2 d k x l Tip
    Bin _ ld lk lx Tip (Bin _ lrd lrk lrx _ _) -> Bin 3 (d+ld+lrd) lrk lrx (Bin 1 (-lrd) lk lx Tip Tip) (Bin 1 (-ld-lrd) k x Tip Tip)
    Bin _ ld lk lx ll@Bin{} Tip -> Bin 3 (d+ld) lk lx ll (Bin 1 (-ld) k x Tip Tip)
    Bin ls ld lk lx ll@(Bin lls _ _ _ _ _) lr@(Bin lrs lrd lrk lrx lrl lrr)
      | lrs < ratio*lls -> Bin (1+ls) (d+ld) lk lx ll (Bin (1+lrs) (-ld) k x (rel ld lr) Tip)
      | otherwise -> Bin (1+ls) (d+ld+lrd) lrk lrx 
        (Bin (1+lls+size lrl) (-lrd) lk lx ll (rel lrd lrl))
        (Bin (1+size lrr) (-lrd-ld) k x (rel (ld+lrd) lrr) Tip)

  Bin rs _ _ _ _ _ -> case l of
    Tip -> Bin (1+rs) d k x Tip r
    Bin ls ld lk lx ll lr
      | ls > delta*rs  -> case (ll, lr) of
        (Bin lls _ _ _ _ _, Bin lrs lrd lrk lrx lrl lrr)
          | lrs < ratio*lls -> Bin (1+ls+rs) (d+ld) lk lx ll (Bin (1+rs+lrs) (-ld) k x (rel ld lr) r)
          | otherwise -> Bin (1+ls+rs) (d+ld+lrd) lrk lrx 
            (Bin (1+lls+size lrl) (-lrd) lk lx ll (rel lrd lrl))
            (Bin (1+rs+size lrr) (-lrd-ld) k x (rel (lrd+ld) lrr) r)
        (_, _) -> error "Coda.Relative.Map.balanceL: failure"
      | otherwise -> Bin (1+ls+rs) d k x l r
{-# noinline balanceL #-}

balanceR :: Delta -> k -> a -> Map k a -> Map k a -> Map k a
balanceR d k x l r = case l of
  Tip -> case r of
    Tip -> Bin 1 d k x Tip Tip
    Bin _ _ _ _ Tip Tip -> Bin 2 d k x Tip r
    Bin _ rd rk rx Tip rr@Bin{} -> Bin 3 (d+rd) rk rx (Bin 1 (-rd) k x Tip Tip) rr
    Bin _ rd rk rx (Bin _ rld rlk rlx _ _) Tip -> Bin 3 (d+rd+rld) rlk rlx (Bin 1 (-rld-rd) k x Tip Tip) (Bin 1 (-rld) rk rx Tip Tip)
    Bin rs rd rk rx rl@(Bin rls rld rlk rlx rll rlr) rr@(Bin rrs _ _ _ _ _)
      | rls < ratio*rrs -> Bin (1+rs) (d+rd) rk rx (Bin (1+rls) (-rd) k x Tip (rel rd rl)) rr
      | otherwise -> Bin (1+rs) (d+rd+rld) rlk rlx
        (Bin (1+size rll) (-rld-rd) k x Tip (rel (rd+rld) rll))
        (Bin (1+rrs+size rlr) (-rld) rk rx (rel rld rlr) rr)

  Bin ls _ _ _ _ _ -> case r of
    Tip -> Bin (1+ls) d k x l Tip
    Bin rs rd rk rx rl rr
      | rs > delta*ls  -> case (rl, rr) of
        (Bin rls rld rlk rlx rll rlr, Bin rrs _ _ _ _ _)
          | rls < ratio*rrs -> Bin (1+ls+rs) (d+rd) rk rx (Bin (1+ls+rls) (-rd) k x l (rel rd rl)) rr
          | otherwise -> Bin (1+ls+rs) (d+rd+rld) rlk rlx
            (Bin (1+ls+size rll) (-rd-rld) k x l (rel (rd+rld) rll))
            (Bin (1+rrs+size rlr) (-rld) rk rx (rel rld rlr) rr)
        (_, _) -> error "Coda.Relative.Map.balanceR: failure"
      | otherwise -> Bin (1+ls+rs) d k x l r
{-# NOINLINE balanceR #-}

data MinView k a = MinView !Delta !k !a !(Map k a)
data MaxView k a = MaxView !Delta !k !a !(Map k a)

glue :: Map k a -> Map k a -> Map k a
glue Tip r = r
glue l Tip = l
glue (Bin sl dl kl xl ll lr) (Bin sr dr kr xr rl rr)
  | sl > sr   = let !(MaxView dm km m l') = maxViewSure dl kl xl ll lr in balanceR dm km m l' (Bin sr (dr-dm) kr xr rl rr)
  | otherwise = let !(MinView dm km m r') = minViewSure dr kr xr rl rr in balanceL dm km m (Bin sl (dl-dm) kl xl ll lr) r'

minViewSure :: Delta -> k -> a -> Map k a -> Map k a -> MinView k a
minViewSure = go where
  go !d k x Tip r = MinView d k x r
  go d k x (Bin _ dl kl xl ll lr) r = case go (d+dl) kl xl ll lr of
    MinView dm km xm l' -> MinView dm km xm (balanceR (-dm) k x (rel dm l') r)

maxViewSure :: Delta -> k -> a -> Map k a -> Map k a -> MaxView k a
maxViewSure = go where
  go !d k x l Tip = MaxView d k x l
  go d k x l (Bin _ dr kr xr rl rr) = case go (d+dr) kr xr rl rr of
    MaxView dm km xm r' -> MaxView dm km xm (balanceL (-dm) k x l (rel dm r'))

delta,ratio :: Int
delta = 3
ratio = 2