Skip to content

Commit 8943433

Browse files
ShimuuarShimuuar
authored
Merge pull request #553 from haskell/simplify-enumFromTo-stream
Add tests for manual specializations and simplification of enumFromTo in vector-streaming
2 parents 8430e8e + 971fb83 commit 8943433

File tree

6 files changed

+159
-174
lines changed

6 files changed

+159
-174
lines changed

‎vector-stream/src/Data/Stream/Monadic.hs‎

Lines changed: 27 additions & 173 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
{-# LANGUAGE MultiParamTypeClasses #-}
88
{-# LANGUAGE RankNTypes #-}
99
{-# LANGUAGE ScopedTypeVariables #-}
10+
{-# LANGUAGE TypeApplications #-}
1011
-- |
1112
-- Module : Data.Stream.Monadic
1213
-- Copyright : (c) Roman Leshchinskiy 2008-2010
@@ -95,22 +96,16 @@ import Prelude
9596
, RealFrac, return, pure, otherwise, seq, error, not, id, show, const, fmap
9697
, (==), (<), (<=), (>), (+), (-), (/), ($), (.), (=<<), (>>=) )
9798

98-
import Data.Int ( Int8, Int16, Int32 )
99+
import Data.Int ( Int8, Int16, Int32, Int64 )
99100
import Data.Word ( Word8, Word16, Word32, Word64 )
100101

101102
import GHC.Stack (HasCallStack)
102103
import GHC.Types ( SPEC(..) )
103104

104-
#include "MachDeps.h"
105-
106105
#define INLINE_FUSED INLINE [1]
107106
#define INLINE_INNER INLINE [0]
108107

109108

110-
#if WORD_SIZE_IN_BITS > 32
111-
import Data.Int ( Int64 )
112-
#endif
113-
114109

115110
-- | Box monad
116111
data Box a = Box { unBox :: a }
@@ -1371,165 +1366,22 @@ enumFromTo :: (Enum a, Monad m) => a -> a -> Stream m a
13711366
{-# INLINE_FUSED enumFromTo #-}
13721367
enumFromTo x y = fromList [x .. y]
13731368

1374-
-- NOTE: We use (x+1) instead of (succ x) below because the latter checks for
1375-
-- overflow which can't happen here.
1376-
1377-
-- FIXME: add "too large" test for Int
1378-
enumFromTo_small :: (Integral a, Monad m) => a -> a -> Stream m a
1379-
{-# INLINE_FUSED enumFromTo_small #-}
1380-
enumFromTo_small x y = x `seq` y `seq` Stream step (Just x)
1381-
where
1382-
{-# INLINE_INNER step #-}
1383-
step Nothing = return $ Done
1384-
step (Just z) | z == y = return $ Yield z Nothing
1385-
| z < y = return $ Yield z (Just (z+1))
1386-
| otherwise = return $ Done
1387-
1388-
{-# RULES
1389-
1390-
"enumFromTo<Int8> [Stream]"
1391-
enumFromTo = enumFromTo_small :: Monad m => Int8 -> Int8 -> Stream m Int8
1392-
1393-
"enumFromTo<Int16> [Stream]"
1394-
enumFromTo = enumFromTo_small :: Monad m => Int16 -> Int16 -> Stream m Int16
1395-
1396-
"enumFromTo<Word8> [Stream]"
1397-
enumFromTo = enumFromTo_small :: Monad m => Word8 -> Word8 -> Stream m Word8
1398-
1399-
"enumFromTo<Word16> [Stream]"
1400-
enumFromTo = enumFromTo_small :: Monad m => Word16 -> Word16 -> Stream m Word16 #-}
1401-
1402-
1403-
#if WORD_SIZE_IN_BITS > 32
1404-
1405-
{-# RULES
1406-
1407-
"enumFromTo<Int32> [Stream]"
1408-
enumFromTo = enumFromTo_small :: Monad m => Int32 -> Int32 -> Stream m Int32
1409-
1410-
"enumFromTo<Word32> [Stream]"
1411-
enumFromTo = enumFromTo_small :: Monad m => Word32 -> Word32 -> Stream m Word32 #-}
1412-
1413-
1414-
#endif
1415-
1416-
-- NOTE: We could implement a generic "too large" test:
1417-
--
1418-
-- len x y | x > y = 0
1419-
-- | n > 0 && n <= fromIntegral (maxBound :: Int) = fromIntegral n
1420-
-- | otherwise = error
1421-
-- where
1422-
-- n = y-x+1
1423-
--
1424-
-- Alas, GHC won't eliminate unnecessary comparisons (such as n >= 0 for
1425-
-- unsigned types). See http://hackage.haskell.org/trac/ghc/ticket/3744
1426-
--
1427-
1428-
enumFromTo_int :: forall m. Monad m => Int -> Int -> Stream m Int
1429-
{-# INLINE_FUSED enumFromTo_int #-}
1430-
enumFromTo_int x y = x `seq` y `seq` Stream step (Just x)
1431-
where
1432-
-- {-# INLINE [0] len #-}
1433-
-- len :: Int -> Int -> Int
1434-
-- len u v | u > v = 0
1435-
-- | otherwise = BOUNDS_CHECK(check) "enumFromTo" "vector too large"
1436-
-- (n > 0)
1437-
-- $ n
1438-
-- where
1439-
-- n = v-u+1
1440-
1441-
{-# INLINE_INNER step #-}
1442-
step Nothing = return $ Done
1443-
step (Just z) | z == y = return $ Yield z Nothing
1444-
| z < y = return $ Yield z (Just (z+1))
1445-
| otherwise = return $ Done
1446-
1447-
1448-
enumFromTo_intlike :: (Integral a, Monad m) => a -> a -> Stream m a
1449-
{-# INLINE_FUSED enumFromTo_intlike #-}
1450-
enumFromTo_intlike x y = x `seq` y `seq` Stream step (Just x)
1451-
where
1452-
{-# INLINE_INNER step #-}
1453-
step Nothing = return $ Done
1454-
step (Just z) | z == y = return $ Yield z Nothing
1455-
| z < y = return $ Yield z (Just (z+1))
1456-
| otherwise = return $ Done
1457-
1458-
{-# RULES
1459-
1460-
"enumFromTo<Int> [Stream]"
1461-
enumFromTo = enumFromTo_int :: Monad m => Int -> Int -> Stream m Int
1462-
1463-
#if WORD_SIZE_IN_BITS > 32
1464-
1465-
"enumFromTo<Int64> [Stream]"
1466-
enumFromTo = enumFromTo_intlike :: Monad m => Int64 -> Int64 -> Stream m Int64 #-}
1467-
1468-
#else
1469-
1470-
"enumFromTo<Int32> [Stream]"
1471-
enumFromTo = enumFromTo_intlike :: Monad m => Int32 -> Int32 -> Stream m Int32 #-}
1472-
1473-
#endif
1474-
1475-
enumFromTo_big_word :: (Integral a, Monad m) => a -> a -> Stream m a
1476-
{-# INLINE_FUSED enumFromTo_big_word #-}
1477-
enumFromTo_big_word x y = x `seq` y `seq` Stream step (Just x)
1478-
where
1479-
{-# INLINE_INNER step #-}
1480-
step Nothing = return $ Done
1481-
step (Just z) | z == y = return $ Yield z Nothing
1482-
| z < y = return $ Yield z (Just (z+1))
1483-
| otherwise = return $ Done
1484-
1485-
{-# RULES
1486-
1487-
"enumFromTo<Word> [Stream]"
1488-
enumFromTo = enumFromTo_big_word :: Monad m => Word -> Word -> Stream m Word
1489-
1490-
"enumFromTo<Word64> [Stream]"
1491-
enumFromTo = enumFromTo_big_word
1492-
:: Monad m => Word64 -> Word64 -> Stream m Word64
1493-
1494-
#if WORD_SIZE_IN_BITS == 32
14951369

1496-
"enumFromTo<Word32> [Stream]"
1497-
enumFromTo = enumFromTo_big_word
1498-
:: Monad m => Word32 -> Word32 -> Stream m Word32
1499-
1500-
#endif
1501-
1502-
"enumFromTo<Integer> [Stream]"
1503-
enumFromTo = enumFromTo_big_word
1504-
:: Monad m => Integer -> Integer -> Stream m Integer #-}
1505-
1506-
1507-
1508-
#if WORD_SIZE_IN_BITS > 32
1509-
1510-
-- FIXME: the "too large" test is totally wrong
1511-
enumFromTo_big_int :: (Integral a, Monad m) => a -> a -> Stream m a
1512-
{-# INLINE_FUSED enumFromTo_big_int #-}
1513-
enumFromTo_big_int x y = x `seq` y `seq` Stream step (Just x)
1370+
enumFromTo_integral :: (Integral a, Monad m) => a -> a -> Stream m a
1371+
{-# INLINE_FUSED enumFromTo_integral #-}
1372+
enumFromTo_integral !x !y = Stream step (Just x)
15141373
where
1374+
-- NOTE: We use (x+1) instead of (succ x) below because the latter
1375+
-- checks for overflow which can't happen here.
15151376
{-# INLINE_INNER step #-}
15161377
step Nothing = return $ Done
15171378
step (Just z) | z == y = return $ Yield z Nothing
15181379
| z < y = return $ Yield z (Just (z+1))
15191380
| otherwise = return $ Done
15201381

1521-
{-# RULES
1522-
1523-
"enumFromTo<Int64> [Stream]"
1524-
enumFromTo = enumFromTo_big_int :: Monad m => Int64 -> Int64 -> Stream m Int64 #-}
1525-
1526-
1527-
1528-
#endif
1529-
15301382
enumFromTo_char :: Monad m => Char -> Char -> Stream m Char
15311383
{-# INLINE_FUSED enumFromTo_char #-}
1532-
enumFromTo_char x y = x `seq` y `seq` Stream step xn
1384+
enumFromTo_char !x !y = Stream step xn
15331385
where
15341386
xn = ord x
15351387
yn = ord y
@@ -1538,21 +1390,10 @@ enumFromTo_char x y = x `seq` y `seq` Stream step xn
15381390
step zn | zn <= yn = return $ Yield (unsafeChr zn) (zn+1)
15391391
| otherwise = return $ Done
15401392

1541-
{-# RULES
1542-
1543-
"enumFromTo<Char> [Stream]"
1544-
enumFromTo = enumFromTo_char #-}
1545-
1546-
1547-
1548-
------------------------------------------------------------------------
1549-
1550-
-- Specialise enumFromTo for Float and Double.
1551-
-- Also, try to do something about pairs?
15521393

15531394
enumFromTo_double :: (Monad m, Ord a, RealFrac a) => a -> a -> Stream m a
15541395
{-# INLINE_FUSED enumFromTo_double #-}
1555-
enumFromTo_double n m = n `seq` m `seq` Stream step ini
1396+
enumFromTo_double !n !m = Stream step ini
15561397
where
15571398
lim = m + 1/2 -- important to float out
15581399
ini = 0
@@ -1561,13 +1402,26 @@ enumFromTo_double n m = n `seq` m `seq` Stream step ini
15611402
where
15621403
x' = x + n
15631404

1564-
{-# RULES
15651405

1566-
"enumFromTo<Double> [Stream]"
1567-
enumFromTo = enumFromTo_double :: Monad m => Double -> Double -> Stream m Double
1406+
{-# RULES
15681407

1569-
"enumFromTo<Float> [Stream]"
1570-
enumFromTo = enumFromTo_double :: Monad m => Float -> Float -> Stream m Float #-}
1408+
"Stream.enumFromTo[Int]" enumFromTo @Int = enumFromTo_integral
1409+
"Stream.enumFromTo[Int8]" enumFromTo @Int8 = enumFromTo_integral
1410+
"Stream.enumFromTo[Int16]" enumFromTo @Int16 = enumFromTo_integral
1411+
"Stream.enumFromTo[Int32]" enumFromTo @Int32 = enumFromTo_integral
1412+
"Stream.enumFromTo[Int64]" enumFromTo @Int64 = enumFromTo_integral
1413+
"Stream.enumFromTo[Word]" enumFromTo @Word = enumFromTo_integral
1414+
"Stream.enumFromTo[Word8]" enumFromTo @Word8 = enumFromTo_integral
1415+
"Stream.enumFromTo[Word16]" enumFromTo @Word16 = enumFromTo_integral
1416+
"Stream.enumFromTo[Word32]" enumFromTo @Word32 = enumFromTo_integral
1417+
"Stream.enumFromTo[Word64]" enumFromTo @Word64 = enumFromTo_integral
1418+
"Stream.enumFromTo[Integer]" enumFromTo @Integer = enumFromTo_integral
1419+
1420+
"Stream.enumFromTo[Float]" enumFromTo @Float = enumFromTo_double
1421+
"Stream.enumFromTo[Double]" enumFromTo @Double = enumFromTo_double
1422+
1423+
"Stream.enumFromTo[Char]" enumFromTo @Char = enumFromTo_char
1424+
#-}
15711425

15721426

15731427

‎vector/tests-inspect/Inspect/Alloc.hs‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,17 @@ tests = testGroup "allocations"
105105
, allocWHNF "test_enumFromTo[Float]" (test_enumFromTo @Float round 0) 100000
106106
, allocWHNF "test_enumFromTo[Double]" (test_enumFromTo @Double round 0) 100000
107107
, allocWHNF "test_enumFromTo[Char]" (test_enumFromTo @Char ord (chr 32)) (chr 8000)
108+
, allocWHNF "test_enumFromToStream[Int]" (test_enumFromToStream @Int fromIntegral 0) 100000
109+
, allocWHNF "test_enumFromToStream[Int64]" (test_enumFromToStream @Int64 fromIntegral 0) 100000
110+
, allocWHNF "test_enumFromToStream[Int32]" (test_enumFromToStream @Int32 fromIntegral 0) 100000
111+
, allocWHNF "test_enumFromToStream[Int16]" (test_enumFromToStream @Int16 fromIntegral 0) maxBound
112+
, allocWHNF "test_enumFromToStream[Word]" (test_enumFromToStream @Word fromIntegral 0) 100000
113+
, allocWHNF "test_enumFromToStream[Word64]" (test_enumFromToStream @Word64 fromIntegral 0) 100000
114+
, allocWHNF "test_enumFromToStream[Word32]" (test_enumFromToStream @Word32 fromIntegral 0) 100000
115+
, allocWHNF "test_enumFromToStream[Word16]" (test_enumFromToStream @Word16 fromIntegral 0) maxBound
116+
, allocWHNF "test_enumFromToStream[Float]" (test_enumFromToStream @Float round 0) 100000
117+
, allocWHNF "test_enumFromToStream[Double]" (test_enumFromToStream @Double round 0) 100000
118+
, allocWHNF "test_enumFromToStream[Char]" (test_enumFromToStream @Char ord (chr 32)) (chr 8000)
108119
-- FIXME: We don't have specializations for enumFromThenTo
109120
--
110121
-- , allocWHNF "test_enumFromThenTo" test_enumFromThenTo size

‎vector/tests-inspect/Inspect/Fusion.hs‎

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,24 @@ module Inspect.Fusion where
66

77
import Test.Tasty
88
-- import Test.Tasty.Inspection
9+
import qualified Data.Vector.Fusion.Bundle.Monadic as B
10+
import Data.Vector.Fusion.Bundle.Size (Size(..))
11+
import qualified Data.Stream.Monadic as S
912
import qualified Data.Vector.Unboxed as VU
1013
import Data.Vector.Unboxed (Vector)
1114
import qualified Data.Vector.Generic as VG
1215
import Data.Vector.Fusion.Util (Box)
1316

1417
import Test.InspectExtra
1518

19+
20+
-- We need to define this function to test rewrite rules in vector-stream.Hv
21+
-- Rewrite rules in Bundle do not reuse rules in vector-stream
22+
enumFromToStream :: (Enum a, VG.Vector v a) => a -> a -> v a
23+
{-# INLINE enumFromToStream #-}
24+
enumFromToStream x y = VG.unstream $ B.fromStream (S.enumFromTo x y) Unknown
25+
26+
1627
-- NOTE: [Fusion tests]
1728
-- ~~~~~~~~~~~~~~~~~~~~
1829
--
@@ -198,7 +209,10 @@ test_enumFromThenTo :: (Enum a, VU.Unbox a) => (a -> Int) -> a -> a -> a -> Int
198209
test_enumFromThenTo fun a b
199210
= goodProducer (VU.map fun . VU.enumFromThenTo a b)
200211

201-
212+
test_enumFromToStream :: (Enum a, VU.Unbox a) => (a -> Int) -> a -> a -> Int
213+
{-# INLINE test_enumFromToStream #-}
214+
test_enumFromToStream fun a
215+
= goodProducer (VU.map fun . enumFromToStream a)
202216

203217
----------------------------------------------------------------
204218
-- Function consuming vectors
@@ -286,6 +300,7 @@ tests = testGroup "Fusion"
286300
, $(inspectFusion 'test_enumFromStepN)
287301
, $(inspectClassyFusion 'test_enumFromTo)
288302
, $(inspectClassyFusion 'test_enumFromThenTo)
303+
, $(inspectClassyFusion 'test_enumFromToStream)
289304
]
290305
, testGroup "consumers"
291306
[ $(inspectFusion 'test_bang)

‎vector/tests/Main.hs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import qualified Tests.Vector.Strict
88
import qualified Tests.Vector.Unboxed
99
import qualified Tests.Bundle
1010
import qualified Tests.Move
11+
import qualified Tests.Specialization
1112
import qualified Tests.Deriving ()
1213

1314
import Test.Tasty (defaultMain,testGroup)
@@ -23,4 +24,5 @@ main = defaultMain $ testGroup "toplevel" $ concat
2324
]
2425
, Tests.Vector.UnitTests.tests
2526
, Tests.Move.tests
27+
, Tests.Specialization.tests
2628
]

0 commit comments

Comments
 (0)