plum

module std

import std/Option

import std/Buffer

import std/Number

import std/Bool

import std/Str

# A node stores the data in a list and contains pointers to the previous and next sibling nodes

enum Node[T] =

  | Node(value: T, prev: Option[Node[T]], next: Option[Node[T]])

# A list is a data structure describing a contiguous section of an array stored separately from the slice variable itself.

# It contains the pointers to the start and end nodes (head, tail) and maintains the size as well

enum List[T: ToStr](ToStr) =

  | List(head: Option[Node[T]], tail: Option[Node[T]], size: Int)

  # `List[T](1, 2, 3)` (or, for an empty list, `List[T]()` — explicit

  # generics required either way, since nothing else in a call this shape

  # pins down `T`) constructs a list via this variadic `init` — see

  # `plum-checker`/`plum-wasm-codegen`'s `ClassCall`/`FnCall` handling, which

  # desugars both call shapes to `List.init(...)`.

  fun init(values: ...T) -> List[T] =

    l := List(head: None, tail: None, size: 0)

    for v := range values

      l.add(v)

    return l

  # gets the element at i'th index of the list

  fun get(self, i: Int) -> Option[T] =

    current := self.head

    index := 0

    while current != None

      match current

        Some(Node(value, _, next)) =>

          if index == i

            return Some(value)

          current = next

          index = index + 1

        None =>

          break

    None

  # sets the element at i'th index of the list, returning the old value (or None if i is out of bounds)

  fun set(self, i: Int, v: T) -> Option[T] =

    current := self.head

    index := 0

    while current != None

      match current

        Some(node) =>

          if index == i

            old := node.value

            node.value = v

            return Some(old)

          current = node.next

          index = index + 1

        None =>

          break

    None

  # returns the no of elements in the list

  fun length(self) -> Int =

    self.size

  # Elementwise comparison via `==` — see the caveat on `expectEq` in

  # `libs/std/testing.plum` for what that means when `T` is a class/enum

  # without its own `equals` (reference identity, not structural equality).

  fun equals(self, other: List[T]) -> Bool =

    if self.length() != other.length()

      return False

    for i := range self.length()

      if self.get(i).unwrap() != other.get(i).unwrap()

        return False

    return True

  # adds the specified elements to the end of the list

  fun add(self, values: ...T) =

    for v := range values

      node := Node(value: v, prev: self.tail, next: None)

      match self.tail

        Some(t) =>

          t.next = Some(node)

        None =>

          self.head = Some(node)

      self.tail = Some(node)

      self.size = self.size + 1

  # unlinks node from the list, patching its neighbors (or head/tail) to close the gap

  fun unlink(self, node: Node[T]) =

    match node.prev

      Some(p) =>

        p.next = node.next

      None =>

        self.head = node.next

    match node.next

      Some(n) =>

        n.prev = node.prev

      None =>

        self.tail = node.prev

    self.size = self.size - 1

  # removes the element at i'th index of the list

  fun removeAt(self, i: Int) =

    current := self.head

    index := 0

    while current != None

      match current

        Some(node) =>

          if index == i

            self.unlink(node)

            return

          current = node.next

          index = index + 1

        None =>

          break

  # removes the first element equal to v from list

  fun remove(self, v: T) =

    current := self.head

    while current != None

      match current

        Some(node) =>

          if node.value == v

            self.unlink(node)

            return

          current = node.next

        None =>

          break

  # removes all objects from this list

  fun clear(self) =

    self.head = None

    self.tail = None

    self.size = 0

  # returns the list with the elements relinked in reverse order.

  fun reverse(self) -> List =

    current := self.head

    while current != None

      match current

        Some(node) =>

          next := node.next

          node.next = node.prev

          node.prev = next

          current = next

        None =>

          break

    old_head := self.head

    self.head = self.tail

    self.tail = old_head

    self

  # Inserts `value` immediately before `node` (which must belong to `self`),

  # patching neighbors (or `head`) to make room — the insertion-point half of

  # `sort`'s pointer surgery, split out since `unlink`'s counterpart already

  # gets its own method.

  fun insertBefore(self, node: Node[T], value: T) -> Unit =

    new_node := Node(value: value, prev: node.prev, next: Some(node))

    match node.prev

      Some(p) =>

        p.next = Some(new_node)

      None =>

        self.head = Some(new_node)

    node.prev = Some(new_node)

    self.size = self.size + 1

  # Inserts `value` into `self` (assumed already sorted by `before`) at the

  # position that keeps it sorted — the first position whose existing element

  # `before(value, existing)` holds for, or the end if none does.

  fun insertSorted(self, value: T, before: fn(T, T) -> Bool) -> Unit =

    current := self.head

    while current != None

      match current

        Some(node) =>

          if before(value, node.value)

            self.insertBefore(node, value)

            return

          current = node.next

        None =>

          break

    self.add(value)

  # Returns a NEW list with `self`'s elements sorted, using `before(a, b)` as

  # the "does a belong before b" comparator. Insertion sort (O(n^2)) — simple

  # and adequate for the list sizes a linked list is already suited to.

  fun sort(self, before: fn(T, T) -> Bool) -> List[T] =

    result := List(head: None, tail: None, size: 0)

    current := self.head

    while current != None

      match current

        Some(Node(value, _, next)) =>

          result.insertSorted(value, before)

          current = next

        None =>

          break

    return result

  # Returns the first element equal to `search`, or `None`.

  fun find(self, search: T) -> Option[T] =

    current := self.head

    while current != None

      match current

        Some(Node(value, _, next)) =>

          if value == search

            return Some(value)

          current = next

        None =>

          break

    return None

  fun contains(self, v: T) -> Bool =

    return self.find(v) != None

  # calls f for each elem in the list

  fun each(self, cb: fn(T)) -> Unit =

    current := self.head

    while current != None

      match current

        Some(Node(value, _, next)) =>

          cb(value)

          current = next

        None =>

          break

  # Returns a new list with each element transformed by `cb`.

  #

  # `U` here is a generic param of `map` ITSELF, separate from `List[T]`'s own

  # `T` — resolved per CALL SITE (from `cb`'s own inferred return type), not

  # when `List[T]` itself is specialized, via `resolveMethodOwnGenerics` in

  # `plum-checker/src/monomorphize.rs`.

  fun map(self, cb: fn(T) -> U) -> List[U] =

    result := List(head: None, tail: None, size: 0)

    current := self.head

    while current != None

      match current

        Some(Node(value, _, next)) =>

          result.add(cb(value))

          current = next

        None =>

          break

    return result

  # returns a new list with each element flat-mapped

  fun flatMap(self) =

    todo

  # Returns a new list with only the elements `predicate` holds for.

  fun retain(self, predicate: fn(T) -> Bool) -> List[T] =

    result := List(head: None, tail: None, size: 0)

    current := self.head

    while current != None

      match current

        Some(Node(value, _, next)) =>

          if predicate(value)

            result.add(value)

          current = next

        None =>

          break

    return result

  # Returns a new list with the elements `predicate` holds for removed.

  fun reject(self, predicate: fn(T) -> Bool) -> List[T] =

    result := List(head: None, tail: None, size: 0)

    current := self.head

    while current != None

      match current

        Some(Node(value, _, next)) =>

          if !predicate(value)

            result.add(value)

          current = next

        None =>

          break

    return result

  fun any(self, predicate: fn(T) -> Bool) -> Bool =

    current := self.head

    while current != None

      match current

        Some(Node(value, _, next)) =>

          if predicate(value)

            return True

          current = next

        None =>

          break

    return False

  fun every(self, predicate: fn(T) -> Bool) -> Bool =

    current := self.head

    while current != None

      match current

        Some(Node(value, _, next)) =>

          if !predicate(value)

            return False

          current = next

        None =>

          break

    return True

  # Folds `self`'s elements into a single value, starting from `acc` and

  # combining each element in with `cb(accumulatedSoFar, element)`.

  #

  # `U` is `reduce`'s OWN generic param (the accumulator's type), separate

  # from `List[T]`'s own `T` — same per-CALL-SITE resolution as `map`'s `U`.

  fun reduce(self, acc: U, cb: fn(U, T) -> U) -> U =

    result := acc

    current := self.head

    while current != None

      match current

        Some(Node(value, _, next)) =>

          result = cb(result, value)

          current = next

        None =>

          break

    return result

  # returns the first element in the list

  fun first(self) -> Option[T] =

    match self.head

      Some(Node(value, _, _)) => Some(value)

      None => None

  # returns the last element in the list

  fun last(self) -> Option[T] =

    match self.tail

      Some(Node(value, _, _)) => Some(value)

      None => None

  # Returns a new list of the elements at index `[start, end)`, clamped into

  # range (a `start`/`end` outside `[0, length]` is clamped, not an error).

  fun sublist(self, start: Int, end: Int) -> List[T] =

    result := List(head: None, tail: None, size: 0)

    s := start < 0 ? 0 : start

    e := end > self.size ? self.size : end

    current := self.head

    i := 0

    while current != None && i < e

      match current

        Some(Node(value, _, next)) =>

          if i >= s

            result.add(value)

          current = next

          i = i + 1

        None =>

          break

    return result

  # The first `n` elements.

  fun take(self, n: Int) -> List[T] =

    return self.sublist(0, n)

  # Every element AFTER the first `n`.

  fun skip(self, n: Int) -> List[T] =

    return self.sublist(n, self.size)

  fun drop(self, n: Int) -> List[T] =

    return self.skip(n)

  # Returns one uniformly random element, or `None` if `self` is empty.

  fun sample(self) -> Option[T] =

    if self.size == 0

      return None

    return self.get(randomInt(self.size))

  # Returns a new list with `self`'s elements in random order (Fisher-Yates).

  fun shuffle(self) -> List[T] =

    result := self.sublist(0, self.size)

    i := result.size - 1

    while i > 0

      j := randomInt(i + 1)

      vi := result.get(i).unwrap()

      vj := result.get(j).unwrap()

      result.set(i, vj)

      result.set(j, vi)

      i = i - 1

    return result

  # `chunk`/`partition` both return `List[List[T]]` — the SAME generic class

  # (`List`) nested inside itself as a method's own return type. This used to

  # be an unconditional stack-overflow landmine for the WHOLE COMPILER (see

  # `plum-checker/src/monomorphize.rs`'s `isSelfNested`/`methodMentionsTemplate`

  # doc comments for the full fix) — that danger is now FIXED (2026-09-02).

  #

  # A second, more contained bug surfaced once past that: two SEPARATE bare

  # `List(head: None, tail: None, size: 0)` constructions in the SAME method

  # body — one for the OUTER `result: List[List[T]]` (matching this

  # method's own declared return type) and one for an INNER per-chunk

  # `piece: List[T]` — are ambiguous to `resolveClassInstantiation`'s

  # `current_return_type` fallback: since `List[T]` and `List[List[T]]` both

  # have exactly ONE generic argument, the fallback (which only ever

  # compares ARITY, not identity) can't tell "this bare construction IS the

  # function's own return value" apart from "this is some OTHER, unrelated

  # nested construction of the same template that happens to need the same

  # arity" — so `piece` was silently resolved to the SAME type as `result`

  # (`List$List$Int` instead of `List$Int`), corrupting its own `add` calls.

  # Worked around by building `piece` via `self.sublist(...)` instead of a

  # second bare construction — `sublist` is an ordinary method already

  # unambiguously resolved through the RECEIVER's own (already-correct)

  # binding, no `current_return_type` guessing involved at all.

  fun chunk(self, size: Int) -> List[List[T]] =

    result := List(head: None, tail: None, size: 0)

    if size <= 0

      return result

    i := 0

    n := self.length()

    while i < n

      result.add(self.sublist(i, i + size))

      i = i + size

    return result

  # Splits into two lists: elements `predicate` holds for, then the rest —

  # in that order. Returns them as a 2-element `List[List[T]]` (`[matched,

  # unmatched]`) since this language has no tuple type. Deliberately built

  # from `self.retain(predicate)`/`self.reject(predicate)` rather than two

  # local `matched`/`unmatched := List(head: None, ...)` bare constructions

  # — see `chunk`'s doc comment above for exactly why that would be

  # ambiguous (`matched`/`unmatched` need `List[T]`, but a bare

  # construction here would resolve via `current_return_type` fallback

  # against THIS method's own `List[List[T]]` return instead, same arity

  # coincidence). `retain`/`reject` are each already unambiguous (single

  # bare construction per method body, matching their own single-level

  # return type).

  fun partition(self, predicate: fn(T) -> Bool) -> List[List[T]] =

    result := List(head: None, tail: None, size: 0)

    result.add(self.retain(predicate))

    result.add(self.reject(predicate))

    return result

  # Grouping by an arbitrary key type would need its own generic param

  # (`groupBy<K>(keyFn: fn(T) -> K) -> Map[K, List[T]]`) resolved from a

  # closure's return type — the same method-level-generics gap documented on

  # `List.map` above, PLUS `Map[K, List[T]]` would make `list.plum` depend on

  # `map.plum`, which already depends on `list.plum`. Not implemented.

  fun groupBy(self) =

    todo

  # Note: NOT written with `each` + a closure-captured "is this the first

  # element" flag — a closure's captured variables are snapshotted at

  # CREATION time (see `libs/std/list.plum`'s own `add`/`each` usage

  # elsewhere), so an assignment to a captured var INSIDE the closure body

  # does not persist across that closure's own repeated invocations from one

  # `each` call. A plain `while`/`match` traversal (same idiom as `unlink`)

  # sidesteps that entirely.

  fun join(self, sep: Str = ",") -> Str =

    res := Buffer()

    current := self.head

    while current != None

      match current

        Some(Node(value, _, next)) =>

          res.write(value.toStr())

          if next != None

            res.write(sep)

          current = next

        None =>

          break

    res.toStr()

  # `List[T: ToStr](ToStr)` declared implementing `ToStr`

  # from the start but never actually defined `toStr` — harmless for a

  # PLAIN `List[Int]`/`List[Str]`/etc (nothing needs to call `.toStr()` on

  # the LIST ITSELF, just on its elements, which `join` above already

  # does), but a genuine gap the moment a `List` is used AS another

  # `List`'s own ELEMENT type (`List[List[T]]`) — the OUTER list's eagerly-

  # compiled `join`/`toStr` calls `.toStr()` on each element, and an inner

  # `List[T]` element had no such method to call, failing to compile with

  # "unknown method 'List$T.toStr'" even for a program that never actually

  # calls `.toStr()`/`.join()` on the outer list (per the "eager compile

  # every method for every specialization" architecture noted elsewhere in

  # this codebase's history).

  fun toStr(self) -> Str =

    return "[" + self.join(",") + "]"

