plum

# Arrays: typed native arrays and arrays of structs.
import std/Int
import std/Float
import std/Array
import std/Grid
import std/Str
import std/Slice

# ==== typed arrays ====

# Typed arrays: unboxed native `Array[Int]` and `Array[Float]` lowering to
# WebAssembly primitive arrays, along with flat 2D `Grid[T]`.

# Tile-grid style data: Array[Int] / Array[Float] / Array[Bool] are native
# unboxed wasm arrays (i64 / f64 / i32 elements), not boxed anyref arrays.
fun sumInts(n: Int) -> Int =
  a := Array.filled(n, 3)
  total := 0
  for i := range n
    total = total + a.get(i)
  return total

fun sumFloats(n: Int) -> Float =
  a := Array.filled(n, 0.5)
  total := 0.0
  for i := range n
    total = total + a.get(i)
  return total

fun countTrue(n: Int) -> Int =
  a := Array.filled(n, false)
  a.set(1, true)
  a.set(3, true)
  c := 0
  for i := range n
    if a.get(i)
      c = c + 1
  return c

# ==== struct arrays ====

# Arrays of structs are stored field by field (struct-of-arrays): `Array[Vec2]`
# is one object holding an `f64` array for `x` and one for `y`, so `get`/`set`
# allocate nothing per element. `Slice`, `Grid` and friends get this for free.

struct Vec2(x: Float, y: Float)

struct Rect(pos: Vec2, size: Vec2) =

  fun area(self) -> Float =
    self.size.x * self.size.y

struct Tile(id: Int, solid: Bool, weight: Float)

struct World(points: Slice[Vec2], name: Str)

struct One(n: Int)

struct Flags(a: Bool, b: Bool, c: Int)

struct L3(r: Rect, k: Int)

struct Holder(arr: Array[Vec2])

struct Named(label: Str, v: Vec2)

fun bump(h: Holder) =
  h.arr.set(0, Vec2(x: 77.0, y: 88.0))

fun first[T](a: Array[T]) -> T =
  a.get(0)

fun mk() -> Array[Vec2] =
  Array[Vec2]().push(Vec2(x: 1.0, y: 2.0))

struct Px(r: Byte, g: Byte, b: Byte, a: Byte)

fun mkPx(v: Int) -> Px =
  Px(r: Byte(v), g: Byte(v + 1), b: Byte(v + 2), a: Byte(255))

# ==== typed arrays tests ====

test "array_int_native: filled arrays hold their values"
  assert sumInts(10) == 30
  assert sumFloats(4) == 2.0
  assert countTrue(5) == 2

# ==== struct arrays tests ====

test "List[Vec2] push/get/set round-trips across growth"
  l := Slice[Vec2]()
  for i := range 100
    l.push(Vec2(x: Float(i), y: Float(i) * 2.0))
  assert l.length() == 100
  v := l.get(37)!
  assert v.x == 37.0
  assert v.y == 74.0
  _ := l.set(37, Vec2(x: v.x + 1.0, y: v.y))!
  assert l.get(37)!.x == 38.0
  assert l.get(37)!.y == 74.0
  assert l.get(36)!.x == 36.0
  assert l.get(99)!.y == 198.0

test "nested struct (Rect of Vec2s) in an Array"
  a := Array[Rect]()
  a = a.push(Rect(pos: Vec2(x: 1.0, y: 2.0), size: Vec2(x: 3.0, y: 4.0)))
  a = a.push(Rect(pos: Vec2(x: 5.0, y: 6.0), size: Vec2(x: 7.0, y: 8.0)))
  assert a.length() == 2
  r := a.get(1)
  assert r.pos.x == 5.0
  assert r.pos.y == 6.0
  assert r.size.x == 7.0
  assert r.size.y == 8.0
  a.set(0, Rect(pos: r.size, size: r.pos))
  assert a.get(0).pos.y == 8.0
  assert a.get(0).size.x == 5.0
  assert a.get(0).area() == 30.0

test "struct with Int, Bool and Float leaves"
  a := Array[Tile]()
  a = a.push(Tile(id: 7, solid: true, weight: 1.5))
  a = a.push(Tile(id: 9, solid: false, weight: 2.5))
  assert a.get(0).id == 7
  assert a.get(0).solid
  assert a.get(0).weight == 1.5
  assert a.get(1).id == 9
  assert !a.get(1).solid
  assert a.get(1).weight == 2.5
  b := a.resize(4)
  assert b.length() == 4
  assert b.get(1).id == 9
  assert b.get(3).id == 0
  assert !b.get(3).solid

test "copy semantics: a value read from the array is a copy until set"
  l := Slice[Vec2]()
  l.push(Vec2(x: 1.0, y: 1.0))
  v := l.get(0)!
  v = Vec2(x: 99.0, y: 99.0)
  assert l.get(0)!.x == 1.0
  w := l.get(0)!
  _ := l.set(0, Vec2(x: 5.0, y: 5.0))!
  assert w.x == 1.0
  assert l.get(0)!.x == 5.0
  assert v.x == 99.0

