实现树的不可变外部迭代器 [blog/iterator/immut-exiter-tree]

Blog 1. 前置知识 [blog/iterator/preorder-traversal]

二叉树的定义

priv enum Tree[A] {
  Nil
  Node(A, Tree[A], Tree[A])
}

先序遍历一棵树

fn Tree::preorder[A](self : Tree[A], f : (A) -> Unit) -> Unit {
  fn dfs(root) {
    match root {
      Nil => ()
      Node(x, left, right) => {
        f(x)
        dfs(left)
        dfs(right)
      }
    }
  }

  dfs(self)
}

我们把 Tree::preorder 改写为手动控制栈的过程 Tree::preorder_stack

fn Tree::preorder_stack[A](self : Tree[A], f : (A) -> Unit) -> Unit {
  let stack = Array::new(capacity=4096)
  stack.push(self)
  while not(stack.is_empty()) {
    let root = stack.unsafe_pop()
    match root {
      Nil => ()
      Node(x, left, right) => {
        f(x)
        stack.push(right) // 先进后出, 先被压入后处理
        stack.push(left) // 先处理左节点
      }
    }
  }
}

下面的部分是用来测试系统栈版本和手动控制栈版本的一致性。

fn Tree::from_n(n : Int) -> Tree[Int] {
  let mut i = 0
  fn dfs() {
    if i < n {
      let x = i
      i += 1
      let res = Node(x, dfs(), dfs())
      res
    } else {
      Nil
    }
  }

  dfs()
}

fn test_preorder(root : Tree[Int]) -> Unit! {
  let b1 = StringBuilder::new()
  let b2 = StringBuilder::new()
  root.preorder(fn(x) { b1.write_object(x) })
  root.preorder_stack(fn(x) { b2.write_object(x) })
  assert_eq!(b1.to_string(), b2.to_string())
}

test "preorder/preorder_stack" {
  let t1 = Node(
    1,
    Node(2, Nil, Nil),
    Node(3, Node(4, Nil, Nil), Node(5, Nil, Nil)),
  )
  let mut sum = 0
  t1.preorder(fn(x) { sum += x })
  inspect!(sum, content="15")
  let t2 = Tree::from_n(15)
  test_preorder!(t2)
}

2. 正文 [blog/iterator/immut-exiter-tree-content]

这里使用不可变单向链表作为栈,每次返回一个 ImmutExIter[A] 需要保存整个迭代上下文,

  • 可变外部迭代器并不会保存,只能运行一次

  • 不可变外部迭代器可以运行多次,但是需要额外保存迭代上下文

typealias Stack[A] = @immut/list.T[A]

fn ImmutExIter::from_tree[A](root : Tree[A]) -> ImmutExIter[A] {

  fn aux(stack : Stack[_]) -> (A,ImmutExIter[A])? {
    match stack {
      Nil => None
      Cons(root,rest_stack) => { // pop root from stack
        match root {
          Nil => None
          Node(x,left,right) => {
            let stack = Stack::Cons(left,Stack::Cons(right,rest_stack))
            // push right into stack
            // push left into stack
            Some((x,fn () { aux(stack)}))
          }
        }
      }
    }
  }
  fn () {
    aux(@immut/list.singleton(root))
  }
}

这里测试一下,不可变外部迭代器和 preorder 的一致性

test "ImmutExIter::from_tree" {
  let t2 = Tree::from_n(15)
  let iter = ImmutExIter::from_tree(t2)
  let b1 = StringBuilder::new()
  loop iter.uncons() {
    None => ()
    Some((x,rest)) => {
      b1.write_string("\{x},")
      continue rest.uncons()
    }
  }
  inspect!(b1, content="0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,")
  let b2 = StringBuilder::new()
  t2.preorder(fn(x) { b2.write_string("\{x},") })
  inspect!(b2, content="0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,")
}

接下来我们来实现 zipWith,这里和普通的List[A]zipWith 别无二致

fn ImmutExIter::zipWith[A,B,C](self : ImmutExIter[A], other : ImmutExIter[B], f : (A,B) -> C) -> ImmutExIter[C] {
  let xs = self
  let ys = other

  match (xs.uncons(),ys.uncons()) {
    (Some((x,xs)),Some((y,ys))) => {
      fn () { Some((f(x,y),  ImmutExIter::zipWith(xs,ys,f)))}
    }
    (_,_) => {
      fn () { None }
    }
  }
}


test {
  let xs = ["apple", "orange", "watermetlon"]
  let ys = Tree::from_n(5)

  let xs = ImmutExIter::from_array(xs)
  let ys = ImmutExIter::from_tree(ys)

  let zs = xs.zipWith(ys,fn (x,y) { (x,y)})

  let b1 = StringBuilder::new()
  loop zs.uncons() {
    None => ()
    Some((x,rest)) => {
      b1.write_string("\{x},")
      continue rest.uncons()
    }
  }
  inspect!(b1, content=
    #|("apple", 0),("orange", 1),("watermetlon", 2),
  )
}