Pure.hs

{- | Pure efficient implementation of the implicit treap data structure with the
segment tree interface.

__NOTE:__ Letter \( d \) in the documentation below means depth of the tree. Real
depth depends on the strategy for creating 'Priority'. If the strategy is poor,
the depth can be linear. However, if priorities are generated randomly, expected
depth is \( O(\log \ n) \).
-}

{-# LANGUAGE DeriveAnyClass    #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies      #-}

module Treap.Pure
       ( -- * Data structure
         Size (..)
       , Priority (..)
       , Treap (..)

         -- * Smart constructors
       , empty
       , one

         -- * Query functions
       , size
       , sizeInt
       , monoid
       , at
       , query

         -- * Cuts and joins
       , splitAt
       , merge
       , take
       , drop
       , rotate

         -- * Modification functions
       , insert
       , delete

         -- * Core internal functions
       , recalculate
       ) where

import Prelude hiding (drop, lookup, splitAt, take)

import Control.DeepSeq (NFData)
import Data.Foldable (foldl')
import Data.Word (Word64)
import GHC.Exts (IsList (..))
import GHC.Generics (Generic)

import Treap.Measured (Measured (..))

-- $setup
-- >>> import Data.Monoid
-- >>> import Treap.Pretty

----------------------------------------------------------------------------
-- Data structure and instances
----------------------------------------------------------------------------

{- | Size of the 'Treap' data structure. Guaranteed to be always non-negative.
-}
newtype Size = Size
    { unSize :: Int
    } deriving stock (Show, Read, Generic)
      deriving newtype (Eq, Ord, Num, NFData)

{- | Priority in the 'Treap' data structure.
-}
newtype Priority = Priority
    { unPriority :: Word64
    } deriving stock (Show, Read, Generic)
      deriving newtype (Eq, Ord, NFData)

-- | 'Treap' data structure.
data Treap m a
    = Node !Size !Priority !m a !(Treap m a) !(Treap m a)
    | Empty
    deriving stock (Show, Read, Eq, Generic, Foldable)
    deriving anyclass (NFData)

{- | (<>) is implemented via 'merge'.
-}
instance Measured m a => Semigroup (Treap m a) where
    (<>) = merge

{- | mempty is implemented via 'empty'.
-}
instance Measured m a => Monoid (Treap m a) where
    mempty = empty

-- | \( O(1) \). Takes cached value from the root.
instance Monoid m => Measured m (Treap m a) where
    measure = monoid
    {-# INLINE measure #-}

-- #if __GLASGOW_HASKELL__ >= 806
-- -- | Safe 'Functor' instance that performs recalculation of monoidal accumulator.
-- instance (forall b . Measured m b) => Functor (Treap m) where
--     fmap :: forall a b . (a -> b) -> Treap m a -> Treap m b
--     fmap f = go
--       where
--         go :: Treap m a -> Treap m b
--         go Empty = Empty
--         go (Node _ p _ a l r) = recalculate $ new p (f a) (go l) (go r)
-- #endif

{- | This instance allows to create 'Treap' from the list of triples. If all
priorities are random then the expected performance of the @fromList@ function
is \( O(n\ \log \ n)\).

__TODO:__ It's possible to implement \( O(n) \) algorithm however.
See issue #15: </~https://github.com/chshersh/treap/issues/15>
-}
instance Measured m a => IsList (Treap m a) where
    type Item (Treap m a) = (Priority, a)

    -- TODO: implement O(n) algorithm
    fromList :: [(Priority, a)] -> Treap m a
    fromList =
        foldl' (\t (i, p, a) -> insert i p a t) Empty
        . zipWith (\i (p, a) -> (i, p, a)) [0..]
    {-# INLINE fromList #-}

    -- TODO: make more efficient
    toList :: Treap m a -> [(Priority, a)]
    toList Empty              = []
    toList (Node _ p _ a l r) = toList l ++ (p, a) : toList r

----------------------------------------------------------------------------
-- Smart constructors
----------------------------------------------------------------------------

-- | \( O(1) \). Creates empty 'Treap'.
empty :: Treap m a
empty = Empty
{-# INLINE empty #-}

-- | \( O(1) \). Creates singleton 'Treap'.
one :: Measured m a => Priority -> a -> Treap m a
one p a = Node (Size 1) p (measure a) a Empty Empty
{-# INLINE one #-}

----------------------------------------------------------------------------
-- Query functions
----------------------------------------------------------------------------

{- | \( O(1) \). Returns the number of the elements in the 'Treap'. Always
non-negative.

__Properties:__

* \( \forall (t\ ::\ \mathrm{Treap}\ m\ a)\ .\ \mathrm{size}\ t \geqslant 0 \)
-}
size :: Treap m a -> Size
size = \case
    Empty -> Size 0
    Node s _ _ _ _ _ -> s
{-# INLINE size #-}

-- | Take size of the 'Treap' as 'Int'.
sizeInt :: Treap m a -> Int
sizeInt = unSize . size
{-# INLINE sizeInt #-}

-- | \( O(1) \). Returns accumulated value in the root of the tree.
monoid :: Monoid m => Treap m a -> m
monoid = \case
    Empty -> mempty
    Node _ _ m _ _ _ -> m
{-# INLINE monoid #-}

-- | \( O(d) \). Lookup a value inside 'Treap' by a given index.
at :: forall m a . Int -> Treap m a -> Maybe a
at i t
    | i < 0          = Nothing
    | i >= sizeInt t = Nothing
    | otherwise      = go i t
  where
    go :: Int -> Treap m a -> Maybe a
    go _ Empty = Nothing
    go k (Node _ _ _ a l r) =
        let lSize = sizeInt l
        in case compare k lSize of
            EQ -> Just a
            LT -> go k l
            GT -> go (k - lSize - 1) r

-- | \( O(d) \). Return value of monoidal accumulator on a segment @[l, r)@.
query :: forall m a . Measured m a => Int -> Int -> Treap m a -> m
query from to t
    | to <= from = mempty
    | otherwise  =
        let (l, _) = splitAt to t
            (_, m) = splitAt from l
        in monoid m

----------------------------------------------------------------------------
-- Cuts and joins
----------------------------------------------------------------------------

-- | Create new 'Node' and recalculate its values.
new :: Measured m a => Priority -> a -> Treap m a -> Treap m a -> Treap m a
new p a l r = recalculate $ Node 0 p mempty a l r
{-# INLINE new #-}

{- | \( O(d) \). @splitAt i t@ splits 'Treap' by the given index into two treaps
@(t1, t2)@ such that the following properties hold:

1. \( \mathrm{size}\ t_1 = i \)
2. \( \mathrm{size}\ t_2 = n - i \)
3. \( \mathrm{merge}\ t_1\ t_2 \equiv t \)

__Special cases:__

1. If \( i \leqslant 0 \) then the result is @('empty', t)@.
2. If \( i \geqslant n \) then the result is @(t, 'empty')@.
-}
splitAt :: forall m a . Measured m a => Int -> Treap m a -> (Treap m a, Treap m a)
splitAt i t
    | i <= 0         = (empty, t)
    | i >= sizeInt t = (t, empty)
    | otherwise      = go i t
  where
    go :: Int -> Treap m a -> (Treap m a, Treap m a)
    go k = \case
        Empty -> (Empty, Empty)
        Node _ p _ a left right ->
            let lSize = sizeInt left
            in case compare k lSize of
                EQ -> (left, new p a Empty right)
                LT ->
                    let (!t1, !newLeft) = go k left
                    in (t1, new p a newLeft right)
                GT ->
                    let (!newRight, !t2) = go (k - lSize - 1) right
                    in (new p a left newRight, t2)

{- | \( O(\max\ d_1\ d_2) \). Merge two 'Treap's into single one.

>>> pone p a = one (Priority p) a :: Treap (Sum Int) Int
>>> prettyPrint $ merge (merge (pone 1 3) (pone 4 5)) (merge (pone 3 0) (pone 5 9))
   4,17:9
     ╱
   3,8:5
     ╱╲
    ╱  ╲
   ╱    ╲
1,3:3 1,0:0
-}
merge :: Measured m a => Treap m a -> Treap m a -> Treap m a
merge Empty r = r
merge l Empty = l
merge l@(Node _ p1 _ a1 l1 r1) r@(Node _ p2 _ a2 l2 r2)
    | p1 > p2   = recalculate $ new p1 a1 l1 (merge r1 r)
    | otherwise = recalculate $ new p2 a2 (merge l l2) r2

{- | \( O(d) \). @'take' n t@ returns 'Treap' that contains first @n@ elements of the given
'Treap' @t@.

__Special cases:__

1. If \( i \leqslant 0 \) then the result is 'empty'.
2. If \( i \geqslant n \) then the result is @t@.
-}
take :: forall m a . Measured m a => Int -> Treap m a -> Treap m a
take n t
    | n <= 0         = Empty
    | n >= sizeInt t = t
    | otherwise      = go n t
  where
    go :: Int -> Treap m a -> Treap m a
    go _ Empty = Empty
    go 0 _     = Empty
    go i (Node _ p _ a l r) =
        let lSize = sizeInt l
        in case compare i lSize of
            LT -> go i l
            EQ -> l
            GT -> new p a l $ go (i - lSize - 1) r

{- | \( O(d) \). @'drop' n t@ returns 'Treap' without first @n@ elements of the given
'Treap' @t@.

__Special cases:__

1. If \( i \leqslant 0 \) then the result is @t@.
2. If \( i \geqslant n \) then the result is 'empty'.
-}
drop :: forall m a . Measured m a => Int -> Treap m a -> Treap m a
drop n t
    | n <= 0         = t
    | n >= sizeInt t = Empty
    | otherwise      = go n t
  where
    go :: Int -> Treap m a -> Treap m a
    go _ Empty = Empty
    go 0 tree  = tree
    go i (Node _ p _ a l r) =
        let lSize = sizeInt l
        in case compare i lSize of
            LT -> new p a (go i l) r
            EQ -> new p a Empty r
            GT -> go (i - lSize - 1) r

{- | \( O(d) \). Rotate a 'Treap' to the right by a given number of elements
modulo treap size. In simple words, @'rotate' n t@ takes first @n@ elements of
@t@ and puts them at the end of @t@ in the same order. If the given index is
negative, then this function rotates left.
-}
rotate :: forall m a . Measured m a => Int -> Treap m a -> Treap m a
rotate n t = case t of
    Empty -> Empty
    _ | n == 0    -> t
      | otherwise -> let (left, right) = splitAt shift t in merge right left
  where
    shift :: Int
    shift = n `mod` sizeInt t

----------------------------------------------------------------------------
-- Modification functions
----------------------------------------------------------------------------

{- | \( O(d) \). Insert a value into 'Treap' with given key and priority.
Updates monoidal accumulator accordingly.
-}
insert :: forall m a . Measured m a => Int -> Priority -> a -> Treap m a -> Treap m a
insert i p a t
    | i < 0          = go 0 t
    | i >= sizeInt t = go (sizeInt t) t
    | otherwise      = go i t
  where
    go :: Int -> Treap m a -> Treap m a
    go _ Empty = one p a
    go k node@(Node _ tp _ ta l r)
      | p <= tp =
          let lSize = sizeInt l
          in if k <= lSize
              then recalculate $ new tp ta (go k l) r
              else recalculate $ new tp ta l (go (k - lSize - 1) r)
      | otherwise =
          let (!newL, !newR) = splitAt k node
          in recalculate $ new p a newL newR

{- | \( O(d) \). Delete element from 'Treap' by the given index. If index is out
of bounds, 'Treap' remains unchanged.
-}
delete :: forall m a . Measured m a => Int -> Treap m a -> Treap m a
delete i t
    | i < 0          = t
    | i >= sizeInt t = t
    | otherwise      = go i t
  where
    go :: Int -> Treap m a -> Treap m a
    go _ Empty = Empty
    go k (Node _ p _ a l r) =
        let lSize = sizeInt l
        in case compare k lSize of
            EQ -> merge l r
            LT -> recalculate $ new p a (go k l) r
            GT -> recalculate $ new p a l (go (k - lSize - 1) r)

----------------------------------------------------------------------------
-- Core internal functions
----------------------------------------------------------------------------

{- | \( O(1) \). Calculate size and the value of the monoidal accumulator
in the given root node. This function doesn't perform any recursive calls and it
assumes that the values in the children are already correct. So use this
function only in bottom-up manner.
-}
recalculate :: Measured m a => Treap m a -> Treap m a
recalculate Empty = Empty
recalculate (Node _ p _ a l r) =
    Node (1 + size l + size r) p (measure l <> measure a <> measure r) a l r