是否有可能使`foldrRanges`和`fol​​drRange2D`一样快？

Question

这个:

foldrRange :: (Int -> t -> t) -> t -> Int -> Int -> t
foldrRange cons nil a b = foldr cons nil [a..b-1]

定义从atil 折叠列表的函数b.这个:

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 #-}

由于适当的严格使用(我们知道最后一个元素,所以我们可以滚动foldl'),是一个快50倍的版本.

这个:

foldrRange2D cons nil (ax,ay) (bx,by) 
    = foldr cons nil 
    $ do
        y <- [ay..by-1]
        x <- [ax..bx-1]
        return (x,y)

是一个2D版本foldrRange,即它可以在2D矩形上工作foldrRange2d (:) [] (0,0) (2,2) == [(0,0),(1,0),(0,1),(1,1)].这个:

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)

再次,由于更严格的使用,定义速度提高了约50倍.写作foldrRange3D,foldrRange4D等等,会很麻烦,所以可以像这样概括它:

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

不幸的是,这个定义的速度要慢大约120倍foldrRange2D,因为可以通过此测试验证:

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]

我可能会用来ST加快速度foldrRanges,但单独使用递归是否可以这样做？

Answer 1

您可以高效地实现算法,该算法在输入维度上是归纳的.幸运的是,你可以在Haskell中做到这一点!

首先,用类型级Nat索引向量替换列表.这给了我们一个进行归纳的类型(它可能用列表来完成......但这更安全).

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

然后你想要的函数和2D情况一样 - 只是推广递归调用:

{-# 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)

虽然当我测试性能时,我很失望地看到它无法跟上2D版本.诀窍似乎更具内线.通过将函数放在一个类中,您可以使它在每个"维度"内联(必须有更好的方法来执行此操作...)

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)

测试如下:

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

和我的系统上的结果

./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