Is it possible to make `foldrRanges` as fast as `foldrRange2D`? - algorithm

Question

This:
foldrRange :: (Int -> t -> t) -> t -> Int -> Int -> t
foldrRange cons nil a b = foldr cons nil [a..b-1]
Defines a function that folds over a list from a til b. This:
foldrRange :: (Int -> t -> t) -> t -> Int -> Int -> t
foldrRange cons nil a b = go (b-1) nil where
go b !r | b < a = r
| otherwise = go (b-1) (cons b r)
{-# INLINE foldrRange #-}
is a ~50x faster version due to proper strictness usage (we know the last element, so we can roll like foldl').
This:
foldrRange2D cons nil (ax,ay) (bx,by)
= foldr cons nil
$ do
y <- [ay..by-1]
x <- [ax..bx-1]
return (x,y)
Is a 2D version of foldrRange, i.e., it works over 2D rectangles so that foldrRange2d (:) [] (0,0) (2,2) == [(0,0),(1,0),(0,1),(1,1)]. This:
foldrRange2D :: ((Int,Int) -> t -> t) -> t -> (Int,Int) -> (Int,Int) -> t
foldrRange2D cons nil (ax,ay) (bx,by) = go (by-1) nil where
go by !r
| by < ay = r
| otherwise = go (by-1) (foldrRange (\ ax -> cons (ax,by)) r ax bx)
Is, again, an ~50x faster definition due to better strictness usage. Writing foldrRange3D, foldrRange4D, etc., would be cumbersome, so one can generalize it like so:
foldrRangeND :: forall t . ([Int] -> t -> t) -> t -> [Int] -> [Int] -> t
foldrRangeND cons nil as bs = foldr co ni (zip as bs) [] nil where
co (a,b) tail lis = foldrRange (\ h t -> tail (h:lis) . t) id a b
ni lis = cons lis
Unfortunately this definition is at around 120 times slower than foldrRange2D, as one can verify with this test:
main = do
let n = 2000
print $ foldrRange2D (\ (a,b) c -> a+b+c) 0 (0,0) (n,n)
print $ foldrRanges (\ [a,b] c -> a+b+c) 0 [0,0] [n,n]
I could probably use ST to get a faster foldrRanges, but is it possible to do so with recursion alone?

Answer 1

You have an efficient implementation of your algorithm which is inductive on the dimension of the input. Fortunately, you can do that in Haskell!
First, replace lists with type level Nat indexed vectors. This gives us a type to do induction on (it could probably be done with lists ... but this is much safer).
data Nat = Z | S Nat
infixl 5 :<
data Vec (n :: Nat) (a :: *) where
Nil :: Vec Z a
(:<) :: Vec n a -> a -> Vec (S n) a
instance Functor (Vec n) where
fmap _ Nil = Nil
fmap f (xs :< x) = fmap f xs :< f x
Then your desired function is just the same as the 2D case - just generalize the recursive call:
{-# INLINE foldrRangeN #-}
foldrRangeN :: (Vec n Int -> t -> t) -> t -> Vec n Int -> Vec n Int -> t
foldrRangeN f x Nil Nil = f Nil x
foldrRangeN cons nil (ax :< ay) (bx :< by) = go (by-1) nil where
go by !r
| by < ay = r
| otherwise = go (by-1) (foldrRangeN (\ ax -> cons (ax :< by)) r ax bx)
Although when I tested the performance, I was disappointed to see it couldn't keep up with the 2D version. The trick seems to be more inlining. By putting the function in a class, you can get it to inline at each 'dimension' (there must be a better way to do this...)
class FoldrRange n where
foldrRangeN' :: (Vec n Int -> t -> t) -> t -> Vec n Int -> Vec n Int -> t
instance FoldrRange Z where
{-# INLINE foldrRangeN' #-}
foldrRangeN' f x Nil Nil = f Nil x
instance FoldrRange n => FoldrRange (S n) where
{-# INLINE foldrRangeN' #-}
foldrRangeN' cons nil (ax :< ay) (bx :< by) = go (by-1) nil where
go by !r
| by < ay = r
| otherwise = go (by-1) (foldrRangeN' (\ ax -> cons (ax :< by)) r ax bx)
Tested as follows:
main = do
i:n':_ <- getArgs
let n = read n' :: Int
rs = [ foldrRange2D (\ (a,b) c -> a+b+c) 0 (0,0) (n,n)
, foldrRangeND (\ [a,b] c -> a+b+c) 0 [0,0] [n,n]
, foldrRangeN (\ (Nil :< a :< b) c -> a+b+c) 0 (Nil :< 0 :< 0) (Nil :< n :< n)
, foldrRangeN' (\ (Nil :< a :< b) c -> a+b+c) 0 (Nil :< 0 :< 0) (Nil :< n :< n)
]
print $ rs !! read i
and the results on my system
./test 0 4000 +RTS -s : 0.02s
./test 1 4000 +RTS -s : 7.63s
./test 2 4000 +RTS -s : 0.59s
./test 3 4000 +RTS -s : 0.03s