Ticket #3181: StrictPair.hs

{-# LANGUAGE TypeOperators #-}

module Main (main) where

import Control.Monad.ST

import Data.Array.Vector

siftByOffset :: (UA e) => (e -> e -> Ordering) -> MUArr e s -> e -> Int -> Int -> Int -> ST s ()
siftByOffset cmp a val off start len = sift val start len
 where
 sift val root len
   | child < len = do (child' :*: ac) <- maximumChild cmp a off child len
                      case cmp val ac of
                        LT -> writeMU a (root + off) ac >> sift val child' len
                        _  -> writeMU a (root + off) val
   | otherwise = writeMU a (root + off) val
  where child = root * 3 + 1
{-# INLINE siftByOffset #-}

maximumChild :: (UA e) => (e -> e -> Ordering) -> MUArr e s -> Int -> Int -> Int -> ST s (Int :*: e)
maximumChild cmp a off child1 len
  | child3 < len = do ac1 <- readMU a (child1 + off)
                      ac2 <- readMU a (child2 + off)
                      ac3 <- readMU a (child3 + off)
                      return $ case cmp ac1 ac2 of
                                 LT -> case cmp ac2 ac3 of
                                         LT -> child3 :*: ac3
                                         _  -> child2 :*: ac2
                                 _  -> case cmp ac1 ac3 of
                                         LT -> child3 :*: ac3
                                         _  -> child1 :*: ac1
  | child2 < len = do ac1 <- readMU a (child1 + off)
                      ac2 <- readMU a (child2 + off)
                      return $ case cmp ac1 ac2 of
                                 LT -> child2 :*: ac2
                                 _  -> child1 :*: ac1
  | otherwise    = do ac1 <- readMU a (child1 + off) ; return (child1 :*: ac1)
 where
 child2 = child1 + 1
 child3 = child1 + 2
{-# INLINE maximumChild #-}

test :: MUArr Int s -> ST s ()
test arr = siftByOffset compare arr len len len len
 where len = lengthMU arr

main = stToIO (newMU 40 >>= test)