tail recursion

Stream: std4

[Mario Carneiro]

@Scott Morrison said:

:light_bulb: Interesting. I wish we had an annotation to require / check tail-recursiveness...

You mean like this?

import Lean
open Lean

partial def visitBody (f : IR.FunId) : IR.FnBody → StateM (NameSet × Bool) Unit
  | .vdecl x _ e b => do
    let tailcall := match b with | .ret (.var y) => x == y | _ => false
    match e with
    | .fap c _ =>
      if tailcall && c == f then modify fun (s, _) => (s, true)
      else modify fun (s, b) => (s.insert c, b)
    | .pap c _ => modify fun (s, b) => (s.insert c, b)
    | _ => pure ()
    visitBody f b
  | .jdecl _ _ v b => do visitBody f v; visitBody f b
  | .case _ _ _ cs => for c in cs do visitBody f c.body
  | e => if e.isTerminal then pure () else visitBody f e.body

elab "#ir_stats " decl:ident : command => do
  let name := decl.getId
  let some decl := (IR.declMapExt.getState (← getEnv)).find? name
    | throwError "IR declaration {name} not found"
  match decl with
  | .extern .. => logInfo m!"{name} is external"
  | .fdecl (f := f) (body := b) .. =>
    let (_, calls, tailRec) := (visitBody f b).run ({}, false)
    if !calls.isEmpty then
      logInfo m!"{name} calls (recursively) {calls.toList}"
    if tailRec then
      logInfo m!"{name} is tail-recursive"
    else if calls.contains f then
      logInfo m!"{name} is properly recursive"
    else
      logInfo m!"{name} is not recursive"

@[specialize] def map {α β} (f : α → β) : List α → List β
  | []    => []
  | a::as => f a :: map f as

@[inline] def mapTR {α β} (f : α → β) (as : List α) : List β :=
  loop as []
where
  @[specialize] loop : List α → List β → List β
  | [],    bs => bs.reverse
  | a::as, bs => loop as (f a :: bs)

#ir_stats map._rarg -- map._rarg is properly recursive
#ir_stats mapTR.loop._rarg -- mapTR.loop._rarg is tail-recursive

[Scott Morrison]

Nice!

Perhaps the name should just be #check_recursive?

[Scott Morrison]

Could this go in Lean.Util.Recursion?

[Scott Morrison]

Would it make sense to allow leaving off the ._rarg?

[Mario Carneiro]

The idea behind this interface was more to give general information about IR declarations, kind of like #print but for IR decls instead of regular decls. For checking for bounded stack space, I was thinking of a more all-in-one solution that traverses the whole dag of function declarations starting at the one you specify, and reports if any of them is properly recursive (or mutually recursive). That way you don't have to put the _rarg. That would also presumably work better as an attribute like you originally suggested.

There is a limitation though: we can only inspect the definitions of declarations in the same file - the IR for other files is not available, only the signatures, so we have to treat any calls to upstream functions as opaque.

[Mario Carneiro]

I didn't show it above but the reason #ir_stats prints the direct calls is so that you can be lead to inspect the _rarg function, I didn't just guess those names above