如何在Lean的ForM.forM定义中递归使用for...in循环？

背景 #

Lean参考手册的《嵌套类型的实例》示例展示了一种实用技巧：在实现类型类方法时，将函数标记为partial，并在函数内部定义局部实例，从而递归使用该类型自身的类型类实例。例如：

inductive NatRoseTree where
  | node (val : Nat) (children : Array NatRoseTree)

partial def NatRoseTree.beq : (tree1 tree2 : NatRoseTree) → Bool
  | .node val1 children1, .node val2 children2 =>
    let _ : BEq NatRoseTree := ⟨NatRoseTree.beq⟩
    val1 == val2 &&
    children1 == children2

-- 后续需定义全局实例：BEq NatRoseTree := ⟨NatRoseTree.beq⟩

即便全局BEq NatRoseTree实例尚未定义，Lean的实例搜索也能找到局部实例，解析函数内的==表达式。

问题描述 #

希望将该技巧应用于ForM/ForIn类型类，在自定义类型的ForM.forM实现中，使用for ... in ... do语法遍历该类型的值。以下是一个示例的“倒计时”类型，目标是在其forM实现中通过递归for循环完成遍历：

structure Countdown where
  val: Nat

partial def Countdown.forM
  [Monad m] (c : Countdown) (action : Nat → m PUnit) : m PUnit
:= do match c with
  | ⟨0⟩ => action 0
  | ⟨n⟩ =>

      -- 此处需添加局部ForIn实例以启用for循环
      -- ↓↓↓ 待填充代码 ↓↓↓

      action n
      for i in (⟨n-1⟩ : Countdown) do
        action i

instance : ForM m Countdown Nat := ⟨Countdown.forM⟩
instance : ForIn m Countdown Nat := ⟨ForM.forIn⟩

def main : IO Unit :=
  for x in (⟨3⟩ : Countdown) do
    IO.println s!"count: {x}"

要求必须使用for循环语法，不能直接递归调用Countdown.forM函数。

已尝试方案 #

1. 直接套用参考手册的单态实例：
   在占位符处添加：

   let _ := (⟨Countdown.forM⟩ : ForM m Countdown Nat)
   let _ := (⟨ForM.forIn⟩ : ForIn m Countdown Nat)
   

   编译报错，提示无法合成复合monad的ForM实例：

   error: Main.lean:10:17: failed to synthesize
     ForM (StateT β✝ (ExceptT β✝ m)) Countdown Nat
   

2. 隐式推导monad类型：
   在占位符处添加：

   let {m'} := (⟨Countdown.forM⟩ : ForM m' Countdown Nat)
   let {m'} := (⟨ForM.forIn⟩ : ForIn m' Countdown Nat)
   

   代码可编译，但仅输出count: 3后崩溃（推测为栈溢出），说明for循环未正确执行。

解决方案 #

问题的核心是：for循环展开后的代码会运行在monad transformer栈中（而非原始的m），因此局部实例需要支持任意Monad类型，而非仅当前的m。在占位符处添加以下多态局部实例定义即可解决：

let instForM {m'} [Monad m'] : ForM m' Countdown Nat := ⟨Countdown.forM⟩
let instForIn {m'} [Monad m'] : ForIn m' Countdown Nat := ⟨ForM.forIn⟩

完整可运行代码 #

structure Countdown where
  val: Nat

partial def Countdown.forM
  [Monad m] (c : Countdown) (action : Nat → m PUnit) : m PUnit
:= do match c with
  | ⟨0⟩ => action 0
  | ⟨n⟩ =>
      -- 添加多态局部实例
      let instForM {m'} [Monad m'] : ForM m' Countdown Nat := ⟨Countdown.forM⟩
      let instForIn {m'} [Monad m'] : ForIn m' Countdown Nat := ⟨ForM.forIn⟩
      
      action n
      for i in (⟨n-1⟩ : Countdown) do
        action i

instance : ForM m Countdown Nat := ⟨Countdown.forM⟩
instance : ForIn m Countdown Nat := ⟨ForM.forIn⟩

def main : IO Unit :=
  for x in (⟨3⟩ : Countdown) do
    IO.println s!"count: {x}"

说明 #

• 定义多态的instForM和instForIn，让实例对任意满足Monad m'的类型都生效，覆盖了for循环展开后可能用到的monad transformer场景。
• partial关键字确保递归调用的合法性，Lean会处理递归终止检查（此处Countdown.val递减至0，满足终止条件）。

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