Skip to content

Add tests for manual specializations and simplification enumFromTo in vector-streaming #553

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Oct 4, 2025
Merged

Add tests for manual specializations and simplification enumFromTo in vector-streaming #553

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4d82bcc
Add tests for custom specializations of functions
Shimuuar Sep 27, 2025
a17001a
Simplify specializations for integral types
Shimuuar Sep 26, 2025
971fb83
Add tests for vector-stream as well
Shimuuar Oct 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
200 changes: 27 additions & 173 deletions vector-stream/src/Data/Stream/Monadic.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
-- |
-- Module : Data.Stream.Monadic
-- Copyright : (c) Roman Leshchinskiy 2008-2010
Expand Down Expand Up @@ -95,22 +96,16 @@ import Prelude
, RealFrac, return, pure, otherwise, seq, error, not, id, show, const, fmap
, (==), (<), (<=), (>), (+), (-), (/), ($), (.), (=<<), (>>=) )

import Data.Int ( Int8, Int16, Int32 )
import Data.Int ( Int8, Int16, Int32, Int64 )
import Data.Word ( Word8, Word16, Word32, Word64 )

import GHC.Stack (HasCallStack)
import GHC.Types ( SPEC(..) )

#include "MachDeps.h"

#define INLINE_FUSED INLINE [1]
#define INLINE_INNER INLINE [0]


#if WORD_SIZE_IN_BITS > 32
import Data.Int ( Int64 )
#endif


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

-- NOTE: We use (x+1) instead of (succ x) below because the latter checks for
-- overflow which can't happen here.

