plum

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

fun each(cb: fn(Int) -> Int) -> Int =
  cb(5)

fun useCapturingClosure() -> Int =
  offset := 100
  cb := |v|
    v + offset
  cb(5)

fun main() -> Int =
  each(|v| v * 3) + useCapturingClosure()

fun useSnapshotClosure() -> Int =
  x = 10
  snapshotCb = |v|
    x + v
  x = 999
  each(snapshotCb)

fun eachF(cb: fn(Float) -> Float) -> Float =
  cb(0.0)

fun useFloatClosure() -> Float =
  offset = 2.5
  floatCb = |v|
    offset + v
  eachF(floatCb)

enum ClosureCat =
  | ClosureCat(age: Int)

fun eachCat(cb: fn(ClosureCat) -> Int) -> Int =
  cb(ClosureCat(age: 7))

fun useClassClosure() -> Int =
  classCb = |c|
    c.age
  eachCat(classCb)

fun identity(value: T) -> T =
  value

fun eachViaGeneric(cb: fn(Int) -> Int) -> Int =
  cb(identity(7))

# `inner` (nested inside `outer`'s body) needs `offset` — a name from
# `useNestedClosure`'s scope, two levels up from `inner` itself, and not
# referenced by `outer` directly — exercising the multi-level capture chain:
# `outer` must itself capture `offset` purely because `inner` needs it.
fun useNestedClosure() -> Int =
  offset = 100
  outer = |v|
    inner = |w|
      w + offset
    inner(v)
  each(outer)

fun doubleForClosureTest(x: Int) -> Int =
  x * 2

fun classifySign(n: Int) -> Str =
  cb = |v| v > 0 ? "pos" : "nonpos"
  cb(n)

test "closures example computes correctly"
  assert main() == 120

# ---- closure regression tests ----

test "non capturing closure passed and called runs correctly"
  assert each(|v| v) == 5

# Regression test: the external scanner used to never emit a dedent for a
# multi-line closure body immediately followed by `)` on the same line as
# the body's last statement, so this shape didn't parse at all before.
test "multi line closure call argument with closing paren on last statement line runs correctly"
  multiLineResult = each(|v|
    bumped = v + 1
    bumped * 2)
  assert multiLineResult == 12

test "capturing closure snapshots value at creation time runs correctly"
  # each calls its callback with 5, so a correct snapshot of x==10 (its value
  # when the closure was created, not 999 — its value when `each(cb)` is
  # actually called) gives 10 + 5 = 15.
  assert useSnapshotClosure() == 15

test "closure assigned then called at float type runs correctly"
  assert useFloatClosure() == 2.5

test "closure assigned then called at class type runs correctly"
  assert useClassClosure() == 7

test "closure passed through already generic higher order function runs correctly"
  assert eachViaGeneric(|v| v * 2) == 14

test "nested closure literal runs correctly"
  assert useNestedClosure() == 105

test "nested closure literal passed directly as call argument runs correctly"
  assert each(|v|
    each(|w| w + v)) == 10

test "named function used as a closure typed value runs correctly"
  assert each(doubleForClosureTest) == 10

test "named function used as a value assigned then called runs correctly"
  f = doubleForClosureTest
  assert f(21) == 42

test "named function used as a value alongside a closure at the same call type runs correctly"
  # Proves the trampoline shares the same call_indirect type as an ordinary
  # closure of the same signature (both must resolve to the same wasm
  # function type, since both flow through the exact same `cb(...)` call site).
  assert each(doubleForClosureTest) + each(|v| v + 1) == 16

test "closure with an inline ternary body runs correctly"
  # A ternary as a closure's own inline body used to parse as an OUTER
  # ternary wrapping the whole closure instead (`(|v| v > 0) ? "pos" :
  # "nonpos"`), a grammar precedence bug — see the closure/PREC.conditional
  # fix in tooling/tree-sitter-plum/grammar.js.
  assert classifySign(5) == "pos"
  assert classifySign(-3) == "nonpos"