iter.scala

/// Effectful, lazy iterators. A reimplementation of Rust's std::iter.
/// Note: This is written in Effekt and should be named iter.effekt.
/// It is only named iter.scala for syntax highlighting.

/// relevant literature surrounding effects, implemented in effekt:
/// effects as capabilities: https://dl.acm.org/doi/10.1145/3428194
/// effects, capabilities, and boxes: https://dl.acm.org/doi/abs/10.1145/3527320
/// from capabilities to regions: https://dl.acm.org/doi/10.1145/3622831
/// handling bidirectional control flow: https://dl.acm.org/doi/10.1145/3428207
///
/// relevant literature, but not to effekt (yet):
/// soundly handling linearity: https://dl.acm.org/doi/10.1145/3632896
/// wasmfx: https://arxiv.org/abs/2308.08347

interface Yield[T] {
  def yield(value: T): Unit
}

def yield[T](value: T) = {
  do yield(value)
}

/**
 * https://doc.rust-lang.org/std/iter/trait.Iterator.html
 * Iterator methods from Rust fall into three categories here:
 * 1. implemented:
 *   next, peek,
 *   fold, reduce, all, any, enumerate, partition, collect,
 *   map, map_while, foreach, inspect, step, nth, last, count, position, unzip, zip,
 *   filter, filter_map, find, find_map, skip, skip_while, take, take_while,
 * 2. unnecessary:
 *   try_, chain, flatten, flat_map, fuse,
 * 3. unimplemented:
 *   rev, rposition, cycle, scan, size_hint,
 *   sum, product, by_ref, cloned, copied, peekable
 *   eq, ne, ge, gt, le, lt, cmp, partial_cmp, is_sorted, is_sorted_by, is_sorted_by_key,
 *   max, max_by, max_by_key, min, min_by, min_by_key
 *   advance_by, array_chunks, collect_into, map_windows, next_chunk,
 *   cmp_by, eq_by, partial_cmp_by, is_partitioned, partition_in_place,
 *   intersperse, intersperse_with,
 **/

// Note: The Yield[T] annotation (and Unit) are both needed here...
def iter[T](self: List[T]): Unit / Yield[T] = {
  self match {
    case Cons(head, tail) => {
      yield(head)
      tail.iter()
    }
    case Nil() => ()
  }
}

def map[T, U]() {self: => Unit / Yield[T]} {fn: T => U} = {
  try self() with Yield[T] {
    def yield(x) = {
      yield(fn(x))
      resume(())
    }
  }
}

// Note: This function is incorrect. Effekt does not seem to have mutable parameters.
// def next[T]() {mut self: => Unit / Yield[T]}: Option[T] = {
//   try {self(); None()} with Yield[T] {
//     def yield(x) = {
//       self = fun() { resume(()) }
//       Some(x)
//     }
//   }
// }

// Note: This function is also incorrect, for the same reason as `.next()`.
// def peek[T]() {mut self: => Unit / Yield[T]}: Option[T] = {
//   try {self(); None()} with Yield[T] {
//     def yield(x) = {
//       self = fun() { yield(x); resume(()) }
//       Some(x)
//     }
//   }
// }

def map_while[T, U]() {self: => Unit / Yield[T]} {fn: T => Option[U]} = {
  try self() with Yield[T] {
    def yield(x) = {
      fn(x) match {
        case Some(x) => {
          yield(x)
          resume(())
        }
        case None() => ()
      }
    }
  }
}

def foreach[T]() {self: => Unit / Yield[T]} {fn: T => Unit} = {
  try self() with Yield[T] {
    def yield(x) = {
      fn(x)
      resume(())
    }
  }
}

def inspect[T]() {self: => Unit / Yield[T]} {fn: T => Unit} = {
  // Note: I don't know how to get anonymous functions to work.
  // `fun (x: T) { fn(x); x }` is incorrect...
  def inspect(x: T) = { fn(x); x }
  map { self } { inspect }
}