# ---- regression tests ----

fun optSumForListTest(o: Option[Int]) -> Int =

  match o

    Some(v) => v

    None => -1000

fun exerciseList() -> Int =

  l = List(1, 2, 3, 4, 5)

  a = l.length()

  b = optSumForListTest(l.get(0))

  c = optSumForListTest(l.get(4))

  oldVal = optSumForListTest(l.set(2, 30))

  d = optSumForListTest(l.get(2))

  l.removeAt(0)

  e = l.length()

  f = optSumForListTest(l.get(0))

  l.remove(30)

  g = l.length()

  l.reverse()

  h = optSumForListTest(l.get(0))

  l.clear()

  i = l.length()

  a + b + c + oldVal + d + e + f + g + h + i

fun removeAllNodesOneAtATime() -> Int =

  l = List(1, 2, 3, 4, 5, 6, 7, 8, 9, 10)

  i = 0

  while i < 10

    l.removeAt(0)

    i = i + 1

  afterLoopLength = l.length()

  isEmpty = optSumForListTest(l.get(0))

  l.add(42)

  afterReAdd = optSumForListTest(l.get(0))

  afterLoopLength + isEmpty + afterReAdd

fun exerciseListInit() -> Int =

  l := List[Int]()

  before := l.length()

  l.add(7)

  before + l.length() + optSumForListTest(l.get(0))

