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