Haskell中Map的链式遍历：实现指定签名的chainTraversal函数

功能说明 #

我们需要实现的函数可以从指定起始key出发，按照处理函数指定的规则链式遍历Map结构，只保留遍历到的、处理成功的键值对，是过滤和遍历能力的结合，函数签名如下：

chainTraversal :: Ord k => k -> (k -> a -> Maybe (k, b)) -> Map k a -> Map k b

注：相比原始签名补充了Ord k约束，这是Haskell中Map操作的必要前置条件。

实现逻辑 #

• 从传入的初始key开始遍历，每次优先查询当前key是否存在于原Map中
• 存在则将当前key和对应值传入处理函数：
  • 处理函数返回Nothing：遍历直接终止
  • 处理函数返回Just (下一个key, 处理后的值)：将当前key和处理后的值存入结果Map，以下一个key为起点继续遍历
• 额外增加已访问key校验逻辑，避免处理函数返回重复key导致死循环，无循环风险的场景可自行删除该逻辑

完整实现代码 #

import qualified Data.Map as M
import Data.Map (Map)

chainTraversal :: Ord k => k -> (k -> a -> Maybe (k, b)) -> Map k a -> Map k b
chainTraversal startKey process originalMap = go startKey M.empty M.empty
  where
    -- 辅助函数参数：当前遍历key、已访问key集合、已生成的结果Map
    go currKey visited result
      -- 遇到已访问的key直接终止，避免死循环
      | M.member currKey visited = result
      | otherwise = case M.lookup currKey originalMap of
          -- 当前key不存在于原Map，终止遍历
          Nothing -> result
          Just currVal -> case process currKey currVal of
            -- 处理函数返回终止信号，结束遍历
            Nothing -> result
            Just (nextKey, processedVal) -> go nextKey
              (M.insert currKey () visited)
              (M.insert currKey processedVal result)

用法示例 #

-- 测试用原Map
testMap :: Map Int Int
testMap = M.fromList [(1, 3), (2, 5), (3, -1), (4, 10)]

-- 测试处理函数：当前值大于0时，下一个遍历key为当前key+1，处理后的值为原值*2；否则终止遍历
testProcess :: Int -> Int -> Maybe (Int, Int)
testProcess k v
  | v > 0 = Just (k + 1, v * 2)
  | otherwise = Nothing

-- 调用示例：从key=1开始遍历
-- 返回结果：fromList [(1,6), (2,10)]
testResult :: Map Int Int
testResult = chainTraversal 1 testProcess testMap

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