def filter[T]() {self: => Unit / Yield[T]} {keep: T => Bool} = {
  try self() with Yield[T] {
    def yield(x) = {
      if (keep(x)) {
        yield(x)
      }
      resume(())
    }
  }
}

def filter_map[T, U]() {self: => Unit / Yield[T]} {keep: T => Option[U]} = {
  try self() with Yield[T] {
    def yield(x) = {
      keep(x) match {
        case Some(x) => yield(x)
        case None() => ()
      }
      resume(())
    }
  }
}

def enumerate[T]() {self: => Unit / Yield[T]} = {
  var n = 0
  try self() with Yield[T] {
    def yield(x) = {
      yield((n, x))
      n = n + 1
      resume(())
    }
  }
}

def skip[T](n: Int) {self: => Unit / Yield[T]} = {
  var count = 0
  try self() with Yield[T] {
    def yield(x) = {
      if (count >= n) {
        yield(x)
        resume(())
      } else {
        count = count + 1
        resume(())
      }
    }
  }
}

def skip_while[T]() {self: => Unit / Yield[T]} {fn: T => Bool} = {
  var flag = false
  try self() with Yield[T] {
    def yield(x) = {
      if (flag) {
        yield(x)
        resume(())
      } else {
        flag = fn(x)
        if (flag) {
          yield(x)
        }
        resume(())
      }
    }
  }
}

def take[T](n: Int) {self: => Unit / Yield[T]} = {
  var count = 0
  try self() with Yield[T] {
    def yield(x) = {
      if (count < n) {
        count = count + 1
        yield(x)
        resume(())
      }
    }
  }
}

def take_while[T]() {self: => Unit / Yield[T]} {fn: T => Bool} = {
  var flag = true
  try self() with Yield[T] {
    def yield(x) = {
      if (flag) {
        flag = fn(x)
        if (not (flag)) {
          yield(x)
          resume(())
        }
      }
    }
  }
}

def find[T]() {self: => Unit / Yield[T]} {fn: T => Bool}: Option[T] = {
  var res = None()
  try self() with Yield[T] {
    def yield(x) = {
      fn(x) match {
        case true => res = Some(x)
        case false => resume(())
      }
    }
  }
  res
}

def find_map[T, U]() {self: => Unit / Yield[T]} {fn: T => Option[U]}: Option[U] = {
  var res = None()
  try self() with Yield[T] {
    def yield(x) = {
      fn(x) match {
        case Some(value) => res = Some(value)
        case None() => resume(())
      }
    }
  }
  res
}

def step[T](n: Int) {self: => Unit / Yield[T]} = {
  var count = 0
  try self() with Yield[T] {
    def yield(x) = {
      if (count == n) {
        count = 0
        yield(x)
      } else {
        count = count + 1
      }
      resume(())
    }
  }
}

def count[T]() {self: => Unit / Yield[T]}: Int = {
  var res = 0
  try self() with Yield[T] {
    def yield(x) = {
      res = res + 1
      resume(())
    }
  }
  res
}

def last[T]() {self: => Unit / Yield[T]}: Option[T] = {
  var res = None()
  try self() with Yield[T] {
    def yield(x) = {
      res = Some(x)
      resume(())
    }
  }
  res
}

def nth[T](n: Int) {self: => Unit / Yield[T]}: Option[T] = {
  var count = 0
  var res = None()
  try self() with Yield[T] {
    def yield(x) = {
      if (count == n) {
        res = Some(x)
      } else {
        count = count + 1
        resume(())
      }
    }
  }
  res
}

def position[T]() {self: => Unit / Yield[T]} {pred: T => Bool}: Option[Int] = {
  var count = 0
  var res = None()
  try self() with Yield[T] {
    def yield(x) = {
      if (pred(x)) {
        res = Some(count)
      } else {
        count = count + 1
        resume(())
      }
    }
  }
  res
}

