Array.val

// Questions:
// - Should `pops` be a requirement of `Collection`?
// - Should we have a `CopyableCollection` trait to require `copies`?
// - Should `ConsumingGenerator` be copyable when `Element: Copyable`?

// # Helpers

// The header of an array.
typealias ArrayHeader = (count: Int, capacity: Int)

// The header of the empty array.
let empty_array_header: ArrayHeader = (count: 0, capacity: 0)

// A dynamically allocated array buffer.
type ArrayBuffer<Element>: Sinkable {

  // The base address of the buffer.
  //
  // - Invariant: The pointer is aligned to `ArrayHeader`.
  var address: MutableRawPointer

  // Creates an empty buffer.
  public init() {
    address = MutableRawPointer(pointer(to: empty_array_header))
  }

  // Creates a new buffer capable of storing `capacity` elements.
  public init(capacity: sink Int) {
    Memory.allocate(
      bytes: Self.payload_offset() + capacity * MemoryLayout<Element>.stride,
      alignment: MemoryLayout<ArrayHeader>.alignment)
    header.initialize(to: (count: 0, capacity: capacity))
  }

  public sink fun deinit() {
    // Deallocate the buffer's memory if necessary.
    if header != MutablePointer(pointer(to: empty_array_header)) { address.deallocate() }
  }

  // A pointer to the buffer's header.
  public property header: MutablePointer<ArrayHeader> {
    address.assume_bound<ArrayHeader>()
  }

  // A pointer to the buffer's payload.
  //
  // The projected pointer is a valid if and only if `header[offset: 0].count > 0`. Otherwise,
  // deferencing it causes undefined behavior.
  public property payload: MutablePointer<Element> {
    address.advanced(by: Self.payload_offset()).assume_bound<Element>()
  }

  // Returns the offset in bytes from the address of the buffer to 
  static fun payload_offset() -> Int {
    let x = MemoryLayout<Element>.alignment - 1
    return (MemoryLayout<ArrayHeader>.size + x) & ~x
  }

}

// # Primary declaration

// An array of homogeneous elements.
public type Array<Element> {

  // The buffer containing the array's elements.
  var source: ArrayBuffer<Element>

  // Creates an empty array.
  public init() {
    source = ArrayBuffer()
  }

  // Creates an empty array with a minimum capacity.
  //
  // - Requires: `minimum_capacity >= 0`
  public init(minimum_capacity: sink Int) {
    source = ArrayBuffer(capacity: minimum_capacity)
  }

  public sink fun deinit() {
    for let i in 0 ..< source.header[unsafe: 0].count {
      source.payload[unsafe: i].deinit()
    }
  }

  // Returns whether the array is empty.
  public fun is_empty() -> Bool { count() == 0 }

  // Returns the number of elements in the array.
  public fun count() -> Int { source.header[unsafe: 0].count.copy() }

  // Returns the capacity of the array, i.e., the number of elements the array can contain without
  // allocating new memory.
  public fun capacity() -> Int { source.header[unsafe: 0].capacity.copy() }

  // Reserves enough space to store `minimum_capacity` elements without allocating new memory.
  public fun reserve(minimum_capacity: Int) {
    inout {
      if capacity() >= minimum_capacity { return }

      // Allocate a new buffer.
      var new_source = ArrayBuffer(capacity: minimum_capacity.copy())

      // Transfer the elements to the new buffer.
      new_source.header[unsafe: 0].count = count()
      new_source.payload.move_initialize(from: source.payload, count: count())
      swap(&source, &new_source)
    }
  }

  // Appends an element, constructing it in-place.
  public fun emplace_back<Args...>(
    _ constructor: inout (set Element, sink ...Args),
    _ arguments: Args
  ) {
    inout {
      if count() + 1 < capacity() { reserve(minimum_capacity: max(1, capacity() * 2)) }
      constructor(&source.payload[unsafe: count()], ...arguments)
      source.header[unsafe: 0].count += 1
    }
  }

  // Returns the array source.
  sink fun take_source() -> ArrayBuffer<Element> { source }

}

// # Conditional API

extension Array where Element: Sinkable {

