Manually transforming tree recursion into tail recursion using a stack

Question

I'm implementing a variation on topological sort (on top of scala-graph) that returns all topological orderings instead of just one. I have a tree recursive implementation that i want to make tail recursive. I don't want to use trampolines, instead, i want to mimic the call stack as described in this answer.

This is the tree recursive version of my algorithm:

import scalax.collection.Graph
import scalax.collection.GraphPredef._
import scalax.collection.GraphEdge._
import scala.collection.Set

def allTopologicalSorts[T](graph: Graph[T, DiEdge]): Unit = {
  val indegree: Map[graph.NodeT, Int] = graph.nodes.map(node => (node, node.inDegree)).toMap

  def isSource(node: graph.NodeT): Boolean = indegree.get(node).get == 0

  def getSources(): Set[graph.NodeT] = graph.nodes.filter(node => isSource(node))

  processSources(getSources(), indegree, List[graph.NodeT](), 0)

  def processSources(sources: Set[graph.NodeT], indegrees: Map[graph.NodeT, Int], topOrder: List[graph.NodeT], cnt: Int): Unit = {
    if (sources.nonEmpty) {
      // `sources` contain all the nodes we can pick
      // --> generate all possibilities
      for (src <- sources) {
        val newTopOrder = src :: topOrder
        var newSources = sources - src

        // Decrease the in-degree of all adjacent nodes
        var newIndegrees = indegrees
        for (adjacent <- src.diSuccessors) {
          val newIndeg = newIndegrees.get(adjacent).get - 1
          newIndegrees = newIndegrees.updated(adjacent, newIndeg)
          // If in-degree becomes zero, add to sources
          if (newIndeg == 0) {
            newSources = newSources + adjacent
          }
        }

        processSources(newSources, newIndegrees, newTopOrder, cnt + 1)
      }
    }
    else if (cnt != graph.nodes.size) {
      println("There is a cycle in the graph.")
    }
    else {
      println(topOrder.reverse)
    }
  }
}

And we can run the algorithm as follows

val graph: Graph[Int, DiEdge] = Graph(2 ~> 4, 2 ~> 7, 4 ~> 5)
allTopologicalSorts(graph)

Which correctly returns

List(2, 7, 4, 5)
List(2, 4, 7, 5)
List(2, 4, 5, 7)

Now, i tried implementing a tail recursive version by manually keeping a stack

import scalax.collection.Graph
import scalax.collection.GraphPredef._
import scalax.collection.GraphEdge._
import scala.collection.Set

def allTopologicalSorts[T](graph: Graph[T, DiEdge]): Unit = { 
  val indegree: Map[graph.NodeT, Int] = graph.nodes.map(node => (node, node.inDegree)).toMap

  def isSource(node: graph.NodeT): Boolean = indegree.get(node).get == 0

  def getSources(): Set[graph.NodeT] = graph.nodes.filter(node => isSource(node))

  def processSources(sources: Set[graph.NodeT], indegrees: Map[graph.NodeT, Int]): Unit = {
    type Order = List[graph.NodeT]
    case class Frame(sources: List[graph.NodeT], indegrees: Map[graph.NodeT, Int], topOrder: Order, cnt: Int)

    def step(stack: List[Frame]): Unit = {
      stack match {
        case Frame(src :: rest, indegrees, topOrder, cnt) :: tail => {
          val onBacktrackingFrame = Frame(rest, indegrees, topOrder, cnt)

          // Process src now and remember to do the rest later
          val newTopOrder = src :: topOrder
          var newSources = rest

          // Decrease the in-degree of all adjacent nodes
          var newIndegrees = indegrees
          for (adjacent <- src.diSuccessors) {
            val newIndeg = newIndegrees.get(adjacent).get - 1
            newIndegrees = newIndegrees.updated(adjacent, newIndeg)
            // If in-degree becomes zero, add to sources
            if (newIndeg == 0) {
              newSources = adjacent :: newSources
            }
          }

          val recursionFrame = Frame(newSources, newIndegrees, newTopOrder, cnt + 1)
          step(recursionFrame :: onBacktrackingFrame :: tail)
        }
        case Frame(Nil, indegrees, topOrder, cnt) :: tail => {
          println(topOrder.reverse)
          step(tail)
        }
        case Nil =>
      }
    }

    step(List(Frame(sources.toList, indegrees, List[graph.NodeT](), 0)))
  }

  processSources(getSources(), indegree)
}

