git.iko.soy/text
on e8ed4fa help sign in
text/Data/Text/Array.hs
272 lines 8.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
30 , MArray
32 -- * Functions
33 , copyM
34 , copyI
35 , empty
36#if defined(ASSERTS)
37 , length
38#endif
39 , run
40 , run2
41 , toList
42 , unsafeFreeze
43 , unsafeIndex
44 , unsafeNew
45 , unsafeWrite
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#, indexWordArray#, newByteArray#,
66 readWord16Array#, readWordArray#, unsafeCoerce#,
67 writeWord16Array#, writeWordArray#)
68import GHC.ST (ST(..), runST)
69import GHC.Word (Word16(..), Word(..))
70import Prelude hiding (length, read)
72-- | Immutable array type.
73data Array = Array {
74 aBA :: ByteArray#
75#if defined(ASSERTS)
76 , aLen :: {-# UNPACK #-} !Int -- length (in units of Word16, not bytes)
77#endif
78 }
80-- | Mutable array type, for use in the ST monad.
81data MArray s = MArray {
82 maBA :: MutableByteArray# s
83#if defined(ASSERTS)
84 , maLen :: {-# UNPACK #-} !Int -- length (in units of Word16, not bytes)
85#endif
86 }
88#if defined(ASSERTS)
89-- | Operations supported by all arrays.
90class IArray a where
91 -- | Return the length of an array.
92 length :: a -> Int
94instance IArray Array where
95 length = aLen
96 {-# INLINE length #-}
98instance IArray (MArray s) where
99 length = maLen
100 {-# INLINE length #-}
101#endif
103-- | Create an uninitialized mutable array.
104unsafeNew :: forall s. Int -> ST s (MArray s)
105unsafeNew n =
106#if defined(ASSERTS)
107 assert (n >= 0) .
108#endif
109 ST $ \s1# ->
110 case bytesInArray n of
111 len@(I# len#) ->
112#if defined(ASSERTS)
113 if len < 0 then error (show ("unsafeNew",len)) else
114#endif
115 case newByteArray# len# s1# of
116 (# s2#, marr# #) -> (# s2#, MArray marr#
117#if defined(ASSERTS)
118 n
119#endif
120 #)
121{-# INLINE unsafeNew #-}
123-- | Freeze a mutable array. Do not mutate the 'MArray' afterwards!
124unsafeFreeze :: MArray s -> ST s Array
125unsafeFreeze MArray{..} = ST $ \s# ->
126 (# s#, Array (unsafeCoerce# maBA)
127#if defined(ASSERTS)
128 maLen
129#endif
130 #)
131{-# INLINE unsafeFreeze #-}
133-- | Indicate how many bytes would be used for an array of the given
134-- size.
135bytesInArray :: Int -> Int
136bytesInArray n = n `shiftL` 1
137{-# INLINE bytesInArray #-}
139-- | Unchecked read of an immutable array. May return garbage or
140-- crash on an out-of-bounds access.
141unsafeIndex :: Array -> Int -> Word16
142unsafeIndex Array{..} i@(I# i#) =
143 CHECK_BOUNDS("unsafeIndex",aLen,i)
144 case indexWord16Array# aBA i# of r# -> (W16# r#)
145{-# INLINE unsafeIndex #-}
147-- | Unchecked read of an immutable array. May return garbage or
148-- crash on an out-of-bounds access.
149unsafeIndexWord :: Array -> Int -> Word
150unsafeIndexWord Array{..} i@(I# i#) =
151 CHECK_BOUNDS("unsafeIndexWord",aLen`div`wordFactor,i)
152 case indexWordArray# aBA i# of r# -> (W# r#)
153{-# INLINE unsafeIndexWord #-}
155-- | Unchecked read of a mutable array. May return garbage or
156-- crash on an out-of-bounds access.
157unsafeRead :: MArray s -> Int -> ST s Word16
158unsafeRead MArray{..} i@(I# i#) = ST $ \s# ->
159 CHECK_BOUNDS("unsafeRead",maLen,i)
160 case readWord16Array# maBA i# s# of
161 (# s2#, r# #) -> (# s2#, W16# r# #)
162{-# INLINE unsafeRead #-}
164-- | Unchecked write of a mutable array. May return garbage or crash
165-- on an out-of-bounds access.
166unsafeWrite :: MArray s -> Int -> Word16 -> ST s ()
167unsafeWrite MArray{..} i@(I# i#) (W16# e#) = ST $ \s1# ->
168 CHECK_BOUNDS("unsafeWrite",maLen,i)
169 case writeWord16Array# maBA i# e# s1# of
170 s2# -> (# s2#, () #)
171{-# INLINE unsafeWrite #-}
173-- | Unchecked read of a mutable array. May return garbage or
174-- crash on an out-of-bounds access.
175unsafeReadWord :: MArray s -> Int -> ST s Word
176unsafeReadWord MArray{..} i@(I# i#) = ST $ \s# ->
177 CHECK_BOUNDS("unsafeRead64",maLen`div`wordFactor,i)
178 case readWordArray# maBA i# s# of
179 (# s2#, r# #) -> (# s2#, W# r# #)
180{-# INLINE unsafeReadWord #-}
182-- | Unchecked write of a mutable array. May return garbage or crash
183-- on an out-of-bounds access.
184unsafeWriteWord :: MArray s -> Int -> Word -> ST s ()
185unsafeWriteWord MArray{..} i@(I# i#) (W# e#) = ST $ \s1# ->
186 CHECK_BOUNDS("unsafeWriteWord",maLen`div`wordFactor,i)
187 case writeWordArray# maBA i# e# s1# of
188 s2# -> (# s2#, () #)
189{-# INLINE unsafeWriteWord #-}
191-- | Convert an immutable array to a list.
192toList :: Array -> Int -> Int -> [Word16]
193toList ary off len = loop 0
194 where loop i | i < len = unsafeIndex ary (off+i) : loop (i+1)
195 | otherwise = []
197-- | An empty immutable array.
198empty :: Array
199empty = runST (unsafeNew 0 >>= unsafeFreeze)
201-- | Run an action in the ST monad and return an immutable array of
202-- its result.
203run :: (forall s. ST s (MArray s)) -> Array
204run k = runST (k >>= unsafeFreeze)
206-- | Run an action in the ST monad and return an immutable array of
207-- its result paired with whatever else the action returns.
208run2 :: (forall s. ST s (MArray s, a)) -> (Array, a)
209run2 k = runST (do
210 (marr,b) <- k
211 arr <- unsafeFreeze marr
212 return (arr,b))
214-- | The amount to divide or multiply by to switch between units of
215-- 'Word16' and units of 'Word'.
216wordFactor :: Int
217wordFactor = SIZEOF_HSWORD `shiftR` 1
219-- | Indicate whether an offset is word-aligned.
220wordAligned :: Int -> Bool
221wordAligned i = i .&. (wordFactor - 1) == 0
223-- | Copy some elements of a mutable array.
224copyM :: MArray s -- ^ Destination
225 -> Int -- ^ Destination offset
226 -> MArray s -- ^ Source
227 -> Int -- ^ Source offset
228 -> Int -- ^ Count
229 -> ST s ()
230copyM dest didx src sidx count =
231#if defined(ASSERTS)
232 assert (sidx + count <= length src) .
233 assert (didx + count <= length dest) $
234#endif
235 if srem == 0 && drem == 0
236 then fast_loop 0
237 else slow_loop 0
238 where
239 (swidx,srem) = sidx `divMod` wordFactor
240 (dwidx,drem) = didx `divMod` wordFactor
241 nwds = count `div` wordFactor
242 fast_loop !i
243 | i >= nwds = slow_loop (i * wordFactor)
244 | otherwise = do w <- unsafeReadWord src (swidx+i)
245 unsafeWriteWord dest (dwidx+i) w
246 fast_loop (i+1)
247 slow_loop !i
248 | i >= count= return ()
249 | otherwise = do unsafeRead src (sidx+i) >>= unsafeWrite dest (didx+i)
250 slow_loop (i+1)
252-- | Copy some elements of an immutable array.
253copyI :: MArray s -- ^ Destination
254 -> Int -- ^ Destination offset
255 -> Array -- ^ Source
256 -> Int -- ^ Source offset
257 -> Int -- ^ First offset in source /not/ to
258 -- copy (i.e. /not/ length)
259 -> ST s ()
260copyI dest i0 src j0 top
261 | wordAligned i0 && wordAligned j0 = fast (i0 `div` wordFactor) (j0 `div` wordFactor)
262 | otherwise = slow i0 j0
263 where
264 topwds = top `div` wordFactor
265 fast !i !j
266 | i >= topwds = slow (i * wordFactor) (j * wordFactor)
267 | otherwise = do unsafeWriteWord dest i (src `unsafeIndexWord` j)
268 fast (i+1) (j+1)
269 slow !i !j
270 | i >= top = return ()
271 | otherwise = do unsafeWrite dest i (src `unsafeIndex` j)
272 slow (i+1) (j+1)