// Note: This is INCREDIBLY DIFFICULT to express reasonably...
def zip[T, U]() {self: => Unit / Yield[T]} {other: => Unit / Yield[U]} = {
  // 1. We create a dummy coroutine so we can mutate it within its definition.
  var coroutine = fun() { None() }
  // 2. We wrap the second computation in a coroutine that:
  //   a) catches its yielded value
  //   b) returns it as Some(value)
  //   c) updates the coroutine wrapper to "freeze" the resume(()) call
  coroutine = fun() {
    try {other(); None() } with Yield[U] {
      def yield(y) = {
        coroutine = fun() { resume(()) }
        Some(y)
      }
    }
  }
  // 3. We then call the first computation and catch its yielded value normally.
  try self() with Yield[T] {
    def yield(x) = {
      // 4. After catching the first yielded value, we run the `coroutine()`
      //   function, which will return its first yielded value and mutate itself.
      //   We then yield that value with our first caught value as a tuple.
      // 5. Future invocations of `coroutine()` will use the mutated `coroutine()`
      //   wrapping `resume()`.
      coroutine() match {
        case Some(y) => yield((x, y))
        case None() => ()
      }
      resume(())
    }
  }
}

def unzip[T, U]() {self: => Unit / Yield[(T, U)]}: (List[T], List[U]) = {
  var a = []
  var b = []
  try self() with Yield[Tuple2[T, U]] {
    def yield(xy) = xy match {
      case (x, y) => {
        a = Cons(x, a)
        b = Cons(y, b)
        resume(())
      }
    }
  }
  (a, b)
}

def fold[T, U](init: U) {self: => Unit / Yield[T]} {fn: (T, U) => U}: U = {
  var acc = init
  try self() with Yield[T] {
    def yield(x) = {
      acc = fn(x, acc)
      resume(())
    }
  }
  acc
}

def reduce[T]() {self: => Unit / Yield[T]} {fn: (T, T) => T}: Option[T] = {
  var acc = None()
  try self() with Yield[T] {
    def yield(x) = {
      acc = acc match {
        case Some(acc) => Some(fn(x, acc))
        case None() => Some(x)
      }
      resume(())
    }
  }
  acc
}

def all[T]() {self: => Unit / Yield[T]} {fn: T => Bool}: Bool = {
  var res = true
  try self() with Yield[T] {
    def yield(x) = {
      if (fn(x)) {
        resume(())
      } else {
        res = false
      }
    }
  }
  res
}

def any[T]() {self: => Unit / Yield[T]} {fn: T => Bool}: Bool = {
  var res = false
  try self() with Yield[T] {
    def yield(x) = {
      if (fn(x)) {
        res = true
      } else {
        resume(())
      }
    }
  }
  res
}

def partition[T]() {self: => Unit / Yield[T]} {fn: T => Bool}: (List[T], List[T]) = {
  var a = []
  var b = []
  try self() with Yield[T] {
    def yield(x) = {
      if (fn(x)) {
        a = Cons(x, a)
      } else {
        b = Cons(x, b)
      }
      resume(())
    }
  }
  (a, b)
}

def collect[T]() {self: => Unit / Yield[T]}: List[T] = {
  var res: List[T] = []
  try self() with Yield[T] {
    def yield(x) = {
      res = Cons(x, res)
      resume(())
    }
  }
  res
}

def fib() = {
  def inner(a: Int, b: Int): Unit / Yield[Int] = {
    yield(a)
    inner(b, a + b)
  }
  inner(0, 1)
}

def genFibs(limit: Int): List[Int] = {
  var count = 0
  var fibs = []
  try fib() with Yield[Int] {
    def yield(x) = {
      if (count < limit) {
        count = count + 1
        fibs = Cons(x, fibs)
        resume(())
      }
    }
  }
  fibs
}

def even(n: Int): Bool = if (n <= 0) true else odd(n - 1)
def odd(n: Int): Bool = if (n <= 0) false else even(n - 1)

def combine(self: (Int, String)): String = show(first(self)) ++ " " ++ second(self)

def main() = {
  println("hello!")
  println(collect[Int] { take[Int](5) { fib } })
  println(collect[Int] { take[Int](5) { filter { fib } { even }}})
  println(collect[String] { map { zip[Int, String] {[1, 2, 3].iter()} {["one", "two", "three"].iter() }} { combine }})
}