git.iko.soy/text
on aeade62 help sign in
text/Data/Text/Array.hs
283 lines 9.7 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 , partialCopyI
36 , empty
37 , run
38 , run2
39 , toList
40 , unsafeFreeze
41 , unsafeIndex
42 , unsafeIndexWord8
43 , unsafeNew
44 , unsafeWrite
45 , unsafeWriteWord8
46 ) where
48#if defined(ASSERTS)
49-- This fugly hack is brought by GHC's apparent reluctance to deal
50-- with MagicHash and UnboxedTuples when inferring types. Eek!
51# define CHECK_BOUNDS(_func_,_len_,_k_) \
52if (_k_) < 0 || (_k_) >= (_len_) then error ("Data.Text.Array." ++ (_func_) ++ ": bounds error, offset " ++ show (_k_) ++ ", length " ++ show (_len_)) else
53#else
54# define CHECK_BOUNDS(_func_,_len_,_k_)
55#endif
57#include "MachDeps.h"
59#if defined(ASSERTS)
60import Control.Exception (assert)
61#endif
62import Data.Bits ((.&.))
63import Data.Text.UnsafeShift (shiftL, shiftR)
64import GHC.Base (ByteArray#, MutableByteArray#, Int(..),
65 indexWord16Array#, indexWord8Array#, indexWordArray#, newByteArray#,
66 readWord16Array#, readWordArray#, unsafeCoerce#,
67 writeWord16Array#, writeWord8Array#, writeWordArray#)
68import GHC.ST (ST(..), runST)
69import GHC.Word (Word16(..), Word8(..), Word(..))
70import Prelude hiding (length, read)
72-- | Immutable array type.
73data Array = Array
74 {-# UNPACK #-} !Int -- length (in units of Word16, not bytes)
75 ByteArray#
77-- | Mutable array type, for use in the ST monad.
78data MArray s = MArray
79 {-# UNPACK #-} !Int -- length (in units of Word16, not bytes)
80 (MutableByteArray# s)
82-- | Operations supported by all arrays.
83class IArray a where
84 -- | Return the length of an array.
85 length :: a -> Int
87instance IArray Array where
88 length (Array len _ba) = len
89 {-# INLINE length #-}
91instance IArray (MArray s) where
92 length (MArray len _ba) = len
93 {-# INLINE length #-}
95-- | Create an uninitialized mutable array.
96unsafeNew :: forall s. Int -> ST s (MArray s)
97unsafeNew n =
98#if defined(ASSERTS)
99 assert (n >= 0) .
100#endif
101 ST $ \s1# ->
102 case bytesInArray n of
103 len@(I# len#) ->
104#if defined(ASSERTS)
105 if len < 0 then error (show ("unsafeNew",len)) else
106#endif
107 case newByteArray# len# s1# of
108 (# s2#, marr# #) -> (# s2#, MArray n marr# #)
109{-# INLINE unsafeNew #-}
111-- | Freeze a mutable array. Do not mutate the 'MArray' afterwards!
112unsafeFreeze :: MArray s -> ST s (Array)
113unsafeFreeze (MArray len mba#) = ST $ \s# ->
114 (# s#, Array len (unsafeCoerce# mba#) #)
115{-# INLINE unsafeFreeze #-}
117-- | Indicate how many bytes would be used for an array of the given
118-- size.
119bytesInArray :: Int -> Int
120bytesInArray n = n `shiftL` 1
121{-# INLINE bytesInArray #-}
123-- | Unchecked read of an immutable array. May return garbage or
124-- crash on an out-of-bounds access.
125unsafeIndex :: Array -> Int -> Word16
126unsafeIndex (Array len ba#) i@(I# i#) =
127 CHECK_BOUNDS("unsafeIndex",len,i)
128 case indexWord16Array# ba# i# of r# -> (W16# r#)
129{-# INLINE unsafeIndex #-}
131-- | Unchecked read of an immutable array. May return garbage or
132-- crash on an out-of-bounds access.
133unsafeIndexWord :: Array -> Int -> Word
134unsafeIndexWord (Array len ba#) i@(I# i#) =
135 CHECK_BOUNDS("unsafeIndexWord",len,i)
136 case indexWordArray# ba# i# of r# -> (W# r#)
137{-# INLINE unsafeIndexWord #-}
139-- | Unchecked read of an immutable array. May return garbage or
140-- crash on an out-of-bounds access.
141unsafeIndexWord8 :: Array -> Int -> Word8
142unsafeIndexWord8 (Array len ba#) i@(I# i#) =
143 CHECK_BOUNDS("unsafeIndexWord8",len,i)
144 case indexWord8Array# ba# i# of r# -> (W8# r#)
145{-# INLINE unsafeIndexWord8 #-}
147-- | Unchecked read of a mutable array. May return garbage or
148-- crash on an out-of-bounds access.
149unsafeRead :: MArray s -> Int -> ST s Word16
150unsafeRead (MArray len mba#) i@(I# i#) = ST $ \s# ->
151 CHECK_BOUNDS("unsafeRead",len,i)
152 case readWord16Array# mba# i# s# of
153 (# s2#, r# #) -> (# s2#, W16# r# #)
154{-# INLINE unsafeRead #-}
156-- | Unchecked write of a mutable array. May return garbage or crash
157-- on an out-of-bounds access.
158unsafeWrite :: MArray s -> Int -> Word16 -> ST s ()
159unsafeWrite (MArray len marr#) i@(I# i#) (W16# e#) = ST $ \s1# ->
160 CHECK_BOUNDS("unsafeWrite",len,i)
161 case writeWord16Array# marr# i# e# s1# of
162 s2# -> (# s2#, () #)
163{-# INLINE unsafeWrite #-}
165-- | Unchecked read of a mutable array. May return garbage or
166-- crash on an out-of-bounds access.
167unsafeReadWord :: MArray s -> Int -> ST s Word
168unsafeReadWord (MArray len mba#) i@(I# i#) = ST $ \s# ->
169 CHECK_BOUNDS("unsafeRead64",len,i)
170 case readWordArray# mba# i# s# of
171 (# s2#, r# #) -> (# s2#, W# r# #)
172{-# INLINE unsafeReadWord #-}
174-- | Unchecked write of a mutable array. May return garbage or crash
175-- on an out-of-bounds access.
176unsafeWriteWord :: MArray s -> Int -> Word -> ST s ()
177unsafeWriteWord (MArray len marr#) i@(I# i#) (W# e#) = ST $ \s1# ->
178 CHECK_BOUNDS("unsafeWriteWord",len,i)
179 case writeWordArray# marr# i# e# s1# of
180 s2# -> (# s2#, () #)
181{-# INLINE unsafeWriteWord #-}
183-- | Unchecked write of a mutable array. May return garbage or crash
184-- on an out-of-bounds access.
185unsafeWriteWord8 :: MArray s -> Int -> Word8 -> ST s ()
186unsafeWriteWord8 (MArray len marr#) i@(I# i#) (W8# e#) = ST $ \s1# ->
187 CHECK_BOUNDS("unsafeWriteWord8",len,i)
188 case writeWord8Array# marr# i# e# s1# of
189 s2# -> (# s2#, () #)
190{-# INLINE unsafeWriteWord8 #-}
192-- | Convert an immutable array to a list.
193toList :: Array -> [Word16]
194toList a = loop 0
195 where loop i | i < len = unsafeIndex a i : loop (i+1)
196 | otherwise = []
197 len = length a
199-- | An empty immutable array.
200empty :: Array
201empty = runST (unsafeNew 0 >>= unsafeFreeze)
203-- | Run an action in the ST monad and return an immutable array of
204-- its result.
205run :: (forall s. ST s (MArray s)) -> Array
206run k = runST (k >>= unsafeFreeze)
208-- | Run an action in the ST monad and return an immutable array of
209-- its result paired with whatever else the action returns.
210run2 :: (forall s. ST s (MArray s, a)) -> (Array, a)
211run2 k = runST (do
212 (marr,b) <- k
213 arr <- unsafeFreeze marr
214 return (arr,b))
216-- | The amount to divide or multiply by to switch between units of
217-- 'Word16' and units of 'Word'.
218wordFactor :: Int
219wordFactor = SIZEOF_HSWORD `shiftR` 1
221-- | Indicate whether an offset is word-aligned.
222wordAligned :: Int -> Bool
223wordAligned i = i .&. (wordFactor - 1) == 0
225-- | Copy an array in its entirety. The destination array must be at
226-- least as big as the source.
227copy :: MArray s -- ^ source array
228 -> MArray s -- ^ destination array
229 -> ST s ()
230copy dest@(MArray dlen _) src@(MArray slen _)
231 | dlen >= slen = fast_loop 0
232 | otherwise = fail "Data.Text.Array.copy: array too small"
233 where
234 nwds = slen `div` wordFactor
235 fast_loop !i
236 | i >= nwds = copy_loop (i * wordFactor)
237 | otherwise = do unsafeReadWord src i >>= unsafeWriteWord dest i
238 fast_loop (i+1)
239 copy_loop !i
240 | i >= slen = return ()
241 | otherwise = do unsafeRead src i >>= unsafeWrite dest i
242 copy_loop (i+1)
244-- | Copy some elements of a mutable array.
245partialCopyM :: MArray s -- ^ Destination
246 -> Int -- ^ Destination offset
247 -> MArray s -- ^ Source
248 -> Int -- ^ Source offset
249 -> Int -- ^ Count
250 -> ST s ()
251partialCopyM dest didx src sidx count =
252#if defined(ASSERTS)
253 assert (sidx + count <= length src) .
254 assert (didx + count <= length dest) $
255#endif
256 copy_loop sidx didx 0
257 where
258 copy_loop !i !j !c
259 | c >= count = return ()
260 | otherwise = do unsafeRead src i >>= unsafeWrite dest j
261 copy_loop (i+1) (j+1) (c+1)
263-- | Copy some elements of an immutable array.
264partialCopyI :: MArray s -- ^ Destination
265 -> Int -- ^ Destination offset
266 -> Array -- ^ Source
267 -> Int -- ^ Source offset
268 -> Int -- ^ First offset in source /not/ to
269 -- copy (i.e. /not/ length)
270 -> ST s ()
271partialCopyI dest i0 src j0 top
272 | wordAligned i0 && wordAligned j0 = fast (i0 `div` wordFactor) (j0 `div` wordFactor)
273 | otherwise = slow i0 j0
274 where
275 topwds = top `div` wordFactor
276 fast !i !j
277 | i >= topwds = slow (i * wordFactor) (j * wordFactor)
278 | otherwise = do unsafeWriteWord dest i (src `unsafeIndexWord` j)
279 fast (i+1) (j+1)
280 slow !i !j
281 | i >= top = return ()
282 | otherwise = do unsafeWrite dest i (src `unsafeIndex` j)
283 slow (i+1) (j+1)