Utils.hs 1.18 KB
Newer Older
1
{-# LANGUAGE BangPatterns #-}
2
3
module Data.Vector.Utils
  ( iforM_
4
5
6
7
  , sort
  , sortBy
  , sortOn
  , hasDuplicates
8
  , imap'
9
10
  ) where

11
import           Data.Ord (comparing)
12
import           Control.Monad (forM_)
13

14
15
import           Data.Vector (Vector)
import qualified Data.Vector as V
16
import qualified Data.Vector.Mutable as VM
17
import qualified Data.Vector.Algorithms.Intro as V
18
19
20

iforM_ :: Monad m => Vector a -> (Int -> a -> m b) -> m ()
iforM_ = flip V.imapM_
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

sort :: Ord a => Vector a -> Vector a
sort = V.modify V.sort


-- FIXME: There should be inline pragmas for sortBy and sortOn, but GHC panics
-- when I do this.

sortBy :: (a -> a -> Ordering) -> Vector a -> Vector a
sortBy predicate = V.modify (V.sortBy predicate)


sortOn :: Ord b => (a -> b) -> Vector a -> Vector a
sortOn f = V.modify (V.sortBy (comparing f))


hasDuplicates :: Eq a => Vector a -> Bool
hasDuplicates v = V.length (V.uniq v) /= V.length v
39
{-# INLINE hasDuplicates #-}
40
41
42
43
44
45
46
47
48
49
50
51
52
53

imap' :: (Int -> a -> b) -> Vector a -> Vector b
imap' !f !v = V.create $ do
  let !len = (length v)
  v' <- VM.new len

  forM_ [0..len-1] $ \i -> do
    let !a = v V.! i
    VM.write v' i (f' i a)

  return v'

  where
    f' !i !a = (f $! i) $! a