1{-# LANGUAGE FlexibleInstances, GADTs #-}
2module Utilities where
3
4import Test.QuickCheck
5
6import Data.Foldable
7import qualified Data.Vector as DV
8import qualified Data.Vector.Generic as DVG
9import qualified Data.Vector.Primitive as DVP
10import qualified Data.Vector.Storable as DVS
11import qualified Data.Vector.Unboxed as DVU
12import qualified Data.Vector.Fusion.Bundle as S
13
14import Control.Monad (foldM, foldM_, zipWithM, zipWithM_)
15import Control.Monad.Trans.Writer
16import Data.Function (on)
17import Data.Functor.Identity
18import Data.List ( sortBy )
19import Data.Monoid
20import Data.Maybe (catMaybes)
21
22instance Show a => Show (S.Bundle v a) where
23    show s = "Data.Vector.Fusion.Bundle.fromList " ++ show (S.toList s)
24
25
26instance Arbitrary a => Arbitrary (DV.Vector a) where
27    arbitrary = fmap DV.fromList arbitrary
28
29instance CoArbitrary a => CoArbitrary (DV.Vector a) where
30    coarbitrary = coarbitrary . DV.toList
31
32instance (Arbitrary a, DVP.Prim a) => Arbitrary (DVP.Vector a) where
33    arbitrary = fmap DVP.fromList arbitrary
34
35instance (CoArbitrary a, DVP.Prim a) => CoArbitrary (DVP.Vector a) where
36    coarbitrary = coarbitrary . DVP.toList
37
38instance (Arbitrary a, DVS.Storable a) => Arbitrary (DVS.Vector a) where
39    arbitrary = fmap DVS.fromList arbitrary
40
41instance (CoArbitrary a, DVS.Storable a) => CoArbitrary (DVS.Vector a) where
42    coarbitrary = coarbitrary . DVS.toList
43
44instance (Arbitrary a, DVU.Unbox a) => Arbitrary (DVU.Vector a) where
45    arbitrary = fmap DVU.fromList arbitrary
46
47instance (CoArbitrary a, DVU.Unbox a) => CoArbitrary (DVU.Vector a) where
48    coarbitrary = coarbitrary . DVU.toList
49
50instance Arbitrary a => Arbitrary (S.Bundle v a) where
51    arbitrary = fmap S.fromList arbitrary
52
53instance CoArbitrary a => CoArbitrary (S.Bundle v a) where
54    coarbitrary = coarbitrary . S.toList
55
56instance (Arbitrary a, Arbitrary b) => Arbitrary (Writer a b) where
57    arbitrary = do b <- arbitrary
58                   a <- arbitrary
59                   return $ writer (b,a)
60
61instance CoArbitrary a => CoArbitrary (Writer a ()) where
62    coarbitrary = coarbitrary . runWriter
63
64class (Testable (EqTest a), Conclusion (EqTest a)) => TestData a where
65  type Model a
66  model :: a -> Model a
67  unmodel :: Model a -> a
68
69  type EqTest a
70  equal :: a -> a -> EqTest a
71
72instance (Eq a, TestData a) => TestData (S.Bundle v a) where
73  type Model (S.Bundle v a) = [Model a]
74  model   = map model  . S.toList
75  unmodel = S.fromList . map unmodel
76
77  type EqTest (S.Bundle v a) = Property
78  equal x y = property (x == y)
79
80instance (Eq a, TestData a) => TestData (DV.Vector a) where
81  type Model (DV.Vector a) = [Model a]
82  model   = map model    . DV.toList
83  unmodel = DV.fromList . map unmodel
84
85  type EqTest (DV.Vector a) = Property
86  equal x y = property (x == y)
87
88instance (Eq a, DVP.Prim a, TestData a) => TestData (DVP.Vector a) where
89  type Model (DVP.Vector a) = [Model a]
90  model   = map model    . DVP.toList
91  unmodel = DVP.fromList . map unmodel
92
93  type EqTest (DVP.Vector a) = Property
94  equal x y = property (x == y)
95
96instance (Eq a, DVS.Storable a, TestData a) => TestData (DVS.Vector a) where
97  type Model (DVS.Vector a) = [Model a]
98  model   = map model    . DVS.toList
99  unmodel = DVS.fromList . map unmodel
100
101  type EqTest (DVS.Vector a) = Property
102  equal x y = property (x == y)
103
104instance (Eq a, DVU.Unbox a, TestData a) => TestData (DVU.Vector a) where
105  type Model (DVU.Vector a) = [Model a]
106  model   = map model    . DVU.toList
107  unmodel = DVU.fromList . map unmodel
108
109  type EqTest (DVU.Vector a) = Property
110  equal x y = property (x == y)
111
112#define id_TestData(ty) \
113instance TestData ty where { \
114  type Model ty = ty;        \
115  model = id;                \
116  unmodel = id;              \
117                             \
118  type EqTest ty = Property; \
119  equal x y = property (x == y) }
120
121id_TestData(())
122id_TestData(Bool)
123id_TestData(Int)
124id_TestData(Float)
125id_TestData(Double)
126id_TestData(Ordering)
127
128bimapEither :: (a -> b) -> (c -> d) -> Either a c -> Either b d
129bimapEither f _ (Left a) = Left (f a)
130bimapEither _ g (Right c) = Right (g c)
131
132-- Functorish models
133-- All of these need UndecidableInstances although they are actually well founded. Oh well.
134instance (Eq a, TestData a) => TestData (Maybe a) where
135  type Model (Maybe a) = Maybe (Model a)
136  model = fmap model
137  unmodel = fmap unmodel
138
139  type EqTest (Maybe a) = Property
140  equal x y = property (x == y)
141
142instance (Eq a, TestData a, Eq b, TestData b) => TestData (Either a b) where
143  type Model (Either a b) = Either (Model a) (Model b)
144  model = bimapEither model model
145  unmodel = bimapEither unmodel unmodel
146
147  type EqTest (Either a b) = Property
148  equal x y = property (x == y)
149
150instance (Eq a, TestData a) => TestData [a] where
151  type Model [a] = [Model a]
152  model = fmap model
153  unmodel = fmap unmodel
154
155  type EqTest [a] = Property
156  equal x y = property (x == y)
157
158instance (Eq a, TestData a) => TestData (Identity a) where
159  type Model (Identity a) = Identity (Model a)
160  model = fmap model
161  unmodel = fmap unmodel
162
163  type EqTest (Identity a) = Property
164  equal = (property .) . on (==) runIdentity
165
166instance (Eq a, TestData a, Eq b, TestData b, Monoid a) => TestData (Writer a b) where
167  type Model (Writer a b) = Writer (Model a) (Model b)
168  model = mapWriter model
169  unmodel = mapWriter unmodel
170
171  type EqTest (Writer a b) = Property
172  equal = (property .) . on (==) runWriter
173
174instance (Eq a, Eq b, TestData a, TestData b) => TestData (a,b) where
175  type Model (a,b) = (Model a, Model b)
176  model (a,b) = (model a, model b)
177  unmodel (a,b) = (unmodel a, unmodel b)
178
179  type EqTest (a,b) = Property
180  equal x y = property (x == y)
181
182instance (Eq a, Eq b, Eq c, TestData a, TestData b, TestData c) => TestData (a,b,c) where
183  type Model (a,b,c) = (Model a, Model b, Model c)
184  model (a,b,c) = (model a, model b, model c)
185  unmodel (a,b,c) = (unmodel a, unmodel b, unmodel c)
186
187  type EqTest (a,b,c) = Property
188  equal x y = property (x == y)
189
190instance (Arbitrary a, Show a, TestData a, TestData b) => TestData (a -> b) where
191  type Model (a -> b) = Model a -> Model b
192  model f = model . f . unmodel
193  unmodel f = unmodel . f . model
194
195  type EqTest (a -> b) = a -> EqTest b
196  equal f g x = equal (f x) (g x)
197
198newtype P a = P { unP :: EqTest a }
199
200instance TestData a => Testable (P a) where
201  property (P a) = property a
202
203infix 4 `eq`
204eq :: TestData a => a -> Model a -> P a
205eq x y = P (equal x (unmodel y))
206
207class Conclusion p where
208  type Predicate p
209
210  predicate :: Predicate p -> p -> p
211
212instance Conclusion Property where
213  type Predicate Property = Bool
214
215  predicate = (==>)
216
217instance Conclusion p => Conclusion (a -> p) where
218  type Predicate (a -> p) = a -> Predicate p
219
220  predicate f p = \x -> predicate (f x) (p x)
221
222infixr 0 ===>
223(===>) :: TestData a => Predicate (EqTest a) -> P a -> P a
224p ===> P a = P (predicate p a)
225
226notNull2 _ xs = not $ DVG.null xs
227notNullS2 _ s = not $ S.null s
228
229-- Generators
230index_value_pairs :: Arbitrary a => Int -> Gen [(Int,a)]
231index_value_pairs 0 = return []
232index_value_pairs m = sized $ \n ->
233  do
234    len <- choose (0,n)
235    is <- sequence [choose (0,m-1) | _i <- [1..len]]
236    xs <- vector len
237    return $ zip is xs
238
239indices :: Int -> Gen [Int]
240indices 0 = return []
241indices m = sized $ \n ->
242  do
243    len <- choose (0,n)
244    sequence [choose (0,m-1) | _i <- [1..len]]
245
246
247-- Additional list functions
248singleton x = [x]
249snoc xs x = xs ++ [x]
250generate n f = [f i | i <- [0 .. n-1]]
251generateM n f = sequence [f i | i <- [0 .. n-1]]
252slice i n xs = take n (drop i xs)
253backpermute xs is = map (xs!!) is
254prescanl f z = init . scanl f z
255postscanl f z = tail . scanl f z
256prescanr f z = tail . scanr f z
257postscanr f z = init . scanr f z
258
259accum :: (a -> b -> a) -> [a] -> [(Int,b)] -> [a]
260accum f xs ps = go xs ps' 0
261  where
262    ps' = sortBy (\p q -> compare (fst p) (fst q)) ps
263
264    go (x:xxs) ((i,y) : pps) j
265      | i == j     = go (f x y : xxs) pps j
266    go (x:xxs) pps j = x : go xxs pps (j+1)
267    go [] _ _      = []
268
269(//) :: [a] -> [(Int, a)] -> [a]
270xs // ps = go xs ps' 0
271  where
272    ps' = sortBy (\p q -> compare (fst p) (fst q)) ps
273
274    go (_x:xxs) ((i,y) : pps) j
275      | i == j     = go (y:xxs) pps j
276    go (x:xxs) pps j = x : go xxs pps (j+1)
277    go [] _ _      = []
278
279
280withIndexFirst m f = m (uncurry f) . zip [0..]
281
282modifyList :: [a] -> (a -> a) -> Int -> [a]
283modifyList xs f i = zipWith merge xs (replicate i Nothing ++ [Just f] ++ repeat Nothing)
284  where
285    merge x Nothing  = x
286    merge x (Just g) = g x
287
288writeList :: [a] -> Int -> a -> [a]
289writeList xs i a = modifyList xs (const a) i
290
291imap :: (Int -> a -> a) -> [a] -> [a]
292imap = withIndexFirst map
293
294imapM :: Monad m => (Int -> a -> m a) -> [a] -> m [a]
295imapM = withIndexFirst mapM
296
297imapM_ :: Monad m => (Int -> a -> m b) -> [a] -> m ()
298imapM_ = withIndexFirst mapM_
299
300izipWith :: (Int -> a -> a -> a) -> [a] -> [a] -> [a]
301izipWith = withIndexFirst zipWith
302
303izipWithM :: Monad m => (Int -> a -> a -> m a) -> [a] -> [a] -> m [a]
304izipWithM = withIndexFirst zipWithM
305
306izipWithM_ :: Monad m => (Int -> a -> a -> m b) -> [a] -> [a] -> m ()
307izipWithM_ = withIndexFirst zipWithM_
308
309izipWith3 :: (Int -> a -> a -> a -> a) -> [a] -> [a] -> [a] -> [a]
310izipWith3 = withIndexFirst zipWith3
311
312ifilter :: (Int -> a -> Bool) -> [a] -> [a]
313ifilter f = map snd . withIndexFirst filter f
314
315mapMaybe :: (a -> Maybe b) -> [a] -> [b]
316mapMaybe f = catMaybes . map f
317
318imapMaybe :: (Int -> a -> Maybe b) -> [a] -> [b]
319imapMaybe f = catMaybes . withIndexFirst map f
320
321indexedLeftFold fld f z = fld (uncurry . f) z . zip [0..]
322
323ifoldl :: (a -> Int -> a -> a) -> a -> [a] -> a
324ifoldl = indexedLeftFold foldl
325
326iscanl :: (Int -> a -> b -> a) -> a -> [b] -> [a]
327iscanl f z = scanl (\a (i, b) -> f i a b) z . zip [0..]
328
329iscanr :: (Int -> a -> b -> b) -> b -> [a] -> [b]
330iscanr f z = scanr (uncurry f) z . zip [0..]
331
332ifoldr :: (Int -> a -> b -> b) -> b -> [a] -> b
333ifoldr f z = foldr (uncurry f) z . zip [0..]
334
335ifoldM :: Monad m => (b -> Int -> a -> m b) -> b -> [a] -> m b
336ifoldM = indexedLeftFold foldM
337
338ifoldrM :: Monad m => (Int -> a -> b -> m b) -> b -> [a] -> m b
339ifoldrM f z xs = foldrM (\(i,a) b -> f i a b) z ([0..] `zip` xs)
340
341ifoldM_ :: Monad m => (b -> Int -> a -> m b) -> b -> [a] -> m ()
342ifoldM_ = indexedLeftFold foldM_
343
344minIndex :: Ord a => [a] -> Int
345minIndex = fst . foldr1 imin . zip [0..]
346  where
347    imin (i,x) (j,y) | x <= y    = (i,x)
348                     | otherwise = (j,y)
349
350maxIndex :: Ord a => [a] -> Int
351maxIndex = fst . foldr1 imax . zip [0..]
352  where
353    imax (i,x) (j,y) | x >  y    = (i,x)
354                     | otherwise = (j,y)
355
356iterateNM :: Monad m => Int -> (a -> m a) -> a -> m [a]
357iterateNM n f x
358    | n <= 0    = return []
359    | n == 1    = return [x]
360    | otherwise =  do x' <- f x
361                      xs <- iterateNM (n-1) f x'
362                      return (x : xs)
363
364unfoldrM :: Monad m => (b -> m (Maybe (a,b))) -> b -> m [a]
365unfoldrM step b0 = do
366    r <- step b0
367    case r of
368      Nothing    -> return []
369      Just (a,b) -> do as <- unfoldrM step b
370                       return (a : as)
371
372
373limitUnfolds f (theirs, ours)
374    | ours >= 0
375    , Just (out, theirs') <- f theirs = Just (out, (theirs', ours - 1))
376    | otherwise                       = Nothing
377