-
Notifications
You must be signed in to change notification settings - Fork 4
/
test.hs
68 lines (59 loc) · 1.9 KB
/
test.hs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE MagicHash #-}
module F where
--import qualified Data.IntMap as IM
import WWFusion as WW
import GHC.Exts(Ptr(..), Addr#, RealWorld, State#)
import GHC.IO(IO(..), unIO)
import Foreign.Ptr
import Foreign.Storable
import Foreign.Marshal.Array
f :: [[Int]] -> Int
f x = foldl' (+) 0 $ WW.concat x
g :: [Int] -> [Int] -> Int
g x y = foldl' (+) 0 $ x WW.++ y
h :: Tree -> Int
h t = WW.foldl (+) 0 $ toList t
-- | Serialize a tree into a buffer. The buffer is assumed to
-- be large enough.
serializeTree :: Tree -> Ptr () -> IO (Ptr ())
serializeTree t buf = foldIO_Ptr write buf $ toList t
where
write p x = poke p' x >> return (castPtr $ advancePtr p' 1)
where
p' :: Ptr Int
p' = castPtr p
data Tree
= Bin !Tree !Tree
| Tip {-# UNPACK #-} !Int
toList :: Tree -> [Int]
toList = \tree -> buildW (toListFB tree)
{-# INLINE toList #-}
toListFB
:: forall r f
. Tree
-> (Wrap f r)
-> (Int -> r -> r) -> r -> r
toListFB root (Wrap wrap unwrap) c n = wrap go root n
where
go :: f Tree
go = unwrap $ \t z -> case t of
Bin l r -> wrap go l (wrap go r z)
Tip x -> c x z
{-# INLINE [0] toListFB #-}
newtype IOInt e = IOInt (e -> Addr# -> State# RealWorld -> (# State# RealWorld, Addr# #))
-- | A version of foldM, specialized for the IO monad and the
-- accumulator type of Ptr (). It uses unboxed ints to pass around
-- intermediate results.
foldIO_Ptr :: (Ptr () -> a -> IO (Ptr ())) -> Ptr () -> [a] -> IO (Ptr ())
foldIO_Ptr f initial = \xs -> foldrW (Wrap wrap unwrap) g return xs initial
where
wrap (IOInt h) e k (Ptr a) = io >>= k
where
io = IO $ \s -> case h e a s of
(# s', a' #) -> (# s', Ptr a' #)
unwrap u = IOInt $ \e a s -> case unIO (u e return (Ptr a)) s of
(# s', Ptr a' #) -> (# s', a' #)
g x next acc = next =<< f acc x