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