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