Ticket #5302: T5302.hs

File T5302.hs, 4.7 KB (added by reinerp, 4 years ago)
Line 
1{-# LANGUAGE ScopedTypeVariables, UnboxedTuples #-}
2module Unboxed (mySplit) where
3
4import Prelude hiding (null, reverse)
5
6import Data.Monoid
7
8data Digit a
9        = One a
10        | Two a a
11        | Three a a a
12        | Four a a a a
13
14newtype Size = Size Int
15
16data FingerTree v a = Empty | Single a | Deep !v !(Digit a) (FingerTree v (Node v a)) !(Digit a)
17data Node v a = Node2 !v a a | Node3 !v a a a
18
19type Split t a = (# t, a, t #)
20 
21
22-- measurements
23myMeasure :: Node v b -> v
24myMeasure (Node2 v _ _) = v
25myMeasure (Node3 v _ _ _) = v
26
27measureFT :: Monoid v => FingerTree v (Node v b) -> v
28measureFT f = case f of
29  Empty -> mempty
30  Single x -> myMeasure x
31  Deep v _ _ _ -> v
32
33measureDigit :: Monoid v => Digit (Node v b) -> v
34measureDigit (One a) = myMeasure a
35measureDigit (Two a b) = myMeasure a `mappend` myMeasure b
36measureDigit (Three a b c) = myMeasure a `mappend` myMeasure b `mappend` myMeasure c
37measureDigit (Four a b c d) = myMeasure a `mappend` myMeasure b `mappend` myMeasure c `mappend` myMeasure d
38
39deep ::  Monoid v => Digit (Node v b) -> FingerTree v (Node v (Node v b)) -> Digit (Node v b) -> FingerTree v (Node v b)
40deep pr m sf = Deep ((measureDigit pr `mappendVal` m) `mappend` measureDigit sf) pr m sf
41
42empty :: FingerTree v b
43empty =  Empty
44 
45mappendVal :: Monoid v => v -> FingerTree v (Node v b) -> v
46mappendVal v t = v `seq` case  t of
47    Empty -> v
48    _ -> v `mappend` measureFT t
49{-# INLINE mappendVal #-}
50
51digitToTree :: Monoid v => Digit (Node v b) -> FingerTree v (Node v b)
52digitToTree (One a) = Single a
53digitToTree (Two a b) = deep (One a) empty (One b)
54digitToTree (Three a b c) = deep (Two a b) empty (One c)
55digitToTree (Four a b c d) = deep (Two a b) empty (Two c d)
56
57------------------------------------------------------------------------------------------
58--                                         Splitting                                    --
59------------------------------------------------------------------------------------------
60-- | /O(log(min(i,n-i)))/. Split a sequence at a point where the predicate
61-- on the accumulated measure changes from 'False' to 'True'.
62split :: forall v c. Monoid v => (v -> Bool) -> FingerTree v (Node v c) -> (FingerTree v (Node v c), FingerTree v (Node v c))
63split p t = case splitTree mempty t of (# l, x, r #) -> (l, undefined)
64--      _ | otherwise -> (t, empty)
65   where
66
67  -- we manually CPR this call, because GHC apparently doesn't want to
68  splitTree :: v -> FingerTree v (Node v b) -> Split (FingerTree v (Node v b)) (Node v b)
69  splitTree i tree = case  tree of
70      Empty -> error "splitTree of empty tree (possible violation of monoid invariant?)"
71      Single x -> i `seq` (# empty, x, empty #)
72      Deep _ pr m sf
73        | p vpr ->  case splitDigit i pr of
74                       (# l, x, r #) -> (# (maybe empty digitToTree l), x, undefined #)
75        | p vm ->  case splitTree vpr m of
76                     (# ml, xs, mr #) -> case splitNode (vpr `mappend` measureFT ml) xs of
77                       (# l, x, r #) -> (# undefined, x, undefined #)
78                         where 
79        | otherwise -> case splitDigit vm sf of
80                          (# l, x, r #) -> (# undefined, x, (maybe empty digitToTree r) #)
81       where 
82         vpr =  i    `mappend`  measureDigit pr
83         vm =  vpr  `mappend` measureFT m
84 
85 
86  splitNode :: v -> Node v (Node v b) -> Split (Maybe (Digit (Node v b))) (Node v b)
87  splitNode i node = error "splitNode"
88
89  splitDigit :: v -> Digit (Node v b) -> Split (Maybe (Digit (Node v b))) (Node v b)
90  splitDigit i (One a) = i `seq` (# Nothing, a, Nothing #)
91  splitDigit i (Two a b)
92    | p va      = (# Nothing, a, (Just (One b)) #)
93    | otherwise = (# (Just (One a)), b, Nothing #)
94    where       va      = i `mappend` myMeasure a
95  splitDigit i (Three a b c)
96    | p va      = (# Nothing, a, (Just (Two b c)) #)
97    | p vab     = (# (Just (One a)), b, (Just (One c)) #)
98    | otherwise = (# (Just (Two a b)), c, Nothing #)
99    where va    = i `mappend` myMeasure a
100          vab   = va `mappend` myMeasure b
101  splitDigit i (Four a b c d)
102    | p va      = (# Nothing, a, (Just (Three b c d)) #)
103    | p vab     = (# (Just (One a)), b, (Just (Two c d)) #)
104    | p vabc    = (# (Just (Two a b)), c, (Just (One d)) #)
105    | otherwise = (# (Just (Three a b c)), d, Nothing #)
106    where va    = i `mappend` myMeasure a
107          vab   = va `mappend` myMeasure b
108          vabc  = vab `mappend` myMeasure c
109
110
111------------------------------------------------------------------------------------------
112
113newtype Seq a = Seq (FingerTree Size (Node Size a))
114
115instance Monoid Size where
116    {-# INLINE mappend #-}
117    mappend (Size a) (Size b) = Size (a + b)
118    {-# INLINE mempty #-}
119    mempty = Size 0
120
121mySplit :: Int -> Seq a -> (Seq a, Seq a)
122mySplit n (Seq a) = n `seq` case split (\(Size s) -> s>n) a of (l, r) -> (Seq l, Seq r)
123