However, this does not work as it results in

List(2, 4, 5, 7)
List(2, 4, 5)
List(2, 4, 7)
List(2, 4)
List(2, 7)
List(2)
List()

There's something off with the stack, but i couldn't find the problem.

Connected question: Tail recursive algorithm for generating all topological orderings in a graph

Have you seen my stack solution here https://stackoverflow.com/a/55786425/5249621 ? (the second one titled "Direct solution") — Dmytro Mitin, May 02 '19 at 08:34
Yes, the question is not about how to do it. I know i need to simulate the stack, which is what i did. However, something is wrong with my implementation and i'm not able to figure it out. I've been spending days on this, but i really don't see where the problem comes from. — Kevin, May 02 '19 at 08:38
Your previous versions took a structure as input and returned a structure as output. Now these versions return `Unit` (usually it's harder to reason about side effects). Normally a frame contains input arguments of the recursive call and **an argument for result**. Even if the result is `Unit` it should be clear if `Unit`s from recursive subcalls have been already calculated or are still being processed. — Dmytro Mitin, May 02 '19 at 09:19
At the moment, the return type is Unit because i print the results as i encounter them. I could easily add an extra argument which accumulates the results. However, for debugging this was easier, and once this works, i can add the argument. Normally, recursive subcalls are done when the `sources` argument is nil. — Kevin, May 02 '19 at 09:22
I mean frame with calculated subcalls and frame with not calculated subcalls should be different frames. In AndreyTyukin's answer it's `todos`, in my answer it's `NotExpanded/Expanded/Calculated`. And it seems you don't watch that. — Dmytro Mitin, May 02 '19 at 09:41
@DmytroMitin The first field of my frames are `sources` which are similar to AndreyTyukin's `todos`. If `sources` is not empty, i.e. `src :: rest` then i construct two different frames, one with the rest of the "todos" `val onBacktrackingFrame = Frame(rest, indegrees, topOrder, cnt)` and one that goes deeper in the graph `val recursionFrame = Frame(newSources, newIndegrees, newTopOrder, cnt + 1)`. And then, i step with the frames in the right order: `step(recursionFrame :: onBacktrackingFrame :: tail)`. — Kevin, May 02 '19 at 11:04

score 1 · Accepted Answer · answered May 02 '19 at 11:20

This solution is tail-recursive AFAICT and works when I run it, though I changed parts of it back to the first version, notably changing some types from List to Set, in order to keep the changes from the original small (I believe changing it back to List again should be relatively straight-forward):

def allTopologicalSortsNew[T](graph: Graph[T, DiEdge]): Unit = { 
  type Order = List[graph.NodeT]
  case class Frame(sources: Set[graph.NodeT], indegrees: Map[graph.NodeT, Int], topOrder: Order, cnt: Int)
  val indegree: Map[graph.NodeT, Int] = graph.nodes.map(node => (node, node.inDegree)).toMap

  def isSource(node: graph.NodeT): Boolean = indegree.get(node).get == 0

  def getSources(): Set[graph.NodeT] = graph.nodes.filter(node => isSource(node))

  def processSources(initialSources: Set[graph.NodeT], initialIndegrees: Map[graph.NodeT, Int]): Unit = {

    def step(stack: List[Frame]): Unit = {
      stack match {
        case Frame(sources, indegrees, topOrder, cnt) :: tail if !sources.isEmpty => {

          val futureFrames = for (src <- sources) yield {
            val newTopOrder = src :: topOrder
            var newSources = sources - src

            // Decrease the in-degree of all adjacent nodes
            var newIndegrees = indegrees
            for (adjacent <- src.diSuccessors) {
              val newIndeg = newIndegrees.get(adjacent).get - 1
              newIndegrees = newIndegrees.updated(adjacent, newIndeg)
              // If in-degree becomes zero, add to sources
              if (newIndeg == 0) {
                newSources = newSources + adjacent
              }
            }

            Frame(newSources, newIndegrees, newTopOrder, cnt + 1)
          }

          step(futureFrames.toList ::: tail)
        }
        case Frame(sources, indegrees, topOrder, cnt) :: tail if sources.isEmpty => {
          println(topOrder.reverse)
          step(tail)
        }
        case Nil =>
      }
    }

    step(List(Frame(initialSources, initialIndegrees, List[graph.NodeT](), 0)))
  }

  processSources(getSources(), indegree)
}

Dmytro Mitin · Answer 2 · 2019-05-02T13:58:21.870

Try

  def allTopologicalSorts[T](graph: Graph[T, DiEdge]): Stream[List[graph.NodeT]] = {
    val indegree: Map[graph.NodeT, Int] = graph.nodes.map(node => (node, node.inDegree)).toMap

    def isSource(node: graph.NodeT): Boolean = indegree.get(node).get == 0
    def getSources(): Set[graph.NodeT] = graph.nodes.filter(node => isSource(node))

    case class Frame(arg: Argument, parentArg: Option[Argument], res: Result)

    case class Argument(sources: Set[graph.NodeT], indegrees: Map[graph.NodeT, Int], topOrder: List[graph.NodeT], cnt: Int)

    sealed trait Result
    case object NotExpanded extends Result
    case object Expanded extends Result
    case class Calculated(value: Output) extends Result

    type Output = Stream[List[graph.NodeT]]

    // extract result from final state of stack
    def processSources(arg: Argument): Output =
      step1(List(Frame(arg, None, NotExpanded))) match {
        case Frame(`arg`, None, Calculated(res)) :: Nil => res
      }

    @tailrec
    // process stack as long as necessary
    def step1(stack: List[Frame]): List[Frame] = {
      val x = step(stack, Nil)
      x match {
        case Frame(arg, None, Calculated(res)) :: Nil => x
        case _ => step1(x)
      }
    }

    // helper method for handling "recursion backward" case in "step"
    def calcFromChildren(stack: List[Frame], parentArg: Argument): Option[(List[Frame], Frame)] = {
      val (childFrames, rest) = stack.span {
        case Frame(_, Some(`parentArg`), Calculated(_)) => true
        case _ => false
      }

      val output: Output = childFrames.map {
        case Frame(arg, Some(`parentArg`), Calculated(res)) => res
      }.toStream.flatten

      rest match {
        case Frame(`parentArg`, parentArg1, Expanded) :: rest1 if parentArg.sources.nonEmpty =>
          Some(rest1, Frame(parentArg, parentArg1, Calculated(output)))
        case _ => None
      }
    }

    @tailrec
    // process stack once
    def step(stack: List[Frame], acc: List[Frame]): List[Frame] = {
      stack match {
          // recursion backward
        case Frame(arg, Some(parentArg), Calculated(res)) :: frames if calcFromChildren(stack, parentArg).isDefined =>
          val (rest1, parentFrame) = calcFromChildren(stack, parentArg).get
          step(rest1, parentFrame :: acc)

          // base
        case Frame(arg, parentArg, _) :: frames if arg.sources.isEmpty && arg.cnt != graph.nodes.size =>
          throw new Error("There is a cycle in the graph.")
        case Frame(arg, parentArg, _) :: frames if arg.sources.isEmpty =>
          val res = arg.topOrder.reverse #:: Stream.empty[List[graph.NodeT]]
          step(frames, Frame(arg, parentArg, Calculated(res)) :: acc)

          // recursion forward
        case Frame(arg, parentArg, NotExpanded) :: frames =>

          val childFrames = arg.sources.toList.map(src => {
            val newTopOrder = src :: arg.topOrder
            var newSources = arg.sources - src

            var newIndegrees = arg.indegrees
            for (adjacent <- src.diSuccessors) {
              val newIndeg = newIndegrees.get(adjacent).get - 1
              newIndegrees = newIndegrees.updated(adjacent, newIndeg)
              if (newIndeg == 0) {
                newSources = newSources + adjacent
              }
            }

            val newArg = Argument(newSources, newIndegrees, newTopOrder, arg.cnt + 1)
            Frame(newArg, Some(arg), NotExpanded)
          })

          step(frames, Frame(arg, parentArg, Expanded) :: childFrames.reverse ::: acc)

          // ignore if not "recursion backward" case
        case Frame(arg, parentArg, Expanded) :: frames => step(frames, Frame(arg, parentArg, Expanded) :: acc)
        case Frame(arg, parentArg, Calculated(res)) :: frames => step(frames, Frame(arg, parentArg, Calculated(res)) :: acc)

          // stack is processed once
        case Nil => acc.reverse
      }
    }

    processSources(Argument(getSources(), indegree, List[graph.NodeT](), 0))
  }

Manually transforming tree recursion into tail recursion using a stack

2 Answers2