{-# language DeriveFoldable #-}
{-# language FlexibleInstances #-}
{-# language FunctionalDependencies #-}
module Data.SplayTree where

import Data.Monoid
import qualified Data.Semigroup as Semigroup

infixr 5 <|
infixl 5 |>

class Monoid v => Measured v a | a -> v where
  measure :: a -> v

data SplayTree v a
  = Leaf
  | Fork (SplayTree v a) a (SplayTree v a) !v -- ^ Cached measure of the whole node
  deriving (Eq, Ord, Show, Foldable)

instance Measured v a => Measured v (SplayTree v a) where
  {-# INLINE measure #-}
  measure Leaf = mempty
  measure (Fork _ _ _ v) = v

instance Measured v a => Semigroup.Semigroup (SplayTree v a) where
  {-# INLINE (<>) #-}
  Leaf <> t = t
  t <> Leaf = t
  Fork l1 a1 r1 lar1 <> Fork l2 a2 r2 lar2
    = Fork l1 a1 (Fork (r1 <> l2) a2 r2 (measure r1 <> lar2)) (lar1 <> lar2)

instance Measured v a => Monoid (SplayTree v a) where
  {-# INLINE mempty #-}
  mempty = Leaf
  {-# INLINE mappend #-}
  mappend = (Semigroup.<>)

-- | Is the splay tree empty?
--
-- @since 0.2.0.0
null :: SplayTree v a -> Bool
null Leaf = True
null Fork {} = False

-------------------------------------------------------------------------------
-- * Construction

{-# INLINE singleton #-}
singleton :: Measured v a => a -> SplayTree v a
singleton a = Fork Leaf a Leaf $ measure a

{-# INLINE (<|) #-}
(<|) :: Measured v a => a -> SplayTree v a -> SplayTree v a
(<|) = fork Leaf

{-# INLINE (|>) #-}
(|>) :: Measured v a => SplayTree v a -> a -> SplayTree v a
(|>) t a = fork t a Leaf

{-# INLINE fork #-}
fork :: Measured v a => SplayTree v a -> a -> SplayTree v a -> SplayTree v a
fork l a r = Fork l a r $ measure l <> measure a <> measure r

-------------------------------------------------------------------------------
-- * Deconstruction

{-# INLINE uncons #-}
uncons :: Measured v a => SplayTree v a -> Maybe (a, SplayTree v a)
uncons Leaf = Nothing
uncons (Fork left el right _) = Just $ go left el right
  where
    go Leaf a r = (a, r)
    go (Fork l a m _) b r = go l a (fork m b r)

{-# INLINE unsnoc #-}
unsnoc :: Measured v a => SplayTree v a -> Maybe (SplayTree v a, a)
unsnoc Leaf = Nothing
unsnoc (Fork left el right _) = Just $ go left el right
  where
    go l a Leaf = (l, a)
    go l a (Fork m b r _) = go (fork l a m) b r

data SplitResult v a
  = Outside
  | Inside (SplayTree v a) a (SplayTree v a)
  deriving (Eq, Ord, Show)

{-# INLINE split #-}
split :: Measured v a => (v -> Bool) -> SplayTree v a -> SplitResult v a
split = go mempty
  where
    go _ _ Leaf = Outside
    go v f (Fork l a r _)
      | f vl = case go v f l of
        Outside -> Outside
        Inside l' a' m -> Inside l' a' $ fork m a r
      | f vla = Inside l a r
      | otherwise = case go vla f r of
        Outside -> Outside
        Inside m a' r' -> Inside (fork l a m) a' r'
      where
        vl = v <> measure l
        vla = vl <> measure a

-------------------------------------------------------------------------------
-- * Maps

{-# INLINE map #-}
map
  :: (Measured v a, Measured w b)
  => (a -> b)
  -> SplayTree v a
  -> SplayTree w b
map _ Leaf = Leaf
map f (Fork l a r _) = fork (Data.SplayTree.map f l) (f a) (Data.SplayTree.map f r)

{-# INLINE mapWithPos #-}
mapWithPos
  :: (Measured v a, Measured w b)
  => (v -> a -> b)
  -> SplayTree v a
  -> SplayTree w b
mapWithPos f = go mempty
  where
    go _ Leaf = Leaf
    go i (Fork l a r _) = fork (go i l) (f il a) (go ila r)
      where
        il = i <> measure l
        ila = il <> measure a

{-# INLINE mapWithContext #-}
mapWithContext
  :: (Measured v a, Measured w b)
  => (v -> a -> v -> b)
  -> SplayTree v a
  -> SplayTree w b
mapWithContext f t = go mempty t mempty
  where
    go _ Leaf _ = Leaf
    go i (Fork l a r _) j = fork (go i l arj) (f il a rj) (go ila r j)
      where
        ma = measure a
        il = i <> measure l
        ila = il <> ma
        rj = measure r <> j
        arj = ma

-------------------------------------------------------------------------------
-- * Traversals

{-# INLINE traverse #-}
traverse
  :: (Measured v a, Measured w b, Applicative f)
  => (a -> f b)
  -> SplayTree v a
  -> f (SplayTree w b)
traverse _ Leaf = pure Leaf
traverse f (Fork l a r _)
  = fork
  <$> Data.SplayTree.traverse f l
  <*> f a
  <*> Data.SplayTree.traverse f r

{-# INLINE traverseWithPos #-}
traverseWithPos
  :: (Measured v a, Measured w b, Applicative f)
  => (v -> a -> f b)
  -> SplayTree v a
  -> f (SplayTree w b)
traverseWithPos f = go mempty
  where
    go _ Leaf = pure Leaf
    go i (Fork l a r _)
      = fork <$> go i l <*> f il a <*> go ila r
      where
        il = i <> measure l
        ila = il <> measure a

{-# INLINE traverseWithContext #-}
traverseWithContext
  :: (Measured v a, Measured w b, Applicative f)
  => (v -> a -> v -> f b)
  -> SplayTree v a
  -> f (SplayTree w b)
traverseWithContext f t = go mempty t mempty
  where
    go _ Leaf _ = pure Leaf
    go i (Fork l a r _) j
      = fork <$> go i l arj <*> f il a rj <*> go ila r j
      where
        ma = measure a
        il = i <> measure l
        ila = il <> ma
        rj = measure r <> j
        arj = ma