@@ -0,0 +1,222 @@
1+{-# LANGUAGE BangPatterns, Rank2Types, UnboxedTuples #-}
2+
3+-----------------------------------------------------------------------------
4+-- |
5+-- Module : Data.Text.Builder
6+-- Copyright : Johan Tibell
7+-- License : BSD3-style (see LICENSE)
8+--
9+-- Maintainer : Johan Tibell <johan.tibell@gmail.com>
10+-- Stability : experimental
11+-- Portability : portable to Hugs and GHC
12+--
13+-- Efficient construction of lazy texts.
14+--
15+-----------------------------------------------------------------------------
16+
17+module Data.Text.Builder
18+ ( -- * The Builder type
19+ Builder
20+ , toLazyText
21+
22+ -- * Constructing Builders
23+ , singleton
24+ , fromText
25+ , fromLazyText
26+
27+ -- * Flushing the buffer state
28+ , flush
29+ ) where
30+
31+import Control.Monad.ST
32+import Data.Bits ((.&.))
33+import Data.Char (ord)
34+import Data.Monoid (Monoid(..))
35+import Data.Text.Internal (Text(..))
36+import Data.Text.Lazy.Internal (defaultChunkSize)
37+import Data.Text.UnsafeShift (shiftR)
38+import Data.Word
39+import GHC.ST (ST(..))
40+import Prelude hiding (map, putChar)
41+
42+import qualified Data.Text as S
43+import qualified Data.Text.Array as A
44+import qualified Data.Text.Lazy as L
45+
46+------------------------------------------------------------------------
47+
48+-- | A 'Builder' is an efficient way to build lazy 'L.Text's. There
49+-- are several functions for constructing 'Builder's, but only one to
50+-- inspect them: to extract any data, you have to turn them into lazy
51+-- 'L.Text's using 'toLazyText'.
52+--
53+-- Internally, a 'Builder' constructs a lazy 'L.Text' by filling byte
54+-- arrays piece by piece. As each buffer is filled, it is \'popped\'
55+-- off, to become a new chunk of the resulting lazy 'L.Text'. All
56+-- this is hidden from the user of the 'Builder'.
57+newtype Builder = Builder {
58+ -- Invariant (from Data.Text.Lazy):
59+ -- The lists include no null Texts.
60+ runBuilder :: forall s. (Buffer s -> ST s [S.Text])
61+ -> Buffer s
62+ -> ST s [S.Text]
63+ }
64+
65+instance Monoid Builder where
66+ mempty = empty
67+ {-# INLINE mempty #-}
68+ mappend = append
69+ {-# INLINE mappend #-}
70+
71+------------------------------------------------------------------------
72+
73+-- | /O(1)./ The empty Builder, satisfying
74+--
75+-- * @'toLazyText' 'empty' = 'L.empty'@
76+--
77+empty :: Builder
78+empty = Builder id
79+{-# INLINE empty #-}
80+
81+-- | /O(1)./ A Builder taking a single character, satisfying
82+--
83+-- * @'toLazyText' ('singleton' c) = 'L.singleton' c@
84+--
85+singleton :: Char -> Builder
86+singleton c = putChar c
87+{-# INLINE singleton #-}
88+
89+------------------------------------------------------------------------
90+
91+-- | /O(1)./ The concatenation of two Builders, an associative
92+-- operation with identity 'empty', satisfying
93+--
94+-- * @'toLazyText' ('append' x y) = 'L.append' ('toLazyText' x) ('toLazyText' y)@
95+--
96+append :: Builder -> Builder -> Builder
97+append (Builder f) (Builder g) = Builder (f . g)
98+{-# INLINE append #-}
99+
100+-- TODO: Experiment to find the right threshold.
101+copyLimit :: Int
102+copyLimit = 128
103+
104+-- This function attempts to merge small Texts instead of treating the
105+-- text as its own chunk. We may not always want this.
106+
107+-- | /O(1)./ A Builder taking a 'S.Text', satisfying
108+--
109+-- * @'toLazyText' ('fromText' t) = 'L.fromChunks' [t]@
110+--
111+fromText :: S.Text -> Builder
112+fromText t@(Text arr off l)
113+ | S.null t = empty
114+ | l <= copyLimit = writeN l $ \marr o ->
115+ sequence_ [A.unsafeWrite marr (o+i) (A.unsafeIndex arr (off+i)) |
116+ i <- [0..l-1]]
117+ | otherwise = flush `append` mapBuilder (t :)
118+{-# INLINE fromText #-}
119+
120+-- | /O(1)./ A Builder taking a lazy 'L.Text', satisfying
121+--
122+-- * @'toLazyText' ('fromLazyText' t) = t@
123+--
124+fromLazyText :: L.Text -> Builder
125+fromLazyText ts = flush `append` mapBuilder (L.toChunks ts ++)
126+{-# INLINE fromLazyText #-}
127+
128+------------------------------------------------------------------------
129+
130+-- Our internal buffer type
131+data Buffer s = Buffer {-# UNPACK #-} !(A.MArray s Word16)
132+ {-# UNPACK #-} !Int -- offset
133+ {-# UNPACK #-} !Int -- used units
134+ {-# UNPACK #-} !Int -- length left
135+
136+------------------------------------------------------------------------
137+
138+-- | /O(n)./ Extract a lazy 'L.T' from a 'Builder'. The construction
139+-- work takes place if and when the relevant part of the lazy 'L.T' is
140+-- demanded.
141+--
142+toLazyText :: Builder -> L.Text
143+toLazyText m = L.fromChunks $ runST $ do
144+ buf <- newBuffer defaultChunkSize
145+ runBuilder (m `append` flush) (const (return [])) buf
146+
147+-- | /O(1)./ Pop the 'S.Text' we have constructed so far, if any,
148+-- yielding a new chunk in the result lazy 'L.Text'.
149+flush :: Builder
150+flush = Builder $ \ k buf@(Buffer p o u l) ->
151+ if u == 0
152+ then k buf
153+ else do arr <- A.unsafeFreeze p
154+ buf' <- inlineInterleaveST (k (Buffer p (o+u) 0 l))
155+ return (Text arr o u : buf')
156+
157+------------------------------------------------------------------------
158+
159+-- | Sequence an ST operation on the buffer
160+unsafeLiftST :: (forall s. Buffer s -> ST s (Buffer s)) -> Builder
161+unsafeLiftST f = Builder $ \k buf -> inlineInterleaveST $ do
162+ buf' <- f buf
163+ k buf'
164+{-# INLINE unsafeLiftST #-}
165+
166+-- | Get the size of the buffer
167+withSize :: (Int -> Builder) -> Builder
168+withSize f = Builder $ \ k buf@(Buffer _ _ _ l) ->
169+ runBuilder (f l) k buf
170+
171+-- | Map the resulting list of texts.
172+mapBuilder :: ([S.Text] -> [S.Text]) -> Builder
173+mapBuilder f = Builder $ \ k buf -> do txt <- k buf
174+ return (f txt)
175+
176+------------------------------------------------------------------------
177+
178+putChar :: Char -> Builder
179+putChar c
180+ | n < 0x10000 = writeN 1 $ \marr o -> A.unsafeWrite marr o (fromIntegral n)
181+ | otherwise = writeN 2 $ \marr o -> do
182+ A.unsafeWrite marr o lo
183+ A.unsafeWrite marr (o+1) hi
184+ where n = ord c
185+ m = n - 0x10000
186+ lo = fromIntegral $ (m `shiftR` 10) + 0xD800
187+ hi = fromIntegral $ (m .&. 0x3FF) + 0xDC00
188+{-# INLINE putChar #-}
189+
190+------------------------------------------------------------------------
191+
192+-- | Ensure that there are at least @n@ many elements available.
193+ensureFree :: Int -> Builder
194+ensureFree n = n `seq` withSize $ \ l ->
195+ if n <= l
196+ then empty
197+ else flush `append` unsafeLiftST (const (newBuffer (max n defaultChunkSize)))
198+{-# INLINE ensureFree #-}
199+
200+-- | Ensure that @n@ many elements are available, and then use @f@ to
201+-- write some elements into the memory.
202+writeN :: Int -> (forall s. A.MArray s Word16 -> Int -> ST s ()) -> Builder
203+writeN n f = ensureFree 1 `append` unsafeLiftST (writeNBuffer n f)
204+{-# INLINE writeN #-}
205+
206+writeNBuffer :: Int -> (A.MArray s Word16 -> Int -> ST s ()) -> (Buffer s)
207+ -> ST s (Buffer s)
208+writeNBuffer n f (Buffer p o u l) = do
209+ f p (o+u)
210+ return (Buffer p o (u+n) (l-n))
211+{-# INLINE writeNBuffer #-}
212+
213+newBuffer :: Int -> ST s (Buffer s)
214+newBuffer size = do
215+ arr <- A.unsafeNew size
216+ return $! Buffer arr 0 0 size
217+{-# INLINE newBuffer #-}
218+
219+inlineInterleaveST :: ST s a -> ST s a
220+inlineInterleaveST (ST m) = ST $ \ s ->
221+ let r = case m s of (# _, res #) -> res in (# s, r #)
222+{-# INLINE inlineInterleaveST #-}