on cf24ead help sign in
text/Data/Text/Array.hs
220 lines 7.2 kB history raw
1{-# LANGUAGE BangPatterns, CPP, MagicHash, Rank2Types, UnboxedTuples #-}
2{-# OPTIONS_GHC -fno-warn-unused-matches #-}
3-- |
4-- Module : Data.Text.Array
5-- Copyright : (c) 2009, 2010 Bryan O'Sullivan
6--
7-- License : BSD-style
8-- Maintainer : bos@serpentine.com, rtomharper@googlemail.com,
9-- duncan@haskell.org
10-- Stability : experimental
11-- Portability : portable
12--
13-- Packed, unboxed, heap-resident arrays. Suitable for performance
14-- critical use, both in terms of large data quantities and high
15-- speed.
16--
17-- This module is intended to be imported @qualified@, to avoid name
18-- clashes with "Prelude" functions, e.g.
19--
20-- > import qualified Data.Text.Array as A
21--
22-- The names in this module resemble those in the 'Data.Array' family
23-- of modules, but are shorter due to the assumption of qualifid
24-- naming.
25module Data.Text.Array
26 (
27 -- * Types
28 IArray(..)
29 , Array
30 , MArray
32 -- * Functions
33 , copy
34 , partialCopyM
35 , empty
36 , run
37 , run2
38 , toList
39 , unsafeFreeze
40 , unsafeIndex
41 , unsafeNew
42 , unsafeWrite
43 ) where
45#if defined(ASSERTS)
46-- This fugly hack is brought by GHC's apparent reluctance to deal
47-- with MagicHash and UnboxedTuples when inferring types. Eek!
48# define CHECK_BOUNDS(_func_,_len_,_k_) \
49if (_k_) < 0 || (_k_) >= (_len_) then error ("Data.Text.Array." ++ (_func_) ++ ": bounds error, offset " ++ show (_k_) ++ ", length " ++ show (_len_)) else
50#else
51# define CHECK_BOUNDS(_func_,_len_,_k_)
52#endif
54#include "MachDeps.h"
56#if defined(ASSERTS)
57import Control.Exception (assert)
58#endif
59import Data.Text.UnsafeShift (shiftL, shiftR)
60import GHC.Base (ByteArray#, MutableByteArray#, Int(..),
61 indexWord16Array#, newByteArray#,
62 readWord16Array#, readWordArray#, unsafeCoerce#,
63 writeWord16Array#, writeWordArray#)
64import GHC.ST (ST(..), runST)
65import GHC.Word (Word16(..), Word(..))
66import Prelude hiding (length, read)
68-- | Immutable array type.
69data Array = Array
70 {-# UNPACK #-} !Int -- length (in units of Word16, not bytes)
71 ByteArray#
73-- | Mutable array type, for use in the ST monad.
74data MArray s = MArray
75 {-# UNPACK #-} !Int -- length (in units of Word16, not bytes)
76 (MutableByteArray# s)
78-- | Operations supported by all arrays.
79class IArray a where
80 -- | Return the length of an array.
81 length :: a -> Int
83instance IArray Array where
84 length (Array len _ba) = len
85 {-# INLINE length #-}
87instance IArray (MArray s) where
88 length (MArray len _ba) = len
89 {-# INLINE length #-}
91-- | Create an uninitialized mutable array.
92unsafeNew :: forall s. Int -> ST s (MArray s)
93unsafeNew n =
94#if defined(ASSERTS)
95 assert (n >= 0) .
96#endif
97 ST $ \s1# ->
98 case bytesInArray n of
99 len@(I# len#) ->
100#if defined(ASSERTS)
101 if len < 0 then error (show ("unsafeNew",len)) else
102#endif
103 case newByteArray# len# s1# of
104 (# s2#, marr# #) -> (# s2#, MArray n marr# #)
105{-# INLINE unsafeNew #-}
107-- | Freeze a mutable array. Do not mutate the 'MArray' afterwards!
108unsafeFreeze :: MArray s -> ST s (Array)
109unsafeFreeze (MArray len mba#) = ST $ \s# ->
110 (# s#, Array len (unsafeCoerce# mba#) #)
111{-# INLINE unsafeFreeze #-}
113-- | Indicate how many bytes would be used for an array of the given
114-- size.
115bytesInArray :: Int -> Int
116bytesInArray n = n `shiftL` 1
117{-# INLINE bytesInArray #-}
119-- | Unchecked read of an immutable array. May return garbage or
120-- crash on an out-of-bounds access.
121unsafeIndex :: Array -> Int -> Word16
122unsafeIndex (Array len ba#) i@(I# i#) =
123 CHECK_BOUNDS("unsafeIndex",len,i)
124 case indexWord16Array# ba# i# of r# -> (W16# r#)
125{-# INLINE unsafeIndex #-}
127-- | Unchecked read of a mutable array. May return garbage or
128-- crash on an out-of-bounds access.
129unsafeRead :: MArray s -> Int -> ST s Word16
130unsafeRead (MArray len mba#) i@(I# i#) = ST $ \s# ->
131 CHECK_BOUNDS("unsafeRead",len,i)
132 case readWord16Array# mba# i# s# of
133 (# s2#, r# #) -> (# s2#, W16# r# #)
134{-# INLINE unsafeRead #-}
136-- | Unchecked write of a mutable array. May return garbage or crash
137-- on an out-of-bounds access.
138unsafeWrite :: MArray s -> Int -> Word16 -> ST s ()
139unsafeWrite (MArray len marr#) i@(I# i#) (W16# e#) = ST $ \s1# ->
140 CHECK_BOUNDS("unsafeWrite",len,i)
141 case writeWord16Array# marr# i# e# s1# of
142 s2# -> (# s2#, () #)
143{-# INLINE unsafeWrite #-}
145-- | Unchecked read of a mutable array. May return garbage or
146-- crash on an out-of-bounds access.
147unsafeReadWord :: MArray s -> Int -> ST s Word
148unsafeReadWord (MArray len mba#) i@(I# i#) = ST $ \s# ->
149 CHECK_BOUNDS("unsafeRead64",len,i)
150 case readWordArray# mba# i# s# of
151 (# s2#, r# #) -> (# s2#, W# r# #)
152{-# INLINE unsafeReadWord #-}
154-- | Unchecked write of a mutable array. May return garbage or crash
155-- on an out-of-bounds access.
156unsafeWriteWord :: MArray s -> Int -> Word -> ST s ()
157unsafeWriteWord (MArray len marr#) i@(I# i#) (W# e#) = ST $ \s1# ->
158 CHECK_BOUNDS("unsafeWriteWord",len,i)
159 case writeWordArray# marr# i# e# s1# of
160 s2# -> (# s2#, () #)
161{-# INLINE unsafeWriteWord #-}
163-- | Convert an immutable array to a list.
164toList :: Array -> [Word16]
165toList a = loop 0
166 where loop i | i < len = unsafeIndex a i : loop (i+1)
167 | otherwise = []
168 len = length a
170-- | An empty immutable array.
171empty :: Array
172empty = runST (unsafeNew 0 >>= unsafeFreeze)
174-- | Run an action in the ST monad and return an immutable array of
175-- its result.
176run :: (forall s. ST s (MArray s)) -> Array
177run k = runST (k >>= unsafeFreeze)
179-- | Run an action in the ST monad and return an immutable array of
180-- its result paired with whatever else the action returns.
181run2 :: (forall s. ST s (MArray s, a)) -> (Array, a)
182run2 k = runST (do
183 (marr,b) <- k
184 arr <- unsafeFreeze marr
185 return (arr,b))
187-- | Copy an array in its entirety. The destination array must be at
188-- least as big as the source.
189copy :: MArray s -- ^ source array
190 -> MArray s -- ^ destination array
191 -> ST s ()
192copy dest@(MArray dlen _) src@(MArray slen _)
193 | dlen >= slen = fast_loop 0
194 | otherwise = fail "Data.Text.Array.copy: array too small"
195 where
196 nwds = slen `div` (SIZEOF_HSWORD `shiftR` 1)
197 fast_loop !i
198 | i >= nwds = copy_loop (i * (SIZEOF_HSWORD `shiftR` 1))
199 | otherwise = do unsafeReadWord src i >>= unsafeWriteWord dest i
200 fast_loop (i+1)
201 copy_loop !i
202 | i >= slen = return ()
203 | otherwise = do unsafeRead src i >>= unsafeWrite dest i
204 copy_loop (i+1)
205{-# INLINE copy #-}
207-- | Copy some elements of a mutable array.
208partialCopyM :: MArray s -> Int -> MArray s -> Int -> Int -> ST s ()
209partialCopyM dest didx src sidx count =
210#if defined(ASSERTS)
211 assert (sidx + count <= length src) .
212 assert (didx + count <= length dest) $
213#endif
214 copy_loop sidx didx 0
215 where
216 copy_loop !i !j !c
217 | c >= count = return ()
218 | otherwise = do unsafeRead src i >>= unsafeWrite dest j
219 copy_loop (i+1) (j+1) (c+1)
220{-# INLINE partialCopyM #-}