Ticket #5302: Unboxed.hs

File Unboxed.hs, 3.7 KB (added by reinerp, 4 years ago)
Line 
1{-# LANGUAGE UnboxedTuples, ScopedTypeVariables #-}
2module Unboxed (mySplit) where
3
4import Prelude hiding (null, reverse)
5import Data.Monoid
6
7data Digit a
8        = One a
9        | Two a a
10        | Three a a a
11        | Four a a a a
12
13data FingerTree v a = Empty | Single a | Deep !v !(Digit a) (FingerTree v (Node v a)) !(Digit a)
14data Node v a = Node2 !v a a | Node3 !v a a a
15
16
17-- measurements
18measureNode :: Node v b -> v
19measureNode (Node2 v _ _) = v
20measureNode (Node3 v _ _ _) = v
21
22measureDigit :: Monoid v => Digit (Node v b) -> v
23measureDigit (One a) = measureNode a
24measureDigit (Two a b) = measureNode a `mappend` measureNode b
25measureDigit (Three a b c) = measureNode a `mappend` measureNode b `mappend` measureNode c
26measureDigit (Four a b c d) = measureNode a `mappend` measureNode b `mappend` measureNode c `mappend` measureNode d
27
28split :: forall v c. Monoid v => (v -> Bool) -> FingerTree v (Node v c) -> (FingerTree v (Node v c), FingerTree v (Node v c))
29split p t = case splitTree mempty t of (# l, _, _ #) -> (l, undefined)
30   where
31
32  splitTree :: v -> FingerTree v (Node v b) -> (# FingerTree v (Node v b), Node v b, FingerTree v (Node v b) #)
33  splitTree i tree = case  tree of
34      Empty -> error "splitTree of Empty tree"
35      Single x -> i `seq` (# Empty,  x, Empty #)
36      Deep _ pr m sf
37        | p vpr ->  case splitDigit i pr of
38                       (# l, x, _ #) -> (# maybe Empty undefined l, x, undefined #)
39        | p vm ->  case splitTree vpr m of
40                     (# _, xs, _ #) -> case splitNode (vpr `mappend` mempty) xs of
41                       (# _, x, _ #) -> (# undefined, x, undefined #)
42                         where 
43        | otherwise -> case splitDigit vm sf of
44                          (# _, x, r #) -> (# undefined, x, maybe Empty undefined r #)
45       where 
46         vpr =  i    `mappend`  measureDigit pr
47         vm =  vpr  `mappend` mempty
48 
49 
50  splitNode :: v -> Node v (Node v b) -> (# Maybe (Digit (Node v b)), Node v b, Maybe (Digit (Node v b)) #)
51  splitNode i node = case node of
52    Node2 _ a b
53      | p va       -> (# Nothing, a, Just (One b) #)
54      | otherwise  -> (# Just (One a), b, Nothing #)
55     where      va      = i `mappend` measureNode a
56    Node3 _ a b c
57      | p va    -> (# Nothing, a, Just (Two b c) #)
58      | p vab   -> (# Just (One a), b, Just (One c) #)
59      | otherwise       -> (# Just (Two a b), c, Nothing #)
60     where va  = i `mappend` measureNode a
61           vab = va `mappend` measureNode b
62   
63  splitDigit :: v -> Digit (Node v b) -> (# Maybe (Digit (Node v b)), Node v b, Maybe (Digit (Node v b)) #)
64  splitDigit i (One a) = i `seq` (# Nothing, a, Nothing #)
65  splitDigit i (Two a b)
66    | p va      = (# Nothing, a, Just (One b) #)
67    | otherwise = (# Just (One a), b, Nothing #)
68    where       va      = i `mappend` measureNode a
69  splitDigit i (Three a b c)
70    | p va      = (# Nothing, a, Just (Two b c) #)
71    | p vab     = (# Just (One a), b, Just (One c) #)
72    | otherwise = (# Just (Two a b), c, Nothing #)
73    where va    = i `mappend` measureNode a
74          vab   = va `mappend` measureNode b
75  splitDigit i (Four a b c d)
76    | p va      = (# Nothing, a, Just (Three b c d) #)
77    | p vab     = (# Just (One a),  b, Just (Two c d) #)
78    | p vabc    = (# Just (Two a b), c, Just (One d) #)
79    | otherwise = (# Just (Three a b c), d, Nothing #)
80    where va    = i `mappend` measureNode a
81          vab   = va `mappend` measureNode b
82          vabc  = vab `mappend` measureNode c
83
84
85------------------------------------------------------------------------------------------
86newtype Size = Size Int
87
88instance Monoid Size where
89    {-# INLINE mappend #-}
90    mappend (Size a) (Size b) = Size (a + b)
91    {-# INLINE mempty #-}
92    mempty = Size 0
93
94mySplit :: Int -> FingerTree Size (Node Size c) -> (FingerTree Size (Node Size c), FingerTree Size (Node Size c))
95mySplit n f = split (\(Size s) -> s > n) f