module Data.RangeMin.Cartesian (buildDepths) where
import Data.RangeMin.Common.Types
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Primitive as PV
import qualified Data.Vector.Fusion.Stream.Monadic as SM
import Data.Vector.Fusion.Stream (Step (..), Stream)
import qualified Data.Vector.Fusion.Stream as S
import Data.RangeMin.Common.Vector
import Data.RangeMin.Cartesian.STInt
import Prelude hiding (read)
type CartesianTree = PV.Vector Int
data IL = IL !Int IL | Nil
mapAccumSM :: Monad m => (b -> a -> m (c, b)) -> b -> SM.Stream m a -> SM.Stream m c
mapAccumSM f z0 (SM.Stream suc s0 n) = SM.Stream suc' (z0, s0) n where
suc' (z, s) = do
step <- suc s
case step of
Done -> return Done
Skip s' -> return (Skip (z, s'))
Yield x s' -> do
(y, z') <- f z x
return (Yield y (z', s'))
mapAccumS :: (b -> a -> (c, b)) -> b -> Stream a -> Stream c
mapAccumS f = mapAccumSM (\ b a -> return (f b a))
buildDepths :: Vector v a => LEq a -> v a -> PV.Vector Int
buildDepths (<=?) xs = makeDepths (makeTree (<=?) (G.length xs) (xs !))
makeDepths :: CartesianTree -> PV.Vector Int
makeDepths !parents = inlineCreate $ do
let !n = PV.length parents
!dest <- newWith n (1)
let depth !i = toSTInt $ do
d0 <- read dest i
case d0 of
1 -> case parents ! i of
1 -> do
write dest i 0
return 0
p -> do
!d' <- runSTInt (depth p)
let !d = d' + 1
write dest i d
return d
_ -> return d0
mapM_ (runSTInt . depth) [0..n1]
return dest
makeTree :: LEq a -> Int -> (Int -> a) -> CartesianTree
makeTree (<=?) !n look = inlineCreate $ do
!dest <- new n
S.mapM_ (\ (IP i j) -> write dest i j) (neighbors $
S.unfoldr (\ !i -> if i == 0 then Nothing else let
!i' = i 1
xi' = look i'
in Just ((i', xi'), i')) n)
let parent (IP i l) = do
r <- read dest i
write dest i $ case (l, r) of
(1, 1) -> 1
(1, r) -> r
(l, 1) -> l
(l, r) -> if look l <=? look r then r else l
S.mapM_ parent (neighbors (S.generate n (\ !i -> (i, look i))))
return dest
where
neighbors xs = mapAccumS suc Nil xs
where suc stk (!i, xi) = let
run Nil = (IP i (1), IL i Nil)
run stk0@(IL j stk)
| not (xi <=? look j)
= (IP i j, IL i stk0)
| otherwise
= run stk
in run stk