{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnboxedTuples #-}
module Std.Data.Vector.Extra (
cons, snoc
, uncons, unsnoc
, headMaybe, tailMayEmpty
, lastMaybe, initMayEmpty
, inits, tails
, take, drop, takeR, dropR
, slice
, splitAt
, takeWhile, takeWhileR, dropWhile, dropWhileR, dropAround
, break, span, breakR, spanR, breakOn
, group, groupBy
, stripPrefix, stripSuffix
, split, splitWith, splitOn
, isPrefixOf, isSuffixOf, isInfixOf
, commonPrefix
, words, lines, unwords, unlines
, padLeft, padRight
, reverse
, intersperse
, intercalate
, intercalateElem
, transpose
, zipWith', unzipWith'
, scanl', scanl1'
, scanr', scanr1'
, rangeCut
, head
, tail
, init
, last
, index
, unsafeHead
, unsafeTail
, unsafeInit
, unsafeLast
, unsafeIndex
, unsafeTake
, unsafeDrop
) where
import Control.Monad.ST
import GHC.Stack
import GHC.Word
import Std.Data.Array
import Std.Data.Vector.Base
import Std.Data.Vector.Search
import Prelude hiding (concat, concatMap,
elem, notElem, null, length, map,
foldl, foldl1, foldr, foldr1,
maximum, minimum, product, sum,
all, any, replicate, traverse,
head, tail, init, last,
take, drop, splitAt,
takeWhile, dropWhile,
break, span, reverse,
words, lines, unwords, unlines)
import qualified Data.List as List
import Data.Bits
import Control.Exception (assert)
cons :: Vec v a => a -> v a -> v a
{-# INLINE cons #-}
cons x (Vec arr s l) = create (l+1) $ \ marr ->
writeArr marr 0 x >> copyArr marr 1 arr s l
snoc :: Vec v a => v a -> a -> v a
{-# INLINE snoc #-}
snoc (Vec arr s l) x = create (l+1) $ \ marr ->
copyArr marr 0 arr s l >> writeArr marr l x
uncons :: Vec v a => v a -> Maybe (a, v a)
{-# INLINE uncons #-}
uncons (Vec arr s l)
| l <= 0 = Nothing
| otherwise = let !v = fromArr arr (s+1) (l-1)
in case indexArr' arr s of (# x #) -> Just (x ,v)
unsnoc :: Vec v a => v a -> Maybe (v a, a)
{-# INLINE unsnoc #-}
unsnoc (Vec arr s l)
| l <= 0 = Nothing
| otherwise = let !v = fromArr arr s (l-1)
in case indexArr' arr (s+l-1) of (# x #) -> Just (v, x)
headMaybe :: Vec v a => v a -> Maybe a
{-# INLINE headMaybe #-}
headMaybe (Vec arr s l)
| l <= 0 = Nothing
| otherwise = indexArrM arr s
tailMayEmpty :: Vec v a => v a -> v a
{-# INLINE tailMayEmpty #-}
tailMayEmpty (Vec arr s l)
| l <= 0 = empty
| otherwise = fromArr arr (s+1) (l-1)
lastMaybe :: Vec v a => v a -> Maybe a
{-# INLINE lastMaybe #-}
lastMaybe (Vec arr s l)
| l <= 0 = Nothing
| otherwise = indexArrM arr (s+l-1)
initMayEmpty :: Vec v a => v a -> v a
{-# INLINE initMayEmpty #-}
initMayEmpty (Vec arr s l)
| l <= 0 = empty
| otherwise = fromArr arr s (l-1)
inits :: Vec v a => v a -> [v a]
{-# INLINE inits #-}
inits (Vec arr s l) = [Vec arr s n | n <- [0..l]]
tails :: Vec v a => v a -> [v a]
{-# INLINE tails #-}
tails (Vec arr s l) = [Vec arr (s+n) (l-n) | n <- [0..l]]
take :: Vec v a => Int -> v a -> v a
{-# INLINE take #-}
take n v@(Vec arr s l)
| n <= 0 = empty
| n >= l = v
| otherwise = fromArr arr s n
takeR :: Vec v a => Int -> v a -> v a
{-# INLINE takeR #-}
takeR n v@(Vec arr s l)
| n <= 0 = empty
| n >= l = v
| otherwise = fromArr arr (s+l-n) n
drop :: Vec v a => Int -> v a -> v a
{-# INLINE drop #-}
drop n v@(Vec arr s l)
| n <= 0 = v
| n >= l = empty
| otherwise = fromArr arr (s+n) (l-n)
dropR :: Vec v a => Int -> v a -> v a
{-# INLINE dropR #-}
dropR n v@(Vec arr s l)
| n <= 0 = v
| n >= l = empty
| otherwise = fromArr arr s (l-n)
slice :: Vec v a => Int
-> Int
-> v a -> v a
{-# INLINE slice #-}
slice x y = drop x . take (x+y)
splitAt :: Vec v a => Int -> v a -> (v a, v a)
{-# INLINE splitAt #-}
splitAt z (Vec arr s l) = let !v1 = fromArr arr s z'
!v2 = fromArr arr (s+z') (l-z')
in (v1, v2)
where z' = rangeCut z 0 l
takeWhile :: Vec v a => (a -> Bool) -> v a -> v a
{-# INLINE takeWhile #-}
takeWhile f v@(Vec arr s l) =
case findIndex (not . f) v of
0 -> empty
i -> Vec arr s i
takeWhileR :: Vec v a => (a -> Bool) -> v a -> v a
{-# INLINE takeWhileR #-}
takeWhileR f v@(Vec arr s l) =
case findIndexR (not . f) v of
-1 -> v
i -> Vec arr (s+i+1) (l-i-1)
dropWhile :: Vec v a => (a -> Bool) -> v a -> v a
{-# INLINE dropWhile #-}
dropWhile f v@(Vec arr s l) =
case findIndex (not . f) v of
i | i == l -> empty
| otherwise -> Vec arr (s+i) (l-i)
dropWhileR :: Vec v a => (a -> Bool) -> v a -> v a
{-# INLINE dropWhileR #-}
dropWhileR f v@(Vec arr s l) =
case findIndexR (not . f) v of
-1 -> empty
i -> Vec arr s (i+1)
dropAround :: Vec v a => (a -> Bool) -> v a -> v a
{-# INLINE dropAround #-}
dropAround f = dropWhile f . dropWhileR f
break :: Vec v a => (a -> Bool) -> v a -> (v a, v a)
{-# INLINE break #-}
break f vs@(Vec arr s l) =
let !n = findIndex f vs
!v1 = Vec arr s n
!v2 = Vec arr (s+n) (l-n)
in (v1, v2)
span :: Vec v a => (a -> Bool) -> v a -> (v a, v a)
{-# INLINE span #-}
span f = break (not . f)
breakR :: Vec v a => (a -> Bool) -> v a -> (v a, v a)
{-# INLINE breakR #-}
breakR f vs@(Vec arr s l) =
let !n = findIndexR f vs
!v1 = Vec arr s (n+1)
!v2 = Vec arr (s+n+1) (l-1-n)
in (v1, v2)
spanR :: Vec v a => (a -> Bool) -> v a -> (v a, v a)
{-# INLINE spanR #-}
spanR f = breakR (not . f)
breakOn :: (Vec v a, Eq a) => v a -> v a -> (v a, v a)
{-# INLINE breakOn #-}
breakOn needle = \ haystack@(Vec arr s l) ->
case search haystack False of
(i:_) -> let !v1 = Vec arr s i
!v2 = Vec arr (s+i) (l-i)
in (v1, v2)
_ -> (haystack, empty)
where search = indices needle
group :: (Vec v a, Eq a) => v a -> [v a]
{-# INLINE group #-}
group = groupBy (==)
groupBy :: Vec v a => (a -> a -> Bool) -> v a -> [v a]
{-# INLINE groupBy #-}
groupBy f (Vec arr s l)
| l == 0 = []
| otherwise = Vec arr s n : groupBy f (Vec arr (s+n) (l-n))
where
n = case indexArr' arr s of
(# x #) -> 1 + findIndex (not . f x) (Vec arr (s+1) (l-1))
stripPrefix :: (Vec v a, Eq (v a))
=> v a
-> v a -> Maybe (v a)
{-# INLINE stripPrefix #-}
stripPrefix v1@(Vec _ _ l1) v2@(Vec arr s l2)
| v1 `isPrefixOf` v2 = Just (Vec arr (s+l1) (l2-l1))
| otherwise = Nothing
isPrefixOf :: (Vec v a, Eq (v a))
=> v a
-> v a -> Bool
{-# INLINE isPrefixOf #-}
isPrefixOf (Vec arrA sA lA) (Vec arrB sB lB)
| lA == 0 = True
| lA > lB = False
| otherwise = Vec arrA sA lA == Vec arrB sB lA
commonPrefix :: (Vec v a, Eq a) => v a -> v a -> (v a, v a, v a)
{-# INLINE commonPrefix #-}
commonPrefix vA@(Vec arrA sA lA) vB@(Vec arrB sB lB) = go sA sB
where
!endA = sA + lA
!endB = sB + lB
go !i !j | i >= endA = let !vB' = fromArr arrB (sB+i-sA) (lB-i+sA) in (vA, empty, vB')
| j >= endB = let !vA' = fromArr arrA (sA+j-sB) (lA-j+sB) in (vB, vA', empty)
| indexArr arrA i == indexArr arrB j = go (i+1) (j+1)
| otherwise =
let !vB' = fromArr arrB (sB+i-sA) (lB-i+sA)
!vA' = fromArr arrA (sA+j-sB) (lA-j+sB)
!vC = fromArr arrA sA (i-sA)
in (vC, vA', vB')
stripSuffix :: (Vec v a, Eq (v a)) => v a -> v a -> Maybe (v a)
{-# INLINE stripSuffix #-}
stripSuffix v1@(Vec _ _ l1) v2@(Vec arr s l2)
| v1 `isSuffixOf` v2 = Just (Vec arr s (l2-l1))
| otherwise = Nothing
isSuffixOf :: (Vec v a, Eq (v a)) => v a -> v a -> Bool
{-# INLINE isSuffixOf #-}
isSuffixOf (Vec arrA sA lA) (Vec arrB sB lB)
| lA == 0 = True
| lA > lB = False
| otherwise = Vec arrA sA lA == Vec arrB (sB+lB-lA) lA
isInfixOf :: (Vec v a, Eq a) => v a -> v a -> Bool
{-# INLINE isInfixOf #-}
isInfixOf needle = \ haystack -> null haystack || search haystack False /= []
where search = indices needle
split :: (Vec v a, Eq a) => a -> v a -> [v a]
{-# INLINE split #-}
split x = splitWith (==x)
splitOn :: (Vec v a, Eq a) => v a -> v a -> [v a]
{-# INLINE splitOn #-}
splitOn needle = splitBySearch
where
splitBySearch haystack@(Vec arr s l) = go s (search haystack False)
where
!l' = length needle
!end = s+l
search = indices needle
go !s' (i:is) = let !v = fromArr arr s' (i+s-s')
in v : go (i+l') is
go !s' _ = let !v = fromArr arr s' (end-s') in [v]
splitWith :: Vec v a => (a -> Bool) -> v a -> [v a]
{-# INLINE splitWith #-}
splitWith f (Vec arr s l) = go s s
where
!end = s + l
go !p !q | q >= end = let !v = Vec arr p (q-p) in [v]
| f x = let !v = Vec arr p (q-p) in v:go (q+1) (q+1)
| otherwise = go p (q+1)
where (# x #) = indexArr' arr q
words :: Bytes -> [Bytes]
{-# INLINE words #-}
words (Vec arr s l) = go s s
where
!end = s + l
go !s' !i | i >= end =
if s' == end
then []
else let !v = fromArr arr s' (end-s') in [v]
| isASCIISpace (indexArr arr i) =
if s' == i
then go (i+1) (i+1)
else
let !v = fromArr arr s' (i-s') in v : go (i+1) (i+1)
| otherwise = go s' (i+1)
lines :: Bytes -> [Bytes]
{-# INLINE lines #-}
lines (Vec arr s l) = go s s
where
!end = s + l
go !p !q | q >= end = if p == q
then []
else let !v = Vec arr p (q-p) in [v]
| indexArr arr q == 10 = let !v = Vec arr p (q-p) in v:go (q+1) (q+1)
| otherwise = go p (q+1)
unwords :: [Bytes] -> Bytes
{-# INLINE unwords #-}
unwords = intercalateElem 32
unlines :: [Bytes] -> Bytes
{-# INLINE unlines #-}
unlines [] = empty
unlines vs = create (len vs 0) (copy 0 vs)
where
len [] !acc = acc
len (Vec _ _ l:vs) !acc = len vs (acc+l+1)
copy !i [] !marr = return ()
copy !i (Vec arr s l:vs) !marr = do
let !i' = i + l
copyArr marr i arr s l
writeArr marr i' 10
copy (i'+1) vs marr
padLeft :: Vec v a => Int -> a -> v a -> v a
{-# INLINE padLeft #-}
padLeft n x v@(Vec arr s l) | n <= l = v
| otherwise = create n (\ marr -> do
setArr marr 0 (n-l) x
copyArr marr (n-l) arr s l)
padRight :: Vec v a => Int -> a -> v a -> v a
{-# INLINE padRight #-}
padRight n x v@(Vec arr s l) | n <= l = v
| otherwise = create n (\ marr -> do
copyArr marr 0 arr s l
setArr marr l (n-l) x)
reverse :: forall v a. (Vec v a) => v a -> v a
{-# INLINE reverse #-}
reverse (Vec arr s l) = create l (go s (l-1))
where
go :: Int -> Int -> MArray v s a -> ST s ()
go !i !j !marr | j < 0 = return ()
| j >= 3 = do
indexArrM arr i >>= writeArr marr j
indexArrM arr (i+1) >>= writeArr marr (j-1)
indexArrM arr (i+2) >>= writeArr marr (j-2)
indexArrM arr (i+3) >>= writeArr marr (j-3)
go (i+4) (j-4) marr
| otherwise = do
indexArrM arr i >>= writeArr marr j
go (i+1) (j-1) marr
intersperse :: forall v a. Vec v a => a -> v a -> v a
{-# INLINE intersperse #-}
intersperse x (Vec _ _ 0) = empty
intersperse x v@(Vec _ _ 1) = v
intersperse x (Vec arr s l) = create (2*l-1) (go s 0)
where
!end = s+l-1
go :: Int
-> Int
-> MArray v s a
-> ST s ()
go !i !j !marr
| i >= end = writeArr marr j =<< indexArrM arr i
| i <= end - 4 = do
writeArr marr j =<< indexArrM arr i
writeArr marr (j+1) x
writeArr marr (j+2) =<< indexArrM arr (i+1)
writeArr marr (j+3) x
writeArr marr (j+4) =<< indexArrM arr (i+2)
writeArr marr (j+5) x
writeArr marr (j+6) =<< indexArrM arr (i+3)
writeArr marr (j+7) x
go (i+4) (j+8) marr
| otherwise = do
writeArr marr j =<< indexArrM arr i
writeArr marr (j+1) x
go (i+1) (j+2) marr
intercalate :: Vec v a => v a -> [v a] -> v a
{-# INLINE intercalate #-}
intercalate s = concat . List.intersperse s
intercalateElem :: Vec v a => a -> [v a] -> v a
{-# INLINE intercalateElem #-}
intercalateElem _ [] = empty
intercalateElem _ [v] = v
intercalateElem w vs = create (len vs 0) (copy 0 vs)
where
len [] !acc = acc
len [Vec _ _ l] !acc = l + acc
len (Vec _ _ l:vs) !acc = len vs (acc+l+1)
copy !i [] !marr = return ()
copy !i (Vec arr s l:[]) !marr = copyArr marr i arr s l
copy !i (Vec arr s l:vs) !marr = do
let !i' = i + l
copyArr marr i arr s l
writeArr marr i' w
copy (i'+1) vs marr
transpose :: Vec v a => [v a] -> [v a]
{-# INLINE transpose #-}
transpose vs =
List.map (packN n) . List.transpose . List.map unpack $ vs
where n = List.length vs
zipWith' :: (Vec v a, Vec u b, Vec w c)
=> (a -> b -> c) -> v a -> u b -> w c
{-# INLINE zipWith' #-}
zipWith' f (Vec arrA sA lA) (Vec arrB sB lB) = create len (go 0)
where
!len = min lA lB
go !i !marr
| i >= len = return ()
| otherwise = case indexArr' arrA (i+sA) of
(# a #) -> case indexArr' arrB (i+sB) of
(# b #) -> do let !c = f a b in writeArr marr i c
go (i+1) marr
unzipWith' :: (Vec v a, Vec u b, Vec w c)
=> (a -> (b, c)) -> v a -> (u b, w c)
{-# INLINE unzipWith' #-}
unzipWith' f (Vec arr s l) = createN2 l l (go 0)
where
go !i !marrB !marrC
| i >= l = return (l,l)
| otherwise = case indexArr' arr (i+s) of
(# a #) -> do let (!b, !c) = f a
writeArr marrB i b
writeArr marrC i c
go (i+1) marrB marrC
scanl' :: forall v u a b. (Vec v a, Vec u b) => (b -> a -> b) -> b -> v a -> u b
{-# INLINE scanl' #-}
scanl' f z (Vec arr s l) =
create (l+1) (\ marr -> writeArr marr 0 z >> go z s 1 marr)
where
go :: b -> Int -> Int -> MArray u s b -> ST s ()
go !acc !i !j !marr
| j > l = return ()
| otherwise = do
x <- indexArrM arr i
let !acc' = acc `f` x
writeArr marr j acc'
go acc' (i+1) (j+1) marr
scanl1' :: forall v a. Vec v a => (a -> a -> a) -> v a -> v a
{-# INLINE scanl1' #-}
scanl1' f (Vec arr s l)
| l <= 0 = empty
| otherwise = case indexArr' arr s of
(# x0 #) -> scanl' f x0 (fromArr arr (s+1) (l-1) :: v a)
scanr' :: forall v u a b. (Vec v a, Vec u b) => (a -> b -> b) -> b -> v a -> u b
{-# INLINE scanr' #-}
scanr' f z (Vec arr s l) =
create (l+1) (\ marr -> writeArr marr l z >> go z (s+l-1) (l-1) marr)
where
go :: b -> Int -> Int -> MArray u s b -> ST s ()
go !acc !i !j !marr
| j < 0 = return ()
| otherwise = do
x <- indexArrM arr i
let !acc' = x `f` acc
writeArr marr j acc'
go acc' (i-1) (j-1) marr
scanr1' :: forall v a. Vec v a => (a -> a -> a) -> v a -> v a
{-# INLINE scanr1' #-}
scanr1' f (Vec arr s l)
| l <= 0 = empty
| otherwise = case indexArr' arr (s+l-1) of
(# x0 #) -> scanr' f x0 (fromArr arr s (l-1) :: v a)
rangeCut :: Int -> Int -> Int -> Int
{-# INLINE rangeCut #-}
rangeCut !r !min !max | r < min = min
| r > max = max
| otherwise = r
isASCIISpace :: Word8 -> Bool
{-# INLINE isASCIISpace #-}
isASCIISpace w = w == 32 || w - 0x9 <= 4 || w == 0xa0
head :: (Vec v a, HasCallStack) => v a -> a
{-# INLINE head #-}
head (Vec arr s l)
| l <= 0 = errorEmptyVector
| otherwise = indexArr arr s
tail :: (Vec v a, HasCallStack) => v a -> v a
{-# INLINE tail #-}
tail (Vec arr s l)
| l <= 0 = errorEmptyVector
| otherwise = fromArr arr (s+1) (l-1)
init :: (Vec v a, HasCallStack) => v a -> v a
{-# INLINE init #-}
init (Vec arr s l)
| l <= 0 = errorEmptyVector
| otherwise = fromArr arr s (l-1)
last :: (Vec v a, HasCallStack) => v a -> a
{-# INLINE last #-}
last (Vec arr s l)
| l <= 0 = errorEmptyVector
| otherwise = indexArr arr (s+l-1)
index :: (Vec v a, HasCallStack) => v a -> Int -> a
{-# INLINE index #-}
index (Vec arr s l) i | i < 0 || i >= l = errorOutRange i
| otherwise = arr `indexArr` (s + i)
unsafeHead :: Vec v a => v a -> a
{-# INLINE unsafeHead #-}
unsafeHead (Vec arr s l) = assert (l > 0) (indexArr arr s)
unsafeTail :: Vec v a => v a -> v a
{-# INLINE unsafeTail #-}
unsafeTail (Vec arr s l) = assert (l > 0) (fromArr arr (s+1) (l-1))
unsafeInit :: Vec v a => v a -> v a
{-# INLINE unsafeInit #-}
unsafeInit (Vec arr s l) = assert (l > 0) (fromArr arr s (l-1))
unsafeLast :: Vec v a => v a -> a
{-# INLINE unsafeLast #-}
unsafeLast (Vec arr s l) = assert (l > 0) (indexArr arr (s+l-1))
unsafeIndex :: Vec v a => v a -> Int -> a
{-# INLINE unsafeIndex #-}
unsafeIndex (Vec arr s l) i = indexArr arr (s + i)
unsafeTake :: Vec v a => Int -> v a -> v a
{-# INLINE unsafeTake #-}
unsafeTake n (Vec arr s l) = assert (0 <= n && n <= l) (fromArr arr s n)
unsafeDrop :: Vec v a => Int -> v a -> v a
{-# INLINE unsafeDrop #-}
unsafeDrop n (Vec arr s l) = assert (0 <= n && n <= l) (fromArr arr (s+n) (l-n))