{-# LANGUAGE Unsafe #-}
{-# LANGUAGE BangPatterns #-}
{-# OPTIONS_HADDOCK not-home #-}
-----------------------------------------------------------------------------
-- |
-- Copyright   :  (C) 2015 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  experimental
-- Portability :  non-portable
--
-----------------------------------------------------------------------------
module Data.Struct.Internal.Order where

import Control.Exception
import Control.Monad
import Control.Monad.Primitive
import Control.Monad.ST
import Data.Bits
import Data.Struct.Internal
import Data.Struct.Internal.Label (Label, Key)
import qualified Data.Struct.Label as Label
import qualified Data.Struct.Internal.Label as Label (key)
import Data.Word

--------------------------------------------------------------------------------
-- * Order Maintenance
--------------------------------------------------------------------------------

-- | This structure maintains an order-maintenance structure as two levels of list-labeling.
--
-- The upper labeling scheme holds @(n / log w)@ elements in a universe of size @w^2@, operating in O(log n) amortized time per operation.
--
-- It is accelerated by an indirection structure where each smaller universe holds O(log w) elements, with total label space @2^log w = w@ and O(1) expected update cost, so we
-- can charge rebuilds to the upper structure to the lower structure.
--
-- Every insert to the upper structure is amortized across @O(log w)@ operations below.
--
-- This means that inserts are O(1) amortized, while comparisons remain O(1) worst-case.

newtype Order a s = Order { forall a s. Order a s -> Object s
runOrder :: Object s }

instance Eq (Order a s) where == :: Order a s -> Order a s -> Bool
(==) = forall (t :: * -> *) s. Struct t => t s -> t s -> Bool
eqStruct

instance Struct (Order a)

key :: Field (Order a) Key
key :: forall a. Field (Order a) Key
key = forall {k} (s :: k) a. Int -> Field s a
field Int
0
{-# INLINE key #-}

value :: Field (Order a) a
value :: forall a. Field (Order a) a
value = forall {k} (s :: k) a. Int -> Field s a
field Int
1
{-# INLINE value #-}

next :: Slot (Order a) (Order a)
next :: forall a. Slot (Order a) (Order a)
next = forall {k1} {k2} (s :: k1) (t :: k2). Int -> Slot s t
slot Int
2
{-# INLINE next #-}

prev :: Slot (Order a) (Order a)
prev :: forall a. Slot (Order a) (Order a)
prev = forall {k1} {k2} (s :: k1) (t :: k2). Int -> Slot s t
slot Int
3
{-# INLINE prev #-}

parent :: Slot (Order a) Label
parent :: forall a. Slot (Order a) Label
parent = forall {k1} {k2} (s :: k1) (t :: k2). Int -> Slot s t
slot Int
4
{-# INLINE parent #-}

makeOrder :: PrimMonad m => Label (PrimState m) -> Key -> a -> Order a (PrimState m) -> Order a (PrimState m) -> m (Order a (PrimState m))
makeOrder :: forall (m :: * -> *) a.
PrimMonad m =>
Label (PrimState m)
-> Key
-> a
-> Order a (PrimState m)
-> Order a (PrimState m)
-> m (Order a (PrimState m))
makeOrder Label (PrimState m)
mom Key
a a
v Order a (PrimState m)
p Order a (PrimState m)
n = forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
st forall a b. (a -> b) -> a -> b
$ do
  Order a (PrimState m)
this <- forall (m :: * -> *) (t :: * -> *).
(PrimMonad m, Struct t) =>
Int -> m (t (PrimState m))
alloc Int
5
  forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> y (PrimState m) -> m ()
set forall a. Slot (Order a) Label
parent Order a (PrimState m)
this Label (PrimState m)
mom
  forall (m :: * -> *) (x :: * -> *) a.
(PrimMonad m, Struct x) =>
Field x a -> x (PrimState m) -> a -> m ()
setField forall a. Field (Order a) Key
key Order a (PrimState m)
this Key
a
  forall (m :: * -> *) (x :: * -> *) a.
(PrimMonad m, Struct x) =>
Field x a -> x (PrimState m) -> a -> m ()
setField forall a. Field (Order a) a
value Order a (PrimState m)
this a
v
  forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> y (PrimState m) -> m ()
set forall a. Slot (Order a) (Order a)
prev Order a (PrimState m)
this Order a (PrimState m)
p
  forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> y (PrimState m) -> m ()
set forall a. Slot (Order a) (Order a)
next Order a (PrimState m)
this Order a (PrimState m)
n
  forall (m :: * -> *) a. Monad m => a -> m a
return Order a (PrimState m)
this
{-# INLINE makeOrder #-}

-- | O(1) compareM, O(1) amortized insert
compareM :: PrimMonad m => Order a (PrimState m) -> Order a (PrimState m) -> m Ordering
compareM :: forall (m :: * -> *) a.
PrimMonad m =>
Order a (PrimState m) -> Order a (PrimState m) -> m Ordering
compareM Order a (PrimState m)
i Order a (PrimState m)
j
  | forall (t :: * -> *) s. Struct t => t s -> Bool
isNil Order a (PrimState m)
i Bool -> Bool -> Bool
|| forall (t :: * -> *) s. Struct t => t s -> Bool
isNil Order a (PrimState m)
j = forall a e. Exception e => e -> a
throw NullPointerException
NullPointerException
  | Bool
otherwise = forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
st forall a b. (a -> b) -> a -> b
$ do
  Label (PrimState m)
ui <- forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> m (y (PrimState m))
get forall a. Slot (Order a) Label
parent Order a (PrimState m)
i
  Label (PrimState m)
uj <- forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> m (y (PrimState m))
get forall a. Slot (Order a) Label
parent Order a (PrimState m)
j
  Ordering
xs <- forall (m :: * -> *).
PrimMonad m =>
Label (PrimState m) -> Label (PrimState m) -> m Ordering
Label.compareM Label (PrimState m)
ui Label (PrimState m)
uj
  case Ordering
xs of
    Ordering
EQ -> forall a. Ord a => a -> a -> Ordering
compare forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) (x :: * -> *) a.
(PrimMonad m, Struct x) =>
Field x a -> x (PrimState m) -> m a
getField forall a. Field (Order a) Key
key Order a (PrimState m)
i forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (m :: * -> *) (x :: * -> *) a.
(PrimMonad m, Struct x) =>
Field x a -> x (PrimState m) -> m a
getField forall a. Field (Order a) Key
key Order a (PrimState m)
j
    Ordering
x -> forall (m :: * -> *) a. Monad m => a -> m a
return Ordering
x
{-# INLINE compareM #-}

insertAfter :: PrimMonad m => Order a (PrimState m) -> a -> m (Order a (PrimState m))
insertAfter :: forall (m :: * -> *) a.
PrimMonad m =>
Order a (PrimState m) -> a -> m (Order a (PrimState m))
insertAfter Order a (PrimState m)
n0 a
a1 = forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
st forall a b. (a -> b) -> a -> b
$ do
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall (t :: * -> *) s. Struct t => t s -> Bool
isNil Order a (PrimState m)
n0) forall a b. (a -> b) -> a -> b
$ forall a e. Exception e => e -> a
throw NullPointerException
NullPointerException
  Label (PrimState m)
mom <- forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> m (y (PrimState m))
get forall a. Slot (Order a) Label
parent Order a (PrimState m)
n0
  Key
k0 <- forall (m :: * -> *) (x :: * -> *) a.
(PrimMonad m, Struct x) =>
Field x a -> x (PrimState m) -> m a
getField forall a. Field (Order a) Key
key Order a (PrimState m)
n0
  Order a (PrimState m)
n2 <- forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> m (y (PrimState m))
get forall a. Slot (Order a) (Order a)
next Order a (PrimState m)
n0
  Key
k2 <- if forall (t :: * -> *) s. Struct t => t s -> Bool
isNil Order a (PrimState m)
n2 then forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Bounded a => a
maxBound else forall (m :: * -> *) (x :: * -> *) a.
(PrimMonad m, Struct x) =>
Field x a -> x (PrimState m) -> m a
getField forall a. Field (Order a) Key
key Order a (PrimState m)
n2
  let !k1 :: Key
k1 = Key
k0 forall a. Num a => a -> a -> a
+ forall a. Bits a => a -> Int -> a
unsafeShiftR (Key
k2 forall a. Num a => a -> a -> a
- Key
k0) Int
1
  Order a (PrimState m)
n1 <- forall (m :: * -> *) a.
PrimMonad m =>
Label (PrimState m)
-> Key
-> a
-> Order a (PrimState m)
-> Order a (PrimState m)
-> m (Order a (PrimState m))
makeOrder Label (PrimState m)
mom Key
k1 a
a1 Order a (PrimState m)
n0 Order a (PrimState m)
n2
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (forall (t :: * -> *) s. Struct t => t s -> Bool
isNil Order a (PrimState m)
n2) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> y (PrimState m) -> m ()
set forall a. Slot (Order a) (Order a)
prev Order a (PrimState m)
n2 Order a (PrimState m)
n1
  forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> y (PrimState m) -> m ()
set forall a. Slot (Order a) (Order a)
next Order a (PrimState m)
n0 Order a (PrimState m)
n1
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Key
k0 forall a. Num a => a -> a -> a
+ Key
1 forall a. Eq a => a -> a -> Bool
== Key
k2) forall a b. (a -> b) -> a -> b
$ forall s a. Label s -> Order a s -> ST s ()
rewind Label (PrimState m)
mom Order a (PrimState m)
n0 -- we have a collision, rebalance
  forall (m :: * -> *) a. Monad m => a -> m a
return Order a (PrimState m)
n1
 where
  -- find the smallest sibling
  rewind :: Label s -> Order a s -> ST s ()
  rewind :: forall s a. Label s -> Order a s -> ST s ()
rewind Label s
mom Order a s
this = do
    Order a s
p <- forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> m (y (PrimState m))
get forall a. Slot (Order a) (Order a)
prev Order a s
this
    if forall (t :: * -> *) s. Struct t => t s -> Bool
isNil Order a s
p then forall s a.
Label s -> Label s -> Order a s -> Key -> Int -> ST s ()
rebalance Label s
mom Label s
mom Order a s
this Key
0 Int
64
    else do
      Label s
dad <- forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> m (y (PrimState m))
get forall a. Slot (Order a) Label
parent Order a s
p
      if Label s
mom forall a. Eq a => a -> a -> Bool
== Label s
dad then forall s a. Label s -> Order a s -> ST s ()
rewind Label s
mom Order a s
p
      else forall s a.
Label s -> Label s -> Order a s -> Key -> Int -> ST s ()
rebalance Label s
mom Label s
mom Order a s
p Key
0 Int
64

  -- break up the family
  rebalance :: Label s -> Label s -> Order a s -> Word64 -> Int -> ST s ()
  rebalance :: forall s a.
Label s -> Label s -> Order a s -> Key -> Int -> ST s ()
rebalance Label s
mom Label s
dad Order a s
this Key
k Int
j = forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (forall (t :: * -> *) s. Struct t => t s -> Bool
isNil Order a s
this) forall a b. (a -> b) -> a -> b
$ do
    Label s
guardian <- forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> m (y (PrimState m))
get forall a. Slot (Order a) Label
parent Order a s
this
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Label s
mom forall a. Eq a => a -> a -> Bool
== Label s
guardian) forall a b. (a -> b) -> a -> b
$ do
      forall (m :: * -> *) (x :: * -> *) a.
(PrimMonad m, Struct x) =>
Field x a -> x (PrimState m) -> a -> m ()
setField forall a. Field (Order a) Key
key Order a s
this Key
k
      forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> y (PrimState m) -> m ()
set forall a. Slot (Order a) Label
parent Order a s
this Label s
dad
      Order a s
n <- forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> m (y (PrimState m))
get forall a. Slot (Order a) (Order a)
next Order a s
this
      if Int
j forall a. Ord a => a -> a -> Bool
> Int
0 then forall s a.
Label s -> Label s -> Order a s -> Key -> Int -> ST s ()
rebalance Label s
mom Label s
dad Order a s
n (Key
k forall a. Num a => a -> a -> a
+ Key
deltaU) (Int
jforall a. Num a => a -> a -> a
-Int
1)
      else do
        Label s
stepdad <- forall (m :: * -> *).
PrimMonad m =>
Label (PrimState m) -> m (Label (PrimState m))
Label.insertAfter Label s
dad
        forall s a.
Label s -> Label s -> Order a s -> Key -> Int -> ST s ()
rebalance Label s
mom Label s
stepdad Order a s
n Key
deltaU Int
logU

delete :: PrimMonad m => Order a (PrimState m) -> m ()
delete :: forall (m :: * -> *) a.
PrimMonad m =>
Order a (PrimState m) -> m ()
delete Order a (PrimState m)
this = forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
st forall a b. (a -> b) -> a -> b
$ do
  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall (t :: * -> *) s. Struct t => t s -> Bool
isNil Order a (PrimState m)
this) forall a b. (a -> b) -> a -> b
$ forall a e. Exception e => e -> a
throw NullPointerException
NullPointerException
  Label (PrimState m)
mom <- forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> m (y (PrimState m))
get forall a. Slot (Order a) Label
parent Order a (PrimState m)
this

  Order a (PrimState m)
p <- forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> m (y (PrimState m))
get forall a. Slot (Order a) (Order a)
prev Order a (PrimState m)
this
  Order a (PrimState m)
n <- forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> m (y (PrimState m))
get forall a. Slot (Order a) (Order a)
next Order a (PrimState m)
this

  forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> y (PrimState m) -> m ()
set forall a. Slot (Order a) (Order a)
prev Order a (PrimState m)
this forall (t :: * -> *) s. Struct t => t s
Nil
  forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> y (PrimState m) -> m ()
set forall a. Slot (Order a) (Order a)
next Order a (PrimState m)
this forall (t :: * -> *) s. Struct t => t s
Nil

  Bool
x <- if forall (t :: * -> *) s. Struct t => t s -> Bool
isNil Order a (PrimState m)
p then forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
  else do
    forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> y (PrimState m) -> m ()
set forall a. Slot (Order a) (Order a)
next Order a (PrimState m)
p Order a (PrimState m)
n
    Label (PrimState m)
pmom <- forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> m (y (PrimState m))
get forall a. Slot (Order a) Label
parent Order a (PrimState m)
p
    forall (m :: * -> *) a. Monad m => a -> m a
return (Label (PrimState m)
mom forall a. Eq a => a -> a -> Bool
== Label (PrimState m)
pmom)

  Bool
y <- if forall (t :: * -> *) s. Struct t => t s -> Bool
isNil Order a (PrimState m)
n then forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
  else do
    forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> y (PrimState m) -> m ()
set forall a. Slot (Order a) (Order a)
prev Order a (PrimState m)
n Order a (PrimState m)
p
    Label (PrimState m)
nmom <- forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> m (y (PrimState m))
get forall a. Slot (Order a) Label
parent Order a (PrimState m)
n
    forall (m :: * -> *) a. Monad m => a -> m a
return (Label (PrimState m)
mom forall a. Eq a => a -> a -> Bool
== Label (PrimState m)
nmom)

  forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Bool
x Bool -> Bool -> Bool
|| Bool
y) forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). PrimMonad m => Label (PrimState m) -> m ()
Label.delete Label (PrimState m)
mom
{-# INLINE delete #-}

logU :: Int
logU :: Int
logU = Int
64

loglogU :: Int
loglogU :: Int
loglogU = Int
6

deltaU :: Key
deltaU :: Key
deltaU = forall a. Bits a => a -> Int -> a
unsafeShiftR forall a. Bounded a => a
maxBound Int
loglogU -- U / log U

new :: PrimMonad m => a -> m (Order a (PrimState m))
new :: forall (m :: * -> *) a.
PrimMonad m =>
a -> m (Order a (PrimState m))
new a
a = forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
st forall a b. (a -> b) -> a -> b
$ do
  Label (PrimState m)
l <- forall (m :: * -> *). PrimMonad m => m (Label (PrimState m))
Label.new
  forall (m :: * -> *) a.
PrimMonad m =>
Label (PrimState m)
-> Key
-> a
-> Order a (PrimState m)
-> Order a (PrimState m)
-> m (Order a (PrimState m))
makeOrder Label (PrimState m)
l (forall a. Bits a => a -> Int -> a
unsafeShiftR forall a. Bounded a => a
maxBound Int
1) a
a forall (t :: * -> *) s. Struct t => t s
Nil forall (t :: * -> *) s. Struct t => t s
Nil
{-# INLINE new #-}

keys :: PrimMonad m => Order a (PrimState m) -> m (Key, Key)
keys :: forall (m :: * -> *) a.
PrimMonad m =>
Order a (PrimState m) -> m (Key, Key)
keys Order a (PrimState m)
this = forall (m :: * -> *) a. PrimMonad m => ST (PrimState m) a -> m a
st forall a b. (a -> b) -> a -> b
$ do
  Label (PrimState m)
mom <- forall (m :: * -> *) (x :: * -> *) (y :: * -> *).
(PrimMonad m, Struct x, Struct y) =>
Slot x y -> x (PrimState m) -> m (y (PrimState m))
get forall a. Slot (Order a) Label
parent Order a (PrimState m)
this
  (,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) (x :: * -> *) a.
(PrimMonad m, Struct x) =>
Field x a -> x (PrimState m) -> m a
getField Field Label Key
Label.key Label (PrimState m)
mom forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall (m :: * -> *) (x :: * -> *) a.
(PrimMonad m, Struct x) =>
Field x a -> x (PrimState m) -> m a
getField forall a. Field (Order a) Key
key Order a (PrimState m)
this
{-# INLINE keys #-}