Breadth-first labeling via queue/zipper in Scala

BFN.scala

import org.junit.Test
import org.junit.Assert._
import scala.util.chaining._
import scala.collection.mutable.ListBuffer
import scala.language.implicitConversions

// Breadth-first labeling -- http://okmij.org/ftp/Algorithms/BFN.html
// Breadth-First Numbering: An Algorithm in Pictures -- https://okasaki.blogspot.com/2008/07/breadth-first-numbering-algorithm-in.html
// The Under-Appreciated Unfold -- https://www.cs.ox.ac.uk/people/jeremy.gibbons/publications/unfold.ps.gz
class BFNTest {
    def [A, B](a: A) |> (f: (A) => B): B = a.pipe(f)

    def map[A, B](f: A => B)(s: List[A]): List[B] = {
        s match {
            case Nil => Nil
            case h::t => 
                val hh: B = f(h) 
                val tt = map(f)(t)
                hh :: tt
        }
    }

    def map_accum[A, B, Z](f: Z => A => (Z, B))(z: Z)(s: List[A]): (Z, List[B]) = {
        s match {
            case Nil => (z, Nil)
            case h::t => 
                val (newState,  hh): (Z, B) = f(z)(h) 
                val (newState2, tt) = map_accum(f)(newState)(t)
                (newState2, hh :: tt)
        }
    }

    def ilabeling[A](z: Int)(x: A): (Int, (Int, A)) = (z + 1, (z, x))

    def list_lab[A](i: Int, l: List[A]): List[(Int, A)] = {
        map_accum[A, (Int, A), Int](ilabeling)(i)(l)._2
    }

    @Test def list_lab_test() = {
        val res = list_lab(101, List(2,3,5,7,11))
        assert(res == List((101,2), (102,3), (103,5), (104,7), (105,11)))
    }

    enum Tree[A] {
        case Leaf(a: A)
        case Node(a: A, left: Tree[A], right: Tree[A])
    }
    import Tree._

    def tree_map[A, B](f: A => B)(t1: Tree[A]): Tree[B] = t1 match {
        case Leaf(a) => Leaf(f(a))
        case Node(a, l, r) => Node(f(a), tree_map(f)(l), tree_map(f)(r))
    }

    // processing the root from a list of trees each time (traversal 1#)
    def bfi_forest1[A](f: A => Unit)(l: List[Tree[A]]): Unit = l match {
        case scala.Nil => () 
        case forest => {
            def iterF(t: Tree[A]): Unit = t match { 
                case Leaf(a) => f(a)
                case Node(a, _, _) => f(a)
            }
            forest.foreach(iterF)

            // prunes the tree level by level (traversal 2#)
            def nodeList(t: Tree[A]) = t match {
                case Leaf(_) => scala.Nil
                case Node(_, l, r) => (l :: r :: scala.Nil)
            }

            // all trees are passed by recursing on the level each time
            forest.map(nodeList).flatten |> bfi_forest1(f) 
        }
    }

    def bfi1[A](f: A => Unit)(l: Tree[A]): Unit = {
        bfi_forest1(f)(List(l))
    }

    // notice the manual fusion that occurs:
    //  the processing with `f` is happening during visitation 
    //  each level is processed sequentially from left to right added on the list in reverse order (LIFO (r))
    def bfi_forest2[A](f: A => Unit)(lp: (List[Tree[A]], List[Tree[A]])): Unit = lp match {
        case (scala.Nil, scala.Nil) => () 
        case (scala.Nil, next) => bfi_forest2(f)((next.reverse, List()))
        case (Leaf(x) :: t, next) => f(x); bfi_forest2(f)((t, next)) // <-- why not rebuilt the tree here?
        case (Node(x, l, r) :: t, next) => f(x); bfi_forest2(f)((t,  r :: l :: next))
    }

    def bfi2[A](f: A => Unit)(l: Tree[A]): Unit = {
        bfi_forest2(f)(List(l), List())
    }

    def bf_fold_left[Z,A](f: Z => A => Z)(z: Z)(t: Tree[A]): Z = {
        var zr = z
        bfi_forest2((x) => f(zr)(x))(List(t), List())
        zr
    }

    def bfa_naive[A, B](f: A => B)(t: Tree[A]): Tree[B] = {
        case class Ref[B](var s: Option[B])
        
        // Map the tree with refcells in a DFS manner
        val tree_s: Tree[(Ref[B], A)] = tree_map[A, (Ref[B], A)](x => (Ref[B](None), x))(t)
        
        // Traverse in BFS manner + mutation
        bfi2[(Ref[B], A)]((rf, x) => rf.s = Some(f(x)))(tree_s)

        // Rebuilding in a DFS manner
        tree_map[(Ref[B], A), B](x => x._1.s.get)(tree_s)
    }

    // revisiting the "fused" version above
    // def bfi_forest2[A](f: A => Unit)(lp: (List[Tree[A]], List[Tree[A]])): Unit 
    type ForestZipper[A] = (List[Tree[A]], List[Tree[A]])

    // The idea is that we will construct the list of trees simultaneously and we will end up with a tree of the same shape
    def bfi_forest3[A, B](f: A => B)(lp: ForestZipper[A]): ForestZipper[B] = lp match {
        case (scala.Nil, scala.Nil) => 
            (scala.Nil, scala.Nil)
        case (scala.Nil, next) => 
            val (next_r, scala.Nil) = bfi_forest3(f)((next.reverse, List()))
            (scala.Nil, next_r.reverse)
        case (Leaf(x) :: t, next) => 
            val v = f(x)
            val (t_, next_) = bfi_forest3(f)((t, next)) 
            (Leaf(v) :: t_, next_)
        case (Node(x, l, r) :: t, next) => 
            val v = f(x)
            val (t_, r_ :: l_ :: next_) = bfi_forest3(f)((t, r :: l :: next)) 
            (Node(v, l_, r_) :: t_, next_)
    }

    def bfa3[A, B](f: A => B)(t: Tree[A]): Tree[B] = {
        bfi_forest3(f)(List(t), List())._1(0)
    }

    val t = Node (1, Node (2, Leaf (3), 
                Node (4, Leaf (5), 
                        Leaf (6))),
        Node (12, Node (14, Leaf (15), 
                            Leaf (16)), 
                Leaf (13)))

    @Test def bfi_forest1_test() = {
        val res = ListBuffer[Int]()
  
        List(t) |> bfi_forest1[Int]((a: Int) => res += a)

        assert(res == List(1, 2, 12, 3, 4, 14, 13, 5, 6, 15, 16))
    }

    @Test def bf_fold_left_test() = {
        val res = bf_fold_left[ListBuffer[Int], Int](z => a => z += a)(ListBuffer[Int]())(t)
           
        assert(res == List(1, 2, 12, 3, 4, 14, 13, 5, 6, 15, 16))
    }

    @Test def bfa3_test() = {
        val res = bfa3[Int, Int](x => x + 1)(t) 
        |> bf_fold_left[ListBuffer[Int], Int](z => a => z += a)(ListBuffer[Int]())

        assert(res == List(2, 3, 13, 4, 5, 15, 14, 6, 7, 16, 17))
    }
}