https://github.com/moonbit-community/community-blog/pull/17

external_iterator.mbt

///|
test "split internal iterator" {
  let xs = [1, 2, 3, 4, 5]
  let iter = xs.iter().drop(4)
  inspect!(xs.iter().last(), content="Some(5)")
  inspect!(iter.head(), content="Some(5)")
}

///|
priv type ImmutExIter[A] () -> (A, ImmutExIter[A])?

///|
fn ImmutExIter::from_array[A](xs : Array[A]) -> ImmutExIter[A] {
  fn aux(i) -> (A, ImmutExIter[A])? {
    if i < xs.length() {
      Some((xs[i], fn() { aux(i + 1) }))
    } else {
      None
    }
  }

  fn() { aux(0) }
}

///|
fn ImmutExIter::uncons[A](self : ImmutExIter[A]) -> (A, ImmutExIter[A])? {
  (self._)()
}

///|
test "ImmutExIter::from_array" {
  let xs = [1, 2, 3, 4, 5]
  let iter = ImmutExIter::from_array(xs)
  let buf = StringBuilder::new()
  loop iter.uncons() {
    None => ()
    Some((x, xs)) => {
      buf.write_object(x)
      continue xs.uncons()
    }
  }
  inspect!(buf, content="12345")
}

///|
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)
}

///|
fn Tree::preorder_stack[A : Show](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)
}

///|
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) =>
        match root {
          Nil => None
          Node(x, left, right) => {
            let stack = Stack::Cons(left, Stack::Cons(right, rest_stack))
            Some((x, fn() { aux(stack) }))
          }
        }
    }
  }

  fn() { aux(@immut/list.singleton(root)) }
}

///|
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,")
}

///|
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),
    ,
  )
}

///|
priv type ExIter[A] () -> A?

///|
fn ExIter::from_array[A](xs : Array[A]) -> ExIter[A] {
  let mut i = 0
  fn() {
    if i < xs.length() {
      let res = xs[i]
      i = i + 1
      Some(res)
    } else {
      None
    }
  }
}

///|
fn ExIter::next[A](self : ExIter[A]) -> A? {
  (self._)()
}

///|
test {
  let xs = [1, 2, 3, 4, 5]
  let iter = ExIter::from_array(xs)
  let mut sum = 0
  loop iter.next() {
    None => ()
    Some(x) => {
      sum += x
      continue iter.next()
    }
  }
  inspect!(sum, content="15")
}

///|
fn ExIter::from_tree[A](root : Tree[A]) -> ExIter[A] {
  let stack = Array::new(capacity=4096)
  stack.push(root)
  fn() {
    if not(stack.is_empty()) {
      let root = stack.unsafe_pop()
      match root {
        Nil => None
        Node(x, left, right) => {
          stack.push(right)
          stack.push(left)
          Some(x)
        }
      }
    } else {
      None
    }
  }
}

///|
fn ExIter::zipWith[A, B, C](
  self : ExIter[A],
  other : ExIter[B],
  f : (A, B) -> C
) -> ExIter[C] {
  let xs = self
  let ys = other
  fn() {
    match (xs.next(), ys.next()) {
      (Some(x), Some(y)) => Some(f(x, y))
      (_, _) => None
    }
  }
}

///|
test "ExIter::zipWith" {
  let xs = ["apple", "orange", "watermetlon"]
  let ys = Tree::from_n(5)
  let xs = ExIter::from_array(xs)
  let ys = ExIter::from_tree(ys)
  let zs = xs.zipWith(ys, fn(x, y) { (x, y) })
  let b1 = StringBuilder::new()
  loop zs.next() {
    None => ()
    Some(x) => {
      b1.write_string("\{x},")
      continue zs.next()
    }
  }
  inspect!(
    b1,
    content=
      #|("apple", 0),("orange", 1),("watermetlon", 2),
    ,
  )
}

///|
test "ExIter::from_tree" {
  let t2 = Tree::from_n(15)
  let iter = ExIter::from_tree(t2)
  let b1 = StringBuilder::new()
  loop iter.next() {
    None => ()
    Some(x) => {
      b1.write_string("\{x},")
      continue iter.next()
    }
  }
  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,")
}