  // A generator that consumes an array to produce its elements.
  public type ConsumingGenerator: Generator, Sinkable {

    // Implementation note:
    // 
    // The generator moves the elements in `source` from the front. Elements are never moved,
    // instead the generator uses an index to point at the next element to output.

    // The buffer from which elements are being consumed.
    var source: ArrayBuffer<Element>

    // The index of the next element to output.
    //
    // - Invariant: `next_element_index <= source.header[unsafe: 0].count`
    var next_element_index = 0

    // Creates a consuming generator that produces the elements in `source`.
    public init(_ source: sink Array) {
      source = source.take_source()
    }

    public sink fun deinit() {
      // Destroy the remaining elements, if any.
      for let i in next_element_index ..< source.header[unsafe: 0].count {
        source.payload[unsafe: i].deinit()
      }
    }

    public fun next() -> Optional<Element> {
      inout {
        if next_element_index == source.header[unsafe: 0].count { return nil }
        defer { next_element_index += 1 }
        return source.payload[unsafe: next_element_index]
      }
    }

  }

  // Appends `new_element` to the array.
  public fun append(_ new_element: sink Element) {
    inout {
      if count() + 1 < capacity() { reserve(minimum_capacity: max(1, capacity() * 2)) }
      source.payload[unsafe: count()] = new_element
      source.header[unsafe: 0].count += 1
    }
  }

  // Appends the elements produced by `source` to the array.
  public fun append<G: Generator where G.Element == Element>(elements_of source: inout G) {
    inout { while let e = source.next() { append(e) } }
  }

  // Removes and returns the last element of the array, if any.
  public fun pop_last() -> Optional<Element> {
    inout {
      if is_empty() { return nil }

      let i = count() - 1
      defer { source.header[unsafe: 0].count = i }
      return source.payload[unsafe: i]
    }
  }

}

extension Array where Element: Copyable {

  // A generator that produces the elements of an array by copying them.
  public type CopyingGenerator: Generator {

    // The buffer from which elements are being copied.
    var source: let ArrayBuffer<Element>

    // The index of the next element to output.
    //
    // - Invariant: `next_element_index <= source.header[unsafe: 0].count`
    var next_element_index = 0

    public fun next() -> Optional<Element> {
      inout {
        if next_element_index == source.header[unsafe: 0].count { return nil }
        defer { next_element_index += 1 }
        return source.payload[unsafe: next_element_index].copy()
      }
    }

  }

  // Projects a generator that produces copies of the array's elements
  public property copies: var CopyingGenerator {
    let { CopyingGenerator(array: self) }
  }

  public fun append(_ new_element: sink Element) {
    let { copy().append(new_element) }
  }

}

// # Conformances

public conformance Array: Sinkable where Element: Sinkable {}

public conformance Array: Copyable where Element: Copyable {

  public fun copy() -> Self {
    if is_empty() { return Array() }
    var other = Array(minimum_capacity: count())
    
    other.source.header[unsafe: 0].count = count()
    for let i in 0 ..< count() {
      other.source.payload[unsafe: i] = source.payload[unsafe: i].copy()
    }
    return other
  }

}

public conformance Array: Equatable where Element: Equatable {

  public infix fun == (_ other: Self) -> Bool {
    if count() != other.count() { return false }
    for let i in 0 ..< count() {
      if source.payload[unsafe: i] != other.source.payload[unsafe: i] { return false }
    }
    return true
  }

}

public conformance Array: Hashable where Element: Hashable {

  public fun hash(into hasher: inout Hasher) {
    for let i in 0 ..< count() { source.payload[unsafe: i].hash(into: &hasher) }
  }

}

public conformance Array: Collection {

  public fun pops() -> ConsumingGenerator {
    sink { ConsumingGenerator(array: self) }
  }

}

public conformance Array: Indexable {

  public typealias ::Indexable.Index = Int

  public typealias ::Indexable.Element = Element

  public fun start_index() -> Int { 0 }

  public fun end_index() -> Int { count() }

  public fun index(after i: Int) -> Int { i + 1 }

  public subscript (_ i: Int): Element {
    let {
      precondition((0 <= i) && (i < count()))
      yield source.payload[unsafe: i]
    }
    inout {
      precondition((0 <= i) && (i < count()))
      yield &source.payload[unsafe: i]
    }
    set {
      precondition((0 <= i) && (i < count()))
      source.payload[unsafe: i] = new_value
    }
    sink {
      precondition((0 <= i) && (i < count()))
      for let j in 0 ..< count() where j != i {
        source.payload[unsafe: j].deinit()
      }
      return source.payload[unsafe: i]
    }
  }

}