iter.rs

/// std::iter: Effectful, lazy iterators.
/// An implementation for a fake effectful language with ownership and Go-like syntax.
/// A corresponding implementation in an actual language (!!), Effekt, is available here:
/// https://gist.github.com/omentic/58b553920b4c8402b04c2873901291df
///
/// relevant literature:
/// effects as capabilities: https://dl.acm.org/doi/10.1145/3428194
/// handling bidirectional control flow: https://dl.acm.org/doi/10.1145/3428207
/// soundly handling linearity: https://dl.acm.org/doi/10.1145/3632896
/// wasmfx: https://arxiv.org/abs/2308.08347

@[magic]
pub Continuation[T] // note: T is only really valid when lent or mut or Copy...
@[magic]
pub func resume[T](self: mut Continuation[T])
pub func resume[T, U](self: mut Continuation[T], with: U)

// Any computation resulting in T can implicitly convert to a Continuation[T].
// We care about computations that **do not** provide a value for the purpose of
// iteration. This is basically like the () -> () thunk type
pub alias Continuation = Continuation[()]

pub effect Yield[T] {
  yield(T)
}

// Inferred effect: Yield[T]
pub func yield[T](self: T) {
  raise yield(self)
}

// alternatively to `raise yield(self)`:
//   `do yield(self)`
//   `suspend yield(self)`
// not sure of the best syntax.
// i like `raise`, but it might get confused with `throw`.
// i also like `do`, but it might get confused with do-notation.
// i also like `suspend`, but it doesn't really evoke the right thing.

example {
  // `try` semantics
  try continuation.resume() {
    effect(x) => {...}
    another_effect(x) => {...}
    // when this is missing, implicitly inserted as value => value
    value => {...}
  }
}

