module Data.Vector.Persistent (
Vector,
empty,
singleton,
snoc,
fromList,
null,
length,
index,
unsafeIndex,
take,
drop,
splitAt,
slice,
shrink,
update,
(//),
foldr,
foldl',
map,
reverse,
filter,
partition
) where
import Prelude hiding ( null, length, tail, take,
drop, map, foldr, reverse,
splitAt, filter
)
import Control.Applicative hiding ( empty )
import Control.DeepSeq
import Data.Bits
import Data.Foldable ( Foldable )
import qualified Data.Foldable as F
import qualified Data.List as L
import Data.Monoid ( Monoid )
import qualified Data.Monoid as M
import Data.Traversable ( Traversable )
import qualified Data.Traversable as T
import Data.Vector.Persistent.Array ( Array )
import qualified Data.Vector.Persistent.Array as A
data Vector a = EmptyVector
| RootNode { vecSize :: !Int
, vecShift :: !Int
, vecOffset :: !Int
, vecCapacity :: !Int
, vecTail :: ![a]
, intVecPtrs :: !(Array (Vector a))
}
| InternalNode { intVecPtrs :: !(Array (Vector a))
}
| DataNode { dataVec :: !(Array a)
}
deriving (Show)
instance (Eq a) => Eq (Vector a) where
(==) = pvEq
instance (Ord a) => Ord (Vector a) where
compare = pvCompare
instance Foldable Vector where
foldr = foldr
instance Functor Vector where
fmap = map
instance Monoid (Vector a) where
mempty = empty
mappend = append
instance Traversable Vector where
traverse = pvTraverse
instance (NFData a) => NFData (Vector a) where
rnf = pvRnf
pvEq :: (Eq a) => Vector a -> Vector a -> Bool
pvEq v1 v2
| length v1 /= length v2 = False
| isNotSliced v1 && isNotSliced v2 = pvSimpleEq v1 v2
| otherwise = F.toList v1 == F.toList v2
pvSimpleEq :: (Eq a) => Vector a -> Vector a -> Bool
pvSimpleEq EmptyVector EmptyVector = True
pvSimpleEq (RootNode sz1 sh1 _ _ t1 v1) (RootNode sz2 sh2 _ _ t2 v2) =
sz1 == sz2 && sh1 == sh2 && t1 == t2 && v1 == v2
pvSimpleEq (DataNode a1) (DataNode a2) = a1 == a2
pvSimpleEq (InternalNode a1) (InternalNode a2) = a1 == a2
pvSimpleEq _ _ = False
pvCompare :: (Ord a) => Vector a -> Vector a -> Ordering
pvCompare v1 v2
| length v1 /= length v2 = compare (length v1) (length v2)
| isNotSliced v1 && isNotSliced v2 = pvSimpleCompare v1 v2
| otherwise = compare (F.toList v1) (F.toList v2)
pvSimpleCompare :: (Ord a) => Vector a -> Vector a -> Ordering
pvSimpleCompare EmptyVector EmptyVector = EQ
pvSimpleCompare (RootNode _ _ _ _ t1 v1) (RootNode _ _ _ _ t2 v2) =
case compare v1 v2 of
EQ -> compare t1 t2
o -> o
pvSimpleCompare (DataNode a1) (DataNode a2) = compare a1 a2
pvSimpleCompare (InternalNode a1) (InternalNode a2) = compare a1 a2
pvSimpleCompare EmptyVector _ = LT
pvSimpleCompare _ EmptyVector = GT
pvSimpleCompare (InternalNode _) (DataNode _) = GT
pvSimpleCompare (DataNode _) (InternalNode _) = LT
pvSimpleCompare _ _ = error "Data.Vector.Persistent.pvSimpleCompare: Unexpected mismatch"
map :: (a -> b) -> Vector a -> Vector b
map f = go
where
go EmptyVector = EmptyVector
go (DataNode v) = DataNode (A.map f v)
go (InternalNode v) = InternalNode (A.map (fmap f) v)
go (RootNode sz sh off cap t v) =
let t' = L.map f t
v' = A.map (fmap f) v
in RootNode sz sh off cap t' v'
foldr :: (a -> b -> b) -> b -> Vector a -> b
foldr _ s0 EmptyVector = s0
foldr f s0 v
| isNotSliced v = sgo v s0
| otherwise =
case go v (s0, max 0 (vecCapacity v vecSize v), length v) of (r, _, _) -> r
where
go EmptyVector seed = seed
go (DataNode a) (seed, nskip, len)
| len <= 0 = (seed, 0, 0)
| nskip == 0 = (A.boundedFoldr f (32 len) 32 seed a, 0, len A.length a)
| nskip >= 32 = (seed, nskip 32, len)
| otherwise =
let end = min (max 0 (32 nskip)) 32
start = 32 (len + nskip)
taken = end max 0 start
in (A.boundedFoldr f start end seed a, 0, len taken)
go (InternalNode as) seed =
A.foldr go seed as
go (RootNode _ _ _ _ t as) (s, nskip, l) =
let tseed = L.foldl' (flip f) s t
seed = (tseed, nskip, l L.length t)
in A.foldr go seed as
sgo EmptyVector seed = seed
sgo (DataNode a) seed = A.foldr f seed a
sgo (InternalNode as) seed = A.foldr sgo seed as
sgo (RootNode _ _ _ _ t as) seed =
let tseed = L.foldl' (flip f) seed t
in A.foldr sgo tseed as
foldl' :: (b -> a -> b) -> b -> Vector a -> b
foldl' _ s0 EmptyVector = s0
foldl' f s0 v
| isNotSliced v = sgo s0 v
| otherwise =
case go (s0, vecOffset v, length v) v of (r, _, _) -> r
where
go seed EmptyVector = seed
go (seed, nskip, len) (DataNode a)
| len <= 0 = (seed, 0, 0)
| nskip == 0 = (A.boundedFoldl' f 0 (min len 32) seed a, 0, len A.length a)
| nskip >= 32 = (seed, nskip 32, len)
| otherwise =
let end = min 32 (len + nskip)
start = nskip
taken = end max 0 start
in (A.boundedFoldl' f start end seed a, 0, len taken)
go seed (InternalNode as) =
A.foldl' go seed as
go (s, nskip, l) (RootNode _ _ _ _ t as) =
let (rseed, _, _) = A.foldl' go (s, nskip, l L.length t) as
in (L.foldr (flip f) rseed t, 0, 0)
sgo seed EmptyVector = seed
sgo seed (DataNode a) = A.foldl' f seed a
sgo seed (InternalNode as) =
A.foldl' sgo seed as
sgo seed (RootNode _ _ _ _ t as) =
let rseed = A.foldl' sgo seed as
in F.foldr (flip f) rseed t
pvTraverse :: (Applicative f) => (a -> f b) -> Vector a -> f (Vector b)
pvTraverse f = go
where
go EmptyVector = pure EmptyVector
go (DataNode a) = DataNode <$> A.traverse f a
go (InternalNode as) = InternalNode <$> A.traverse go as
go (RootNode sz sh off cap t as) =
RootNode sz sh off cap <$> T.traverse f t <*> A.traverse go as
append :: Vector a -> Vector a -> Vector a
append EmptyVector v = v
append v EmptyVector = v
append v1 v2 = F.foldl' snoc v1 (F.toList v2)
pvRnf :: (NFData a) => Vector a -> ()
pvRnf = F.foldr deepseq ()
empty :: Vector a
empty = EmptyVector
null :: Vector a -> Bool
null EmptyVector = True
null _ = False
length :: Vector a -> Int
length EmptyVector = 0
length RootNode { vecSize = s, vecOffset = off } = s off
length InternalNode {} = error "Data.Vector.Persistent.length: Internal nodes should not be exposed"
length DataNode {} = error "Data.Vector.Persistent.length: Data nodes should not be exposed"
index :: Vector a -> Int -> Maybe a
index v ix
| length v > ix = Just $ unsafeIndex v ix
| otherwise = Nothing
unsafeIndex :: Vector a -> Int -> a
unsafeIndex vec userIndex
| ix >= tailOffset vec && vecCapacity vec < vecSize vec =
L.reverse (vecTail vec) !! (ix .&. 0x1f)
| otherwise = go (vecShift vec) vec
where
ix = vecOffset vec + userIndex
go level v
| level == 0 = A.index (dataVec v) (ix .&. 0x1f)
| otherwise =
let nextVecIx = (ix `shiftR` level) .&. 0x1f
v' = intVecPtrs v
in go (level 5) (A.index v' nextVecIx)
singleton :: a -> Vector a
singleton elt =
RootNode { vecSize = 1
, vecShift = 5
, vecOffset = 0
, vecCapacity = 0
, vecTail = [elt]
, intVecPtrs = A.fromList 0 []
}
arraySnoc :: Array a -> a -> Array a
arraySnoc a elt = A.run $ do
let alen = A.length a
a' <- A.new_ (1 + alen)
A.copy a 0 a' 0 alen
A.write a' alen elt
return a'
snoc :: Vector a -> a -> Vector a
snoc EmptyVector elt = singleton elt
snoc v@RootNode { vecSize = sz, vecShift = sh, vecTail = t } elt
| vecCapacity v > sz =
let v' = update sz elt v
in v' { vecSize = vecSize v' + 1 }
| sz .&. 0x1f /= 0 = v { vecTail = elt : t, vecSize = sz + 1 }
| sz `shiftR` 5 > 1 `shiftL` sh =
RootNode { vecSize = sz + 1
, vecShift = sh + 5
, vecOffset = vecOffset v
, vecCapacity = vecCapacity v + 32
, vecTail = [elt]
, intVecPtrs = A.fromList 2 [ InternalNode (intVecPtrs v)
, newPath sh t
]
}
| otherwise =
RootNode { vecSize = sz + 1
, vecShift = sh
, vecOffset = vecOffset v
, vecCapacity = vecCapacity v + 32
, vecTail = [elt]
, intVecPtrs = pushTail sz t sh (intVecPtrs v)
}
snoc _ _ = error "Data.Vector.Persistent.snoc: Internal nodes should not be exposed to the user"
pushTail :: Int -> [a] -> Int -> Array (Vector a) -> Array (Vector a)
pushTail cnt t = go
where
go level parent
| level == 5 = arraySnoc parent (DataNode (A.fromList 32 (L.reverse t)))
| subIdx < A.length parent =
let nextVec = A.index parent subIdx
toInsert = go (level 5) (intVecPtrs nextVec)
in A.update parent subIdx (InternalNode toInsert)
| otherwise = arraySnoc parent (newPath (level 5) t)
where
subIdx = ((cnt 1) `shiftR` level) .&. 0x1f
newPath :: Int -> [a] -> Vector a
newPath level t
| level == 0 = DataNode (A.fromList 32 (L.reverse t))
| otherwise = InternalNode $ A.fromList 1 $ [newPath (level 5) t]
update :: Int -> a -> Vector a -> Vector a
update ix elt = (// [(ix, elt)])
(//) :: Vector a -> [(Int, a)] -> Vector a
(//) = L.foldr replaceElement
replaceElement :: (Int, a) -> Vector a -> Vector a
replaceElement _ EmptyVector = EmptyVector
replaceElement (userIndex, elt) v@(RootNode { vecSize = sz, vecShift = sh, vecTail = t })
| sz <= ix || ix < 0 = v
| ix >= toff && vecCapacity v < sz =
let tix = sz 1 ix
(keepHead, _:keepTail) = L.splitAt tix t
in v { vecTail = keepHead ++ (elt : keepTail) }
| otherwise = v { intVecPtrs = go sh (intVecPtrs v) }
where
ix = userIndex + vecOffset v
toff = tailOffset v
go level vec
| level == 5 =
let dnode = DataNode $ A.update (dataVec vec') (ix .&. 0x1f) elt
in A.update vec vix dnode
| otherwise =
let rnode = go (level 5) (intVecPtrs vec')
in A.update vec vix (InternalNode rnode)
where
vix = (ix `shiftR` level) .&. 0x1f
vec' = A.index vec vix
replaceElement _ _ = error "Data.Vector.Persistent.replaceElement: should not see internal nodes"
slice :: Int -> Int -> Vector a -> Vector a
slice _ _ EmptyVector = EmptyVector
slice start userLen v@RootNode { vecSize = sz, vecOffset = off, vecCapacity = cap, vecTail = t }
| len <= 0 = EmptyVector
| start < 0 =
let eltsRetained = min (len + start) sz
in v { vecSize = eltsRetained
, vecTail = L.drop (sz eltsRetained) t
}
| otherwise =
let newOff = off + start
newSize = min (newOff + len) sz
ntake = max 0 (start cap)
t' = L.drop (sz newSize) t
in v { vecOffset = newOff
, vecSize = newSize
, vecTail = L.take (L.length t' ntake) t'
}
where
len = max 0 (min userLen (sz start))
slice _ _ _ = error "Data.Vector.Persistent.slice: Internal node"
take :: Int -> Vector a -> Vector a
take = slice 0
drop :: Int -> Vector a -> Vector a
drop i v = slice i (length v) v
splitAt :: Int -> Vector a -> (Vector a, Vector a)
splitAt ix v = (take ix v, drop ix v)
shrink :: Vector a -> Vector a
shrink EmptyVector = EmptyVector
shrink v
| isNotSliced v = v
| otherwise = fromList $ F.toList v
reverse :: Vector a -> Vector a
reverse = fromList . foldl' (flip (:)) []
filter :: (a -> Bool) -> Vector a -> Vector a
filter p = foldl' go empty
where
go acc e = if p e then snoc acc e else acc
partition :: (a -> Bool) -> Vector a -> (Vector a, Vector a)
partition p = foldl' go (empty, empty)
where
go (atrue, afalse) e =
if p e then (snoc atrue e, afalse) else (atrue, snoc afalse e)
fromList :: [a] -> Vector a
fromList = F.foldl' snoc empty
tailOffset :: Vector a -> Int
tailOffset v
| len < 32 = 0
| otherwise = (len 1) `shiftR` 5 `shiftL` 5
where
len = length v
isNotSliced :: Vector a -> Bool
isNotSliced v = vecOffset v == 0 && vecCapacity v < vecSize v