探究Haskell zipWithAll函数性能差异的原因与优化方案

问题背景 #

我需要实现一个Haskell函数，它能通过指定函数合并两个列表，但和标准zipWith不同，它不会在较短列表结束时停止，而是会处理到较长列表的末尾，对较短列表的剩余位置使用默认值。我先后实现了三个版本：

实现版本 #

1. 显式递归版本 #

zipWithAll :: (a -> b -> c) -> a -> b -> [a] -> [b] -> [c]
zipWithAll f x y = go
    where go []     []     = []
          go (a:as) []     = f a y : go as []
          go []     (b:bs) = f x b : go [] bs
          go (a:as) (b:bs) = f a b : go as bs

2. 基于标准库高阶函数的版本 #

zipWithAll' :: (a -> b -> c) -> a -> b -> [a] -> [b] -> [c]
zipWithAll' f x y xs ys = zipWith f xs' ys'
    where n   = max (length xs) (length ys)
          xs' = take n $ xs ++ repeat x
          ys' = take n $ ys ++ repeat y

3. 传入列表长度参数的优化版本 #

zipWithAll'' :: (a -> b -> c) -> Int -> Int -> a -> b -> [a] -> [b] -> [c]
zipWithAll'' f n m x y xs ys = zipWith f xs' ys'
    where k   = max n m
          xs' = take k $ xs ++ repeat x
          ys' = take k $ ys ++ repeat y

性能测试结果 #

对长度均为1000000的随机Int列表测试后，得到如下性能数据：

疑问 #

这与我的性能直觉不符，现提出以下问题：

• 为何显式递归版本性能最差？是否仅因标准zipWith的优化？能否让递归版性能接近zipWith？
• 我认为第二个版本执行3n次操作、第三个版本执行n次操作的假设是否错误？为何二者性能差异极小？
• 能否在未知列表长度的前提下，利用标准库优化实现更高效的zipWithAll？

基准测试代码 #

{-# LANGUAGE BangPatterns #-}

import Control.Monad (replicateM, forM_)
import Data.Foldable (foldl')
import Data.Time (diffUTCTime, getCurrentTime, NominalDiffTime)
import Numeric (showEFloat, showFFloat)
import Test.QuickCheck

main = do
    let n  = 1000000
        fs = [ ("zipWithAll", uncurry4 $ zipWithAll (+))
             , ("zipWithAll'", uncurry4 $ zipWithAll' (+))
             , ("zipWithAll''", uncurry4 $ zipWithAll'' (+) n n)]

    xs <- generate (vectorOf n arbitrary :: Gen [Int])
    ys <- generate (vectorOf n arbitrary :: Gen [Int])
    benchmark fs (0, 0, xs, ys) 30

uncurry4 :: (a -> b -> c -> d -> e) -> (a,b,c,d) -> e
uncurry4 f (a,b,c,d) = f a b c d

-- | Measure and print the average time it takes for each function in the list to return.
benchmark :: (Show a, Show b) => [(String, (a -> b))] -> a -> Int -> IO ()
benchmark fs x rep = do
    force x
    forM_ fs $ \(name, f) -> do
        ts <- replicateM rep (measureTime f x)
        putStrLn $ "function: " ++ name ++ ", time = " ++ (showSignificant 2 $ average ts)

-- | Get the time measurement for a function applied to an arguemnt
measureTime :: Show b => (a -> b) -> a -> IO NominalDiffTime
measureTime f x = do
    t1 <- getCurrentTime
    force (f x)
    t2 <- getCurrentTime
    return $ diffUTCTime t2 t1

-- | Force the computation of a value
force :: Show a => a -> IO ()
force a = maximum (show a) `seq` return ()


-- | Show a time difference using @n@ significant figures
showSignificant :: Int -> NominalDiffTime -> String
showSignificant n a = showFFloat Nothing b "s"
  where
    ae = showEFloat (Just (n-1)) (fromRational (toRational a)) ""
    b  = read ae :: Double

-- | Take the average of the elements in a foldable data structure
average :: (Foldable t, Fractional a) => t a -> a
average = uncurry (/) . foldl' f (0,0)
    where f (s,l) x = (s', l')
            where !s' = x + s
                  !l' = 1 + l

内容的提问来源于stack exchange，提问作者user20726076