1{-# LANGUAGE BangPatterns, CPP, MagicHash, Rank2Types, UnboxedTuples, ScopedTypeVariables #-} 2{-# OPTIONS_GHC -fno-full-laziness -funbox-strict-fields #-} 3{-# OPTIONS_HADDOCK not-home #-} 4 5-- | = WARNING 6-- 7-- This module is considered __internal__. 8-- 9-- The Package Versioning Policy __does not apply__. 10-- 11-- The contents of this module may change __in any way whatsoever__ 12-- and __without any warning__ between minor versions of this package. 13-- 14-- Authors importing this module are expected to track development 15-- closely. 16-- 17-- = Description 18-- 19-- Zero based arrays. 20-- 21-- Note that no bounds checking are performed. 22module Data.HashMap.Internal.Array 23 ( Array 24 , MArray 25 26 -- * Creation 27 , new 28 , new_ 29 , singleton 30 , singletonM 31 , pair 32 33 -- * Basic interface 34 , length 35 , lengthM 36 , read 37 , write 38 , index 39 , indexM 40 , index# 41 , update 42 , updateWith' 43 , unsafeUpdateM 44 , insert 45 , insertM 46 , delete 47 , sameArray1 48 , trim 49 50 , unsafeFreeze 51 , unsafeThaw 52 , unsafeSameArray 53 , run 54 , copy 55 , copyM 56 57 -- * Folds 58 , foldl 59 , foldl' 60 , foldr 61 , foldr' 62 , foldMap 63 , all 64 65 , thaw 66 , map 67 , map' 68 , traverse 69 , traverse' 70 , toList 71 , fromList 72 ) where 73 74#if !MIN_VERSION_base(4,8,0) 75import Control.Applicative (Applicative (..), (<$>)) 76#endif 77import Control.Applicative (liftA2) 78import Control.DeepSeq (NFData (..)) 79import GHC.Exts(Int(..), Int#, reallyUnsafePtrEquality#, tagToEnum#, unsafeCoerce#, State#) 80import GHC.ST (ST(..)) 81import Control.Monad.ST (stToIO) 82 83#if __GLASGOW_HASKELL__ >= 709 84import Prelude hiding (filter, foldMap, foldr, foldl, length, map, read, traverse, all) 85#else 86import Prelude hiding (filter, foldr, foldl, length, map, read, all) 87#endif 88 89#if __GLASGOW_HASKELL__ >= 710 90import GHC.Exts (SmallArray#, newSmallArray#, readSmallArray#, writeSmallArray#, 91 indexSmallArray#, unsafeFreezeSmallArray#, unsafeThawSmallArray#, 92 SmallMutableArray#, sizeofSmallArray#, copySmallArray#, thawSmallArray#, 93 sizeofSmallMutableArray#, copySmallMutableArray#, cloneSmallMutableArray#) 94 95#else 96import GHC.Exts (Array#, newArray#, readArray#, writeArray#, 97 indexArray#, unsafeFreezeArray#, unsafeThawArray#, 98 MutableArray#, sizeofArray#, copyArray#, thawArray#, 99 sizeofMutableArray#, copyMutableArray#, cloneMutableArray#) 100import Data.Monoid (Monoid (..)) 101#endif 102 103#if defined(ASSERTS) 104import qualified Prelude 105#endif 106 107#if MIN_VERSION_deepseq(1,4,3) 108import qualified Control.DeepSeq as NF 109#endif 110 111import Data.HashMap.Internal.Unsafe (runST) 112import Control.Monad ((>=>)) 113 114 115#if __GLASGOW_HASKELL__ >= 710 116type Array# a = SmallArray# a 117type MutableArray# a = SmallMutableArray# a 118 119newArray# :: Int# -> a -> State# d -> (# State# d, SmallMutableArray# d a #) 120newArray# = newSmallArray# 121 122unsafeFreezeArray# :: SmallMutableArray# d a 123 -> State# d -> (# State# d, SmallArray# a #) 124unsafeFreezeArray# = unsafeFreezeSmallArray# 125 126readArray# :: SmallMutableArray# d a 127 -> Int# -> State# d -> (# State# d, a #) 128readArray# = readSmallArray# 129 130writeArray# :: SmallMutableArray# d a 131 -> Int# -> a -> State# d -> State# d 132writeArray# = writeSmallArray# 133 134indexArray# :: SmallArray# a -> Int# -> (# a #) 135indexArray# = indexSmallArray# 136 137unsafeThawArray# :: SmallArray# a 138 -> State# d -> (# State# d, SmallMutableArray# d a #) 139unsafeThawArray# = unsafeThawSmallArray# 140 141sizeofArray# :: SmallArray# a -> Int# 142sizeofArray# = sizeofSmallArray# 143 144copyArray# :: SmallArray# a 145 -> Int# 146 -> SmallMutableArray# d a 147 -> Int# 148 -> Int# 149 -> State# d 150 -> State# d 151copyArray# = copySmallArray# 152 153cloneMutableArray# :: SmallMutableArray# s a 154 -> Int# 155 -> Int# 156 -> State# s 157 -> (# State# s, SmallMutableArray# s a #) 158cloneMutableArray# = cloneSmallMutableArray# 159 160thawArray# :: SmallArray# a 161 -> Int# 162 -> Int# 163 -> State# d 164 -> (# State# d, SmallMutableArray# d a #) 165thawArray# = thawSmallArray# 166 167sizeofMutableArray# :: SmallMutableArray# s a -> Int# 168sizeofMutableArray# = sizeofSmallMutableArray# 169 170copyMutableArray# :: SmallMutableArray# d a 171 -> Int# 172 -> SmallMutableArray# d a 173 -> Int# 174 -> Int# 175 -> State# d 176 -> State# d 177copyMutableArray# = copySmallMutableArray# 178#endif 179 180------------------------------------------------------------------------ 181 182#if defined(ASSERTS) 183-- This fugly hack is brought by GHC's apparent reluctance to deal 184-- with MagicHash and UnboxedTuples when inferring types. Eek! 185# define CHECK_BOUNDS(_func_,_len_,_k_) \ 186if (_k_) < 0 || (_k_) >= (_len_) then error ("Data.HashMap.Internal.Array." ++ (_func_) ++ ": bounds error, offset " ++ show (_k_) ++ ", length " ++ show (_len_)) else 187# define CHECK_OP(_func_,_op_,_lhs_,_rhs_) \ 188if not ((_lhs_) _op_ (_rhs_)) then error ("Data.HashMap.Internal.Array." ++ (_func_) ++ ": Check failed: _lhs_ _op_ _rhs_ (" ++ show (_lhs_) ++ " vs. " ++ show (_rhs_) ++ ")") else 189# define CHECK_GT(_func_,_lhs_,_rhs_) CHECK_OP(_func_,>,_lhs_,_rhs_) 190# define CHECK_LE(_func_,_lhs_,_rhs_) CHECK_OP(_func_,<=,_lhs_,_rhs_) 191# define CHECK_EQ(_func_,_lhs_,_rhs_) CHECK_OP(_func_,==,_lhs_,_rhs_) 192#else 193# define CHECK_BOUNDS(_func_,_len_,_k_) 194# define CHECK_OP(_func_,_op_,_lhs_,_rhs_) 195# define CHECK_GT(_func_,_lhs_,_rhs_) 196# define CHECK_LE(_func_,_lhs_,_rhs_) 197# define CHECK_EQ(_func_,_lhs_,_rhs_) 198#endif 199 200data Array a = Array { 201 unArray :: !(Array# a) 202 } 203 204instance Show a => Show (Array a) where 205 show = show . toList 206 207-- Determines whether two arrays have the same memory address. 208-- This is more reliable than testing pointer equality on the 209-- Array wrappers, but it's still slightly bogus. 210unsafeSameArray :: Array a -> Array b -> Bool 211unsafeSameArray (Array xs) (Array ys) = 212 tagToEnum# (unsafeCoerce# reallyUnsafePtrEquality# xs ys) 213 214sameArray1 :: (a -> b -> Bool) -> Array a -> Array b -> Bool 215sameArray1 eq !xs0 !ys0 216 | lenxs /= lenys = False 217 | otherwise = go 0 xs0 ys0 218 where 219 go !k !xs !ys 220 | k == lenxs = True 221 | (# x #) <- index# xs k 222 , (# y #) <- index# ys k 223 = eq x y && go (k + 1) xs ys 224 225 !lenxs = length xs0 226 !lenys = length ys0 227 228length :: Array a -> Int 229length ary = I# (sizeofArray# (unArray ary)) 230{-# INLINE length #-} 231 232data MArray s a = MArray { 233 unMArray :: !(MutableArray# s a) 234 } 235 236lengthM :: MArray s a -> Int 237lengthM mary = I# (sizeofMutableArray# (unMArray mary)) 238{-# INLINE lengthM #-} 239 240------------------------------------------------------------------------ 241 242instance NFData a => NFData (Array a) where 243 rnf = rnfArray 244 245rnfArray :: NFData a => Array a -> () 246rnfArray ary0 = go ary0 n0 0 247 where 248 n0 = length ary0 249 go !ary !n !i 250 | i >= n = () 251 | (# x #) <- index# ary i 252 = rnf x `seq` go ary n (i+1) 253-- We use index# just in case GHC can't see that the 254-- relevant rnf is strict, or in case it actually isn't. 255{-# INLINE rnfArray #-} 256 257#if MIN_VERSION_deepseq(1,4,3) 258-- | @since 0.2.14.0 259instance NF.NFData1 Array where 260 liftRnf = liftRnfArray 261 262liftRnfArray :: (a -> ()) -> Array a -> () 263liftRnfArray rnf0 ary0 = go ary0 n0 0 264 where 265 n0 = length ary0 266 go !ary !n !i 267 | i >= n = () 268 | (# x #) <- index# ary i 269 = rnf0 x `seq` go ary n (i+1) 270{-# INLINE liftRnfArray #-} 271#endif 272 273-- | Create a new mutable array of specified size, in the specified 274-- state thread, with each element containing the specified initial 275-- value. 276new :: Int -> a -> ST s (MArray s a) 277new (I# n#) b = 278 CHECK_GT("new",n,(0 :: Int)) 279 ST $ \s -> 280 case newArray# n# b s of 281 (# s', ary #) -> (# s', MArray ary #) 282{-# INLINE new #-} 283 284new_ :: Int -> ST s (MArray s a) 285new_ n = new n undefinedElem 286 287singleton :: a -> Array a 288singleton x = runST (singletonM x) 289{-# INLINE singleton #-} 290 291singletonM :: a -> ST s (Array a) 292singletonM x = new 1 x >>= unsafeFreeze 293{-# INLINE singletonM #-} 294 295pair :: a -> a -> Array a 296pair x y = run $ do 297 ary <- new 2 x 298 write ary 1 y 299 return ary 300{-# INLINE pair #-} 301 302read :: MArray s a -> Int -> ST s a 303read ary _i@(I# i#) = ST $ \ s -> 304 CHECK_BOUNDS("read", lengthM ary, _i) 305 readArray# (unMArray ary) i# s 306{-# INLINE read #-} 307 308write :: MArray s a -> Int -> a -> ST s () 309write ary _i@(I# i#) b = ST $ \ s -> 310 CHECK_BOUNDS("write", lengthM ary, _i) 311 case writeArray# (unMArray ary) i# b s of 312 s' -> (# s' , () #) 313{-# INLINE write #-} 314 315index :: Array a -> Int -> a 316index ary _i@(I# i#) = 317 CHECK_BOUNDS("index", length ary, _i) 318 case indexArray# (unArray ary) i# of (# b #) -> b 319{-# INLINE index #-} 320 321index# :: Array a -> Int -> (# a #) 322index# ary _i@(I# i#) = 323 CHECK_BOUNDS("index#", length ary, _i) 324 indexArray# (unArray ary) i# 325{-# INLINE index# #-} 326 327indexM :: Array a -> Int -> ST s a 328indexM ary _i@(I# i#) = 329 CHECK_BOUNDS("indexM", length ary, _i) 330 case indexArray# (unArray ary) i# of (# b #) -> return b 331{-# INLINE indexM #-} 332 333unsafeFreeze :: MArray s a -> ST s (Array a) 334unsafeFreeze mary 335 = ST $ \s -> case unsafeFreezeArray# (unMArray mary) s of 336 (# s', ary #) -> (# s', Array ary #) 337{-# INLINE unsafeFreeze #-} 338 339unsafeThaw :: Array a -> ST s (MArray s a) 340unsafeThaw ary 341 = ST $ \s -> case unsafeThawArray# (unArray ary) s of 342 (# s', mary #) -> (# s', MArray mary #) 343{-# INLINE unsafeThaw #-} 344 345run :: (forall s . ST s (MArray s e)) -> Array e 346run act = runST $ act >>= unsafeFreeze 347{-# INLINE run #-} 348 349-- | Unsafely copy the elements of an array. Array bounds are not checked. 350copy :: Array e -> Int -> MArray s e -> Int -> Int -> ST s () 351copy !src !_sidx@(I# sidx#) !dst !_didx@(I# didx#) _n@(I# n#) = 352 CHECK_LE("copy", _sidx + _n, length src) 353 CHECK_LE("copy", _didx + _n, lengthM dst) 354 ST $ \ s# -> 355 case copyArray# (unArray src) sidx# (unMArray dst) didx# n# s# of 356 s2 -> (# s2, () #) 357 358-- | Unsafely copy the elements of an array. Array bounds are not checked. 359copyM :: MArray s e -> Int -> MArray s e -> Int -> Int -> ST s () 360copyM !src !_sidx@(I# sidx#) !dst !_didx@(I# didx#) _n@(I# n#) = 361 CHECK_BOUNDS("copyM: src", lengthM src, _sidx + _n - 1) 362 CHECK_BOUNDS("copyM: dst", lengthM dst, _didx + _n - 1) 363 ST $ \ s# -> 364 case copyMutableArray# (unMArray src) sidx# (unMArray dst) didx# n# s# of 365 s2 -> (# s2, () #) 366 367cloneM :: MArray s a -> Int -> Int -> ST s (MArray s a) 368cloneM _mary@(MArray mary#) _off@(I# off#) _len@(I# len#) = 369 CHECK_BOUNDS("cloneM_off", lengthM _mary, _off - 1) 370 CHECK_BOUNDS("cloneM_end", lengthM _mary, _off + _len - 1) 371 ST $ \ s -> 372 case cloneMutableArray# mary# off# len# s of 373 (# s', mary'# #) -> (# s', MArray mary'# #) 374 375-- | Create a new array of the @n@ first elements of @mary@. 376trim :: MArray s a -> Int -> ST s (Array a) 377trim mary n = cloneM mary 0 n >>= unsafeFreeze 378{-# INLINE trim #-} 379 380-- | /O(n)/ Insert an element at the given position in this array, 381-- increasing its size by one. 382insert :: Array e -> Int -> e -> Array e 383insert ary idx b = runST (insertM ary idx b) 384{-# INLINE insert #-} 385 386-- | /O(n)/ Insert an element at the given position in this array, 387-- increasing its size by one. 388insertM :: Array e -> Int -> e -> ST s (Array e) 389insertM ary idx b = 390 CHECK_BOUNDS("insertM", count + 1, idx) 391 do mary <- new_ (count+1) 392 copy ary 0 mary 0 idx 393 write mary idx b 394 copy ary idx mary (idx+1) (count-idx) 395 unsafeFreeze mary 396 where !count = length ary 397{-# INLINE insertM #-} 398 399-- | /O(n)/ Update the element at the given position in this array. 400update :: Array e -> Int -> e -> Array e 401update ary idx b = runST (updateM ary idx b) 402{-# INLINE update #-} 403 404-- | /O(n)/ Update the element at the given position in this array. 405updateM :: Array e -> Int -> e -> ST s (Array e) 406updateM ary idx b = 407 CHECK_BOUNDS("updateM", count, idx) 408 do mary <- thaw ary 0 count 409 write mary idx b 410 unsafeFreeze mary 411 where !count = length ary 412{-# INLINE updateM #-} 413 414-- | /O(n)/ Update the element at the given positio in this array, by 415-- applying a function to it. Evaluates the element to WHNF before 416-- inserting it into the array. 417updateWith' :: Array e -> Int -> (e -> e) -> Array e 418updateWith' ary idx f 419 | (# x #) <- index# ary idx 420 = update ary idx $! f x 421{-# INLINE updateWith' #-} 422 423-- | /O(1)/ Update the element at the given position in this array, 424-- without copying. 425unsafeUpdateM :: Array e -> Int -> e -> ST s () 426unsafeUpdateM ary idx b = 427 CHECK_BOUNDS("unsafeUpdateM", length ary, idx) 428 do mary <- unsafeThaw ary 429 write mary idx b 430 _ <- unsafeFreeze mary 431 return () 432{-# INLINE unsafeUpdateM #-} 433 434foldl' :: (b -> a -> b) -> b -> Array a -> b 435foldl' f = \ z0 ary0 -> go ary0 (length ary0) 0 z0 436 where 437 go ary n i !z 438 | i >= n = z 439 | otherwise 440 = case index# ary i of 441 (# x #) -> go ary n (i+1) (f z x) 442{-# INLINE foldl' #-} 443 444foldr' :: (a -> b -> b) -> b -> Array a -> b 445foldr' f = \ z0 ary0 -> go ary0 (length ary0 - 1) z0 446 where 447 go !_ary (-1) z = z 448 go !ary i !z 449 | (# x #) <- index# ary i 450 = go ary (i - 1) (f x z) 451{-# INLINE foldr' #-} 452 453foldr :: (a -> b -> b) -> b -> Array a -> b 454foldr f = \ z0 ary0 -> go ary0 (length ary0) 0 z0 455 where 456 go ary n i z 457 | i >= n = z 458 | otherwise 459 = case index# ary i of 460 (# x #) -> f x (go ary n (i+1) z) 461{-# INLINE foldr #-} 462 463foldl :: (b -> a -> b) -> b -> Array a -> b 464foldl f = \ z0 ary0 -> go ary0 (length ary0 - 1) z0 465 where 466 go _ary (-1) z = z 467 go ary i z 468 | (# x #) <- index# ary i 469 = f (go ary (i - 1) z) x 470{-# INLINE foldl #-} 471 472-- We go to a bit of trouble here to avoid appending an extra mempty. 473-- The below implementation is by Mateusz Kowalczyk, who indicates that 474-- benchmarks show it to be faster than one that avoids lifting out 475-- lst. 476foldMap :: Monoid m => (a -> m) -> Array a -> m 477foldMap f = \ary0 -> case length ary0 of 478 0 -> mempty 479 len -> 480 let !lst = len - 1 481 go i | (# x #) <- index# ary0 i, let fx = f x = 482 if i == lst then fx else fx `mappend` go (i + 1) 483 in go 0 484{-# INLINE foldMap #-} 485 486-- | Verifies that a predicate holds for all elements of an array. 487all :: (a -> Bool) -> Array a -> Bool 488all p = foldr (\a acc -> p a && acc) True 489{-# INLINE all #-} 490 491undefinedElem :: a 492undefinedElem = error "Data.HashMap.Internal.Array: Undefined element" 493{-# NOINLINE undefinedElem #-} 494 495thaw :: Array e -> Int -> Int -> ST s (MArray s e) 496thaw !ary !_o@(I# o#) (I# n#) = 497 CHECK_LE("thaw", _o + n, length ary) 498 ST $ \ s -> case thawArray# (unArray ary) o# n# s of 499 (# s2, mary# #) -> (# s2, MArray mary# #) 500{-# INLINE thaw #-} 501 502-- | /O(n)/ Delete an element at the given position in this array, 503-- decreasing its size by one. 504delete :: Array e -> Int -> Array e 505delete ary idx = runST (deleteM ary idx) 506{-# INLINE delete #-} 507 508-- | /O(n)/ Delete an element at the given position in this array, 509-- decreasing its size by one. 510deleteM :: Array e -> Int -> ST s (Array e) 511deleteM ary idx = do 512 CHECK_BOUNDS("deleteM", count, idx) 513 do mary <- new_ (count-1) 514 copy ary 0 mary 0 idx 515 copy ary (idx+1) mary idx (count-(idx+1)) 516 unsafeFreeze mary 517 where !count = length ary 518{-# INLINE deleteM #-} 519 520map :: (a -> b) -> Array a -> Array b 521map f = \ ary -> 522 let !n = length ary 523 in run $ do 524 mary <- new_ n 525 go ary mary 0 n 526 where 527 go ary mary i n 528 | i >= n = return mary 529 | otherwise = do 530 x <- indexM ary i 531 write mary i $ f x 532 go ary mary (i+1) n 533{-# INLINE map #-} 534 535-- | Strict version of 'map'. 536map' :: (a -> b) -> Array a -> Array b 537map' f = \ ary -> 538 let !n = length ary 539 in run $ do 540 mary <- new_ n 541 go ary mary 0 n 542 where 543 go ary mary i n 544 | i >= n = return mary 545 | otherwise = do 546 x <- indexM ary i 547 write mary i $! f x 548 go ary mary (i+1) n 549{-# INLINE map' #-} 550 551fromList :: Int -> [a] -> Array a 552fromList n xs0 = 553 CHECK_EQ("fromList", n, Prelude.length xs0) 554 run $ do 555 mary <- new_ n 556 go xs0 mary 0 557 where 558 go [] !mary !_ = return mary 559 go (x:xs) mary i = do write mary i x 560 go xs mary (i+1) 561 562toList :: Array a -> [a] 563toList = foldr (:) [] 564 565newtype STA a = STA {_runSTA :: forall s. MutableArray# s a -> ST s (Array a)} 566 567runSTA :: Int -> STA a -> Array a 568runSTA !n (STA m) = runST $ new_ n >>= \ (MArray ar) -> m ar 569 570traverse :: Applicative f => (a -> f b) -> Array a -> f (Array b) 571traverse f = \ !ary -> 572 let 573 !len = length ary 574 go !i 575 | i == len = pure $ STA $ \mary -> unsafeFreeze (MArray mary) 576 | (# x #) <- index# ary i 577 = liftA2 (\b (STA m) -> STA $ \mary -> 578 write (MArray mary) i b >> m mary) 579 (f x) (go (i + 1)) 580 in runSTA len <$> go 0 581{-# INLINE [1] traverse #-} 582 583-- TODO: Would it be better to just use a lazy traversal 584-- and then force the elements of the result? My guess is 585-- yes. 586traverse' :: Applicative f => (a -> f b) -> Array a -> f (Array b) 587traverse' f = \ !ary -> 588 let 589 !len = length ary 590 go !i 591 | i == len = pure $ STA $ \mary -> unsafeFreeze (MArray mary) 592 | (# x #) <- index# ary i 593 = liftA2 (\ !b (STA m) -> STA $ \mary -> 594 write (MArray mary) i b >> m mary) 595 (f x) (go (i + 1)) 596 in runSTA len <$> go 0 597{-# INLINE [1] traverse' #-} 598 599-- Traversing in ST, we don't need to get fancy; we 600-- can just do it directly. 601traverseST :: (a -> ST s b) -> Array a -> ST s (Array b) 602traverseST f = \ ary0 -> 603 let 604 !len = length ary0 605 go k !mary 606 | k == len = return mary 607 | otherwise = do 608 x <- indexM ary0 k 609 y <- f x 610 write mary k y 611 go (k + 1) mary 612 in new_ len >>= (go 0 >=> unsafeFreeze) 613{-# INLINE traverseST #-} 614 615traverseIO :: (a -> IO b) -> Array a -> IO (Array b) 616traverseIO f = \ ary0 -> 617 let 618 !len = length ary0 619 go k !mary 620 | k == len = return mary 621 | otherwise = do 622 x <- stToIO $ indexM ary0 k 623 y <- f x 624 stToIO $ write mary k y 625 go (k + 1) mary 626 in stToIO (new_ len) >>= (go 0 >=> stToIO . unsafeFreeze) 627{-# INLINE traverseIO #-} 628 629 630-- Why don't we have similar RULES for traverse'? The efficient 631-- way to traverse strictly in IO or ST is to force results as 632-- they come in, which leads to different semantics. In particular, 633-- we need to ensure that 634-- 635-- traverse' (\x -> print x *> pure undefined) xs 636-- 637-- will actually print all the values and then return undefined. 638-- We could add a strict mapMWithIndex, operating in an arbitrary 639-- Monad, that supported such rules, but we don't have that right now. 640{-# RULES 641"traverse/ST" forall f. traverse f = traverseST f 642"traverse/IO" forall f. traverse f = traverseIO f 643 #-} 644