-----------------------------------------------------------------------------
-- |
-- Module      :  Data.LCA.Online.Naive
-- Copyright   :  (C) 2011-2012 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  experimental
-- Portability :  portable
--
-- Naive online calculation of the the lowest common ancestor in /O(h)/
----------------------------------------------------------------------------
module Data.LCA.Online.Naive
  ( Path
  , empty
  , cons
  , uncons
  , view
  , null
  , length
  , isAncestorOf
  , lca
  , keep
  , drop
  , traverseWithKey
  , toList
  , fromList
  , (~=)
  ) where

import Control.Applicative hiding (empty)
import Data.Foldable hiding (toList)
import Data.Traversable
import Prelude hiding (length, null, drop)
import qualified Prelude
import Data.LCA.View

data Path a = Path {-# UNPACK #-} !Int [(Int,a)]
  deriving (Show, Read)

toList :: Path a -> [(Int,a)]
toList (Path _ xs) = xs
{-# INLINE toList #-}

fromList :: [(Int,a)] -> Path a
fromList xs = Path (Prelude.length xs) xs
{-# INLINE fromList #-}

instance Functor Path where
  fmap f (Path n xs) = Path n [ (k, f a) | (k,a) <- xs]

instance Foldable Path where
  foldMap f (Path _ xs) = foldMap (f . snd) xs

instance Traversable Path where
  traverse f (Path n xs) = Path n <$> traverse (\(k,a) -> (,) k <$> f a) xs

traverseWithKey :: Applicative f => (Int -> a -> f b) -> Path a -> f (Path b)
traverseWithKey f (Path n xs) = Path n <$> traverse (\(k,a) -> (,) k <$> f k a) xs
{-# INLINE traverseWithKey #-}

-- | The empty path
empty :: Path a
empty = Path 0 []

-- | /O(1)/
length :: Path a -> Int
length (Path n _) = n
{-# INLINE length #-}

-- | /O(1)/
null :: Path a -> Bool
null (Path n _) = n == 0
{-# INLINE null #-}

-- | /O(1)/ Invariant: most operations assume that the keys @k@ are globally unique
cons :: Int -> a -> Path a -> Path a
cons k a (Path n xs) = Path (n + 1) $ (k,a):xs
{-# INLINE cons #-}

uncons :: Path a -> Maybe (Int, a, Path a)
uncons (Path _ []) = Nothing
uncons (Path n ((k,a):xs)) = Just (k,a,Path (n - 1) xs)
{-# INLINE uncons #-}

view :: Path a -> View Path a
view (Path _ []) = Root
view (Path n ((k,a):xs)) = Node k a (Path (n - 1) xs)
{-# INLINE view #-}

-- | /O(h - k)/ to @keep k@ elements of path of height @h@
keep :: Int -> Path a -> Path a
keep k p@(Path n xs)
  | k >= n    = p
  | otherwise = Path k $ Prelude.drop (n - k) xs
{-# INLINE keep #-}

-- | /O(k)/ to @drop k@ elements from a path
drop :: Int -> Path a -> Path a
drop k (Path n xs)
  | k >= n    = empty
  | otherwise = Path (n - k) (Prelude.drop k xs)
{-# INLINE drop #-}

-- /O(h)/ @xs `isAncestorOf` ys@ holds when @xs@ is a prefix starting at the root of path @ys@.
isAncestorOf :: Path a -> Path b -> Bool
isAncestorOf xs ys = xs ~= keep (length xs) ys
{-# INLINE isAncestorOf #-}

infix 4 ~=
-- | /O(1)/ Compare to see if two trees have the same leaf key
(~=) :: Path a -> Path b -> Bool
Path _ []        ~= Path _ []        = True
Path _ ((i,_):_) ~= Path _ ((j,_):_) = i == j
_                ~= _                = False
{-# INLINE (~=) #-}

-- | /O(h)/ Compute the lowest common ancestor
lca :: Path a -> Path b -> Path a
lca xs0 ys0 = case compare nxs nys of
    LT -> go nxs (toList xs0) (toList (keep nxs ys0))
    EQ -> go nxs (toList xs0) (toList ys0)
    GT -> go nys (toList (keep nys xs0)) (toList ys0)
  where
    nxs = length xs0
    nys = length ys0
    go k xss@((i,_):xs) ((j,_):ys)
      | i == j    = Path k xss
      | otherwise = go (k - 1) xs ys
    go _ _ _ = empty
{-# INLINE lca #-}