{-# LANGUAGE CPP #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.LCA.Online.Monoidal
-- Copyright   :  (C) 2012-2015 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
--
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  experimental
-- Portability :  portable
--
-- Provides online calculation of the the lowest common ancestor in /O(log h)/
-- by compressing the spine of the paths using a skew-binary random access
-- list.
--
-- This library implements the technique described in my talk
--
-- <http://www.slideshare.net/ekmett/skewbinary-online-lowest-common-ancestor-search>
--
-- to improve the known asymptotic bounds on both online lowest common ancestor search
--
-- <http://en.wikipedia.org/wiki/Lowest_common_ancestor>
--
-- and the online level ancestor problem:
--
-- <http://en.wikipedia.org/wiki/Level_ancestor_problem>
--
-- Algorithms used here assume that the key values chosen for @k@ are
-- globally unique.
--
-- This version provides access to a monoidal \"summary\" of the
-- elided path for many operations.
--
----------------------------------------------------------------------------
module Data.LCA.Online.Monoidal
  ( Path
  , toList, fromList
  , map, mapHom, mapWithKey
  , traverse, traverseWithKey
  , empty
  , cons
  , uncons, view
  , null
  , length
  , measure
  , isAncestorOf
  , keep, mkeep
  , drop, mdrop
  , (~=)
  , lca, mlca
  ) where

import Control.Applicative hiding (empty)
import qualified Data.Foldable as F

#if __GLASGOW_HASKELL__ < 710
import Data.Monoid (Monoid(..))
#endif

import Prelude hiding
  ( drop
  , map
  , length
  , null
#if __GLASGOW_HASKELL__ < 710
#else
  , traverse
#endif
#if MIN_VERSION_base(4,11,0)
  , (<>)
#endif
  )
import Data.LCA.View

-- $setup
-- >>> let length = Data.LCA.Online.Monoidal.length

infixl 6 <>
(<>) :: Monoid a => a -> a -> a
<> :: a -> a -> a
(<>) = a -> a -> a
forall a. Monoid a => a -> a -> a
mappend
{-# INLINE (<>) #-}

-- | Complete binary trees
--
-- NB: we could ensure the complete tree invariant
data Tree a
  = Bin a {-# UNPACK #-} !Int a (Tree a) (Tree a)
  | Tip {-# UNPACK #-} !Int a
  deriving (Int -> Tree a -> ShowS
[Tree a] -> ShowS
Tree a -> String
(Int -> Tree a -> ShowS)
-> (Tree a -> String) -> ([Tree a] -> ShowS) -> Show (Tree a)
forall a. Show a => Int -> Tree a -> ShowS
forall a. Show a => [Tree a] -> ShowS
forall a. Show a => Tree a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Tree a] -> ShowS
$cshowList :: forall a. Show a => [Tree a] -> ShowS
show :: Tree a -> String
$cshow :: forall a. Show a => Tree a -> String
showsPrec :: Int -> Tree a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Tree a -> ShowS
Show, ReadPrec [Tree a]
ReadPrec (Tree a)
Int -> ReadS (Tree a)
ReadS [Tree a]
(Int -> ReadS (Tree a))
-> ReadS [Tree a]
-> ReadPrec (Tree a)
-> ReadPrec [Tree a]
-> Read (Tree a)
forall a. Read a => ReadPrec [Tree a]
forall a. Read a => ReadPrec (Tree a)
forall a. Read a => Int -> ReadS (Tree a)
forall a. Read a => ReadS [Tree a]
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [Tree a]
$creadListPrec :: forall a. Read a => ReadPrec [Tree a]
readPrec :: ReadPrec (Tree a)
$creadPrec :: forall a. Read a => ReadPrec (Tree a)
readList :: ReadS [Tree a]
$creadList :: forall a. Read a => ReadS [Tree a]
readsPrec :: Int -> ReadS (Tree a)
$creadsPrec :: forall a. Read a => Int -> ReadS (Tree a)
Read)

instance F.Foldable Tree where
  foldMap :: (a -> m) -> Tree a -> m
foldMap a -> m
f (Tip Int
_ a
a) = a -> m
f a
a
  foldMap a -> m
f (Bin a
_ Int
_ a
a Tree a
l Tree a
r) = a -> m
f a
a m -> m -> m
forall a. Monoid a => a -> a -> a
<> (a -> m) -> Tree a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
F.foldMap a -> m
f Tree a
l m -> m -> m
forall a. Monoid a => a -> a -> a
<> (a -> m) -> Tree a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
F.foldMap a -> m
f Tree a
r

measureT :: Tree a -> a
measureT :: Tree a -> a
measureT (Tip Int
_ a
a)       = a
a
measureT (Bin a
a Int
_ a
_ Tree a
_ Tree a
_) = a
a

bin :: Monoid a => Int -> a -> Tree a -> Tree a -> Tree a
bin :: Int -> a -> Tree a -> Tree a -> Tree a
bin Int
n a
a Tree a
l Tree a
r = a -> Int -> a -> Tree a -> Tree a -> Tree a
forall a. a -> Int -> a -> Tree a -> Tree a -> Tree a
Bin (a
a a -> a -> a
forall a. Monoid a => a -> a -> a
<> Tree a -> a
forall a. Tree a -> a
measureT Tree a
l a -> a -> a
forall a. Monoid a => a -> a -> a
<> Tree a -> a
forall a. Tree a -> a
measureT Tree a
r) Int
n a
a Tree a
l Tree a
r

sameT :: Tree a -> Tree b -> Bool
sameT :: Tree a -> Tree b -> Bool
sameT Tree a
xs Tree b
ys = Tree a -> Int
forall a. Tree a -> Int
root Tree a
xs Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Tree b -> Int
forall a. Tree a -> Int
root Tree b
ys where
  root :: Tree a -> Int
root (Tip Int
k a
_)     = Int
k
  root (Bin a
_ Int
k a
_ Tree a
_ Tree a
_) = Int
k

-- | A compressed 'Path' as a skew binary random access list
data Path a
  = Nil
  | Cons a
         {-# UNPACK #-} !Int -- the number of elements @n@ in this entire skew list
         {-# UNPACK #-} !Int -- the number of elements @w@ in this binary tree node
         (Tree a)          -- a complete binary tree @t@ of with @w@ elements
         (Path a)          -- @n - w@ elements in a linked list @ts@, of complete trees in ascending order by size
  deriving (Int -> Path a -> ShowS
[Path a] -> ShowS
Path a -> String
(Int -> Path a -> ShowS)
-> (Path a -> String) -> ([Path a] -> ShowS) -> Show (Path a)
forall a. Show a => Int -> Path a -> ShowS
forall a. Show a => [Path a] -> ShowS
forall a. Show a => Path a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Path a] -> ShowS
$cshowList :: forall a. Show a => [Path a] -> ShowS
show :: Path a -> String
$cshow :: forall a. Show a => Path a -> String
showsPrec :: Int -> Path a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Path a -> ShowS
Show, ReadPrec [Path a]
ReadPrec (Path a)
Int -> ReadS (Path a)
ReadS [Path a]
(Int -> ReadS (Path a))
-> ReadS [Path a]
-> ReadPrec (Path a)
-> ReadPrec [Path a]
-> Read (Path a)
forall a. Read a => ReadPrec [Path a]
forall a. Read a => ReadPrec (Path a)
forall a. Read a => Int -> ReadS (Path a)
forall a. Read a => ReadS [Path a]
forall a.
(Int -> ReadS a)
-> ReadS [a] -> ReadPrec a -> ReadPrec [a] -> Read a
readListPrec :: ReadPrec [Path a]
$creadListPrec :: forall a. Read a => ReadPrec [Path a]
readPrec :: ReadPrec (Path a)
$creadPrec :: forall a. Read a => ReadPrec (Path a)
readList :: ReadS [Path a]
$creadList :: forall a. Read a => ReadS [Path a]
readsPrec :: Int -> ReadS (Path a)
$creadsPrec :: forall a. Read a => Int -> ReadS (Path a)
Read)

instance F.Foldable Path where
  foldMap :: (a -> m) -> Path a -> m
foldMap a -> m
_ Path a
Nil = m
forall a. Monoid a => a
mempty
  foldMap a -> m
f (Cons a
_ Int
_ Int
_ Tree a
t Path a
ts) = (a -> m) -> Tree a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
F.foldMap a -> m
f Tree a
t m -> m -> m
forall a. Monoid a => a -> a -> a
<> (a -> m) -> Path a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
F.foldMap a -> m
f Path a
ts
#if __GLASGOW_HASKELL__ >= 710
  length :: Path a -> Int
length = Path a -> Int
forall a. Path a -> Int
length
  null :: Path a -> Bool
null   = Path a -> Bool
forall a. Path a -> Bool
null
#endif

-- | /O(1)/ Determine the 'length' of a 'Path'.
length :: Path a -> Int
length :: Path a -> Int
length Path a
Nil = Int
0
length (Cons a
_ Int
n Int
_ Tree a
_ Path a
_) = Int
n
{-# INLINE length #-}

-- | /O(1)/ Returns 'True' iff the path is 'empty'.
null :: Path a -> Bool
null :: Path a -> Bool
null Path a
Nil = Bool
True
null Path a
_ = Bool
False
{-# INLINE null #-}

-- | Extract a monoidal summary of a 'Path'.
measure :: Monoid a => Path a -> a
measure :: Path a -> a
measure Path a
Nil = a
forall a. Monoid a => a
mempty
measure (Cons a
a Int
_ Int
_ Tree a
_ Path a
_) = a
a

consT :: Monoid a => Int -> Tree a -> Path a -> Path a
consT :: Int -> Tree a -> Path a -> Path a
consT Int
w Tree a
t Path a
ts = a -> Int -> Int -> Tree a -> Path a -> Path a
forall a. a -> Int -> Int -> Tree a -> Path a -> Path a
Cons (Tree a -> a
forall a. Tree a -> a
measureT Tree a
t a -> a -> a
forall a. Monoid a => a -> a -> a
<> Path a -> a
forall m. Monoid m => Path m -> m
measure Path a
ts) (Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Path a -> Int
forall a. Path a -> Int
length Path a
ts) Int
w Tree a
t Path a
ts

consN :: Monoid a => Int -> Int -> Tree a -> Path a -> Path a
consN :: Int -> Int -> Tree a -> Path a -> Path a
consN Int
n Int
w Tree a
t Path a
ts = a -> Int -> Int -> Tree a -> Path a -> Path a
forall a. a -> Int -> Int -> Tree a -> Path a -> Path a
Cons (Tree a -> a
forall a. Tree a -> a
measureT Tree a
t a -> a -> a
forall a. Monoid a => a -> a -> a
<> Path a -> a
forall m. Monoid m => Path m -> m
measure Path a
ts) Int
n Int
w Tree a
t Path a
ts

-- | /O(n)/ Re-annotate a 'Path' full of monoidal values using a different 'Monoid'.
map :: Monoid b => (a -> b) -> Path a -> Path b
map :: (a -> b) -> Path a -> Path b
map a -> b
f = Path a -> Path b
go where
  go :: Path a -> Path b
go Path a
Nil = Path b
forall a. Path a
Nil
  go (Cons a
_ Int
n Int
k Tree a
t Path a
ts) = Int -> Int -> Tree b -> Path b -> Path b
forall a. Monoid a => Int -> Int -> Tree a -> Path a -> Path a
consN Int
n Int
k (Tree a -> Tree b
goT Tree a
t) (Path a -> Path b
go Path a
ts)
  goT :: Tree a -> Tree b
goT (Tip Int
k a
a) = Int -> b -> Tree b
forall a. Int -> a -> Tree a
Tip Int
k (a -> b
f a
a)
  goT (Bin a
_ Int
k a
a Tree a
l Tree a
r) = Int -> b -> Tree b -> Tree b -> Tree b
forall a. Monoid a => Int -> a -> Tree a -> Tree a -> Tree a
bin Int
k (a -> b
f a
a) (Tree a -> Tree b
goT Tree a
l) (Tree a -> Tree b
goT Tree a
r)
{-# INLINE map #-}

-- | /O(n)/ Re-annotate a 'Path' full of monoidal values with access to the key.
mapWithKey :: Monoid b => (Int -> a -> b) -> Path a -> Path b
mapWithKey :: (Int -> a -> b) -> Path a -> Path b
mapWithKey Int -> a -> b
f = Path a -> Path b
go where
  go :: Path a -> Path b
go Path a
Nil = Path b
forall a. Path a
Nil
  go (Cons a
_ Int
n Int
k Tree a
t Path a
ts) = Int -> Int -> Tree b -> Path b -> Path b
forall a. Monoid a => Int -> Int -> Tree a -> Path a -> Path a
consN Int
n Int
k (Tree a -> Tree b
goT Tree a
t) (Path a -> Path b
go Path a
ts)
  goT :: Tree a -> Tree b
goT (Tip Int
k a
a) = Int -> b -> Tree b
forall a. Int -> a -> Tree a
Tip Int
k (Int -> a -> b
f Int
k a
a)
  goT (Bin a
_ Int
k a
a Tree a
l Tree a
r) = Int -> b -> Tree b -> Tree b -> Tree b
forall a. Monoid a => Int -> a -> Tree a -> Tree a -> Tree a
bin Int
k (Int -> a -> b
f Int
k a
a) (Tree a -> Tree b
goT Tree a
l) (Tree a -> Tree b
goT Tree a
r)
{-# INLINE mapWithKey #-}

-- | /O(n)/ Re-annotate a 'Path' full of monoidal values/
--
-- Unlike 'map', @'mapHom' f@ assumes that @f@ is a 'Monoid' homomorphism, that is to say you must ensure
--
-- @
-- f a `'mappend'` f b = f (a `'mappend'` b)
-- f 'mempty' = 'mempty'
-- @
mapHom :: (a -> b) -> Path a -> Path b
mapHom :: (a -> b) -> Path a -> Path b
mapHom a -> b
f = Path a -> Path b
go where
  go :: Path a -> Path b
go Path a
Nil               = Path b
forall a. Path a
Nil
  go (Cons a
a Int
n Int
k Tree a
t Path a
ts) = b -> Int -> Int -> Tree b -> Path b -> Path b
forall a. a -> Int -> Int -> Tree a -> Path a -> Path a
Cons (a -> b
f a
a) Int
n Int
k (Tree a -> Tree b
goT Tree a
t) (Path a -> Path b
go Path a
ts)
  goT :: Tree a -> Tree b
goT (Tip Int
k a
a)        = Int -> b -> Tree b
forall a. Int -> a -> Tree a
Tip Int
k (a -> b
f a
a)
  goT (Bin a
m Int
k a
a Tree a
l Tree a
r)  = b -> Int -> b -> Tree b -> Tree b -> Tree b
forall a. a -> Int -> a -> Tree a -> Tree a -> Tree a
Bin (a -> b
f a
m) Int
k (a -> b
f a
a) (Tree a -> Tree b
goT Tree a
l) (Tree a -> Tree b
goT Tree a
r)
{-# INLINE mapHom #-}

-- | Convert a 'Path' to a list of @(ID, value)@ pairs.
toList :: Path a -> [(Int,a)]
toList :: Path a -> [(Int, a)]
toList Path a
Nil = []
toList (Cons a
_ Int
_ Int
_ Tree a
t Path a
ts)  = Tree a -> [(Int, a)] -> [(Int, a)]
forall b. Tree b -> [(Int, b)] -> [(Int, b)]
go Tree a
t (Path a -> [(Int, a)]
forall a. Path a -> [(Int, a)]
toList Path a
ts) where
  go :: Tree b -> [(Int, b)] -> [(Int, b)]
go (Tip Int
k b
a) [(Int, b)]
xs       = (Int
k,b
a) (Int, b) -> [(Int, b)] -> [(Int, b)]
forall a. a -> [a] -> [a]
: [(Int, b)]
xs
  go (Bin b
_ Int
k b
a Tree b
l Tree b
r) [(Int, b)]
xs = (Int
k,b
a) (Int, b) -> [(Int, b)] -> [(Int, b)]
forall a. a -> [a] -> [a]
: Tree b -> [(Int, b)] -> [(Int, b)]
go Tree b
l (Tree b -> [(Int, b)] -> [(Int, b)]
go Tree b
r [(Int, b)]
xs)

-- | Build a 'Path' from a list of @(ID, value)@ pairs.
fromList :: Monoid a => [(Int,a)] -> Path a
fromList :: [(Int, a)] -> Path a
fromList [] = Path a
forall a. Path a
Nil
fromList ((Int
k,a
a):[(Int, a)]
xs) = Int -> a -> Path a -> Path a
forall a. Monoid a => Int -> a -> Path a -> Path a
cons Int
k a
a ([(Int, a)] -> Path a
forall a. Monoid a => [(Int, a)] -> Path a
fromList [(Int, a)]
xs)

-- | Traverse a 'Path' with access to the node IDs.
traverseWithKey :: (Applicative f, Monoid b) => (Int -> a -> f b) -> Path a -> f (Path b)
traverseWithKey :: (Int -> a -> f b) -> Path a -> f (Path b)
traverseWithKey Int -> a -> f b
f = Path a -> f (Path b)
go where
  go :: Path a -> f (Path b)
go Path a
Nil = Path b -> f (Path b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Path b
forall a. Path a
Nil
  go (Cons a
_ Int
n Int
k Tree a
t Path a
ts) = Int -> Int -> Tree b -> Path b -> Path b
forall a. Monoid a => Int -> Int -> Tree a -> Path a -> Path a
consN Int
n Int
k (Tree b -> Path b -> Path b) -> f (Tree b) -> f (Path b -> Path b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Tree a -> f (Tree b)
goT Tree a
t f (Path b -> Path b) -> f (Path b) -> f (Path b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Path a -> f (Path b)
go Path a
ts
  goT :: Tree a -> f (Tree b)
goT (Tip Int
k a
a)        = Int -> b -> Tree b
forall a. Int -> a -> Tree a
Tip Int
k (b -> Tree b) -> f b -> f (Tree b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> a -> f b
f Int
k a
a
  goT (Bin a
_ Int
k a
a Tree a
l Tree a
r)  = Int -> b -> Tree b -> Tree b -> Tree b
forall a. Monoid a => Int -> a -> Tree a -> Tree a -> Tree a
bin Int
k (b -> Tree b -> Tree b -> Tree b)
-> f b -> f (Tree b -> Tree b -> Tree b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> a -> f b
f Int
k a
a f (Tree b -> Tree b -> Tree b)
-> f (Tree b) -> f (Tree b -> Tree b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Tree a -> f (Tree b)
goT Tree a
l f (Tree b -> Tree b) -> f (Tree b) -> f (Tree b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Tree a -> f (Tree b)
goT Tree a
r
{-# INLINE traverseWithKey #-}

-- | Traverse a 'Path' yielding a new monoidal annotation.
traverse :: (Applicative f, Monoid b) => (a -> f b) -> Path a -> f (Path b)
traverse :: (a -> f b) -> Path a -> f (Path b)
traverse a -> f b
f = Path a -> f (Path b)
go where
  go :: Path a -> f (Path b)
go Path a
Nil = Path b -> f (Path b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Path b
forall a. Path a
Nil
  go (Cons a
_ Int
n Int
k Tree a
t Path a
ts) = Int -> Int -> Tree b -> Path b -> Path b
forall a. Monoid a => Int -> Int -> Tree a -> Path a -> Path a
consN Int
n Int
k (Tree b -> Path b -> Path b) -> f (Tree b) -> f (Path b -> Path b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Tree a -> f (Tree b)
goT Tree a
t f (Path b -> Path b) -> f (Path b) -> f (Path b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Path a -> f (Path b)
go Path a
ts
  goT :: Tree a -> f (Tree b)
goT (Tip Int
k a
a)        = Int -> b -> Tree b
forall a. Int -> a -> Tree a
Tip Int
k (b -> Tree b) -> f b -> f (Tree b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
f a
a
  goT (Bin a
_ Int
k a
a Tree a
l Tree a
r)  = Int -> b -> Tree b -> Tree b -> Tree b
forall a. Monoid a => Int -> a -> Tree a -> Tree a -> Tree a
bin Int
k (b -> Tree b -> Tree b -> Tree b)
-> f b -> f (Tree b -> Tree b -> Tree b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f b
f a
a f (Tree b -> Tree b -> Tree b)
-> f (Tree b) -> f (Tree b -> Tree b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Tree a -> f (Tree b)
goT Tree a
l f (Tree b -> Tree b) -> f (Tree b) -> f (Tree b)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Tree a -> f (Tree b)
goT Tree a
r
{-# INLINE traverse #-}

-- | The empty 'Path'
empty :: Path a
empty :: Path a
empty = Path a
forall a. Path a
Nil
{-# INLINE empty #-}

-- | /O(1)/ Invariant: most operations assume that the keys @k@ are globally unique
--
-- Extend the 'Path' with a new node ID and value.
cons :: Monoid a => Int -> a -> Path a -> Path a
cons :: Int -> a -> Path a -> Path a
cons Int
k a
a (Cons a
m Int
n Int
w Tree a
t (Cons a
_ Int
_ Int
w' Tree a
t2 Path a
ts)) | Int
w Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
w' = a -> Int -> Int -> Tree a -> Path a -> Path a
forall a. a -> Int -> Int -> Tree a -> Path a -> Path a
Cons (a
a a -> a -> a
forall a. Monoid a => a -> a -> a
<> a
m) (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
2 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> a -> Tree a -> Tree a -> Tree a
forall a. Monoid a => Int -> a -> Tree a -> Tree a -> Tree a
bin Int
k a
a Tree a
t Tree a
t2) Path a
ts
cons Int
k a
a Path a
ts = a -> Int -> Int -> Tree a -> Path a -> Path a
forall a. a -> Int -> Int -> Tree a -> Path a -> Path a
Cons (a
a a -> a -> a
forall a. Monoid a => a -> a -> a
<> Path a -> a
forall m. Monoid m => Path m -> m
measure Path a
ts) (Path a -> Int
forall a. Path a -> Int
length Path a
ts Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Int
1 (Int -> a -> Tree a
forall a. Int -> a -> Tree a
Tip Int
k a
a) Path a
ts
{-# INLINE cons #-}

-- | /O(1)/ Extract the node ID and value from the newest node on the 'Path'.
uncons :: Monoid a => Path a -> Maybe (Int, a, Path a)
uncons :: Path a -> Maybe (Int, a, Path a)
uncons Path a
Nil = Maybe (Int, a, Path a)
forall a. Maybe a
Nothing
uncons (Cons a
_ Int
_ Int
_ (Tip Int
k a
a) Path a
ts) = (Int, a, Path a) -> Maybe (Int, a, Path a)
forall a. a -> Maybe a
Just (Int
k, a
a, Path a
ts)
uncons (Cons a
_ Int
_ Int
w (Bin a
_ Int
k a
a Tree a
l Tree a
r) Path a
ts) = (Int, a, Path a) -> Maybe (Int, a, Path a)
forall a. a -> Maybe a
Just (Int
k, a
a, Int -> Tree a -> Path a -> Path a
forall a. Monoid a => Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
l (Int -> Tree a -> Path a -> Path a
forall a. Monoid a => Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
r Path a
ts)) where w2 :: Int
w2 = Int -> Int -> Int
forall a. Integral a => a -> a -> a
div Int
w Int
2
{-# INLINE uncons #-}

-- | /O(1)/ Extract the node ID and value from the newest node on the 'Path', slightly faster than 'uncons'.
view :: Monoid a => Path a -> View Path a
view :: Path a -> View Path a
view Path a
Nil = View Path a
forall (f :: * -> *) a. View f a
Root
view (Cons a
_ Int
_ Int
_ (Tip Int
k a
a) Path a
ts) = Int -> a -> Path a -> View Path a
forall (f :: * -> *) a. Int -> a -> f a -> View f a
Node Int
k a
a Path a
ts
view (Cons a
_ Int
_ Int
w (Bin a
_ Int
k a
a Tree a
l Tree a
r) Path a
ts) = Int -> a -> Path a -> View Path a
forall (f :: * -> *) a. Int -> a -> f a -> View f a
Node Int
k a
a (Int -> Tree a -> Path a -> Path a
forall a. Monoid a => Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
l (Int -> Tree a -> Path a -> Path a
forall a. Monoid a => Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
r Path a
ts)) where w2 :: Int
w2 = Int -> Int -> Int
forall a. Integral a => a -> a -> a
div Int
w Int
2
{-# INLINE view #-}

-- | /O(log (h - k))/ to keep @k@ elements of 'Path' of 'length' @h@, and provide a monoidal summary of the dropped elements
-- using a supplied monoid homomorphism.
--
mkeep :: (Monoid a, Monoid b) => (a -> b) -> Int -> Path a -> (b, Path a)
mkeep :: (a -> b) -> Int -> Path a -> (b, Path a)
mkeep a -> b
f = b -> Int -> Path a -> (b, Path a)
go b
forall a. Monoid a => a
mempty where
  go :: b -> Int -> Path a -> (b, Path a)
go b
as Int
_ Path a
Nil = (b
as, Path a
forall a. Path a
Nil)
  go b
as Int
k xs :: Path a
xs@(Cons a
_ Int
n Int
w Tree a
t Path a
ts)
    | Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n    = (b
as, Path a
xs)
    | Bool
otherwise = case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
k (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
w) of
      Ordering
GT -> b -> Int -> Int -> Tree a -> Path a -> (b, Path a)
goT b
as (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
w) Int
w Tree a
t Path a
ts
      Ordering
EQ -> (b
as b -> b -> b
forall a. Monoid a => a -> a -> a
<> a -> b
f (Tree a -> a
forall a. Tree a -> a
measureT Tree a
t), Path a
ts)
      Ordering
LT -> b -> Int -> Path a -> (b, Path a)
go (b
as b -> b -> b
forall a. Monoid a => a -> a -> a
<> a -> b
f (Tree a -> a
forall a. Tree a -> a
measureT Tree a
t)) Int
k Path a
ts
  -- goT :: Monoid a => Int -> Int -> Tree a -> Path a -> Path a
  goT :: b -> Int -> Int -> Tree a -> Path a -> (b, Path a)
goT b
as Int
n Int
w (Bin a
_ Int
_ a
a Tree a
l Tree a
r) Path a
ts = case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
n Int
w2 of
    Ordering
LT              -> b -> Int -> Int -> Tree a -> Path a -> (b, Path a)
goT (b
as b -> b -> b
forall a. Monoid a => a -> a -> a
<> a -> b
f a
a b -> b -> b
forall a. Monoid a => a -> a -> a
<> a -> b
f (Tree a -> a
forall a. Tree a -> a
measureT Tree a
l)) Int
n Int
w2 Tree a
r Path a
ts
    Ordering
EQ              -> (b
as b -> b -> b
forall a. Monoid a => a -> a -> a
<> a -> b
f a
a b -> b -> b
forall a. Monoid a => a -> a -> a
<> a -> b
f (Tree a -> a
forall a. Tree a -> a
measureT Tree a
l), Int -> Tree a -> Path a -> Path a
forall a. Monoid a => Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
r Path a
ts)
    Ordering
GT | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
w Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 -> (b
as b -> b -> b
forall a. Monoid a => a -> a -> a
<> a -> b
f a
a, Int -> Tree a -> Path a -> Path a
forall a. Monoid a => Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
l (Int -> Tree a -> Path a -> Path a
forall a. Monoid a => Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
r Path a
ts))
       | Bool
otherwise  -> b -> Int -> Int -> Tree a -> Path a -> (b, Path a)
goT (b
as b -> b -> b
forall a. Monoid a => a -> a -> a
<> a -> b
f a
a) (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
w2) Int
w2 Tree a
l (Int -> Tree a -> Path a -> Path a
forall a. Monoid a => Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
r Path a
ts)
    where w2 :: Int
w2 = Int -> Int -> Int
forall a. Integral a => a -> a -> a
div Int
w Int
2
  goT b
as Int
_ Int
_ Tree a
_ Path a
ts = (b
as, Path a
ts)
{-# INLINE mkeep #-}

-- | /O(log (h - k))/ to @'keep' k@ elements of 'Path' of 'length' @h@
--
-- This solves the online version of the \"level ancestor problem\" with no preprocessing in /O(log h)/ time,
-- improving known complexity bounds.
--
-- <http://en.wikipedia.org/wiki/Level_ancestor_problem>
keep :: Monoid a => Int -> Path a -> Path a
keep :: Int -> Path a -> Path a
keep Int
k Path a
xs = ((), Path a) -> Path a
forall a b. (a, b) -> b
snd ((a -> ()) -> Int -> Path a -> ((), Path a)
forall a b.
(Monoid a, Monoid b) =>
(a -> b) -> Int -> Path a -> (b, Path a)
mkeep (\a
_ -> ()) Int
k Path a
xs)
{-# INLINE keep #-}

-- | /O(log k)/ to @'drop' k@ elements from a 'Path'
drop :: Monoid a => Int -> Path a -> Path a
drop :: Int -> Path a -> Path a
drop Int
k Path a
xs = ((), Path a) -> Path a
forall a b. (a, b) -> b
snd ((a -> ()) -> Int -> Path a -> ((), Path a)
forall a b.
(Monoid a, Monoid b) =>
(a -> b) -> Int -> Path a -> (b, Path a)
mdrop (\a
_ -> ()) Int
k Path a
xs)
{-# INLINE drop #-}

-- | /O(log k)/ to drop @k@ elements from a 'Path' and provide a monoidal summary of the dropped elements
-- using a suplied monoid homomorphism
mdrop :: (Monoid a, Monoid b) => (a -> b) -> Int -> Path a -> (b, Path a)
mdrop :: (a -> b) -> Int -> Path a -> (b, Path a)
mdrop a -> b
f Int
k Path a
xs = (a -> b) -> Int -> Path a -> (b, Path a)
forall a b.
(Monoid a, Monoid b) =>
(a -> b) -> Int -> Path a -> (b, Path a)
mkeep a -> b
f (Path a -> Int
forall a. Path a -> Int
length Path a
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
k) Path a
xs
{-# INLINE mdrop #-}

-- | /O(log h)/ @xs `'isAncestorOf'` ys@ holds when @xs@ is a prefix starting at the root of path @ys@.
isAncestorOf :: Monoid b => Path a -> Path b -> Bool
isAncestorOf :: Path a -> Path b -> Bool
isAncestorOf Path a
xs Path b
ys = Path a
xs Path a -> Path b -> Bool
forall a b. Path a -> Path b -> Bool
~= Int -> Path b -> Path b
forall a. Monoid a => Int -> Path a -> Path a
keep (Path a -> Int
forall a. Path a -> Int
length Path a
xs) Path b
ys

infix 4 ~=
-- | /O(1)/ Compare to see if two trees have the same root key
(~=) :: Path a -> Path b -> Bool
Path a
Nil            ~= :: Path a -> Path b -> Bool
~= Path b
Nil            = Bool
True
Cons a
_ Int
_ Int
_ Tree a
s Path a
_ ~= Cons b
_ Int
_ Int
_ Tree b
t Path b
_ = Tree a -> Tree b -> Bool
forall a b. Tree a -> Tree b -> Bool
sameT Tree a
s Tree b
t
Path a
_              ~= Path b
_              = Bool
False

-- | /O(log h)/ Compute the lowest common ancestor of two paths
--
-- >>> let fromList' = fromList . fmap (flip (,) ())
-- >>> length (lca (fromList' [1, 2, 3, 4, 5, 6]) (fromList' [7, 8, 3, 4, 5, 6]))
-- 4
--
lca :: (Monoid a, Monoid b) => Path a -> Path b -> Path a
lca :: Path a -> Path b -> Path a
lca Path a
xs Path b
ys = Path a
zs where (()
_, Path a
zs, ()
_, Path b
_) = (a -> ())
-> (b -> ()) -> Path a -> Path b -> ((), Path a, (), Path b)
forall a b c d.
(Monoid a, Monoid b, Monoid c, Monoid d) =>
(a -> c) -> (b -> d) -> Path a -> Path b -> (c, Path a, d, Path b)
mlca (\a
_ -> ()) (\b
_ -> ()) Path a
xs Path b
ys

-- | /O(log h)/ Compute the lowest common ancestor of two paths along with a monoidal summary of their respective tails using
-- the supplied monoid homomorphisms.
mlca :: (Monoid a, Monoid b, Monoid c, Monoid d) => (a -> c) -> (b -> d) -> Path a -> Path b -> (c, Path a, d, Path b)
mlca :: (a -> c) -> (b -> d) -> Path a -> Path b -> (c, Path a, d, Path b)
mlca a -> c
f b -> d
g Path a
xs0 Path b
ys0 = case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
nxs Int
nys of
  Ordering
LT -> let (d
bs, Path b
ys) = (b -> d) -> Int -> Path b -> (d, Path b)
forall a b.
(Monoid a, Monoid b) =>
(a -> b) -> Int -> Path a -> (b, Path a)
mkeep b -> d
g Int
nxs Path b
ys0 in c -> d -> Path a -> Path b -> (c, Path a, d, Path b)
go c
forall a. Monoid a => a
mempty d
bs Path a
xs0 Path b
ys
  Ordering
EQ -> c -> d -> Path a -> Path b -> (c, Path a, d, Path b)
go c
forall a. Monoid a => a
mempty d
forall a. Monoid a => a
mempty Path a
xs0 Path b
ys0
  Ordering
GT -> let (c
as, Path a
xs) = (a -> c) -> Int -> Path a -> (c, Path a)
forall a b.
(Monoid a, Monoid b) =>
(a -> b) -> Int -> Path a -> (b, Path a)
mkeep a -> c
f Int
nys Path a
xs0 in c -> d -> Path a -> Path b -> (c, Path a, d, Path b)
go c
as d
forall a. Monoid a => a
mempty Path a
xs Path b
ys0
  where
    nxs :: Int
nxs = Path a -> Int
forall a. Path a -> Int
length Path a
xs0
    nys :: Int
nys = Path b -> Int
forall a. Path a -> Int
length Path b
ys0

    go :: c -> d -> Path a -> Path b -> (c, Path a, d, Path b)
go c
as d
bs pa :: Path a
pa@(Cons a
_ Int
_ Int
w Tree a
x Path a
xs) pb :: Path b
pb@(Cons b
_ Int
_ Int
_ Tree b
y Path b
ys)
      | Tree a -> Tree b -> Bool
forall a b. Tree a -> Tree b -> Bool
sameT Tree a
x Tree b
y = (c
as, Path a
pa, d
bs, Path b
pb)
      | Path a
xs Path a -> Path b -> Bool
forall a b. Path a -> Path b -> Bool
~= Path b
ys  = c
-> d
-> Int
-> Tree a
-> Tree b
-> Path a
-> Path b
-> (c, Path a, d, Path b)
goT c
as d
bs Int
w Tree a
x Tree b
y Path a
xs Path b
ys
      | Bool
otherwise = c -> d -> Path a -> Path b -> (c, Path a, d, Path b)
go (c
as c -> c -> c
forall a. Monoid a => a -> a -> a
<> a -> c
f (Tree a -> a
forall a. Tree a -> a
measureT Tree a
x)) (d
bs d -> d -> d
forall a. Monoid a => a -> a -> a
<> b -> d
g (Tree b -> b
forall a. Tree a -> a
measureT Tree b
y)) Path a
xs Path b
ys
    go c
as d
bs Path a
_ Path b
_ = (c
as, Path a
forall a. Path a
Nil, d
bs, Path b
forall a. Path a
Nil)

    goT :: c
-> d
-> Int
-> Tree a
-> Tree b
-> Path a
-> Path b
-> (c, Path a, d, Path b)
goT c
as d
bs Int
w (Bin a
_ Int
_ a
a Tree a
la Tree a
ra) (Bin b
_ Int
_ b
b Tree b
lb Tree b
rb) Path a
pa Path b
pb
      | Tree a -> Tree b -> Bool
forall a b. Tree a -> Tree b -> Bool
sameT Tree a
la Tree b
lb = (c
as c -> c -> c
forall a. Monoid a => a -> a -> a
<> a -> c
f a
a, Int -> Tree a -> Path a -> Path a
forall a. Monoid a => Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
la (Int -> Tree a -> Path a -> Path a
forall a. Monoid a => Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
ra Path a
pa), d
bs d -> d -> d
forall a. Monoid a => a -> a -> a
<> b -> d
g b
b, Int -> Tree b -> Path b -> Path b
forall a. Monoid a => Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree b
lb (Int -> Tree b -> Path b -> Path b
forall a. Monoid a => Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree b
rb Path b
pb))
      | Tree a -> Tree b -> Bool
forall a b. Tree a -> Tree b -> Bool
sameT Tree a
ra Tree b
rb = c
-> d
-> Int
-> Tree a
-> Tree b
-> Path a
-> Path b
-> (c, Path a, d, Path b)
goT (c
as c -> c -> c
forall a. Monoid a => a -> a -> a
<> a -> c
f a
a) (d
bs d -> d -> d
forall a. Monoid a => a -> a -> a
<> b -> d
g b
b) Int
w2 Tree a
la Tree b
lb (Int -> Tree a -> Path a -> Path a
forall a. Monoid a => Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree a
ra Path a
pa) (Int -> Tree b -> Path b -> Path b
forall a. Monoid a => Int -> Tree a -> Path a -> Path a
consT Int
w2 Tree b
rb Path b
pb)
      | Bool
otherwise   = c
-> d
-> Int
-> Tree a
-> Tree b
-> Path a
-> Path b
-> (c, Path a, d, Path b)
goT (c
as c -> c -> c
forall a. Monoid a => a -> a -> a
<> a -> c
f a
a c -> c -> c
forall a. Monoid a => a -> a -> a
<> a -> c
f (Tree a -> a
forall a. Tree a -> a
measureT Tree a
la)) (d
bs d -> d -> d
forall a. Monoid a => a -> a -> a
<> b -> d
g b
b d -> d -> d
forall a. Monoid a => a -> a -> a
<> b -> d
g (Tree b -> b
forall a. Tree a -> a
measureT Tree b
lb)) Int
w2 Tree a
ra Tree b
rb Path a
pa Path b
pb
      where w2 :: Int
w2 = Int -> Int -> Int
forall a. Integral a => a -> a -> a
div Int
w Int
2
    goT c
as d
bs Int
_ Tree a
ta Tree b
tb Path a
pa Path b
pb = (c
as c -> c -> c
forall a. Monoid a => a -> a -> a
<> a -> c
f (Tree a -> a
forall a. Tree a -> a
measureT Tree a
ta), Path a
pa, d
bs d -> d -> d
forall a. Monoid a => a -> a -> a
<> b -> d
g (Tree b -> b
forall a. Tree a -> a
measureT Tree b
tb), Path b
pb)
{-# INLINE mlca #-}