plum

module std

import std/Result
import std/Str
import std/Bool
import std/Number

# Option[T] represents a value that may or may not be present — Plum's
# counterpart to Rust's `Option`/Go's "zero value or ok bool" idiom.
enum Option[T] =
  | Some(T)
  | None

  fun isSome(self) -> Bool =
    match self
      Some(_) => True
      None => False

  fun isNone(self) -> Bool =
    match self
      Some(_) => False
      None => True

  # Returns the wrapped value, or traps if `self` is `None`.
  fun unwrap(self) -> T =
    match self
      Some(v) =>
        return v
      None =>
        todo

  # Returns the wrapped value, or traps if `self` is `None`. `msg` documents
  # the expectation at the call site (matching Rust's `Option::expect`) but
  # isn't surfaced anywhere at runtime — there's no host-independent "print
  # this and trap" primitive to build that on.
  fun expect(self, msg: Str) -> T =
    match self
      Some(v) =>
        return v
      None =>
        todo

  fun unwrapOr(self, default: T) -> T =
    match self
      Some(v) =>
        return v
      None =>
        return default

  # Like `unwrapOr`, but the fallback is computed lazily (only when `self` is
  # `None`) via `cb` — useful when producing the default is expensive.
  fun unwrapOrElse(self, cb: fn() -> T) -> T =
    match self
      Some(v) =>
        return v
      None =>
        return cb()

  # Keeps `self` only if it's `Some` AND `predicate` holds for its value;
  # otherwise returns `None`.
  fun filter(self, predicate: fn(T) -> Bool) -> Option[T] =
    match self
      Some(v) =>
        if predicate(v)
          return Some(v)
        return None
      None =>
        return None

  # Transforms the wrapped value with `cb`, leaving `None` as `None`. `U` is
  # `map`'s OWN generic param (the transformed value's type), separate from
  # `Option[T]`'s own `T` — resolved per CALL SITE, not when `Option[T]`
  # itself is specialized (see `resolveMethodOwnGenerics`).
  fun map(self, cb: fn(T) -> U) -> Option[U] =
    match self
      Some(v) =>
        return Some(cb(v))
      None =>
        return None

  # Like `map`, but `cb` itself returns an `Option[U]` (rather than a bare
  # `U`) and the result isn't re-wrapped — useful for chaining together
  # several fallible steps without nesting (`Option[Option[U]]`). `U` here is
  # `andThen`'s OWN generic param, nested inside `cb`'s declared return type
  # `Option[U]` rather than being `cb`'s bare return type — resolved from the
  # closure's actual inferred return type via `bindGenericArg` in
  # `plum-checker/src/monomorphize.rs`'s `resolveMethodOwnGenerics`.
  fun andThen(self, cb: fn(T) -> Option[U]) -> Option[U] =
    match self
      Some(v) =>
        return cb(v)
      None =>
        return None

  # Converts to a `Result`, using `err` as the failure value if `self` is
  # `None`. `E` is `okOr`'s OWN generic param, inferred directly from `err`'s
  # own value (same mechanism as `List.reduce`'s accumulator param).
  fun okOr(self, err: E) -> Result[T, E] =
    match self
      Some(v) =>
        return Ok(v)
      None =>
        return Err(err)

  # No `toStr(self) -> Str` here — every method on a generic type gets
  # compiled for EVERY concrete specialization that type is ever used at
  # ANYWHERE in the whole program, whether or not that particular method is
  # actually called for that specialization (there's no lazy/on-demand
  # method compilation). `libs/std/list.plum`'s `List[T]` internally uses
  # `Option[Node[T]]` for its own `head`/`tail` fields — and `Node` has no
  # `toStr` — so an `Option.toStr` calling `.toStr()` on the wrapped value
  # would fail to compile for THAT specialization even though nothing ever
  # actually calls `.toStr()` on an `Option[Node[Int]]`. Use `toOptionStr`
  # below (an explicit stringifier callback, not an implicit `T: ToStr`
  # bound) wherever printing an `Option` is needed.
  fun toOptionStr(self, valueToStr: fn(T) -> Str) -> Str =
    match self
      Some(v) =>
        return "Some({valueToStr(v)})"
      None =>
        return "None"

fun makeNoneIntForOptionTest() -> Option[Int] =
  return None

test "filter keeps a Some value only when the predicate holds"
  assert Some(4).filter(|v| v > 2).isSome()
  assert Some(1).filter(|v| v > 2).isNone()
  assert makeNoneIntForOptionTest().filter(|v| v > 2).isNone()

test "andThen chains fallible steps without double-wrapping"
  half := |n| n % 2 == 0
  a := Some(8).andThen(|n| Some(n / 2)).andThen(|n| Some(n / 2))
  assert a.isSome()
  assert a.unwrap() == 2
  assert half(a.unwrap())
  b := makeNoneIntForOptionTest().andThen(|n| Some(n / 2))
  assert b.isNone()

test "okOr converts to a Result using the given error on None"
  assert Some(5).okOr("missing").isOk()
  assert Some(5).okOr("missing").unwrap() == 5
  assert makeNoneIntForOptionTest().okOr("missing").isErr()
  assert makeNoneIntForOptionTest().okOr("missing").unwrapErr() == "missing"