fun exerciseListInitWithValues() -> Int =

  l := List(10, 20, 30)

  l.length() + optSumForListTest(l.get(0)) + optSumForListTest(l.get(2))

test "list add set remove at remove clear reverse all work correctly"

  # add(1,2,3,4,5): a=length=5, b=get(0)=1, c=get(4)=5

  # set(2,30): oldVal=3, d=get(2)=30 -> list [1,2,30,4,5]

  # removeAt(0): e=length=4, f=get(0)=2 -> list [2,30,4,5]

  # remove(30): g=length=3 -> list [2,4,5]

  # reverse(): h=get(0)=5 -> list [5,4,2]

  # clear(): i=length=0

  # 5+1+5+3+30+4+2+3+5+0 = 58

  assert exerciseList() == 58

test "removing every node in a loop leaves an empty correctly functioning list"

  # afterLoopLength=0, isEmpty(get(0) on empty list)=-1000, afterReAdd=42

  # 0 + -1000 + 42 = -958

  assert removeAllNodesOneAtATime() == -958

test "List[Int]() with no args runs init() and produces a working empty list"

  assert exerciseListInit() == 8

test "List(values) with positional args runs init(values) directly, infers T from args"

  assert exerciseListInitWithValues() == 43

test "list map and reduce work correctly"

  l = List(1, 2, 3, 4)

  doubled = l.map(|v| v * 2)

  total = doubled.reduce(0, |acc, v| acc + v)

  assert total == 20

