Splay.hs

{-# LANGUAGE ViewPatterns #-}

import Control.Monad (forM_)
import Control.Monad.ST (ST)
import Data.Functor ((<&>))
import Data.Monoid (Endo (..))
import Data.STRef (STRef, newSTRef)
import Data.Vector qualified as V
import Data.Vector.Generic.Mutable qualified as GV
import Data.Vector.Mutable qualified as MV
import System.Random (Random (..), getStdGen, getStdRandom)
import System.Random.Stateful (RandomGen, STGenM, applySTGen, runSTGen, runSTGen_)

data Node a = Node {value :: a, left :: Maybe (Node a), right :: Maybe (Node a)}

deriving instance (Show a) => Show (Node a)

rotLeft, rotRight :: Node a -> Node a
rotLeft (Node {value = x, left = a, right = Just (Node {value = y, left = b, right = c})}) =
  Node {value = y, left = Just (Node {value = x, left = a, right = b}), right = c}
rotRight (Node {value = x, left = Just (Node {value = y, left = a, right = b}), right = c}) =
  Node {value = y, left = a, right = Just (Node {value = x, left = b, right = c})}

splay :: (Ord a) => a -> Node a -> Node a
splay key n = case compare key (value n) of
  EQ -> n
  LT -> case left n of
    Nothing -> n
    Just m -> case compare key (value m) of
      EQ -> rotRight n
      LT -> case left m of
        Nothing -> rotRight n
        Just o ->
          rotRight (rotRight (n {left = Just (m {left = Just (splay key o)})}))
      GT -> case right m of
        Nothing -> rotRight n
        Just o ->
          rotRight (n {left = Just (rotLeft (m {right = Just (splay key o)}))})
  GT -> case right n of
    Nothing -> n
    Just m -> case compare key (value m) of
      EQ -> rotLeft n
      LT -> case left m of
        Nothing -> rotLeft n
        Just o ->
          rotLeft (n {right = Just (rotRight (m {left = Just (splay key o)}))})
      GT -> case right m of
        Nothing -> rotLeft n
        Just o ->
          rotLeft (rotLeft (n {right = Just (m {right = Just (splay key o)})}))

splitPresplayed :: (Ord a) => a -> Node a -> (Maybe (Node a), Maybe (Node a))
splitPresplayed key n =
  if key <= value n
    then
      (left n, Just (n {left = Nothing}))
    else
      (Just (n {right = Nothing}), right n)

split :: (Ord a) => a -> Node a -> (Maybe (Node a), Maybe (Node a))
split key (splay key -> n) = splitPresplayed key n

insert :: (Ord a) => a -> Node a -> Node a
insert key (splay key -> n) = case splitPresplayed key n of
  (_, Just r) | key == value r -> n
  (l, r) -> Node {value = key, left = l, right = r}

singleton :: (Ord a) => a -> Node a
singleton key = Node {value = key, left = Nothing, right = Nothing}

contains :: (Ord a) => a -> Node a -> (Bool, Node a)
contains a (splay a -> n) = (value n == a, n)

fromList :: (Ord a) => [a] -> Node a
fromList (x : xs) = foldl (flip insert) (singleton x) xs

foldMapT :: (Monoid m) => (a -> m) -> Node a -> m
foldMapT f (Node {value, left, right}) = foldMap (foldMapT f) left <> f value <> foldMap (foldMapT f) right

sfold :: b -> (a -> b -> b -> b) -> Maybe (Node a) -> b
sfold nothing node Nothing = nothing
sfold nothing node (Just (Node {value, left, right})) = node value (sfold nothing node left) (sfold nothing node right)

sfold1 :: b -> (a -> b -> b -> b) -> Node a -> b
sfold1 nothing node = sfold nothing node . Just

height :: Node a -> Int
height = sfold1 0 (\_ l r -> max l r + 1)

toList :: Node a -> [a]
toList = flip appEndo [] . foldMapT (Endo . (:))

shuffle :: forall g s a. (RandomGen g) => [a] -> STGenM g s -> ST s [a]
shuffle list gen = do
  let initialVector = V.fromList list
  a <- V.thaw initialVector
  let n = length list
  forM_ [0 .. n - 1] $ \i -> do
    pos <- applySTGen (randomR (i, n - 1)) gen
    MV.swap a i pos
  a' <- V.freeze a
  pure $ V.toList a'

shuffleIO :: [a] -> IO [a]
shuffleIO list = getStdRandom $ \gen -> runSTGen gen (shuffle list)