Haskell中Nat类型Enum实例处理infinity时丢失首元素问题排查

问题背景 #

已定义自然数类型：

data Nat
  = Zero
  | Succ Nat

需求为其实现Enum类型类，要求最大化惰性——由于Nat多用于表示列表长度，完全求值Nat往往意味着完全求值列表，例如f x = [2 .. natLength x]应无需求值整个列表即可获取首元素，仅需前两个元素。

用infinity = Succ infinity做惰性测试，若能正确处理infinity则说明惰性足够（因infinity无法被完全求值）。

当前仅关注enumFromThenTo方法，实现代码如下：

instance Enum Nat
  where
    enumFromThenTo =
      go Zero
      where
        go n (Succ x) (Succ y) i =
          go (Succ n) x y i
        go n Zero y i = do
          j <- i - y
          (y +) <$> (n : go n Zero y j)
        go n y Zero Zero =
          n : do
            m <- n - y
            go m y Zero Zero
        go n y Zero i = do
          v <- (n + y) : go n y Zero Zero
          (+ i) <$> (v - i)

        x - Zero =
          [x]
        Succ x - Succ y =
          x - y
        Zero - Succ _ =
          []

异常现象 #

有限Nat值的表现符合预期，但处理infinity时丢失首元素：

ghci> take 14 [1, 9 .. 100]  -- 有限值正常（假设已实现Nat与Int的转换）
[1,9,17,25,33,41,49,57,65,73,81,89,97]
ghci> take 14 [1, 9 .. infinity]  -- infinity场景丢失首元素1
[9,17,25,33,41,49,57,65,73,81,89,97,105,113]

问题分析 #

当调用enumFromThenTo 1 9 infinity时：

1. 初始进入go Zero (Succ Zero) (Succ^8 Zero) infinity，触发第一个分支递归一次，变为go (Succ Zero) Zero (Succ^7 Zero) infinity（此时n为原start值1，y为步长8）。
2. 进入第二个分支go n Zero y i，执行i - y得到infinity（因infinity减有限值仍为infinity）。
3. 原代码中(y +) <$> (n : go n Zero y j)会将整个列表的所有元素（包括首元素n）都加上步长y，导致首元素从1变为1+8=9，直接丢失了原start值。

有限值场景下可能因后续截断逻辑让结果看似正常，但infinity是无限序列，首元素丢失的问题直接暴露。

修复方案 #

修改第二个分支，将首元素单独保留，仅对递归生成的后续元素施加步长加法：

go n Zero y i = do
  j <- i - y
  n : ((y +) <$> go n Zero y j)

修改后，enumFromThenTo 1 9 infinity生成的序列首元素为1，后续依次为9、17等，符合预期。

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