Quick implementation of a segment tree with fast range mconcat

FRM.hs

{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE NoMonoLocalBinds    #-}
{-# LANGUAGE RankNTypes          #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE ViewPatterns        #-}
module FRM
       ( FastRangeMconcat
       , fromVector
       , toVector
       , (!), (!?)
       , mconcatRange, mconcatRange'
       , update
       ) where

import qualified Control.Monad               as CM
import           Control.Monad.ST
import           Data.Bits
import           Data.Foldable
import           Data.Function
import           Data.Maybe
import           Data.Monoid
import qualified Data.Vector                 as V
import qualified Data.Vector.Generic.Mutable as VGM
import qualified Data.Vector.Mutable         as VM

newtype FastRangeMconcat a = FRM (V.Vector a)

instance Eq a => Eq (FastRangeMconcat a) where
  (==) = (==) `on` toVector

instance Ord a => Ord (FastRangeMconcat a) where
  compare = compare `on` toVector

instance Show a => Show (FastRangeMconcat a) where
  show = show . toVector

instance Foldable FastRangeMconcat where
  fold v = mconcatRange 0 (length v) v -- O(1)
  null (FRM v) = V.null v -- O(1)
  length (FRM v) = V.length v `div` 2 -- O(1)

  foldMap f = foldMap f . toVector
  foldr f i = foldr f i . toVector
  foldr' f i = foldr' f i . toVector
  foldl f i = foldl f i . toVector
  foldl' f i = foldl' f i . toVector
  foldr1 f = foldr1 f . toVector
  foldl1 f = foldl1 f . toVector
  toList = V.toList . toVector
  elem a = elem a . toVector
  maximum = maximum . toVector
  minimum = minimum . toVector
  sum = sum . toVector
  product = product . toVector


fromVector :: (Monoid a) => V.Vector a -> FastRangeMconcat a
fromVector arr = FRM $ V.constructrN (n * 2) go
  where n = V.length arr
        go v | V.length v < n     = arr V.! (n - 1 - V.length v)
             | V.length v < n * 2 = let i = n * 2 - V.length v - 1
                                        lc = i * 2
                                        rc = lc + 1
                                        lc' = lc - (n * 2 - V.length v)
                                        rc' = rc - (n * 2 - V.length v)
                                    in v V.! lc' <> v V.! rc'
             | otherwise          = mempty

fromVector' :: forall a. (Monoid a) => V.Vector a -> FastRangeMconcat a
fromVector' arr = FRM (V.create act)
  where act :: ST s (VM.MVector s a)
        act = do
          arr' <- V.thaw arr
          tree <- VGM.growFront arr' (V.length arr)
          build tree
            where build tree = loop (V.length arr - 1) >> return tree
                    where loop !i = CM.when (i > 0) $ do
                            lc <- VGM.read tree (i * 2)
                            rc <- VGM.read tree (i * 2 + 1)
                            VGM.write tree i (lc <> rc)
                            loop (i - 1)

toVector :: FastRangeMconcat a -> V.Vector a
toVector (FRM v) = V.slice n n v
  where n = V.length v `div` 2
{-# INLINE toVector #-}

(!) :: FastRangeMconcat a -> Int -> a
(!) t i = fromMaybe (error "index out of range") (t !? i)
{-# INLINE (!) #-}

(!?) :: FastRangeMconcat a -> Int -> Maybe a
(FRM v) !? i
  | 0 <= i && i < n = Just (v V.! (i + n))
  | otherwise       = Nothing
  where n = V.length v `div` 2
{-# INLINE (!?) #-}

mconcatRange :: (Monoid a) => Int -> Int -> FastRangeMconcat a -> a
mconcatRange l r t = fromMaybe (error "index out of range") (mconcatRange' l r t)
{-# INLINE mconcatRange #-}

mconcatRange' :: (Monoid a) => Int -> Int -> FastRangeMconcat a -> Maybe a
mconcatRange' l r (FRM v)
  | 0 <= l && l < n && 0 <= r && r < n = mconcatRangeImpl v (l + n) (r + n) mempty mempty
  | otherwise = Nothing
  where n = V.length v `div` 2
{-# INLINE mconcatRange' #-}

mconcatRangeImpl :: (Monoid a) => V.Vector a -> Int -> Int -> a -> a -> Maybe a
mconcatRangeImpl !v = loop
  where loop !l !r !rvL !rvR =
          if l < r then do
            let (!l', !rvL') = if testBit l 0 then ((l + 1) `div` 2, rvL <> (v V.! l))       else (l `div` 2, rvL)
            let (!r', !rvR') = if testBit r 0 then ((r - 1) `div` 2, (v V.! (r - 1)) <> rvR) else (r `div` 2, rvR)
            loop l' r' rvL' rvR'
          else return $! (rvL <> rvR)
{-# INLINE mconcatRangeImpl #-}

update :: (Monoid a) => Int -> a -> FastRangeMconcat a -> FastRangeMconcat a
update i a (FRM v) | 0 <= i && i < n = FRM (V.modify action v)
                   | otherwise = FRM v
  where n = V.length v `div` 2
        action tree = do
          VGM.write tree (i + n) a
          loop (i + n)
            where loop !p = CM.when (p > 1) $ do
                    lc <- VGM.read tree p
                    rc <- VGM.read tree (p `xor` 1)
                    VGM.write tree (p `div` 2) (lc <> rc)
                    loop (p `div` 2)
{-# INLINE update #-}

----------------------------
-- Properties for Testing --
----------------------------

prop_canRoundTripToVectorFromVector :: [String] -> Bool
prop_canRoundTripToVectorFromVector (V.fromList -> v) = v == toVector (fromVector v)

prop_constructrNSameAsMutable :: [String] -> Bool
prop_constructrNSameAsMutable (V.fromList -> v) = fromVector v == fromVector' v

prop_rangeMconcatSameAsSlow :: [String] -> Int -> Int -> Bool
prop_rangeMconcatSameAsSlow (V.fromList -> v) i j = let t = fromVector v in naiveMconcatRange i j == mconcatRange' i j t
  where naiveMconcatRange l r | 0 <= l && l < n && 0 <= r && r < n && l < r = Just (mconcat (V.toList (V.slice i (j - i) v)))
                              | 0 <= l && l < n && 0 <= r && r < n = Just mempty
                              | otherwise = Nothing
          where n = V.length v

prop_canUpdateFRM :: [String] -> Int -> String -> Bool
prop_canUpdateFRM (V.fromList -> v) i a = let t = fromVector v in fromVector vecUpdated == update i a t
  where vecUpdated | 0 <= i && i < n = v V.// [(i, a)]
                   | otherwise = v
          where n = V.length v