test "for x := range over Array[Vec2]"
  a := Array[Vec2]().push(Vec2(x: 1.0, y: 10.0)).push(Vec2(x: 2.0, y: 20.0)).push(Vec2(x: 3.0, y: 30.0))
  sx := 0.0
  sy := 0.0
  for v := range a
    sx = sx + v.x
    sy = sy + v.y
  assert sx == 6.0
  assert sy == 60.0

test "Array.filled with a struct"
  a := Array.filled(5, Vec2(x: 2.0, y: 3.0))
  assert a.length() == 5
  assert a.get(4).y == 3.0

test "Grid[Tile] get/set"
  g := Grid.filled(3, 2, Tile(id: 1, solid: false, weight: 0.5))
  g.set(2, 1, Tile(id: 42, solid: true, weight: 9.0))
  assert g.get(2, 1).id == 42
  assert g.get(2, 1).solid
  assert g.get(0, 0).id == 1
  assert g.get(1, 1).weight == 0.5

test "a List[Vec2] inside a record field"
  w := World(points: Slice[Vec2](), name: "w")
  w.points.push(Vec2(x: 1.0, y: 2.0))
  w.points.push(Vec2(x: 3.0, y: 4.0))
  assert w.points.length() == 2
  _ := w.points.set(1, Vec2(x: 30.0, y: 40.0))!
  assert w.points.get(1)!.x == 30.0
  assert w.points.get(0)!.y == 2.0

test "probe push growth, shrink, single leaf"
  a := Array[One]()
  for i := range 50
    a = a.push(One(n: i))
  assert a.length() == 50
  assert a.get(49).n == 49
  b := a.resize(10)
  assert b.length() == 10
  assert b.get(9).n == 9
  c := b.resize(0)
  assert c.length() == 0

test "probe Byte/Bool leaves"
  a := Array[Flags]()
  a = a.push(Flags(a: false, b: true, c: 3))
  assert a.get(0).b
  assert a.get(0).c == 3

test "probe three-level nested"
  a := Array[L3]()
  a = a.push(L3(r: Rect(pos: Vec2(x: 1.0, y: 2.0), size: Vec2(x: 3.0, y: 4.0)), k: 9))
  assert a.get(0).r.size.y == 4.0
  assert a.get(0).k == 9
  assert a.get(0).r.area() == 12.0

test "probe array field reference semantics"
  h := Holder(arr: Array[Vec2]().push(Vec2(x: 1.0, y: 1.0)))
  bump(h)
  assert h.arr.get(0).x == 77.0
  alias := h.arr
  alias.set(0, Vec2(x: 5.0, y: 5.0))
  assert h.arr.get(0).x == 5.0

test "probe generic first, equality"
  a := Array[Vec2]().push(Vec2(x: 1.0, y: 2.0)).push(Vec2(x: 1.0, y: 2.0))
  assert first(a).y == 2.0
  eq := a.get(0) == a.get(1)
  assert eq

test "probe option and list"
  n: Option[Array[Vec2]] := Some(Array[Vec2]().push(Vec2(x: 3.0, y: 3.0)))
  assert n!.get(0).x == 3.0
  l := Slice[Array[Vec2]]()
  l.push(Array[Vec2]().push(Vec2(x: 4.0, y: 4.0)))
  assert l.get(0)!.get(0).x == 4.0
  r := mk()
  assert r.get(0).y == 2.0

test "probe break/continue"
  a := Array[Vec2]().push(Vec2(x: 1.0, y: 0.0)).push(Vec2(x: 2.0, y: 0.0)).push(Vec2(x: 3.0, y: 0.0)).push(Vec2(x: 4.0, y: 0.0))
  s := 0.0
  for v := range a
    if v.x == 2.0
      continue
    if v.x == 4.0
      break
    s = s + v.x
  assert s == 4.0

test "probe Str leaf"
  a := Array[Named]()
  a = a.push(Named(label: "hi", v: Vec2(x: 1.0, y: 2.0)))
  assert a.get(0).label == "hi"

test "Byte struct fields: copy, compare, store"
  p := mkPx(10)
  q := p
  assert p == q
  assert p.r == Byte(10)
  assert p.a == Byte(255)
  assert p != mkPx(11)
  a := Array[Px]().push(p).push(mkPx(20))
  assert a.get(1).g == Byte(21)
  assert a.get(0) == q
  l := Slice[Px]()
  for i := range 50
    l.push(mkPx(i))
  assert l.get(37)!.b == Byte(39)
  assert l.get(49)! == mkPx(49)

test "assert == on structs without toStr"
  assert Vec2(x: 1.0, y: 2.0) == Vec2(x: 1.0, y: 2.0)