How to avoid stack overflow in this F# program (recursive tree search)?

Question

I've got a discriminated union tree like this:

type rbtree =
 | LeafB of int
 | LeafR of int
 | Node of int*rbtree*rbtree

And what I have to do is to search for every LeafB present in the tree, so I came with a this recursive function:

let rec searchB (tree:rbtree) : rbtree list = 
 match tree with
 | LeafB(n) -> LeafB(n)::searchB tree
 | LeafR(n) -> []
 | Node(n,left,right) -> List.append (searchB left) (searchB right)

But when I try to test it I get stack overflow exception and I have no idea how to modify it to work properly.

Answer 1

As @kvb says your updated version isn't truely tail-rec and might cause a stackoverflow as well.

What you can do is using continuations essentially using heap space instead of stack space.

let searchB_ tree =
 let rec tail results continuation tree =
 match tree with
 | LeafB v -> continuation (v::results)
 | LeafR _ -> continuation results
 | Node (_, lt, rt) -> tail results (fun leftResults -> tail leftResults continuation rt) lt
 tail [] id tree |> List.rev

If we looks at the generated code in ILSpy it looks essentially like this:

internal static a tail@13<a>(FSharpList<int> results, FSharpFunc<FSharpList<int>, a> continuation, Program.rbtree tree)
{
 while (true)
 {
 Program.rbtree rbtree = tree;
 if (rbtree is Program.rbtree.LeafR)
 {
 goto IL_34;
 }
 if (!(rbtree is Program.rbtree.Node))
 {
 break;
 }
 Program.rbtree.Node node = (Program.rbtree.Node)tree;
 Program.rbtree rt = node.item3;
 FSharpList<int> arg_5E_0 = results;
 FSharpFunc<FSharpList<int>, a> arg_5C_0 = new Program<a>.tail@17-1(continuation, rt);
 tree = node.item2;
 continuation = arg_5C_0;
 results = arg_5E_0;
 }
 Program.rbtree.LeafB leafB = (Program.rbtree.LeafB)tree;
 int v = leafB.item;
 return continuation.Invoke(FSharpList<int>.Cons(v, results));
 IL_34:
 return continuation.Invoke(results);
}

So as expected with tail recursive functions in F# it is tranformed into a while loop. If we look at the non-tail recursive function:

// Program
public static FSharpList<int> searchB(Program.rbtree tree)
{
 if (tree is Program.rbtree.LeafR)
 {
 return FSharpList<int>.Empty;
 }
 if (!(tree is Program.rbtree.Node))
 {
 Program.rbtree.LeafB leafB = (Program.rbtree.LeafB)tree;
 return FSharpList<int>.Cons(leafB.item, FSharpList<int>.Empty);
 }
 Program.rbtree.Node node = (Program.rbtree.Node)tree;
 Program.rbtree right = node.item3;
 Program.rbtree left = node.item2;
 return Operators.op_Append<int>(Program.searchB(left), Program.searchB(right));
}

We see the recursive call at the end of the function Operators.op_Append<int>(Program.searchB(left), Program.searchB(right));

So the tail-recursive function allocates continuations functions instead of creating a new stack frame. We can still run out of heap but there's lot more heap than stack.

Full example demonstrating a stackoverflow:

type rbtree =
 | LeafB of int
 | LeafR of int
 | Node of int*rbtree*rbtree
let rec searchB tree = 
 match tree with
 | LeafB(n) -> n::[]
 | LeafR(n) -> []
 | Node(n,left,right) -> List.append (searchB left) (searchB right)
let searchB_ tree =
 let rec tail results continuation tree =
 match tree with
 | LeafB v -> continuation (v::results)
 | LeafR _ -> continuation results
 | Node (_, lt, rt) -> tail results (fun leftResults -> tail leftResults continuation rt) lt
 tail [] id tree |> List.rev
let rec genTree n =
 let rec loop i t =
 if i > 0 then
 loop (i - 1) (Node (i, t, LeafB i))
 else
 t
 loop n (LeafB n)
[<EntryPoint>]
let main argv =
 printfn "generate left leaning tree..."
 let tree = genTree 100000
 printfn "tail rec"
 let s = searchB_ tree
 printfn "rec"
 let f = searchB tree
 printfn "Is equal? %A" (f = s)
 0

Answer 2

Oh, I might came with an solution:

let rec searchB (tree:rbtree) : rbtree list = 
match tree with
| LeafB(n) -> LeafB(n)::[]
| LeafR(n) -> []
| Node(n,left,right) -> List.append (searchB left) (searchB right)

Now it looks like working properly when I try it.

• kvb: As long as your tree isn't too deep this will work; however, note that the recursive calls to searchB will cause the stack to grow so with a very deep tree it's still possible to get a stack overflow.