What tools can help in realizing tail recursion?

Question

I had nice discussions with Leonid and Rojo that got me interested in tail recursion. Tail recursion is not always easy to realize with Mathematica, so it would be nice to have some tools to help with this. The question is the title: What tools can help in realizing tail recursion?

Answer 1

The idea: using continuations

I was preparing my answer when I saw the one by Mr.Wizard, which turns out to be a special case of what I am to offer (and which is generally a common trick that many of us used many times for some recursive problems). I still decided to post it however, since I believe it adds some value. What I will describe here is a version of the transformation to the so-called continuation-passing style, albeit limited in certain ways. The basic requirement of the recursive function would be to have a separate rule for the base case, and only a single call to itself inside its body.

The idea will be to make the function accept a second argument, which will be a so-called continuation - a function that encodes the rest of the program's execution. The recursive function will build the continuation as it proceeeds, and then the continuation will effectively unwind the stack.

The implementation

Without further ado, here is the code-generating macro to perform such a conversion:

ClearAll[def];
SetAttributes[def, HoldAll];
def[SetDelayed[f_[args___], rhs_]] :=
  With[{heldBody = Hold[rhs]},
    With[{innerCalls = 
       Cases[
          heldBody, 
          p : _f :>
            (Hold[p] /. Hold[f[x___]] :>
               (heldBody /. HoldPattern[p] -> # /. Hold[code_] :> 
                   Hold[f[x, Composition[fun, Function[code]]]])), 
          Infinity, 
          1
       ]
    },
    Module[{},
       First[innerCalls] /. Hold[code_] :>
          SetDelayed @@ Prepend[Hold@code, Unevaluated[f[args, fun_]]];
       f[p : PatternSequence[args]] := f[p, Identity]
    ] /; Length[innerCalls] == 1
  ]];

def[SetDelayed[f_[args___], rhs_]] :=
   f[args, fun_: Identity] := fun[rhs];

