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