git.iko.soy/text
on cedee36 help sign in
text/Data/Text/Array.hs
267 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#if defined(ASSERTS)
37 , length
38#endif
39 , run
40 , run2
41 , toList
42 , unsafeFreeze
43 , unsafeIndex
44 , new
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 ((.&.), xor)
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.
104new :: forall s. Int -> ST s (MArray s)
105new n
106 | n < 0 || n .&. highBit /= 0 = error $ "Data.Text.Array.new: size overflow"
107 | otherwise = ST $ \s1# ->
108 case newByteArray# len# s1# of
109 (# s2#, marr# #) -> (# s2#, MArray marr#
110#if defined(ASSERTS)
111 n
112#endif
113 #)
114 where !(I# len#) = bytesInArray n
115 highBit = maxBound `xor` (maxBound `shiftR` 1)
116{-# INLINE new #-}
118-- | Freeze a mutable array. Do not mutate the 'MArray' afterwards!
119unsafeFreeze :: MArray s -> ST s Array
120unsafeFreeze MArray{..} = ST $ \s# ->
121 (# s#, Array (unsafeCoerce# maBA)
122#if defined(ASSERTS)
123 maLen
124#endif
125 #)
126{-# INLINE unsafeFreeze #-}
128-- | Indicate how many bytes would be used for an array of the given
129-- size.
130bytesInArray :: Int -> Int
131bytesInArray n = n `shiftL` 1
132{-# INLINE bytesInArray #-}
134-- | Unchecked read of an immutable array. May return garbage or
135-- crash on an out-of-bounds access.
136unsafeIndex :: Array -> Int -> Word16
137unsafeIndex Array{..} i@(I# i#) =
138 CHECK_BOUNDS("unsafeIndex",aLen,i)
139 case indexWord16Array# aBA i# of r# -> (W16# r#)
140{-# INLINE unsafeIndex #-}
142-- | Unchecked read of an immutable array. May return garbage or
143-- crash on an out-of-bounds access.
144unsafeIndexWord :: Array -> Int -> Word
145unsafeIndexWord Array{..} i@(I# i#) =
146 CHECK_BOUNDS("unsafeIndexWord",aLen`div`wordFactor,i)
147 case indexWordArray# aBA i# of r# -> (W# r#)
148{-# INLINE unsafeIndexWord #-}
150-- | Unchecked read of a mutable array. May return garbage or
151-- crash on an out-of-bounds access.
152unsafeRead :: MArray s -> Int -> ST s Word16
153unsafeRead MArray{..} i@(I# i#) = ST $ \s# ->
154 CHECK_BOUNDS("unsafeRead",maLen,i)
155 case readWord16Array# maBA i# s# of
156 (# s2#, r# #) -> (# s2#, W16# r# #)
157{-# INLINE unsafeRead #-}
159-- | Unchecked write of a mutable array. May return garbage or crash
160-- on an out-of-bounds access.
161unsafeWrite :: MArray s -> Int -> Word16 -> ST s ()
162unsafeWrite MArray{..} i@(I# i#) (W16# e#) = ST $ \s1# ->
163 CHECK_BOUNDS("unsafeWrite",maLen,i)
164 case writeWord16Array# maBA i# e# s1# of
165 s2# -> (# s2#, () #)
166{-# INLINE unsafeWrite #-}
168-- | Unchecked read of a mutable array. May return garbage or
169-- crash on an out-of-bounds access.
170unsafeReadWord :: MArray s -> Int -> ST s Word
171unsafeReadWord MArray{..} i@(I# i#) = ST $ \s# ->
172 CHECK_BOUNDS("unsafeRead64",maLen`div`wordFactor,i)
173 case readWordArray# maBA i# s# of
174 (# s2#, r# #) -> (# s2#, W# r# #)
175{-# INLINE unsafeReadWord #-}
177-- | Unchecked write of a mutable array. May return garbage or crash
178-- on an out-of-bounds access.
179unsafeWriteWord :: MArray s -> Int -> Word -> ST s ()
180unsafeWriteWord MArray{..} i@(I# i#) (W# e#) = ST $ \s1# ->
181 CHECK_BOUNDS("unsafeWriteWord",maLen`div`wordFactor,i)
182 case writeWordArray# maBA i# e# s1# of
183 s2# -> (# s2#, () #)
184{-# INLINE unsafeWriteWord #-}
186-- | Convert an immutable array to a list.
187toList :: Array -> Int -> Int -> [Word16]
188toList ary off len = loop 0
189 where loop i | i < len = unsafeIndex ary (off+i) : loop (i+1)
190 | otherwise = []
192-- | An empty immutable array.
193empty :: Array
194empty = runST (new 0 >>= unsafeFreeze)
196-- | Run an action in the ST monad and return an immutable array of
197-- its result.
198run :: (forall s. ST s (MArray s)) -> Array
199run k = runST (k >>= unsafeFreeze)
201-- | Run an action in the ST monad and return an immutable array of
202-- its result paired with whatever else the action returns.
203run2 :: (forall s. ST s (MArray s, a)) -> (Array, a)
204run2 k = runST (do
205 (marr,b) <- k
206 arr <- unsafeFreeze marr
207 return (arr,b))
209-- | The amount to divide or multiply by to switch between units of
210-- 'Word16' and units of 'Word'.
211wordFactor :: Int
212wordFactor = SIZEOF_HSWORD `shiftR` 1
214-- | Indicate whether an offset is word-aligned.
215wordAligned :: Int -> Bool
216wordAligned i = i .&. (wordFactor - 1) == 0
218-- | Copy some elements of a mutable array.
219copyM :: MArray s -- ^ Destination
220 -> Int -- ^ Destination offset
221 -> MArray s -- ^ Source
222 -> Int -- ^ Source offset
223 -> Int -- ^ Count
224 -> ST s ()
225copyM dest didx src sidx count =
226#if defined(ASSERTS)
227 assert (sidx + count <= length src) .
228 assert (didx + count <= length dest) $
229#endif
230 if srem == 0 && drem == 0
231 then fast_loop 0
232 else slow_loop 0
233 where
234 (swidx,srem) = sidx `divMod` wordFactor
235 (dwidx,drem) = didx `divMod` wordFactor
236 nwds = count `div` wordFactor
237 fast_loop !i
238 | i >= nwds = slow_loop (i * wordFactor)
239 | otherwise = do w <- unsafeReadWord src (swidx+i)
240 unsafeWriteWord dest (dwidx+i) w
241 fast_loop (i+1)
242 slow_loop !i
243 | i >= count= return ()
244 | otherwise = do unsafeRead src (sidx+i) >>= unsafeWrite dest (didx+i)
245 slow_loop (i+1)
247-- | Copy some elements of an immutable array.
248copyI :: MArray s -- ^ Destination
249 -> Int -- ^ Destination offset
250 -> Array -- ^ Source
251 -> Int -- ^ Source offset
252 -> Int -- ^ First offset in source /not/ to
253 -- copy (i.e. /not/ length)
254 -> ST s ()
255copyI dest i0 src j0 top
256 | wordAligned i0 && wordAligned j0 = fast (i0 `div` wordFactor) (j0 `div` wordFactor)
257 | otherwise = slow i0 j0
258 where
259 topwds = top `div` wordFactor
260 fast !i !j
261 | i >= topwds = slow (i * wordFactor) (j * wordFactor)
262 | otherwise = do unsafeWriteWord dest i (src `unsafeIndexWord` j)
263 fast (i+1) (j+1)
264 slow !i !j
265 | i >= top = return ()
266 | otherwise = do unsafeWrite dest i (src `unsafeIndex` j)
267 slow (i+1) (j+1)