The code of this macro is quite involved, and uses a nested injector pattern, but what it does is simple to describe. It parses the definition-to-be-made, and hunts for calls of f[x___] in the body (r.h.s.). If the body has the form g[someCode__,f[x__]] (can be more nested too), it constructs the continuation, which is g[someCode,#]&, and passes that into f as a last argument. So, instead of having

f[x_]:=g[someCode___,f[something]]

we now have

f[x_,fun_]:= f[something,Composition[ g[someCode,#]&, fun]]

which is the essence of the CPS transformation in this case.

Fixing some speed problems

As discovered by Mr.Wizard, the above version can be quite slow for large number of iterations (deep recursion). I analyzed this and found that the culrpit is Composition with its Flat attribute. As the number of functions in the chain grows, the Flat attribute leads to that array of functione re-evaluated at every step, which leads to an overall quadratic complexity of the solution, in much the same way as when one uses Append or Prepend to build a large list.

The way out is to have our own version of Composition, which will be free of this flaw. Here it is:

ClearAll[composition];
SetAttributes[composition, HoldAllComplete];
composition[f_, g_][x_] := f[g[x]] 

With this one, my version of Mr.Wizard's benchmarks, on my system:

Block[{$IterationLimit = 1*^6, Composition   = composition },
  f[100000]
] // AbsoluteTiming

(* {0.362305, 5000050000} *)

whereas

Block[{$IterationLimit = 1*^6}, g[100000]] // AbsoluteTiming

(* {0.143555, 5000050001} *)

where f and g are defined as

ClearAll[f];
def[f[1] := 1]
def[f[n_] := n + f[n - 1]]

and

ClearAll[g]
g[0, total_] := total;
g[n_, total_] := g[n - 1, total + n]
g[n_] := g[n, 1]

So, we have a factor of 2 here, more or less as I predicted. And this is for computationally not very intensive Plus. Should we look at Times (the factorial), and one can find that the time for large n is roughly the same.

Examples

Building a simple linked list

The first example is a simple one, concerned with recursive building of linked list. Very natural application, isn't it? Here is a toy function:

Clear[lim, f];
f[x_ /; x < lim] := {x, f[x + 1]}
f[x_] := x

We can test it:

lim=5;
res =f[1]

(* {1,{2,{3,{4,5}}}} *)

However, it is not tail-recursive:

lim=2000;
f[1];

$RecursionLimit::reclim: Recursion depth of 1024 exceeded. >>

Now, we use def:

Clear[f];
def[f[x_ /; x < lim] := {x, f[x + 1]}]
def[f[x_] := x]

We can test that the simple case is fine:

lim=5;
res =f[1]

(* {1,{2,{3,{4,5}}}} *)

but we also have no problem with a general case:

lim=2000;
res =f[1];
Flatten[res]//Short

(* {1,2,3,4,5,6,7,<<1986>>,1994,1995,1996,1997,1998,1999,2000}  *)

It makes sense to inspect the actual definitions for f:

?f

Global`f
f[x_/;x<lim,fun_]:=f[x+1,Composition[fun,{x,#1}&]] 
f[p$:PatternSequence[x_/;x<lim]]:=f[p$,Identity] 
f[x_,fun$_:Identity]:=fun$[x]

Recursive addition

Here is the recursive addition function:

ClearAll[ff];
def[ff[0] := 1];
def[ff[n_] := n + ff[n - 1]]

Now, its true definitions are

?ff
Global`ff
ff[0,fun$_:Identity]:=fun$[1] 
ff[n_,fun_]:=ff[n-1,Composition[fun,n+#1&]] 
ff[p$:PatternSequence[n_]]:=ff[p$,Identity]

and again, no problem with $RecursionLimit:

ff[3000]

(* 4501501 *)

Recusrive select

I will take this example from here, and modify it so that it satisfies our requirement to have a separate base case:

Clear[toLinkedList, test, selrecBad, sel]
toLinkedList[x_List] := Fold[{#2, #1} &, {}, Reverse[x]];
sel[x_List, testF_] := 
   Block[{test = testF}, Flatten[selrecBad @@ toLinkedList[x]]]

and the function proper:

selrecBad[fst_, {}] := If[test[fst], {fst}, {}]
selrecBad[fst_?test, rest_List] := 
  {fst, If[rest === {}, {}, selrecBad[Sequence @@ rest]]};
selrecBad[fst_, rest_List] := selrecBad[Sequence @@ rest];

Note that If is not really necessary, it is there just to show that the function call can be deeper into the body of the function. We can see that it works:

sel[Range[30], EvenQ]

(* {2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30}  *)

However, it is not tail-recursive:

sel[Range[3000], EvenQ];

$RecursionLimit::reclim: Recursion depth of 1024 exceeded. >>

We will now perform our magic:

Clear[selrecBad];
def[selrecBad[fst_?test, {}] := If[test[fst], {fst}, {}]];
def[selrecBad[fst_?test, rest_List] := {fst, 
    If[rest === {}, {}, selrecBad[Sequence @@ rest]]}];
def[selrecBad[fst_, rest_List] := selrecBad[Sequence @@ rest]];

and now:

sel[Range[3000],EvenQ]//Short

(* {2,4,6,8,10,12,14,<<1486>>,2988,2990,2992,2994,2996,2998,3000} *)

Note that we only wrapped def around previous definitions, nothing else changed. Again, it is instructive to inspect actual definitions for selrecBad (which is not so bad any more):

?selrecBad

Global`selrecBad
selrecBad[fst_?test,{},fun$_:Identity]:=fun$[If[test[fst],{fst},{}]]

selrecBad[fst_?test,rest_List,fun_]:=
   selrecBad[Sequence@@rest,Composition[fun,{fst,If[rest==={},{},#1]}&]]

selrecBad[p$:PatternSequence[fst_?test,rest_List]]:=selrecBad[p$,Identity]

selrecBad[fst_,rest_List,fun_]:=selrecBad[Sequence@@rest,Composition[fun,#1&]]

selrecBad[p$:PatternSequence[fst_,rest_List]]:=selrecBad[p$,Identity]

These definitions allow one to understand what really happens when the function is called.

Remarks and limitations

While it is interesting to see how continuations can help us here, one thing to note is that this method won't be very efficient. It will likely be a constant factor of 2 to 10 slower than hand-tuned tail-recursive functions. The value of this construct seems to be more in the concept it illustrates, and the generality of the approach.

It is also important to note that the result of this transformation can be not innocent if the body of the function contains mutable code, since the order of evaluation here will generally change.

The def macro itself has certain purely technical limitations, in that it scopes lexically, and will miss function calls such as f @@ {args}.

Finally, it is not clear to me yet whether the procedure I described above can be called a real use of continuations or using trampolines. Likely the latter (with Composition acting as a trampolining device), but I like the word "continuations" more ;).

Answer 2

For issues with CompoundExpression, I refer to my answer here

https://stackoverflow.com/questions/4481301/tail-call-optimization-in-mathematica/15292525#15292525

In this answer, I define a function called wrapper, acting as a replacement for CompoundExpression. Also follow this link for a nice answer on tail recursion by Leonid.

I have also made a function that can make algorithms in which Times occurs tail recursive, or rather, make tail recursive implementations possible. The function jTimes below is intended to replace Times in such cases. Note, however, that this is a work in progress. Naively combining the function below with its analog for Plus will not yield a tail recursion.

jTimes is defined as follows.

SetAttributes[jTimes, HoldAll];
Module[{functionToken1},

 (*initialisation step*)

 jTimes[arg1_, arg2_] :=
  Block[
   {jTimes},
   SetAttributes[jTimes, HoldAll];
   functionToken1[Evaluate[arg1], Evaluate[arg2]]
   ];

 SetAttributes[{functionToken1}, HoldAll];

 functionToken1[arg1_, jTimes[argd1First_, argd1Last_]] := 
  functionToken1[Evaluate[Times[arg1, argd1First]], 
   Evaluate[argd1Last]];

 (*end step*)

 functionToken1[arg1_, arg2_] := Times[arg1, arg2]
 ]

Then if we set

Clear[f];
f[0] = 1;
f[n_] := jTimes[n, f[n - 1]]

We have

Block[{timing},
 {
  timing = Block[
     {$IterationLimit = 30000},
     f[5000]

     ] // Timing;
  Last[timing] == 5000!
  ,
  timing // First
  ,
  5000! // Timing // First
  }
 ]

-> {True, 0.026255, 0.000311}

The timing is just an indication of the speed, I compare it with the timing of 5000! only to be fair, as our recursive formula is simply not the most efficient way to calculate a large factorial. Note that we do not reach $RecursionLimit, and that including the Block statement scaling up $IterationLimit is necessary.

Note that if we need to multiply more than two things, we can just let the rest of the arguments be handled by Times as follows.

f[n_] := jTimes[Times[n,n], f[n - 1]]

The reason we can do this without a problem is that normally most work in expression like jTimes[Times[n,n],f[1000]] will be in calculating f[1000]. jTimes puts emphasis on its last argument, evaluating this as late as possible.

Note that jTimes does not handle more than two arguments. With this in mind, the following function may be a little more user friendly, though in the case of two arguments it must be slower. I will probably make a function that refers to either jTimes or jTimesMultiple below, based on the number of arguments provided. But note again that the function below is really not necessary if you know about the trick above. jTimesMultiple is defined as follows.

SetAttributes[jTimesMultiple, HoldAll];

Module[{token, functionToken1},

 (*initialisation step*)
 jTimesMultiple[args__, last_] := 
  Block[{jTimesMultiple},
   SetAttributes[jTimesMultiple, HoldAll];
   functionToken1[Evaluate[Times[args]], Evaluate[last]]
   ];

 SetAttributes[functionToken1, HoldAll];

 functionToken1[arg1_, jTimesMultiple_[argd1Firsts__, argd1Last_]] := 
  functionToken1[Evaluate[Times[arg1, argd1Firsts]], 
   Evaluate[argd1Last]];

 (*end step*)

 functionToken1[arg1_, arg2_] := Times[arg1, arg2]

 ]

We then have, for

Clear[f2, f3, f4];
f2[0] = 1;
f2[n_] := jTimesMultiple[n*n, f2[n - 1]]

f3[0] = 1;
f3[n_] := jTimes[n*n, f3[n - 1]]

f4[0] = 1;
f4[n_] := jTimesMultiple[n, n, f2[n - 1]]

That

Block[
 {$IterationLimit = 30000},
 f2[5000] == f3[5000] == f4[5000]
 ]

( -> True)

Combination with wrapper

Notice that wrapper is the function that solves some issues with CompoundExpression in the answer I mentioned above. It is defined as follows

SetAttributes[wrapper, HoldRest];
wrapper[first_, fin_] := fin
wrapper[first_, rest__] := wrapper[rest]

Then, if we set

Clear[fff];
fff[0] = 1;
fff[n_] := wrapper[c++, jTimes[n, fff[n - 1]]];

We have

Block[{$IterationLimit = 100000}
 ,
 c = 0;
 fff[5000] == 5000! && c == 5000
 ]

(-> True)

More tools: If and With

Note that If works great with tail recursion. Setting

Clear[ff]
ff[n_] := If[True, ff[n - 1]]
ff[0] = 0;

We have

Block[{$IterationLimit = 20000},
 ff[5000]==0
]

(-> True)

We could do the same thing with With.

Plus

I suppose there is nothing stopping us from doing exactly the same for plus. And probably many other functions. I may make a nice piece of code emphasizing the similarities between the code for Plus and that for Times.

About how jTimes works

Of course jTimes should evaluate to something. But it is a bit of an issue to let for example f[1000] evaluate to jTimes[a,b] and then stop evaluation there. The reason we want to stop evaluating jTimes[a,b] is that we will have something like jTimes[c,jTimes[a,b]] on "the stack" and if evaluation of jTimes[a,b] leads to another nested jTimes etc, our stack explodes. The trick I used to stop this is to use Block, to make Mathematica forget the definition attached to jTimes for a while. The evaluation stops and we can return to a function that is "higher up" on the stack to massage our expression.

Remarks

Note that it is necessary to give jTimes the attribute a second time when it is considered local. That was confusing to me.

I have tried to make my function evaluate like Times in case of tricky input like

Clear[ggg, f5, f6]
ggg := If[c <=   10, c++; ggg, c];
f5[n_] := jTimes[ggg, f5[n - 1]];
f5[0] = 1;
f6[n_] := Times[ggg, f6[n - 1]];
f6[0] = 1;

With these definitions, the following gives

c = 0;
{f5[11], c}
c = 0;
{f6[11], c}

-> {285311670611, 11} -> {285311670611, 11}

The function jTimes seems to be a bit of a beast, using some (possibly) unconventional techniques. Possibly there are much more elegant ways of doing things. Any feedback is welcome :).

Answer 3

As a point of reference it's not hard to make Times iterative, you just need to put everything inside a single function call:

ClearAll[g]
g[0, total_] := total;
g[n_, total_] := g[n - 1, total*n]
g[n_] := g[n, 1]

Block[{$IterationLimit = 30000},
 g[5000] === 5000!
]

True

Answer 4

Well, here's a new idea. I've put it here mainly to discuss it. It is not really practical yet. I really should have played with Leonids function first, but I couldn't resist making this function.

This approach requires the user to put the symbol token around the last call of the recursive function. In my newest versions, it is also needed to set a state variable in the recursive definitions. I guess this can easily automized (by a new assignment operator). The idea is similar to the idea in my answer above.

The code below defines my newest version of initializer, that is specific for the case where we have only one recursive call.

ClearAll[initializer];
SetAttributes[initializer, HoldAll];
SetAttributes[token, HoldAll];

Module[{evaluator}
  ,
  SetAttributes[evaluator, HoldAll];
  initializer::tokenWarning = "no token present in expression";
  initializer[expr_] :=
   Block[{stateBool = True},

    With[
     {newExpr = expr
      ,

      replacer =
       Function[
        ReplaceRepeated[#,
         evaluator[replacerPos_, replacerExpr_] :> 
          With[
           {
            replacerNewExpr =
             Expand[
              Delete[
               Unevaluated[replacerExpr],
               replacerPos
               ]
              ]
            }
           ,
           If[
            stateBool
            ,
            evaluator[
             replacerPos,
             replacerNewExpr
             ]
            ,
            replacerNewExpr
            ]
           ]
         ]
        ]

      },
     With[{pos = Position[Unevaluated[newExpr], token]}
      ,
      If[
       pos == {},
       Message[initializer::tokenWarning],

       replacer[evaluator[pos, newExpr]]
       ]
      ]

     ]
    ]
  ];

Now, setting

ClearAll[fRep];
fRep[n_] := Plus[n, Times[n, token[fRep[n - 1]]]];
fRep[0] := (stateBool = False; 1);
ClearAll[ggg2]
ggg2[nnn_, total_] := ggg2[nnn + 1, Plus[nnn + 1, (nnn + 1)*total]];
ggg2[5000, total_] := total;

We have

Block[{$IterationLimit = 30000, a, b},
 a = Timing[initializer[fRep[5000]]];
 b = Timing[ggg2[0, 1]];
 {Last[a] == Last[b], First[a], First[b]}
 ]

-> {True, 0.191861, 0.027366}

If we use a version that automatically clears and later sets the attributes Flat and Orderless of Plus and Times, the result is

-> {True, 0.166841, 0.027617}

If I can say so myself, it is pretty magical :). But it is a factor 7 slower, so experienced users may want to keep doing their own recursions. In the section below, all the ideas that go into this function are touched.

Note that I may have been wrong in thinking that ReplaceRepeated would make things faster. Especially, consider the following

Block[{$IterationLimit = 30000, a, b},
 a = Timing@
   ReplaceRepeated[{0, 
     1}, {{5000, total_} :> 
      total, {nnn_, total_} :> {nnn + 1, 
       Plus[nnn + 1, (nnn + 1)*total]}}];
 b = Timing@ ggg2[0, 1];
 {Last[a] == Last[b], First[a], First[b]}
 ]

-> {True, 0.029521, 0.027656}

Note that this works for quite a lot of functions. For example, if we set

ClearAll[fRep];
fRep[n_] :=
  Module[{},
   Block[{d},
    CompoundExpression[d = c++,
     d + Plus[n, Times[n, token[fRep[n - 1]]]]]]];
fRep[0] := (stateBool = False; 1);

We have

{c = 0; N[Timing[initializer[fRep[5000]]]], c}

-> {{0.218156, 3.633367163547326*10^16329}, 5000}

Older, more stabile version

Below is an older version. The idea is more evident here. We block evaluation with some symbol and have a function at the top of our expression, in this case evaluator, to massage it. We let Mathematica get rid of Times and Plus(es) on the stack using Expand. Of course, this use of Expand makes this a specific solution for cases where Times and Plus are used in a recursive definition. However, note that other regular evaluations are also made, in fact I added Expand later as I thought Mathematica would simplify expressions like 1+2(3+token) automatically.

We define initializer and evaluator as follows

ClearAll[evaluator];
SetAttributes[evaluator, HoldAll];
ClearAll[initializer];
SetAttributes[initializer, HoldAll];
SetAttributes[token, HoldAll];

initializer::tokenWarning = "no token present in expression";
initializer[expr_] :=
  With[{newExpr = expr},
   With[{pos = Position[Unevaluated[newExpr], token]}
    ,
    If[
     pos == {},
     Message[initializer::tokenWarning],
     evaluator[pos, newExpr]
     ]

    ]
   ];

evaluator[pos_, expr_] := 
 With[
  {newExpr = Expand[Delete[Unevaluated[expr], pos]]},
  With[
   {newPos = Position[Unevaluated[newExpr], token]}
   ,
   evaluator[
    Evaluate[
     If[
      newPos == {},
      Null,
      newPos
      ]
     ]
    ,
    newExpr
    ]
   ]
  ]

(*end step*)
evaluator[Null, expr_] := expr

Now, setting

ClearAll[f, ff];
f[n_] := Plus[n, Times[n, token[f[n - 1]]]];
f[0] = 0;
ff[n_] := Plus[n, Times[n, ff[n - 1]]];
ff[0] = 0;
ClearAll[ggg2]
ggg2[nnn_, total_] := ggg2[nnn + 1, Plus[nnn + 1, (nnn + 1)*total]];
ggg2[5000, total_] := total;

We have

Block[
 {$IterationLimit = 30000, a, b},
 a = Timing[ggg2[0, 0]];
 b = Timing[initializer[f[5000]]];
 {Last[a] == Last[b], First[a], First[b], N[Last[a]]}
 ]

-> {True, 0.027325, 0.253952, 1.149446653811488*10^16326}

So this approach is 10 times slower than necessary.

I've also tried to make a version where multiple tokens can be provided, but it doesn't work yet. I have also made a version which first checks if at the specified position we do indeed have a token, and if not, it finds it using position.

Note that this also solves the problem for CompoundExpression. Setting

ClearAll[f3]
f3[n_] := CompoundExpression[c++, Plus[n, Times[n, token[f3[n - 1]]]]];
f3[0] = 0;

We find

Block[{$IterationLimit = 20000},
 {c = 0; N[Timing[initializer[f3[5000]]]] , c}
 ]

-> {{0.296686, 1.149446653811488*10^16326}, 5000}

Note that in this way, this approach may still have additional value, despite Leonids fantastic answer. Leonids answer does not seem to handle CompoundExpression well. If we set, using Leonids definitions

ClearAll[f1];
f1[1] = 1;
f1[n_] := CompoundExpression[c++, c + n (n + f1[n - 1])]
ClearAll[f2];
def[f2[1] := 1]
def[f2[n_] := CompoundExpression[c++, c + n (n + f2[n - 1])]]
ClearAll[f3];
f3[1] = 1;
f3[n_] := CompoundExpression[c++, c + n (n + token[f3[n - 1]])]

We have

Block[{a1, a2, a3},
 a1 = (c = 0; f2[100]);
 a2 = (c = 0; f[100]);
 a3 = (c = 0; initializer[f3[100]]);

 {a1 == a2, a1 == a3}
 ]

-> {False, True}

Where a1 is the correct answer, so that indeed Leonids method fails and the method describes here evaluates to the desired result.

Conclusion

I guess this may be a strong candidate, especially if we can make a nice assignment operator that keeps track of which rules generate new tokens. Maybe it would then be nice to generate warnings if certain functions are used in the recursive definition handled by our new assignment operator. Of course things like

f[n_]:= {c++, Sequence@@ConstantArray[Null,c],token[f[n-1]]}

Are impossible to handle in this way, but then again, this example is really useless anyway and will reach $RecursionLimit because of nested Lists.

So.. again, feedback is welcome!