Improve the enumeration APIs, fix overflow, improve perf

benchmark/Streamly/Benchmark/Data/Scanl/Container.hs

@@ -15,14 +15,14 @@
 -- Benchmarks for operations exported from Streamly.Internal.Data.Scanl.Container.
 module Scanl.Container (benchmarks) where
 
--- import Data.IntSet (IntSet)
+import Data.IntSet (IntSet)
 import Data.IORef (IORef, newIORef, readIORef, modifyIORef)
 import Data.Map.Strict (Map)
--- import Data.Set (Set)
--- import Data.STRef (STRef)
-import Streamly.Internal.Data.Scanl (Scanl(..))
+import Data.Set (Set)
+import Data.STRef (STRef)
+import Streamly.Internal.Data.Scanl (Scanl(..), Tuple'Fused)
 import Streamly.Internal.Data.Stream (Stream)
--- import Streamly.Internal.Data.Tuple.Strict (Tuple', Tuple3')
+import Streamly.Internal.Data.Tuple.Strict (Tuple', Tuple3')
 import System.IO.Unsafe (unsafePerformIO)
 import System.Random (randomRIO)
 
@@ -83,42 +83,59 @@ getScanl k = do
 -- Set operations
 -------------------------------------------------------------------------------
 
--- {-# ANN toSet (PermitTypes [''Int,''Set,''STRef,''(,)]) #-}
+{-# ANN toSet (PermitTypes [''Int,''Set,''STRef,''(,)]) #-}
 {-# ANN toSet (PermitTypeClasses []) #-}
 {-# ANN toSet (MaxCoreSize 1000) #-}
 {-# NOINLINE toSet #-}
 toSet :: Int -> IO ()
 toSet n = withPostscanl n Scanl.toSet
 
--- {-# ANN toIntSet (PermitTypes [''Int,''STRef,''IntSet,''(,)]) #-}
+{-# ANN toIntSet (PermitTypes [''Int,''STRef,''IntSet,''(,)]) #-}
 {-# ANN toIntSet (PermitTypeClasses []) #-}
 {-# ANN toIntSet (MaxCoreSize 1000) #-}
 {-# NOINLINE toIntSet #-}
 toIntSet :: Int -> IO ()
 toIntSet n = withPostscanl n Scanl.toIntSet
 
--- {-# ANN countDistinct (PermitTypes [''Int,''STRef,''Set,''(,)]) #-}
+{-# ANN countDistinct (PermitTypes [''Int,''STRef,''Set,''(,)]) #-}
 {-# ANN countDistinct (PermitTypeClasses []) #-}
 {-# ANN countDistinct (MaxCoreSize 1000) #-}
 {-# NOINLINE countDistinct #-}
 countDistinct :: Int -> IO ()
 countDistinct n = withPostscanl n Scanl.countDistinct
 
--- {-# ANN countDistinctInt (PermitTypes [''Int,''STRef,''IntSet,''(,)]) #-}
+{-# ANN countDistinctInt (PermitTypes [''Int,''STRef,''IntSet,''(,)]) #-}
 {-# ANN countDistinctInt (PermitTypeClasses []) #-}
 {-# ANN countDistinctInt (MaxCoreSize 1000) #-}
 {-# NOINLINE countDistinctInt #-}
 countDistinctInt :: Int -> IO ()
 countDistinctInt n = withPostscanl n Scanl.countDistinctInt
 
--- {-# ANN nub (PermitTypes [''Int,''STRef,''Set,''Maybe,''Tuple',''(,)]) #-}
+{-# ANN nub (PermitTypes
+   [ ''Int
+   , ''STRef
+   , ''Set
+   , ''Maybe
+   , ''Tuple'
+   , ''(,)
+   -- For GHC-9.14
+   , ''Scanl.Step
+   ]) #-}
 {-# ANN nub (PermitTypeClasses []) #-}
 {-# ANN nub (MaxCoreSize 1000) #-}
 {-# NOINLINE nub #-}
 nub :: Int -> IO ()
 nub n = withPostscanl n Scanl.nub
 
--- {-# ANN nubInt (PermitTypes [''Int,''STRef,''IntSet,''(,)]) #-}
+{-# ANN nubInt (PermitTypes
+   [ ''Int
+   , ''STRef
+   , ''IntSet
+   , ''(,)
+   -- For GHC-9.14
+   , ''Tuple'
+   , ''Scanl.Step
+   ]) #-}
 {-# ANN nubInt (PermitTypeClasses []) #-}
 {-# ANN nubInt (MaxCoreSize 1000) #-}
 {-# NOINLINE nubInt #-}
@@ -129,8 +146,19 @@ nubInt n = withPostscanl n Scanl.nubInt
 -- Demultiplexing
 -------------------------------------------------------------------------------
 
--- {-# ANN demuxIOOneShot (PermitTypes [''Int,''STRef,''IORef,''Map,''Scanl,
---   ''FL.Step, ''Scanl.Step,''(,)]) #-}
+{-# ANN demuxIOOneShot (PermitTypes
+   [ ''Int
+   , ''STRef
+   , ''IORef
+   , ''Map
+   , ''Scanl
+   , ''FL.Step
+   , ''Scanl.Step
+   , ''(,)
+   -- For GHC-9.14
+   , ''Tuple'
+   , ''Maybe
+   ]) #-}
 {-# ANN demuxIOOneShot (PermitTypeClasses []) #-}
 {-# ANN demuxIOOneShot (MaxCoreSize 2000) #-}
 {-# NOINLINE demuxIOOneShot #-}
@@ -140,10 +168,21 @@ demuxIOOneShot len =
         Stream.fold FL.drain
         . Stream.postscanl (Scanl.demuxIO (getKey 64) getScanl)
 
--- {-# ANN demuxIOSum (PermitTypes
---   [''Int,''STRef,''IORef,''Map,''Scanl,''FL.Step,''Scanl.Step,''(,)]) #-}
+{-# ANN demuxIOSum (PermitTypes
+   [ ''Int
+   , ''STRef
+   , ''IORef
+   , ''Map
+   , ''Scanl
+   , ''FL.Step
+   , ''Scanl.Step
+   , ''(,)
+   -- For GHC-9.14
+   , ''Tuple'
+   , ''Maybe
+   ]) #-}
 {-# ANN demuxIOSum (PermitTypeClasses []) #-}
-{-# ANN demuxIOSum (MaxCoreSize 1000) #-}
+{-# ANN demuxIOSum (MaxCoreSize 2000) #-}
 {-# NOINLINE demuxIOSum #-}
 demuxIOSum :: Int -> IO ()
 demuxIOSum len =
@@ -152,9 +191,21 @@ demuxIOSum len =
         . Stream.postscanl
             (Scanl.demuxIO (getKey 64) (const (pure (Just Scanl.sum))))
 
--- {-# ANN demuxSum (PermitTypes [''Int,''STRef,''IO,''Map,''Scanl,''FL.Step,''Scanl.Step,''(,)]) #-}
+{-# ANN demuxSum (PermitTypes
+   [ ''Int
+   , ''STRef
+   , ''IO
+   , ''Map
+   , ''Scanl
+   , ''FL.Step
+   , ''Scanl.Step
+   , ''(,)
+   -- For GHC-9.14
+   , ''Tuple'
+   , ''Maybe
+   ]) #-}
 {-# ANN demuxSum (PermitTypeClasses []) #-}
-{-# ANN demuxSum (MaxCoreSize 1000) #-}
+{-# ANN demuxSum (MaxCoreSize 2000) #-}
 {-# NOINLINE demuxSum #-}
 demuxSum :: Int -> IO ()
 demuxSum len =
@@ -163,10 +214,10 @@ demuxSum len =
         . Stream.postscanl
             (Scanl.demux (getKey 64) (const (pure (Just Scanl.sum))))
 
--- {-# ANN demuxGenericSum (PermitTypes
---   [''Int,''STRef,''Map,''Maybe,''Scanl,''FL.Step,''Scanl.Step,''(,)]) #-}
+{-# ANN demuxGenericSum (PermitTypes
+   [''Int,''STRef,''Map,''Maybe,''Scanl,''FL.Step,''Scanl.Step,''(,)]) #-}
 {-# ANN demuxGenericSum (PermitTypeClasses []) #-}
-{-# ANN demuxGenericSum (MaxCoreSize 1000) #-}
+{-# ANN demuxGenericSum (MaxCoreSize 2000) #-}
 {-# NOINLINE demuxGenericSum #-}
 demuxGenericSum :: Int -> IO ()
 demuxGenericSum len =
@@ -176,10 +227,10 @@ demuxGenericSum len =
             (Scanl.demuxGeneric (getKey 64) (const (pure (Just Scanl.sum)))
                 :: Scanl IO Int (IO (Map Int Int), Maybe (Int, Int)))
 
--- {-# ANN demuxGenericIOSum (PermitTypes
---   [''Int,''STRef,''Map,''Maybe,''Scanl,''FL.Step,''Scanl.Step,''(,)]) #-}
+{-# ANN demuxGenericIOSum (PermitTypes
+   [''Int,''STRef,''Map,''Maybe,''Scanl,''FL.Step,''Scanl.Step,''(,)]) #-}
 {-# ANN demuxGenericIOSum (PermitTypeClasses []) #-}
-{-# ANN demuxGenericIOSum (MaxCoreSize 1000) #-}
+{-# ANN demuxGenericIOSum (MaxCoreSize 2000) #-}
 {-# NOINLINE demuxGenericIOSum #-}
 demuxGenericIOSum :: Int -> IO ()
 demuxGenericIOSum len =
@@ -193,17 +244,42 @@ demuxGenericIOSum len =
 -- Classifying
 -------------------------------------------------------------------------------
 
--- {-# ANN classifyLimitedSum (PermitTypes [''Int,''STRef,''IORef,''Map,''Set,''Tuple',''(,)]) #-}
+{-# ANN classifyLimitedSum (PermitTypes
+   [ ''Int
+   , ''STRef
+   , ''IORef
+   , ''Map
+   , ''Set
+   , ''Tuple'
+   , ''Tuple'Fused
+   , ''(,)
+   -- For GHC-9.14
+   , ''Tuple3'
+   , ''Scanl.Step
+   , ''Maybe
+   ]) #-}
 {-# ANN classifyLimitedSum (PermitTypeClasses []) #-}
-{-# ANN classifyLimitedSum (MaxCoreSize 1000) #-}
+{-# ANN classifyLimitedSum (MaxCoreSize 2000) #-}
 {-# NOINLINE classifyLimitedSum #-}
 classifyLimitedSum :: Int -> IO ()
 classifyLimitedSum len =
     withStream len $
         Stream.fold FL.drain
         . Stream.postscanl (Scanl.classifyIO (getKey 64) (limitedSum 100))
 
--- {-# ANN classifyIOSum (PermitTypes [''Int,''STRef,''IORef,''Map,''Tuple',''(,)]) #-}
+{-# ANN classifyIOSum (PermitTypes
+   [ ''Int
+   , ''STRef
+   , ''IORef
+   , ''Map
+   , ''Tuple'
+   , ''(,)
+   -- For GHC-9.14
+   , ''Set
+   , ''Tuple3'
+   , ''Scanl.Step
+   , ''Maybe
+   ]) #-}
 {-# ANN classifyIOSum (PermitTypeClasses []) #-}
 {-# ANN classifyIOSum (MaxCoreSize 1000) #-}
 {-# NOINLINE classifyIOSum #-}
@@ -213,7 +289,18 @@ classifyIOSum len =
         Stream.fold FL.drain
         . Stream.postscanl (Scanl.classifyIO (getKey 64) Scanl.sum)
 
--- {-# ANN classifySum (PermitTypes [''Int,''STRef,''Map,''Tuple',''(,)]) #-}
+{-# ANN classifySum (PermitTypes
+   [ ''Int
+   , ''STRef
+   , ''Map
+   , ''Tuple'
+   , ''(,)
+   -- For GHC-9.14
+   , ''Set
+   , ''Tuple3'
+   , ''Scanl.Step
+   , ''Maybe
+   ]) #-}
 {-# ANN classifySum (PermitTypeClasses []) #-}
 {-# ANN classifySum (MaxCoreSize 1000) #-}
 {-# NOINLINE classifySum #-}
@@ -223,9 +310,9 @@ classifySum len =
         Stream.fold FL.drain
         . Stream.postscanl (Scanl.classify (getKey 64) Scanl.sum)
 
--- {-# ANN classifyGenericSum (PermitTypes [''Int,''STRef,''Map,''Maybe,''Set,''Tuple',''Tuple3',''(,)]) #-}
+{-# ANN classifyGenericSum (PermitTypes [''Int,''STRef,''Map,''Maybe,''Set,''Tuple',''Tuple3',''(,)]) #-}
 {-# ANN classifyGenericSum (PermitTypeClasses []) #-}
-{-# ANN classifyGenericSum (MaxCoreSize 1000) #-}
+{-# ANN classifyGenericSum (MaxCoreSize 1500) #-}
 {-# NOINLINE classifyGenericSum #-}
 classifyGenericSum :: Int -> IO ()
 classifyGenericSum len =
@@ -235,9 +322,9 @@ classifyGenericSum len =
             (Scanl.classifyGeneric (getKey 64) Scanl.sum
                 :: Scanl IO Int (IO (Map Int Int), Maybe (Int, Int)))
 
--- {-# ANN classifyGenericIOSum (PermitTypes [''Int,''STRef,''Map,''Maybe,''Set,''Tuple',''Tuple3',''(,)]) #-}
+{-# ANN classifyGenericIOSum (PermitTypes [''Int,''STRef,''Map,''Maybe,''Set,''Tuple',''Tuple3',''(,)]) #-}
 {-# ANN classifyGenericIOSum (PermitTypeClasses []) #-}
-{-# ANN classifyGenericIOSum (MaxCoreSize 1000) #-}
+{-# ANN classifyGenericIOSum (MaxCoreSize 1500) #-}
 {-# NOINLINE classifyGenericIOSum #-}
 classifyGenericIOSum :: Int -> IO ()
 classifyGenericIOSum len =

benchmark/Streamly/Benchmark/Data/Stream/Generate.hs

@@ -24,6 +24,7 @@ module Stream.Generate (benchmarks) where
 import GHC.Types (SPEC(..))
 import Test.Inspection
 import qualified Streamly.Internal.Data.Fold as Fold
+import qualified Streamly.Internal.Data.Producer as Producer
 #endif
 
 import Control.Monad.IO.Class (MonadIO)
@@ -89,7 +90,7 @@ sourceIntFromThenTo value = withDrain $ \n ->
 #ifdef INSPECTION
 inspect $ hasNoTypeClasses 'sourceIntFromThenTo
 inspect $ 'sourceIntFromThenTo `hasNoType` ''Stream.Step
-inspect $ 'sourceIntFromThenTo `hasNoType` ''Stream.EnumState
+inspect $ 'sourceIntFromThenTo `hasNoType` ''Producer.EnumState
 inspect $ 'sourceIntFromThenTo `hasNoType` ''Fold.Step
 inspect $ 'sourceIntFromThenTo `hasNoType` ''SPEC
 #endif

benchmark/Streamly/Benchmark/Data/Stream/Nesting/LogicUnfold.hs

@@ -6,11 +6,13 @@
 
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE FlexibleContexts #-}
-{-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TemplateHaskellQuotes #-}
 
 module Stream.Nesting.LogicUnfold (benchmarks) where
 
+import Streamly.Internal.Data.Producer (InterleaveState(..))
 import Streamly.Data.Stream (Stream)
 import Streamly.Data.Unfold (Unfold)
 
@@ -22,6 +24,7 @@ import Stream.Type (benchIO, withRandomIntIO)
 import Streamly.Benchmark.Common
 import qualified Stream.Type as Type
 import Prelude hiding (concatMap, zipWith)
+import Fusion.Plugin.Types
 
 -------------------------------------------------------------------------------
 -- Monad
@@ -32,7 +35,7 @@ infiniteIntsUnfold :: Monad m => Int -> Int -> Unfold m ((), ()) Int
 infiniteIntsUnfold _ _ =
     Unfold.interleave
         (Unfold.supply (0 :: Int) Unfold.enumerateFrom)
-        (Unfold.supply (-1, -2) Unfold.enumerateFromThen)
+        (Unfold.supply (-1) Unfold.enumerateDownFromNum)
 
 {-# INLINE unfoldCrossEqn #-}
 unfoldCrossEqn :: Monad m => Int -> Unfold m ((), ()) Int -> m ()
@@ -72,6 +75,11 @@ fairUnfoldSchedEqn maxVal input ints =
         $ Stream.mapM (Type.checkPair maxVal)
         $ Stream.fairUnfoldSched intu ints
 
+-- GHC 9.14.1 cannot fuse InterleaveState, though 9.10 can
+{-# ANN unfoldCrossBounded (PermitTypes [''Int,''Maybe,''InterleaveState]) #-}
+{-# ANN unfoldCrossBounded (PermitTypeClasses []) #-}
+{-# ANN unfoldCrossBounded (MaxCoreSize 20000) #-}
+{-# NOINLINE unfoldCrossBounded #-}
 unfoldCrossBounded :: Int -> IO ()
 unfoldCrossBounded maxVal = unfoldCrossEqn maxVal (Type.boundedIntsUnfold maxVal 0)
 

benchmark/Streamly/Benchmark/Data/Stream/Type/Logic.hs

@@ -39,21 +39,21 @@ boundedInts :: Monad m => Int -> Int -> Stream m Int
 boundedInts n _ =
     Stream.interleave
         (Stream.enumerateFromTo (0 :: Int) n)
-        (Stream.enumerateFromThenTo (-1) (-2) (-n))
+        (Stream.enumerateDownFromToNum (-1) (-n))
 
 {-# INLINE infiniteInts #-}
 infiniteInts :: Monad m => Int -> Int -> Stream m Int
 infiniteInts _ _ =
     Stream.interleave
         (Stream.enumerateFrom (0 :: Int))
-        (Stream.enumerateFromThen (-1) (-2))
+        (Stream.enumerateDownFromNum (-1))
 
 {-# INLINE boundedIntsUnfold #-}
 boundedIntsUnfold :: Monad m => Int -> Int -> Unfold m ((), ()) Int
 boundedIntsUnfold n _ =
     Unfold.interleave
         (Unfold.supply (0 :: Int, n) Unfold.enumerateFromTo)
-        (Unfold.supply (-1, -2, -n) Unfold.enumerateFromThenTo)
+        (Unfold.supply (-1, -n) Unfold.enumerateDownFromToNum)
 
 {-# INLINE checkStream #-}
 checkStream :: Applicative m =>