Splay.hs

{-# LANGUAGE GHC2021 #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE NoFieldSelectors #-}

import Data.Function ((&))
import Data.Functor.Identity (Identity (..))
import Data.Monoid (Endo (..), Sum (..))

data Node m a = Node
  { left :: ONode m a,
    value :: a,
    right :: ONode m a,
    metadata :: m
  }
  deriving (Show, Eq, Foldable, Functor, Traversable)

class Metadata a m where
  calc :: ONode m a -> a -> ONode m a -> m

instance Metadata a () where
  calc _ _ _ = ()

instance Metadata Int (Sum Int) where
  calc l x r = maybe 0 (.metadata) l <> Sum x <> maybe 0 (.metadata) r

fixup :: (Metadata a m) => Node m a -> Node m a
fixup n@(Node {left, value, right}) = n {metadata = calc left value right}

mkNode :: (Metadata a m) => ONode m a -> a -> ONode m a -> Node m a
mkNode left value right =
  Node
    { left,
      value,
      right,
      metadata = calc left value right
    }

type ONode m a = Maybe (Node m a)

data Dir = L | R deriving (Show)

mapL, mapR :: (Metadata a m) => (Node m a -> Node m a) -> (Node m a -> Node m a)
mapL f n = fixup $ n {left = f <$> n.left}
mapR f n = fixup $ n {right = f <$> n.right}

splayByPath :: (Metadata a m) => [Dir] -> Node m a -> Node m a
splayByPath [] n = n
splayByPath [L] n = rotR n
splayByPath [R] n = rotL n
splayByPath (L : L : d) n = n & (mapL . mapL) (splayByPath d) & rotR & rotR
splayByPath (R : R : d) n = n & (mapR . mapR) (splayByPath d) & rotL & rotL
splayByPath (L : R : d) n = n & (mapL . mapR) (splayByPath d) & mapL rotL & rotR
splayByPath (R : L : d) n = n & (mapR . mapL) (splayByPath d) & mapR rotR & rotL

splay :: (Metadata a m, Ord a) => a -> Node m a -> Node m a
splay key n = splayByPath (locateByKey key (Just n)) n

rotL, rotR :: (Metadata a m) => Node m a -> Node m a
rotL (Node a x (Just (Node b y c _)) _) = mkNode (Just (mkNode a x b)) y c
rotL n = n
rotR (Node (Just (Node a x b _)) y c _) = mkNode a x (Just (mkNode b y c))
rotR n = n

locateM :: (Ord a, Monad m) => (Node n a -> m (Maybe Dir)) -> ONode n a -> m [Dir]
locateM _f Nothing = pure []
locateM f (Just n) =
  f n >>= \case
    Nothing -> pure []
    Just L -> (L :) <$> locateM f n.left
    Just R -> (R :) <$> locateM f n.right

locate :: (Ord a) => (Node n a -> Maybe Dir) -> ONode n a -> [Dir]
locate f = runIdentity . locateM (Identity . f)

locateByKey :: (Ord a) => a -> ONode m a -> [Dir]
locateByKey key = locate $ \x -> case compare key x.value of
  EQ -> Nothing
  LT -> Just L
  GT -> Just R

split :: (Metadata a m, Ord a) => a -> Node m a -> (ONode m a, ONode m a)
split key n =
  let n' = splay key n
   in if key <= n'.value
        then (n'.left, Just (fixup (n' {left = Nothing})))
        else (Just (fixup (n' {right = Nothing})), n'.right)

insert :: (Metadata a m, Ord a) => a -> Node m a -> Node m a
insert key (split key -> (l, r)) = mkNode l key r

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

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

-- locateByKeyAndAlsoLog :: Ord a => a -> ONode a -> IO [Dir]
-- locateByKeyAndAlsoLog key = locateM $ \x -> do
--   let result = compare key x.value
--   putStrLn $ "Comparison result: " <> show result
--   pure $ case result of
--     EQ -> Nothing
--     LT -> Just L
--     GT -> Just R

singleton :: (Metadata a m) => a -> Node m a
singleton x = mkNode Nothing x Nothing