codegen: ref-count instruction destroying tail recursion

Prerequisites

• Check that your issue is not already filed:
  https://github.com/leanprover/lean4/issues
• Reduce the issue to a minimal, self-contained, reproducible test case.
  Avoid dependencies to Mathlib or Batteries.
• Test your test case against the latest nightly release, for example on
  https://live.lean-lang.org/#project=lean-nightly
  (You can also use the settings there to switch to “Lean nightly”)

Description

import Lean
import Std

open Lean Meta

@[inline] def forInNewArray {α : Type u} {σ β : Type v} {m : Type v → Type w} [Inhabited α]
    (as : Array α) (s : σ) (kcons : (a : α) → (σ → m β) → σ → m β) (knil : σ → m β) : m β :=
  let sz := as.size
  let rec @[specialize] loop (i : Nat) (s : σ) : m β :=
    if i < sz then
      let a := as[i]!
      kcons a (loop (i+1)) s
    else
      knil s
  loop 0 s

@[specialize] def forInNewAssoc {m : Type v → Type w} (as : Std.DHashMap.Internal.AssocList α β) (init : σ) (kcons : (a : α) → β a → (σ → m δ) → σ → m δ) (knil : σ → m δ) : m δ :=
  match as with
  | .nil => knil init
  | .cons k v t => kcons k v (forInNewAssoc t · kcons knil) init

set_option trace.Compiler.saveBase true in
def size : MetaM Nat := do
  let env ← getEnv
  forInNewArray env.constants.map₁.inner.inner.buckets ()
      (fun b kcontinue _ =>
        forInNewAssoc b ()
          (fun _ _ kcontinue _ => do
            pure ()
            kcontinue ())
          fun _ => kcontinue ())
      fun _ => pure 0

run_meta size

Context

This is code generated for nested for loops by the new do elaborator.

Steps to Reproduce

Run the program.

Expected behavior: run_meta size terminates successfully

Actual behavior: stack overflow

Observe the output of size after specialization.

Expected behavior: No redundant code

Actual behavior: forInNewAssoc._at_.size.spec_0 is redundant and forInNewAssoc._at_.size.spec_0._at_.forInNewArray.loop._at_.size.spec_1.spec_1 is just peeling off one layer of the loop forInNewAssoc._at_.forInNewAssoc._at_.size.spec_0._at_.forInNewArray.loop._at_.size.spec_1.spec_1.spec_2. As far as I can tell, every call to the former could be replaced by a call to the latter.

Versions

Lean 4.28.0-nightly-2025-12-24
Target: x86_64-unknown-linux-gnu

Additional Information

I did a bit of diagnosis. Note that the LCNF output has all recursive calls in tail position, so I would tail call elimination to trigger. Alas, it doesn't. Here's the last few items of the stack trace:

#1 forInNewArray.loop._at_.size.spec_1._redArg
#2 forInNewAssoc._at_.forInNewAssoc._at_.size.spec_0._at_.forInNewArray.loop._at_.size.spec_1.spec_1.spec_2._redArg
#3 forInNewAssoc._at_.size.spec_0._at_.forInNewArray.loop._at_.size.spec_1.spec_1._redArg
#4 forInNewArray.loop._at_.size.spec_1._redArg
#5 forInNewAssoc._at_.size.spec_0._at_.forInNewArray.loop._at_.size.spec_1.spec_1._redArg
#6 forInNewArray.loop._at_.size.spec_1._redArg
#7 forInNewAssoc._at_.size.spec_0._at_.forInNewArray.loop._at_.size.spec_1.spec_1._redArg
#8 forInNewArray.loop._at_.size.spec_1._redArg
#9 forInNewAssoc._at_.size.spec_0._at_.forInNewArray.loop._at_.size.spec_1.spec_1._redArg
#10 forInNewArray.loop._at_.size.spec_1._redArg
#11 forInNewAssoc._at_.size.spec_0._at_.forInNewArray.loop._at_.size.spec_1.spec_1._redArg
#12 forInNewArray.loop._at_.size.spec_1._redArg
#13 forInNewAssoc._at_.forInNewAssoc._at_.size.spec_0._at_.forInNewArray.loop._at_.size.spec_1.spec_1.spec_2._redArg

One reason that the calls aren't eliminated is the supposed tail call in forInNewArray.loop._at_.size.spec_1._redArg:

        let x_16 : tobj := Nat.add x_5 x_15;
        let x_17 : obj := forInNewAssoc._at_.size.spec_0._at_.forInNewArray.loop._at_.size.spec_1.spec_1._redArg x_1 x_2 x_3 x_4 x_16 x_1 x_14 x_6 x_7 x_8 x_9 ◾;
        dec x_16;
        ret x_17

Why is there a dec x_16? x_16 is passed as the 5th arg to forInNewAssoc._at_.size.spec_0._at_.forInNewArray.loop._at_.size.spec_1.spec_1._redArg, and that function takes its 5th arg by borrow. This borrow has been introduced by the borrow pass. Ouch!! It appears that borrow inference may destroy tail calls. I don't think it should do that!

Impact

Add 👍 to issues you consider important. If others are impacted by this issue, please ask them to add 👍 to it.

Start by running the minimal reproducer and inspecting the specialized output and stack trace shown in the issue. Trace the borrow pass around the generated dec/ret sequence in forInNewArray.loop and its forInNewAssoc caller; done means the redundant wrapper is removed or tail calls are preserved, so run_meta size terminates without stack overflow.