haskell 编写更高效的Data.Map.findWithDefault版本

cqoc49vn  于 2023-03-03  发布在  其他
关注(0)|答案(1)|浏览(125)

Data.Map.findWithDefault是这样实现的:

findWithDefault :: Ord k => a -> k -> Map k a -> a
findWithDefault def k m = case lookup k m of
    Nothing -> def
    Just x  -> x

这是非常愚蠢的,因为没有使用Maybe类型。
我想写我自己的查找:

lookup k = k `seq` go
   where
      go Data.Map.Tip = (0,0,0,0)
      go (Data.Map.Bin _ kx x l r) =
         case compare k kx of
            LT -> go l
            GT -> go r
            EQ -> x

但是Data.Map不导出Data.Map.TipData.Map.Bin,除非启用TESTING,我不知道如何启用TESTING,请参见:
https://hackage.haskell.org/package/containers-0.4.0.0/docs/src/Data-Map.html

bxfogqkk

bxfogqkk1#

编译findWithDefault时,对lookup的调用将被内联,这将导致Maybe被删除。使用containers包中的实际代码来证明这一点有点困难,但您可以在以下简化的、自包含的示例中看到这一点:

module MyMap where

-- a really terrible `Map`
newtype Map k v = Map [(k, v)]

myLookup :: (Eq k) => k -> Map k v -> Maybe v
myLookup k (Map m) = go m
  where go [] = Nothing
        go ((k', v): rest) | k == k' = Just v
                           | otherwise = go rest

myFindWithDefault :: (Eq k) => v -> k -> Map k v -> v
myFindWithDefault def k m = case myLookup k m of
  Nothing -> def
  Just x -> x

如果使用ghc -O2 -fforce-recomp -ddump-simpl -dsuppress-all -dsuppress-uniques MyMap.hs编译该代码,您会发现为myFindWithDefault生成了以下内核:

myFindWithDefault
  = \ @k @v $dEq def k1 m ->
      joinrec {
        go ds
          = case ds of {
              [] -> def;
              : ds1 rest ->
                case ds1 of { (k', v1) ->
                case == $dEq k1 k' of {
                  False -> jump go rest;
                  True -> v1
                }
                }
            }; } in
      jump go (m `cast` <Co:3>)

注意,这里没有调用myLookup,也没有涉及JustNothingMaybe

myOptimizedFindWithDefault :: (Eq k) => v -> k -> Map k v -> v
myOptimizedFindWithDefault def k (Map m) = go m
  where go [] = def
        go ((k', v): rest) | k == k' = v
                           | otherwise = go rest

并使用相同的标志重新编译,您可能会惊讶地发现,针对myOptimizedFindWithDefault的GHC编译代码如下所示:

myOptimizedFindWithDefault = myFindWithDefault

GHC不仅自动优化原始的myFindWithDefault,而且它足够聪明,能够确定“难以置信的愚蠢”和手动优化的定义是等价的,并且它消除了重复的代码。
因此,即使Data.Map中的findWithDefault定义“愚蠢得令人难以置信”,但聪明得令人难以置信的GHC编译器弥补了这一点,而且没有特别的理由去尝试优化它。
然而,回到你的问题上,当你试图修改一个包时,作者不明智地决定将内部结构保留在内部,你会陷入困境。看看containers-0.4.0.0的源代码,启用TESTING的唯一方法是使用-DTESTING标志重新编译Data.Map。如果你无论如何都要重新编译Data.Map,你也可以直接修改findWithDefault的定义,你实际安装修改后的包的方法取决于你如何管理你的Haskell安装和Haskell项目。
幸运的是,正如评论中所指出的,containers的最新版本在适当命名的模块中导出内部组件,因此这可能是最好的方法。

import Data.Map.Internal (Map(..))

findWithDefault' :: Ord k => a -> k -> Map k a -> a
findWithDefault' def k = k `seq` go
  where go Tip = def
        go (Bin _ kx x l r) =
          case compare k kx of
            LT -> go l
            GT -> go r
            EQ -> x

不过,正如注解中 * 也 * 指出的那样,在容器的后续版本中,findWithDefault的定义已经被重写,大致如下所示:

-- actual definition form containers-0.6.5.1
findWithDefault :: Ord k => a -> k -> Map k a -> a
findWithDefault = go
  where
    go def !_ Tip = def
    go def k (Bin _ kx x l r) = case compare k kx of
      LT -> go def k l
      GT -> go def k r
      EQ -> x

相关问题