-- FIXME: add "too large" test for Int
enumFromTo_small :: (Integral a, Monad m) => a -> a -> Stream m a
{-# INLINE_FUSED enumFromTo_small #-}
enumFromTo_small x y = x `seq` y `seq` Stream step (Just x)
where
{-# INLINE_INNER step #-}
step Nothing = return $ Done
step (Just z) | z == y = return $ Yield z Nothing
| z < y = return $ Yield z (Just (z+1))
| otherwise = return $ Done

{-# RULES

"enumFromTo<Int8> [Stream]"
enumFromTo = enumFromTo_small :: Monad m => Int8 -> Int8 -> Stream m Int8

"enumFromTo<Int16> [Stream]"
enumFromTo = enumFromTo_small :: Monad m => Int16 -> Int16 -> Stream m Int16

"enumFromTo<Word8> [Stream]"
enumFromTo = enumFromTo_small :: Monad m => Word8 -> Word8 -> Stream m Word8

"enumFromTo<Word16> [Stream]"
enumFromTo = enumFromTo_small :: Monad m => Word16 -> Word16 -> Stream m Word16 #-}


#if WORD_SIZE_IN_BITS > 32

{-# RULES

"enumFromTo<Int32> [Stream]"
enumFromTo = enumFromTo_small :: Monad m => Int32 -> Int32 -> Stream m Int32

"enumFromTo<Word32> [Stream]"
enumFromTo = enumFromTo_small :: Monad m => Word32 -> Word32 -> Stream m Word32 #-}


#endif

-- NOTE: We could implement a generic "too large" test:
--
-- len x y | x > y = 0
-- | n > 0 && n <= fromIntegral (maxBound :: Int) = fromIntegral n
-- | otherwise = error
-- where
-- n = y-x+1
--
-- Alas, GHC won't eliminate unnecessary comparisons (such as n >= 0 for
-- unsigned types). See http://hackage.haskell.org/trac/ghc/ticket/3744
--

enumFromTo_int :: forall m. Monad m => Int -> Int -> Stream m Int
{-# INLINE_FUSED enumFromTo_int #-}
enumFromTo_int x y = x `seq` y `seq` Stream step (Just x)
where
-- {-# INLINE [0] len #-}
-- len :: Int -> Int -> Int
-- len u v | u > v = 0
-- | otherwise = BOUNDS_CHECK(check) "enumFromTo" "vector too large"
-- (n > 0)
-- $ n
-- where
-- n = v-u+1

{-# INLINE_INNER step #-}
step Nothing = return $ Done
step (Just z) | z == y = return $ Yield z Nothing
| z < y = return $ Yield z (Just (z+1))
| otherwise = return $ Done


enumFromTo_intlike :: (Integral a, Monad m) => a -> a -> Stream m a
{-# INLINE_FUSED enumFromTo_intlike #-}
enumFromTo_intlike x y = x `seq` y `seq` Stream step (Just x)
where
{-# INLINE_INNER step #-}
step Nothing = return $ Done
step (Just z) | z == y = return $ Yield z Nothing
| z < y = return $ Yield z (Just (z+1))
| otherwise = return $ Done

{-# RULES

"enumFromTo<Int> [Stream]"
enumFromTo = enumFromTo_int :: Monad m => Int -> Int -> Stream m Int

#if WORD_SIZE_IN_BITS > 32

"enumFromTo<Int64> [Stream]"
enumFromTo = enumFromTo_intlike :: Monad m => Int64 -> Int64 -> Stream m Int64 #-}

#else

"enumFromTo<Int32> [Stream]"
enumFromTo = enumFromTo_intlike :: Monad m => Int32 -> Int32 -> Stream m Int32 #-}

#endif

enumFromTo_big_word :: (Integral a, Monad m) => a -> a -> Stream m a
{-# INLINE_FUSED enumFromTo_big_word #-}
enumFromTo_big_word x y = x `seq` y `seq` Stream step (Just x)
where
{-# INLINE_INNER step #-}
step Nothing = return $ Done
step (Just z) | z == y = return $ Yield z Nothing
| z < y = return $ Yield z (Just (z+1))
| otherwise = return $ Done

{-# RULES

"enumFromTo<Word> [Stream]"
enumFromTo = enumFromTo_big_word :: Monad m => Word -> Word -> Stream m Word

"enumFromTo<Word64> [Stream]"
enumFromTo = enumFromTo_big_word
:: Monad m => Word64 -> Word64 -> Stream m Word64

#if WORD_SIZE_IN_BITS == 32

"enumFromTo<Word32> [Stream]"
enumFromTo = enumFromTo_big_word
:: Monad m => Word32 -> Word32 -> Stream m Word32

#endif

"enumFromTo<Integer> [Stream]"
enumFromTo = enumFromTo_big_word
:: Monad m => Integer -> Integer -> Stream m Integer #-}



#if WORD_SIZE_IN_BITS > 32

-- FIXME: the "too large" test is totally wrong
enumFromTo_big_int :: (Integral a, Monad m) => a -> a -> Stream m a
{-# INLINE_FUSED enumFromTo_big_int #-}
enumFromTo_big_int x y = x `seq` y `seq` Stream step (Just x)
enumFromTo_integral :: (Integral a, Monad m) => a -> a -> Stream m a
{-# INLINE_FUSED enumFromTo_integral #-}
enumFromTo_integral !x !y = Stream step (Just x)
where
-- NOTE: We use (x+1) instead of (succ x) below because the latter
-- checks for overflow which can't happen here.
{-# INLINE_INNER step #-}
step Nothing = return $ Done
step (Just z) | z == y = return $ Yield z Nothing
| z < y = return $ Yield z (Just (z+1))
| otherwise = return $ Done

{-# RULES

"enumFromTo<Int64> [Stream]"
enumFromTo = enumFromTo_big_int :: Monad m => Int64 -> Int64 -> Stream m Int64 #-}



#endif

enumFromTo_char :: Monad m => Char -> Char -> Stream m Char
{-# INLINE_FUSED enumFromTo_char #-}
enumFromTo_char x y = x `seq` y `seq` Stream step xn
enumFromTo_char !x !y = Stream step xn
where
xn = ord x
yn = ord y
Expand All @@ -1538,21 +1390,10 @@ enumFromTo_char x y = x `seq` y `seq` Stream step xn
step zn | zn <= yn = return $ Yield (unsafeChr zn) (zn+1)
| otherwise = return $ Done

{-# RULES

"enumFromTo<Char> [Stream]"
enumFromTo = enumFromTo_char #-}



------------------------------------------------------------------------

-- Specialise enumFromTo for Float and Double.
-- Also, try to do something about pairs?

enumFromTo_double :: (Monad m, Ord a, RealFrac a) => a -> a -> Stream m a
{-# INLINE_FUSED enumFromTo_double #-}
enumFromTo_double n m = n `seq` m `seq` Stream step ini
enumFromTo_double !n !m = Stream step ini
where
lim = m + 1/2 -- important to float out
ini = 0
Expand All @@ -1561,13 +1402,26 @@ enumFromTo_double n m = n `seq` m `seq` Stream step ini
where
x' = x + n

{-# RULES

"enumFromTo<Double> [Stream]"
enumFromTo = enumFromTo_double :: Monad m => Double -> Double -> Stream m Double
{-# RULES

"enumFromTo<Float> [Stream]"
enumFromTo = enumFromTo_double :: Monad m => Float -> Float -> Stream m Float #-}
"Stream.enumFromTo[Int]" enumFromTo @Int = enumFromTo_integral
Copy link
Contributor

@lehins lehins Oct 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need to add SPECIALIZE pragma on enumFromTo_integral in order for such specialization to actually work?

Copy link
Contributor Author

@Shimuuar Shimuuar Oct 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In my understanding RHS of rewrite rule will get inlined afterward. But it's always better to measure than to guess. Speaking of which I only test rewrite rules from Bundle and not ones from vector-stream. I need to think how to write them

Copy link
Contributor Author

@Shimuuar Shimuuar Oct 4, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ive checked. With these rewrite rules GHC can produce nonallocating loop and fails to do so if they're disabled.

"Stream.enumFromTo[Int8]" enumFromTo @Int8 = enumFromTo_integral
"Stream.enumFromTo[Int16]" enumFromTo @Int16 = enumFromTo_integral
"Stream.enumFromTo[Int32]" enumFromTo @Int32 = enumFromTo_integral
"Stream.enumFromTo[Int64]" enumFromTo @Int64 = enumFromTo_integral
"Stream.enumFromTo[Word]" enumFromTo @Word = enumFromTo_integral
"Stream.enumFromTo[Word8]" enumFromTo @Word8 = enumFromTo_integral
"Stream.enumFromTo[Word16]" enumFromTo @Word16 = enumFromTo_integral
"Stream.enumFromTo[Word32]" enumFromTo @Word32 = enumFromTo_integral
"Stream.enumFromTo[Word64]" enumFromTo @Word64 = enumFromTo_integral
"Stream.enumFromTo[Integer]" enumFromTo @Integer = enumFromTo_integral

"Stream.enumFromTo[Float]" enumFromTo @Float = enumFromTo_double
"Stream.enumFromTo[Double]" enumFromTo @Double = enumFromTo_double

"Stream.enumFromTo[Char]" enumFromTo @Char = enumFromTo_char
#-}



Expand Down
11 changes: 11 additions & 0 deletions vector/tests-inspect/Inspect/Alloc.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,17 @@ tests = testGroup "allocations"
, allocWHNF "test_enumFromTo[Float]" (test_enumFromTo @Float round 0) 100000
, allocWHNF "test_enumFromTo[Double]" (test_enumFromTo @Double round 0) 100000
, allocWHNF "test_enumFromTo[Char]" (test_enumFromTo @Char ord (chr 32)) (chr 8000)
, allocWHNF "test_enumFromToStream[Int]" (test_enumFromToStream @Int fromIntegral 0) 100000
, allocWHNF "test_enumFromToStream[Int64]" (test_enumFromToStream @Int64 fromIntegral 0) 100000
, allocWHNF "test_enumFromToStream[Int32]" (test_enumFromToStream @Int32 fromIntegral 0) 100000
, allocWHNF "test_enumFromToStream[Int16]" (test_enumFromToStream @Int16 fromIntegral 0) maxBound
, allocWHNF "test_enumFromToStream[Word]" (test_enumFromToStream @Word fromIntegral 0) 100000
, allocWHNF "test_enumFromToStream[Word64]" (test_enumFromToStream @Word64 fromIntegral 0) 100000
, allocWHNF "test_enumFromToStream[Word32]" (test_enumFromToStream @Word32 fromIntegral 0) 100000
, allocWHNF "test_enumFromToStream[Word16]" (test_enumFromToStream @Word16 fromIntegral 0) maxBound
, allocWHNF "test_enumFromToStream[Float]" (test_enumFromToStream @Float round 0) 100000
, allocWHNF "test_enumFromToStream[Double]" (test_enumFromToStream @Double round 0) 100000
, allocWHNF "test_enumFromToStream[Char]" (test_enumFromToStream @Char ord (chr 32)) (chr 8000)
-- FIXME: We don't have specializations for enumFromThenTo
--
-- , allocWHNF "test_enumFromThenTo" test_enumFromThenTo size
Expand Down
17 changes: 16 additions & 1 deletion vector/tests-inspect/Inspect/Fusion.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,24 @@ module Inspect.Fusion where

import Test.Tasty
-- import Test.Tasty.Inspection
import qualified Data.Vector.Fusion.Bundle.Monadic as B
import Data.Vector.Fusion.Bundle.Size (Size(..))
import qualified Data.Stream.Monadic as S
import qualified Data.Vector.Unboxed as VU
import Data.Vector.Unboxed (Vector)
import qualified Data.Vector.Generic as VG
import Data.Vector.Fusion.Util (Box)

import Test.InspectExtra


-- We need to define this function to test rewrite rules in vector-stream.Hv
-- Rewrite rules in Bundle do not reuse rules in vector-stream
enumFromToStream :: (Enum a, VG.Vector v a) => a -> a -> v a
{-# INLINE enumFromToStream #-}
enumFromToStream x y = VG.unstream $ B.fromStream (S.enumFromTo x y) Unknown


-- NOTE: [Fusion tests]
-- ~~~~~~~~~~~~~~~~~~~~
--
Expand Down Expand Up @@ -198,7 +209,10 @@ test_enumFromThenTo :: (Enum a, VU.Unbox a) => (a -> Int) -> a -> a -> a -> Int
test_enumFromThenTo fun a b
= goodProducer (VU.map fun . VU.enumFromThenTo a b)


test_enumFromToStream :: (Enum a, VU.Unbox a) => (a -> Int) -> a -> a -> Int
{-# INLINE test_enumFromToStream #-}
test_enumFromToStream fun a
= goodProducer (VU.map fun . enumFromToStream a)

----------------------------------------------------------------
-- Function consuming vectors
Expand Down Expand Up @@ -286,6 +300,7 @@ tests = testGroup "Fusion"
, $(inspectFusion 'test_enumFromStepN)
, $(inspectClassyFusion 'test_enumFromTo)
, $(inspectClassyFusion 'test_enumFromThenTo)
, $(inspectClassyFusion 'test_enumFromToStream)
]
, testGroup "consumers"
[ $(inspectFusion 'test_bang)
Expand Down
2 changes: 2 additions & 0 deletions vector/tests/Main.hs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ import qualified Tests.Vector.Strict
import qualified Tests.Vector.Unboxed
import qualified Tests.Bundle
import qualified Tests.Move
import qualified Tests.Specialization
import qualified Tests.Deriving ()

import Test.Tasty (defaultMain,testGroup)
Expand All @@ -23,4 +24,5 @@ main = defaultMain $ testGroup "toplevel" $ concat
]
, Tests.Vector.UnitTests.tests
, Tests.Move.tests
, Tests.Specialization.tests
]
Loading
Loading