Allow higher-order patterns to match on constructors
Summary
I am trying to implement concatMap
fusion with the new higher order patterns in rewrite rules, but I now discover a quite serious limitation: higher order patterns cannot currently match on data constructors.
Steps to reproduce
Here's a short example which doesn't fuse:
module ConcatMap where
data Stream a = forall s. Stream (s -> Step s a) !s
data Step s a = Yield a !s | Skip !s | Done
data Tuple a b = a :!: b
data Option a = None | Some !a
concatMapS :: (a -> Stream b) -> Stream a -> Stream b
concatMapS f (Stream next0 s0) = Stream next (s0 :!: None)
where
{-# INLINE next #-}
next (s :!: None) = case next0 s of
Done -> Done
Skip s' -> Skip (s' :!: None)
Yield x s' -> Skip (s' :!: Some (f x))
next (s :!: Some (Stream g t)) = case g t of
Done -> Skip (s :!: None)
Skip t' -> Skip (s :!: Some (Stream g t'))
Yield x t' -> Yield x (s :!: Some (Stream g t'))
{-# INLINE [1] concatMapS #-}
concatMapS' :: (s -> Step s b) -> (a -> s) -> Stream a -> Stream b
concatMapS' next2 f (Stream next1 s0) = Stream next (s0 :!: None)
where
{-# INLINE next #-}
next (s :!: None) = case next1 s of
Done -> Done
Skip s' -> Skip (s' :!: None)
Yield x s' -> Skip (s' :!: Some (f x))
next (s :!: Some t) = case next2 t of
Done -> Skip (s :!: None)
Skip t' -> Skip (s :!: Some t')
Yield x t' -> Yield x (s :!: Some t')
{-# INLINE concatMapS' #-}
{-# RULES "concatMap" forall step f. concatMapS (\x -> Stream step (f x)) = concatMapS' step f #-}
replicateStep1 :: Tuple Int a -> Step (Tuple Int a) a
replicateStep1 (0 :!: _) = Done
replicateStep1 (n :!: x) = Yield x ((n - 1) :!: x)
replicateS1 :: Int -> a -> Stream a
replicateS1 n x = Stream replicateStep1 (n :!: x) -- rewrite rule does not match
{-# INLINE replicateS1 #-}
foo1 :: Stream Int -> Stream Int
foo1 = concatMapS (replicateS1 2)
I've tried two variations, one where we flip the elements of the tuple:
replicateStep2 :: Tuple a Int -> Step (Tuple a Int) a
replicateStep2 (_ :!: 0) = Done
replicateStep2 (x :!: n) = Yield x (x :!: (n - 1))
replicateS2 :: Int -> a -> Stream a
replicateS2 n x = Stream replicateStep2 (x :!: n) -- rewrite rule matches
{-# INLINE replicateS2 #-}
foo2 :: Stream Int -> Stream Int
foo2 = concatMapS (replicateS2 2)
That does work, which suggests that it fails when the local variable x
must be the last argument.
And I've tried another variant with a NOINLINE function in between:
pair :: a -> b -> Tuple a b
pair x y = x :!: y
{-# NOINLINE pair #-}
replicateS3 :: Int -> a -> Stream a
replicateS3 n x = Stream replicateStep1 (pair n x) -- rewrite rule matches
{-# INLINE replicateS3 #-}
foo3 :: Stream Int -> Stream Int
foo3 = concatMapS (replicateS3 2)
Here the rewrite rule also does work, which suggests that the problematic case really has to do with constructors.
Expected behavior
The higher order matching should also work on constructors.
Environment
- GHC version used: 9.8.2