example {
  // Inferred effect: Yield[Nat]
  func fib() {
    func inner(a: Nat, b: Nat) {
      yield a
      inner(b, a + b)
    }
    inner(0, 1)
  }

  // Inferred effect: Total
  func evenFibs(limit: Int): List[Int] {
    fib.filter(x => x mod 2 == 0).take(limit).collect()
  }

  // Inferred effect: Total
  // Note that use of `try` is PROHIBITED outside of a module.
  // This is just an example so it's fine...
  @[breaks_encapsulation("example")]
  func genFibs(limit: Int): List[Int] {
    let mut count = 0
    let mut fibs = []
    let mut fib = fib.continuify()
    try fib.resume() {
      yield(x) => {
        if count < limit {
          count += 1
          fibs.push(x)
          fib.resume()
        }
    }
    fibs
  }
}


impl Functor {
  // Inferred effects: Yield[U] (+ other effects of self)
  // this is never over-applied b/c the restriction that the continuation yields T
  pub func map[T, U](self: mut Continuation / Yield[T], fn: T -> U) {
    try self.resume() {
      yield(x) => {
        raise yield(fn(x))
        self.resume()
      }
    }
  }
}

/**
 * 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,
 **/

impl {
  // Inferred effect: at least Yield[T]
  pub func iter[T](self: List[T]) {
    if let Cons(head, tail) = self {
      raise yield(head)
      tail.iter()
    }
  }

  pub func next[T](self: mut Continuation / Yield[T]): T? {
    try self.resume() {
      yield(x) => some x
      _ => none
    }
  }

  pub func peek[T](self: mut Continuation / Yield[T]): T? {
    try self.resume() {
      yield(x) => {
        // this might not work wrt. ownership and also efficiency
        self = () => { raise yield(x); self.resume() }
        some x
      }
      _ => none
    }
  }

  // Inferred effect: at least Yield[T]
  pub func map_while[T, U](self: mut Continuation / Yield[T], fn: T -> U?) {
    try self.resume() {
      yield(x) => {
        if let some x = fn(x) {
          raise yield(x)
          self.resume()
        }
      }
    }
  }

  /// The `for/in` syntax is sugar atop this.
  pub func foreach[T](self: mut Continuation / Yield[T], fn: T -> ()) {
    try self.resume() {
      yield(x) => {
        fn(x)
        self.resume()
      }
    }
  }

  // Inferred effect: at least Yield[T]
  pub func inspect[T](self: mut Continuation / Yield[T], fn: T -> ()) {
    self.map(x => fn(x); x)
  }

  // Inferred effect: at least Yield[T]
  pub func filter[T](self: mut Continuation / Yield[T], keep: T -> Bool) {
    try self.resume() {
      yield(x) => {
        // If keeping the value, act like nothing happened.
        // raise the Yield and resume afterwards.
        if keep(x) {
          raise yield(x)
        }
        // Otherwise, do not yield the value
        // and immediately resume the iterator
        self.resume()
      }
    }
  }

  // Inferred effect: at least Yield[U]
  pub func filter_map[T, U](self: mut Continuation / Yield[T], keep: T -> U?) {
    try self.resume() {
      yield(x) => {
        if let some x = keep(x) {
          raise yield(x)
        }
        self.resume()
      }
    }
  }

  // Inferred effect: at least Yield[(Nat, T)]
  pub func enumerate[T](self: mut Continuation / Yield[T]) {
    let mut n = 0
    try self.resume() {
      yield(x) => {
        raise yield((n, x))
        n += 1
        self.resume()
      }
    }
  }

  // Inferred effect: at least Yield[T]
  pub func skip[T](self: mut Continuation / Yield[T], n: Nat) {
    let mut count = 0
    try self.resume() {
      yield(x) => {
        if count >= n {
          raise yield(x)
          self.resume()
        } else {
          count += 1
          self.resume()
        }
      }
    }
  }

  // Inferred effect: at least Yield[T]
  pub func skip_while[T](self: mut Continuation / Yield[T], fn: T -> Bool) {
    let mut flag = false
    try self.resume() {
      yield(x) => {
        if flag {
          raise yield(x)
          self.resume()
        } else {
          flag = fn(x)
          if flag {
            raise yield(x)
          }
          self.resume()
        }
      }
    }
  }

  // Inferred effect: at least Yield[T]
  pub func take[T](self: mut Continuation / Yield[T], n: Nat) {
    let mut count = 0
    try self.resume() {
      yield(x) => {
        if count < n {
          raise yield(x)
          self.resume()
        } else {
          count += 1
          self.resume()
        }
      }
    }
  }

  // Inferred effect: at least Yield[T]
  pub func take_while[T](self: mut Continuation / Yield[T], fn: T -> Bool) {
    let mut flag = true
    try self.resume() {
      yield(x) => {
        if flag {
          flag = fn(x)
          if not flag {
            raise yield(x)
          }
          self.resume()
        } else {
          self.resume()
        }
      }
    }
  }

  pub func find[T](self: mut Continuation / Yield[T], fn: T -> Bool): T? {
    try self.resume() {
      yield(x) => {
        if fn(x) {
          return some x
        }
        self.resume()
      }
    }
    none
  }

  pub func find_map[T, U](self: mut Continuation / Yield[T], fn: T -> U?): U? {
    try self.resume() {
      yield(x) => {
        let res = fn(x)
        match res {
          some _ => return res
          none => self.resume()
        }
      }
    }
    none
  }

  // Inferred effect: at least Yield[T]
  pub func step[T](self: mut Continuation / Yield[T], n: Nat) {
    let mut count = 0
    try self.resume() {
      yield(x) => {
        if count == n {
          count = 0
          raise yield(x)
          self.resume()
        } else {
          count += 1
          self.resume()
        }
      }
    }
  }

  pub func count[T](self: mut Continuation / Yield[T]): Nat {
    let mut res = 0
    try self.resume() {
      yield(x) => {
        res += 1
        self.resume()
      }
    }
    res
  }

  pub func last[T](self: mut Continuation / Yield[T]): T? {
    let mut res = none
    try self.resume() {
      yield(x) => {
        res = some x
        self.resume()
      }
    }
    res
  }

  pub func nth[T](self: mut Continuation / Yield[T], n: Nat): T? {
    let mut count = 0
    try self.resume() {
      yield(x) => {
        if count == n {
          return some(x)
        } else {
          count += 1
          self.resume()
        }
      }
    }
    // If the iterator runs out, return none.
    none
  }

  pub func position[T](self: mut Continuation / Yield[T], pred: T -> Bool): Nat? {
    let mut count = 0
    try self.resume() {
      yield(x) => {
        if pred(x) {
          return some(count)
        } else {
          count += 1
          self.resume()
        }
      }
    }
    none
  }

  // Inferred effect: at least Yield[(T, U)]
  // Note: This is INCREDIBLY DIFFICULT to express reasonably...
  pub func zip[T, U](self: mut Continuation / Yield[T], other: mut Continuation / Yield[U]) {
    try self.resume() {
      // 1. Upon the Yield catch, `self` is modified to be the rest of the computation.
      yield(x) => {
        // 4. Thus, the modified `other` is called to resume again.
        try other.resume() {
          // 2. Upon the Yield handler, `other` is modified to be the rest of the computation.
          Yield[U](y) => {
            raise yield((x, y))
            // 3. Thus, resuming the modified `self` puts us back in the handlers.
            self.resume()
          }
        }
      }
    }
  }

  pub func unzip[T, U](self: mut Continuation / Yield[(T, U)]): (List[T], List[U]) {
    let mut (a, b) = ([], [])
    try self.resume() {
      yield((x, y)) => {
        a.push(x)
        b.push(y)
        self.resume()
      }
    }
    (a, b)
  }

  pub func fold[T, U](self: mut Continuation / Yield[T], init: U, fn: (T, U) -> U): U {
    let mut acc = init
    try self.resume() {
      yield(x) => {
        acc = fn(x, acc)
        self.resume()
      }
    }
    acc
  }

  pub func reduce[T](self: mut Continuation / Yield[T], fn: (T, T) -> T): T? {
    let mut acc = none
    try self.resume() {
      yield(x) => {
        acc = match acc {
          some acc => some fn(x, acc)
          none => some x
        }
        self.resume()
      }
    }
    acc
  }

  pub func all[T](self: mut Continuation / Yield[T], fn: A -> Bool): Bool {
    try self.resume() {
      yield(x) => {
        if fn(x) {
          self.resume()
        } else {
          return false
        }
      }
    }
    true
  }

  pub func any[T](self: mut Continuation / Yield[T], fn: A -> Bool): Bool {
    try self.resume() {
      yield(x) => {
        if fn(x) {
          return true
        } else {
          self.resume()
        }
      }
    }
    false
  }

  pub func partition[T](self: mut Continuation / Yield[T], fn: T -> Bool): (List[T], List[T]) {
    let mut a = []
    let mut b = []
    try self.resume() {
      yield(x) => {
        if fn(x) {
          a.push(x)
        } else {
          b.push(x)
        }
        self.resume()
      }
    }
    (a, b)
  }

  pub func collect[T](self: mut Continuation / Yield[T], n: Int): List[T] {
    let mut res: List[T] = []
    let mut count = 0
    try self.resume() {
      yield(x) => {
        if count < n {
          res.push(x)
          count += 1
          self.resume()
        }
      }
    }
    res
  }
}