test "list equals compares elementwise"

  a = List(1, 2, 3)

  b = List(1, 2, 3)

  c = List(1, 2)

  assert a.equals(b)

  assert !a.equals(c)

test "list sort orders elements correctly"

  l = List(3, 1, 4, 1, 5)

  sorted = l.sort(|a, b| a < b)

  assert optSumForListTest(sorted.get(0)) == 1

  assert optSumForListTest(sorted.get(4)) == 5

test "list join renders elements correctly with separator"

  l = List(1, 2, 3)

  assert l.join(",") == "1,2,3"

test "list toStr formats a plain list and a nested list of lists"

  l := List(1, 2, 3)

  assert l.toStr() == "[1,2,3]"

  chunks := l.chunk(2)

  assert chunks.toStr() == "[[1,2],[3]]"

test "chunk splits into sublists of at most size elements each"

  l := List(1, 2, 3, 4, 5)

  chunks := l.chunk(2)

  assert chunks.length() == 3

  assert chunks.get(0).unwrap().join(",") == "1,2"

  assert chunks.get(1).unwrap().join(",") == "3,4"

  assert chunks.get(2).unwrap().join(",") == "5"

  assert l.chunk(0).length() == 0

test "partition splits into elements matching a predicate, then the rest"

  l := List(1, 2, 3, 4, 5, 6)

  parts := l.partition(|v| v % 2 == 0)

  assert parts.length() == 2

  assert parts.get(0).unwrap().join(",") == "2,4,6"

  assert parts.get(1).unwrap().join(",") == "1,3,5"