git.iko.soy/text
on builder help sign in
text/Data/Text/Array.hs
296 lines 9.9 kB history raw
1{-# LANGUAGE BangPatterns, CPP, MagicHash, Rank2Types, RecordWildCards,
2 UnboxedTuples #-}
3{-# OPTIONS_GHC -fno-warn-unused-matches #-}
4-- |
5-- Module : Data.Text.Array
6-- Copyright : (c) 2009, 2010 Bryan O'Sullivan
7--
8-- License : BSD-style
9-- Maintainer : bos@serpentine.com, rtomharper@googlemail.com,
10-- duncan@haskell.org
11-- Stability : experimental
12-- Portability : portable
13--
14-- Packed, unboxed, heap-resident arrays. Suitable for performance
15-- critical use, both in terms of large data quantities and high
16-- speed.
17--
18-- This module is intended to be imported @qualified@, to avoid name
19-- clashes with "Prelude" functions, e.g.
20--
21-- > import qualified Data.Text.Array as A
22--
23-- The names in this module resemble those in the 'Data.Array' family
24-- of modules, but are shorter due to the assumption of qualifid
25-- naming.
26module Data.Text.Array
27 (
28 -- * Types
29 Array(aBA)
30 , MArray(maBA)
32 -- * Functions
33 , copyM
34 , copyI
35 , empty
36 , equal
37#if defined(ASSERTS)
38 , length
39#endif
40 , run
41 , run2
42 , toList
43 , unsafeFreeze
44 , unsafeIndex
45 , new
46 , unsafeWrite
47 ) where
49#if defined(ASSERTS)
50-- This fugly hack is brought by GHC's apparent reluctance to deal
51-- with MagicHash and UnboxedTuples when inferring types. Eek!
52# define CHECK_BOUNDS(_func_,_len_,_k_) \
53if (_k_) < 0 || (_k_) >= (_len_) then error ("Data.Text.Array." ++ (_func_) ++ ": bounds error, offset " ++ show (_k_) ++ ", length " ++ show (_len_)) else
54#else
55# define CHECK_BOUNDS(_func_,_len_,_k_)
56#endif
58#include "MachDeps.h"
60#if defined(ASSERTS)
61import Control.Exception (assert)
62#endif
63import Data.Bits ((.&.), xor)
64import Data.Text.UnsafeShift (shiftL, shiftR)
65import GHC.Base (ByteArray#, MutableByteArray#, Int(..),
66 indexWord16Array#, indexWordArray#, newByteArray#,
67 readWord16Array#, readWordArray#, unsafeCoerce#,
68 writeWord16Array#, writeWordArray#)
69import GHC.ST (ST(..), runST)
70import GHC.Word (Word16(..), Word(..))
71import Prelude hiding (length, read)
73-- | Immutable array type.
74data Array = Array {
75 aBA :: ByteArray#
76#if defined(ASSERTS)
77 , aLen :: {-# UNPACK #-} !Int -- length (in units of Word16, not bytes)
78#endif
79 }
81-- | Mutable array type, for use in the ST monad.
82data MArray s = MArray {
83 maBA :: MutableByteArray# s
84#if defined(ASSERTS)
85 , maLen :: {-# UNPACK #-} !Int -- length (in units of Word16, not bytes)
86#endif
87 }
89#if defined(ASSERTS)
90-- | Operations supported by all arrays.
91class IArray a where
92 -- | Return the length of an array.
93 length :: a -> Int
95instance IArray Array where
96 length = aLen
97 {-# INLINE length #-}
99instance IArray (MArray s) where
100 length = maLen
101 {-# INLINE length #-}
102#endif
104-- | Create an uninitialized mutable array.
105new :: forall s. Int -> ST s (MArray s)
106new n
107 | n < 0 || n .&. highBit /= 0 = error $ "Data.Text.Array.new: size overflow"
108 | otherwise = ST $ \s1# ->
109 case newByteArray# len# s1# of
110 (# s2#, marr# #) -> (# s2#, MArray marr#
111#if defined(ASSERTS)
112 n
113#endif
114 #)
115 where !(I# len#) = bytesInArray n
116 highBit = maxBound `xor` (maxBound `shiftR` 1)
117{-# INLINE new #-}
119-- | Freeze a mutable array. Do not mutate the 'MArray' afterwards!
120unsafeFreeze :: MArray s -> ST s Array
121unsafeFreeze MArray{..} = ST $ \s# ->
122 (# s#, Array (unsafeCoerce# maBA)
123#if defined(ASSERTS)
124 maLen
125#endif
126 #)
127{-# INLINE unsafeFreeze #-}
129-- | Indicate how many bytes would be used for an array of the given
130-- size.
131bytesInArray :: Int -> Int
132bytesInArray n = n `shiftL` 1
133{-# INLINE bytesInArray #-}
135-- | Unchecked read of an immutable array. May return garbage or
136-- crash on an out-of-bounds access.
137unsafeIndex :: Array -> Int -> Word16
138unsafeIndex Array{..} i@(I# i#) =
139 CHECK_BOUNDS("unsafeIndex",aLen,i)
140 case indexWord16Array# aBA i# of r# -> (W16# r#)
141{-# INLINE unsafeIndex #-}
143-- | Unchecked read of an immutable array. May return garbage or
144-- crash on an out-of-bounds access.
145unsafeIndexWord :: Array -> Int -> Word
146unsafeIndexWord Array{..} i@(I# i#) =
147 CHECK_BOUNDS("unsafeIndexWord",aLen`div`wordFactor,i)
148 case indexWordArray# aBA i# of r# -> (W# r#)
149{-# INLINE unsafeIndexWord #-}
151-- | Unchecked read of a mutable array. May return garbage or
152-- crash on an out-of-bounds access.
153unsafeRead :: MArray s -> Int -> ST s Word16
154unsafeRead MArray{..} i@(I# i#) = ST $ \s# ->
155 CHECK_BOUNDS("unsafeRead",maLen,i)
156 case readWord16Array# maBA i# s# of
157 (# s2#, r# #) -> (# s2#, W16# r# #)
158{-# INLINE unsafeRead #-}
160-- | Unchecked write of a mutable array. May return garbage or crash
161-- on an out-of-bounds access.
162unsafeWrite :: MArray s -> Int -> Word16 -> ST s ()
163unsafeWrite MArray{..} i@(I# i#) (W16# e#) = ST $ \s1# ->
164 CHECK_BOUNDS("unsafeWrite",maLen,i)
165 case writeWord16Array# maBA i# e# s1# of
166 s2# -> (# s2#, () #)
167{-# INLINE unsafeWrite #-}
169-- | Unchecked read of a mutable array. May return garbage or
170-- crash on an out-of-bounds access.
171unsafeReadWord :: MArray s -> Int -> ST s Word
172unsafeReadWord MArray{..} i@(I# i#) = ST $ \s# ->
173 CHECK_BOUNDS("unsafeRead64",maLen`div`wordFactor,i)
174 case readWordArray# maBA i# s# of
175 (# s2#, r# #) -> (# s2#, W# r# #)
176{-# INLINE unsafeReadWord #-}
178-- | Unchecked write of a mutable array. May return garbage or crash
179-- on an out-of-bounds access.
180unsafeWriteWord :: MArray s -> Int -> Word -> ST s ()
181unsafeWriteWord MArray{..} i@(I# i#) (W# e#) = ST $ \s1# ->
182 CHECK_BOUNDS("unsafeWriteWord",maLen`div`wordFactor,i)
183 case writeWordArray# maBA i# e# s1# of
184 s2# -> (# s2#, () #)
185{-# INLINE unsafeWriteWord #-}
187-- | Convert an immutable array to a list.
188toList :: Array -> Int -> Int -> [Word16]
189toList ary off len = loop 0
190 where loop i | i < len = unsafeIndex ary (off+i) : loop (i+1)
191 | otherwise = []
193-- | An empty immutable array.
194empty :: Array
195empty = runST (new 0 >>= unsafeFreeze)
197-- | Run an action in the ST monad and return an immutable array of
198-- its result.
199run :: (forall s. ST s (MArray s)) -> Array
200run k = runST (k >>= unsafeFreeze)
202-- | Run an action in the ST monad and return an immutable array of
203-- its result paired with whatever else the action returns.
204run2 :: (forall s. ST s (MArray s, a)) -> (Array, a)
205run2 k = runST (do
206 (marr,b) <- k
207 arr <- unsafeFreeze marr
208 return (arr,b))
210-- | The amount to divide or multiply by to switch between units of
211-- 'Word16' and units of 'Word'.
212wordFactor :: Int
213wordFactor = SIZEOF_HSWORD `shiftR` 1
215-- | Indicate whether an offset is word-aligned.
216wordAligned :: Int -> Bool
217wordAligned i = i .&. (wordFactor - 1) == 0
219-- | Copy some elements of a mutable array.
220copyM :: MArray s -- ^ Destination
221 -> Int -- ^ Destination offset
222 -> MArray s -- ^ Source
223 -> Int -- ^ Source offset
224 -> Int -- ^ Count
225 -> ST s ()
226copyM dest didx src sidx count =
227#if defined(ASSERTS)
228 assert (sidx + count <= length src) .
229 assert (didx + count <= length dest) $
230#endif
231 if srem == 0 && drem == 0
232 then fast_loop 0
233 else slow_loop 0
234 where
235 (swidx,srem) = sidx `divMod` wordFactor
236 (dwidx,drem) = didx `divMod` wordFactor
237 nwds = count `div` wordFactor
238 fast_loop !i
239 | i >= nwds = slow_loop (i * wordFactor)
240 | otherwise = do w <- unsafeReadWord src (swidx+i)
241 unsafeWriteWord dest (dwidx+i) w
242 fast_loop (i+1)
243 slow_loop !i
244 | i >= count= return ()
245 | otherwise = do unsafeRead src (sidx+i) >>= unsafeWrite dest (didx+i)
246 slow_loop (i+1)
248-- | Copy some elements of an immutable array.
249copyI :: MArray s -- ^ Destination
250 -> Int -- ^ Destination offset
251 -> Array -- ^ Source
252 -> Int -- ^ Source offset
253 -> Int -- ^ First offset in source /not/ to
254 -- copy (i.e. /not/ length)
255 -> ST s ()
256copyI dest i0 src j0 top
257 | wordAligned i0 && wordAligned j0 = fast (i0 `div` wordFactor) (j0 `div` wordFactor)
258 | otherwise = slow i0 j0
259 where
260 topwds = top `div` wordFactor
261 fast !i !j
262 | i >= topwds = slow (i * wordFactor) (j * wordFactor)
263 | otherwise = do unsafeWriteWord dest i (src `unsafeIndexWord` j)
264 fast (i+1) (j+1)
265 slow !i !j
266 | i >= top = return ()
267 | otherwise = do unsafeWrite dest i (src `unsafeIndex` j)
268 slow (i+1) (j+1)
270-- | Compare portions of two arrays for equality. No bounds checking
271-- is performed.
272equal :: Array -- ^ First
273 -> Int -- ^ Offset into first
274 -> Array -- ^ Second
275 -> Int -- ^ Offset into second
276 -> Int -- ^ Count
277 -> Bool
278equal arrA offA arrB offB count
279 | wordAligned offA && wordAligned offB = fast 0
280 | otherwise = slow 0
281 where
282 countWords = count `div` wordFactor
283 fast !i
284 | i >= countWords = slow (i * wordFactor)
285 | a /= b = False
286 | otherwise = fast (i+1)
287 where a = unsafeIndexWord arrA (offAW+i)
288 b = unsafeIndexWord arrB (offBW+i)
289 offAW = offA `div` wordFactor
290 offBW = offB `div` wordFactor
291 slow !i
292 | i >= count = True
293 | a /= b = False
294 | otherwise = slow (i+1)
295 where a = unsafeIndex arrA (offA+i)
296 b = unsafeIndex arrB (offB+i)