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,
takeWhile,
dropWhile
) where
import Prelude hiding ( null, length, tail, take,
drop, map, foldr, reverse,
splitAt, filter, takeWhile, dropWhile
)
import qualified Control.Applicative as Ap
import Control.DeepSeq
import Data.Bits
import qualified Data.Foldable as F
import qualified Data.List as L
import qualified Data.Monoid as M
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 F.Foldable Vector where
foldr = foldr
instance Functor Vector where
fmap = map
instance M.Monoid (Vector a) where
mempty = empty
mappend = append
instance T.Traversable Vector where
traverse = pvTraverse
instance (NFData a) => NFData (Vector a) where
rnf = pvRnf
pvEq :: (Eq a) => Vector a -> Vector a -> Bool
pvEq EmptyVector EmptyVector = True
pvEq v1@RootNode { } v2@RootNode { }
| length v1 /= length v2 = False
| isNotSliced v1 && isNotSliced v2 = pvSimpleEq v1 v2
| otherwise = F.toList v1 == F.toList v2
pvEq (DataNode a1) (DataNode a2) = a1 == a2
pvEq (InternalNode a1) (InternalNode a2) = a1 == a2
pvEq _ _ = False
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 EmptyVector EmptyVector = EQ
pvCompare (DataNode a1) (DataNode a2) = compare a1 a2
pvCompare (InternalNode a1) (InternalNode a2) = compare a1 a2
pvCompare v1@RootNode { vecSize = s1 } v2@RootNode { vecSize = s2 }
| s1 /= s2 = compare s1 s2
| isNotSliced v1 && isNotSliced v2 = pvSimpleCompare v1 v2
| otherwise = compare (F.toList v1) (F.toList v2)
pvCompare EmptyVector _ = LT
pvCompare _ EmptyVector = GT
pvCompare (DataNode _) (InternalNode _) = LT
pvCompare (InternalNode _) (DataNode _) = GT
pvCompare _ _ = error "Data.Vector.Persistent: unexpected root node"
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'
data FoldInfo a = FI a !Int !Int
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 (FI s0 (max 0 (vecCapacity v vecSize v)) (length v)) of (FI r _ _) -> r
where
go EmptyVector seed = seed
go (DataNode a) s@(FI seed nskip len)
| len <= 0 = s
| nskip == 0 = FI (A.boundedFoldr f (32 len) 32 seed a) 0 (len A.length a)
| nskip >= 32 = FI seed (nskip 32) len
| otherwise =
let end = min (max 0 (32 nskip)) 32
start = 32 (len + nskip)
taken = end max 0 start
in FI (A.boundedFoldr f start end seed a) 0 (len taken)
go (InternalNode as) seed =
A.foldr go seed as
go (RootNode _ _ _ _ t as) (FI s nskip l) =
let tseed = L.foldl' (flip f) s t
seed = FI 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 (FI s0 (vecOffset v) (length v)) v of (FI r _ _) -> r
where
go seed EmptyVector = seed
go s@(FI seed nskip len) (DataNode a)
| len <= 0 = s
| nskip == 0 = FI (A.boundedFoldl' f 0 (min len 32) seed a) 0 (len A.length a)
| nskip >= 32 = FI seed (nskip 32) len
| otherwise =
let end = min 32 (len + nskip)
start = nskip
taken = end max 0 start
in FI (A.boundedFoldl' f start end seed a) 0 (len taken)
go seed (InternalNode as) =
A.foldl' go seed as
go (FI s nskip l) (RootNode _ _ _ _ t as) =
let FI rseed _ _ = A.foldl' go (FI s nskip (l L.length t)) as
in FI (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 :: (Ap.Applicative f) => (a -> f b) -> Vector a -> f (Vector b)
pvTraverse f = go
where
go EmptyVector = Ap.pure EmptyVector
go (DataNode a) = DataNode Ap.<$> A.traverseArray f a
go (InternalNode as) = InternalNode Ap.<$> A.traverseArray go as
go (RootNode sz sh off cap t as) =
RootNode sz sh off cap Ap.<$> T.traverse f t Ap.<*> A.traverseArray go as
append :: Vector a -> Vector a -> Vector a
append EmptyVector v = v
append v EmptyVector = v
append v1 v2 = foldl' snoc v1 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 =
let sh = vecShift vec
in go (sh 5) (A.index (intVecPtrs vec) (ix `shiftR` sh))
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, vecOffset = off, vecTail = t } elt
| vecCapacity v >= sz =
let v' = v { vecSize = sz + 1 }
in update (sz off) elt v'
| 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 =
case t of
[] -> v { vecTail = [elt] }
_ ->
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
| toff < start =
let t' = L.reverse $ L.take userLen $ L.drop (start toff) $ L.reverse t
in v { vecOffset = 0
, vecCapacity = 0
, intVecPtrs = A.fromList 0 []
, vecSize = L.length t'
, vecTail = t'
}
| 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
toff = tailOffset v
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 = spToPair . foldl' go (SP empty empty)
where
go (SP atrue afalse) e =
if p e then SP (snoc atrue e) afalse else SP atrue (snoc afalse e)
fromList :: [a] -> Vector a
fromList = F.foldl' snoc empty
data StrictPair a b = SP !a !b
spSnd :: StrictPair a b -> b
spSnd (SP _ v) = v
spToPair :: StrictPair a b -> (a, b)
spToPair (SP a b) = (a, b)
takeWhile :: (a -> Bool) -> Vector a -> Vector a
takeWhile p = spSnd . foldl' f (SP True empty)
where
f (SP True v) e =
if p e then SP True (snoc v e)
else SP False v
f a _ = a
dropWhile :: (a -> Bool) -> Vector a -> Vector a
dropWhile p = spSnd . foldl' f (SP True empty)
where
f a@(SP True v) e =
if p e then a
else SP False (snoc v e)
f (SP False v) e = SP False (snoc v e)
tailOffset :: Vector a -> Int
tailOffset EmptyVector = 0
tailOffset v
| len < 32 = 0
| otherwise = (len 1) `shiftR` 5 `shiftL` 5
where
len = vecSize v
isNotSliced :: Vector a -> Bool
isNotSliced v = vecOffset v == 0 && vecCapacity v < vecSize v