Ticket #5302: T5302_full.hs

File T5302_full.hs, 12.8 KB (added by reinerp, 4 years ago)
Line 
1{-# LANGUAGE   MultiParamTypeClasses,
2                 FunctionalDependencies,
3                 FlexibleInstances,
4                 UndecidableInstances,
5                 ScopedTypeVariables,
6                 UnboxedTuples #-}
7
8module T5302_full (
9        mySplit,
10        ) where
11
12import Prelude hiding (null, reverse)
13
14import Data.Monoid
15import GHC.Magic(inline)
16
17data Digit a
18        = One a
19        | Two a a
20        | Three a a a
21        | Four a a a a
22
23-----------------------------------------------------------------------------------------------
24-- Template Haskell part
25-----------------------------------------------------------------------------------------------
26
27newtype Size = Size { unSize :: Int }
28
29data FingerTree v a = Empty | Single a | Deep !v !(Digit a) (FingerTree v (Node v a)) !(Digit a)
30data Node v a = Node2 !v a a | Node3 !v a a a
31
32
33----------------------------------------------------------------------------------------------------
34--                                                   Types                                        --
35----------------------------------------------------------------------------------------------------
36-- | View of the left end of a sequence.
37data ViewL s a
38        = EmptyL        -- ^ empty sequence
39        | a :< s a      -- ^ leftmost element and the rest of the sequence
40
41
42-- | View of the right end of a sequence.
43data ViewR s a
44        = EmptyR        -- ^ empty sequence
45        | s a :> a      -- ^ the sequence minus the rightmost element,
46                        -- and the rightmost element
47
48
49instance (Measured v a) => Measured v (FingerTree v a) where
50        measure f = case f of
51            Empty -> mempty
52            Single x -> measure x
53            Deep v _ _ _ -> v
54        {-# INLINABLE measure #-}
55
56--data Split t a = Split t a t
57type Split t a = (# t, a, t #)
58mkSplit :: t -> a -> t -> Split t a
59mkSplit a b c = (# a, b, c #)
60 
61----------------------------------------------------------------------------------------------------
62--                                         Some cheap typeclasses                                 --
63----------------------------------------------------------------------------------------------------
64-- | Things that can be measured.
65class (Monoid v) => Measured v a | a -> v where
66        measure :: a -> v
67
68instance (Monoid v) => Measured v (Node v a) where
69    measure node = case node of
70        Node2 v _ _ -> v
71        Node3 v _ _ _ -> v
72    {-# INLINABLE measure #-}
73
74myMeasure :: Node v b -> v
75myMeasure (Node2 v _ _) = v
76myMeasure (Node3 v _ _ _) = v
77
78measureFT :: Monoid v => FingerTree v (Node v b) -> v
79measureFT f = case f of
80  Empty -> mempty
81  Single x -> myMeasure x
82  Deep v _ _ _ -> v
83
84measureDigit :: Monoid v => Digit (Node v b) -> v
85measureDigit (One a) = myMeasure a
86measureDigit (Two a b) = myMeasure a `mappend` myMeasure b
87measureDigit (Three a b c) = myMeasure a `mappend` myMeasure b `mappend` myMeasure c
88measureDigit (Four a b c d) = myMeasure a `mappend` myMeasure b `mappend` myMeasure c `mappend` myMeasure d
89------------------------------------------------------------------------------------------
90--                                         Smart constructors                           --
91------------------------------------------------------------------------------------------
92node2        ::  Monoid v => Node v b -> Node v b -> Node v (Node v b)
93node2 a b    =   Node2 (myMeasure a `mappend` myMeasure b) a b
94
95node3        ::  Monoid v => Node v b -> Node v b -> Node v b -> Node v (Node v b)
96node3 a b c  =   Node3 (myMeasure a `mappend` myMeasure b `mappend` myMeasure c) a b c
97
98nodeToDigit :: Node v b -> Digit b
99nodeToDigit node = case node of
100  Node2 _ a b -> Two a b
101  Node3 _ a b c -> Three a b c
102
103deep ::  Monoid v => Digit (Node v b) -> FingerTree v (Node v (Node v b)) -> Digit (Node v b) -> FingerTree v (Node v b)
104deep pr m sf = Deep ((measureDigit pr `mappendVal` m) `mappend` measureDigit sf) pr m sf
105
106empty :: FingerTree v b
107empty =  Empty
108 
109empty' :: FingerTree v b
110empty' =  Empty
111
112mappendVal :: Monoid v => v -> FingerTree v (Node v b) -> v
113mappendVal v t = v `seq` case  t of
114    Empty -> v
115    _ -> v `mappend` measureFT t
116{-# INLINE mappendVal #-}
117
118------------------------------------------------------------------------------------------
119--                                    Empty and singleton                               --
120------------------------------------------------------------------------------------------
121-- | /O(1)/. The empty sequence.
122emptyD :: FingerTree v a
123emptyD = empty'
124
125-- | /O(1)/. A singleton sequence.
126singletonD :: a -> FingerTree v a
127singletonD = Single
128
129------------------------------------------------------------------------------------------
130--                                         Cons                                         --
131------------------------------------------------------------------------------------------
132-- | /O(1)/. Add an element to the left end of a sequence.
133-- Mnemonic: a triangle with the single element at the pointy end.
134consD :: Monoid v => Node v b -> FingerTree v (Node v b) -> FingerTree v (Node v b)
135consD a t = case  t of
136      Empty -> Single a
137      Single b -> deep (One a) empty' (One b)
138      Deep v (Four b c d e) m sf -> m `seq`
139          (Deep (myMeasure a `mappend` v) (Two a b) (consD (node3 c d e) m) sf)
140      Deep v pr m sf -> Deep (myMeasure a `mappend` v) (consDigit a pr) m sf
141
142consDigit :: b -> Digit b -> Digit b
143consDigit a (One b) = Two a b
144consDigit a (Two b c) = Three a b c
145consDigit a (Three b c d) = Four a b c d
146
147------------------------------------------------------------------------------------------
148--                                         Snoc                                         --
149------------------------------------------------------------------------------------------
150-- | /O(Add an element to the right end of a sequence.
151-- Mnemonic: a triangle with the single element at the pointy end.
152snocD :: Monoid v => FingerTree v (Node v b) -> Node v b -> FingerTree v (Node v b)
153snocD t l = case  t of
154      Empty -> Single l
155      Single a -> deep (One a) empty' (One l)
156      Deep v pr m (Four a b c d) -> m `seq`
157         (Deep (v `mappend` myMeasure l) pr (snocD m (node3 a b c)) (Two d l))
158      Deep v pr m sf -> Deep (v `mappend` myMeasure l) pr m (snocDigit sf l)
159
160snocDigit :: Digit b -> b -> Digit b
161snocDigit (One a) b = Two a b
162snocDigit (Two a b) c = Three a b c
163snocDigit (Three a b c) d = Four a b c d
164
165------------------------------------------------------------------------------------------
166--                                         viewl                                        --
167------------------------------------------------------------------------------------------
168-- | /O(1)/. Analyse the left end of a sequence.
169viewlD :: Monoid v => FingerTree v (Node v b) -> ViewL (FingerTree v) (Node v b)
170viewlD f = case  f of
171  Empty -> EmptyL
172  Single x -> x :< empty'
173  Deep _ pr m sf -> case viewlDigit pr of
174     Left x -> x :< rotL m sf
175     Right (x, pr') -> x :< deep pr' m sf
176
177rotL :: Monoid v => FingerTree v (Node v (Node v b)) -> Digit (Node v b) -> FingerTree v (Node v b)
178rotL m sf      =   case viewlD m of
179        EmptyL  ->  digitToTree sf
180        node :< m' ->  case node of
181          Node2 _ a b -> Deep (measureFT m `mappend` measureDigit sf) (Two a b) m' sf
182          Node3 _ a b c ->  Deep (measureFT m `mappend` measureDigit sf) (Three a b c) m' sf
183
184
185viewlDigit :: Digit b -> Either b (b, Digit b)
186viewlDigit (One a) = Left a
187viewlDigit (Two a b) = Right (a, One b)
188viewlDigit (Three a b c) = Right (a, Two b c)
189viewlDigit (Four a b c d) = Right (a, Three b c d)
190
191digitToTree :: Monoid v => Digit (Node v b) -> FingerTree v (Node v b)
192digitToTree (One a) = Single a
193digitToTree (Two a b) = deep (One a) empty' (One b)
194digitToTree (Three a b c) = deep (Two a b) empty' (One c)
195digitToTree (Four a b c d) = deep (Two a b) empty' (Two c d)
196
197
198------------------------------------------------------------------------------------------
199--                                         viewr                                        --
200------------------------------------------------------------------------------------------
201-- | /O(1)/. Analyse the right end of a sequence.
202viewrD :: Monoid v => FingerTree v (Node v b) -> ViewR (FingerTree v) (Node v b)
203viewrD f = case f of
204    Empty -> EmptyR
205    Single x -> empty' :> x
206    Deep _ pr m sf -> case viewrDigit sf of
207        Left x -> rotR pr m :> x
208        Right (sf', x) -> deep pr m sf' :> x
209
210
211rotR :: Monoid v => Digit (Node v b) -> FingerTree v (Node v (Node v b)) -> FingerTree v (Node v b)
212rotR pr m = case viewrD m of
213    EmptyR      ->  digitToTree pr
214    m' :> node -> case node of
215        Node2 _ a b -> Deep (measureDigit pr `mappend` measureFT m) pr m' (Two a b)
216        Node3 _ a b c -> Deep (measureDigit pr `mappend` measureFT m) pr m' (Three a b c)
217
218
219viewrDigit :: Digit b -> Either b (Digit b, b)
220viewrDigit (One a) = Left a
221viewrDigit (Two a b) = Right (One a, b)
222viewrDigit (Three a b c) = Right (Two a b, c)
223viewrDigit (Four a b c d) = Right (Three a b c, d)
224
225------------------------------------------------------------------------------------------
226--                                         Splitting                                    --
227------------------------------------------------------------------------------------------
228-- | /O(log(min(i,n-i)))/. Split a sequence at a point where the predicate
229-- on the accumulated measure changes from 'False' to 'True'.
230split :: forall v c. Monoid v => (v -> Bool) -> FingerTree v (Node v c) -> (FingerTree v (Node v c), FingerTree v (Node v c))
231split p t = case t of
232      Empty -> (empty, empty)
233      _ -> case splitTree mempty t of (# l, x, r #) -> (l, consD x r)
234   where
235
236  -- we manually CPR this call, because GHC apparently doesn't want to
237  splitTree :: v -> FingerTree v (Node v b) -> Split (FingerTree v (Node v b)) (Node v b)
238  splitTree i tree = case  tree of
239      Empty -> error "splitTree of empty tree (possible violation of monoid invariant?)"
240      Single x -> i `seq` mkSplit empty x empty
241      Deep _ pr m sf
242        | p vpr ->  case splitDigit i pr of
243                       (# l, x, r #) -> mkSplit (maybe empty digitToTree l) x (deepL r m sf)
244        | p vm ->  case splitTree vpr m of
245                     (# ml, xs, mr #) -> case splitNode (vpr `mappend` measureFT ml) xs of
246                       (# l, x, r #) -> mkSplit (deepR pr  ml l) x (deepL r mr sf)
247                         where 
248        | otherwise -> case splitDigit vm sf of
249                          (# l, x, r #) -> mkSplit (deepR pr  m  l) x (maybe empty digitToTree r)
250       where 
251         vpr =  i    `mappend`  measureDigit pr
252         vm =  vpr  `mappend` measureFT m
253 
254 
255  splitNode :: v -> Node v (Node v b) -> Split (Maybe (Digit (Node v b))) (Node v b)
256  splitNode i node = case node of
257    Node2 _ a b
258      | p va       -> mkSplit Nothing a (Just (One b))
259      | otherwise  -> mkSplit (Just (One a)) b Nothing
260     where      va      = i `mappend` myMeasure a
261    Node3 _ a b c
262      | p va    -> mkSplit Nothing a (Just (Two b c))
263      | p vab   -> mkSplit (Just (One a)) b (Just (One c))
264      | otherwise       -> mkSplit (Just (Two a b)) c Nothing
265     where va  = i `mappend` myMeasure a
266           vab = va `mappend` myMeasure b
267   
268  splitDigit :: v -> Digit (Node v b) -> Split (Maybe (Digit (Node v b))) (Node v b)
269  splitDigit i (One a) = i `seq` mkSplit Nothing a Nothing
270  splitDigit i (Two a b)
271    | p va      = mkSplit Nothing a (Just (One b))
272    | otherwise = mkSplit (Just (One a)) b Nothing
273    where       va      = i `mappend` myMeasure a
274  splitDigit i (Three a b c)
275    | p va      = mkSplit Nothing a (Just (Two b c))
276    | p vab     = mkSplit (Just (One a)) b (Just (One c))
277    | otherwise = mkSplit (Just (Two a b)) c Nothing
278    where va    = i `mappend` myMeasure a
279          vab   = va `mappend` myMeasure b
280  splitDigit i (Four a b c d)
281    | p va      = mkSplit Nothing a (Just (Three b c d))
282    | p vab     = mkSplit (Just (One a)) b (Just (Two c d))
283    | p vabc    = mkSplit (Just (Two a b)) c (Just (One d))
284    | otherwise = mkSplit (Just (Three a b c)) d Nothing
285    where va    = i `mappend` myMeasure a
286          vab   = va `mappend` myMeasure b
287          vabc  = vab `mappend` myMeasure c
288
289
290deepL :: Monoid v => Maybe (Digit (Node v b)) -> FingerTree v (Node v (Node v b)) -> Digit (Node v b) -> FingerTree v (Node v b)
291deepL Nothing m sf      =   rotL m sf
292deepL (Just pr) m sf    =   deep pr m sf
293
294 
295deepR :: Monoid v => Digit (Node v b) -> FingerTree v (Node v (Node v b)) -> Maybe (Digit (Node v b)) -> FingerTree v (Node v b)
296deepR pr m Nothing      =   rotR pr m
297deepR pr m (Just sf)    =   deep pr m sf
298------------------------------------------------------------------------------------------
299
300newtype Elem a = Elem { unElem :: a }
301newtype Seq a = Seq { unSeq :: FingerTree Size (Node Size a) }
302
303instance Monoid Size where
304    {-# INLINE mappend #-}
305    mappend (Size a) (Size b) = Size (a + b)
306    {-# INLINE mempty #-}
307    mempty = Size 0
308
309mySplit :: Int -> Seq a -> (Seq a, Seq a)
310mySplit n (Seq a) = n `seq` case split (\(Size s) -> s>n) a of (l, r) -> (Seq l, Seq r)
311