Vector.milone

module rec Std.Vector

open Std.Block
open Std.Own
open Std.Ptr
open Std.StdError

module VectorBase = Std.VectorBase

// For memcpy.
__nativeDecl "#include <string.h>"

/// Owned, mutable and growable array structure.
///
/// - len and capacity represents the number of items.
///
/// Invariants:
///
/// - If capacity isn't zero, ptr isn't null and can be disposed with a particular method.
///     If capacity is zero, ptr is null.
/// - len is less than or equal to capacity.
/// - For all i in {0, ..., len - 1}, ptr[i] must be valid as 'T type.
type Vector<'T> = private Vector of ptr: Own<nativeptr<'T>> * len: int * capacity: int

module Vector =
  /// Gets an empty vector.
  let empty () : Vector<'T> = Vector(Own.acquire Ptr.nullPtr, 0, 0)

  /// Allocates an empty vector with the specified capacity.
  let alloc (capacity: int) : Vector<'T> =
    if capacity > 0 then
      let ptr = VectorBase.alloc capacity sizeof<'T>

      Vector(Ptr.cast ptr, 0, capacity)
    else
      empty ()

  /// Disposes a vector.
  let dispose (v: Vector<'T>) : unit =
    let (Vector(ptr, _, _)) = v
    VectorBase.unsafeFree (Ptr.cast ptr)

  /// Gets the length.
  let length (v: Vector<'T>) : int * Vector<'T> =
    let (Vector(ptr, length, capacity)) = v
    length, Vector(ptr, length, capacity)

  /// Gets an item at the position. Out of bound causes crash.
  let forceGet (i: int) (v: Vector<'T>) : 'T * Vector<'T> =
    let (Vector(ptr, length, capacity)) = v
    assert (uint i < uint length)

    let ptr = Own.release ptr
    let item = Ptr.read (Ptr.cast ptr: InPtr<'T>).[i]
    item, Vector(Own.acquire ptr, length, capacity)

  /// Sets an item at the position. Out of bound causes crash.
  let forceSet (i: int) (item: 'T) (v: Vector<'T>) : Vector<'T> =
    let (Vector(ptr, length, capacity)) = v
    assert (uint i < uint length)

    let ptr = Own.release ptr
    Ptr.write ptr.[i] item
    Vector(Own.acquire ptr, length, capacity)

  /// Appends an item to the end of vector.
  let push (item: 'T) (v: Vector<'T>) : Vector<'T> =
    let (Vector(ptr, length, capacity)) = v

    let ptr, capacity =
      let itemPtr: VoidInPtr = __nativeExpr ("&({0})", item)
      VectorBase.unsafePush itemPtr (Ptr.cast ptr) length capacity sizeof<'T>

    Vector(Ptr.cast ptr, length + 1, capacity)

  /// Removes the last item from the vector unless empty.
  ///
  /// Remark: Allocated memory doesn't get freed.
  let pop (v: Vector<'T>) : 'T option * Vector<'T> =
    let (Vector(ptr, length, capacity)) = v

    if length <> 0 then
      let ptr = Own.release ptr
      let item = Ptr.read (Ptr.cast ptr: InPtr<'T>).[length - 1]
      Some item, Vector(Own.acquire ptr, length - 1, capacity)
    else
      None, Vector(ptr, length, capacity)

  /// Appends all items in a list to the vector.
  let extendFromList (xs: 'T list) (v: Vector<'T>) : Vector<'T> =
    match xs with
    | [] -> v

    | x :: xs ->
      let (Vector(ptr, length, capacity)) = v

      let addition = 1 + List.length xs

      let ptr, capacity =
        VectorBase.unsafeReserve addition (Ptr.cast ptr) length capacity sizeof<'T>

      let ptr: nativeptr<'T> = Ptr.cast (Own.release ptr)

      let rec extendFromListLoop i x xs =
        Ptr.write ptr.[i] x

        match xs with
        | [] -> ()
        | x :: xs -> extendFromListLoop (i + 1) x xs

      extendFromListLoop length x xs
      Vector(Own.acquire ptr, length + addition, capacity)

  /// Creates a vector from a list.
  let ofList (xs: 'T list) : Vector<'T> = extendFromList xs (empty ())

  /// Creates a list from a range of the vector.
  let sliceToList (start: int) (endIndex: int) (v: Vector<'T>) : 'T list * Vector<'T> =
    let (Vector(ptr, length, capacity)) = v
    let ptr = Own.release ptr

    assert (uint endIndex <= uint length)
    assert (uint start <= uint endIndex)

    let xs =
      let ptr: InPtr<'T> = Ptr.cast ptr

      let rec toListLoop acc i =
        let acc = Ptr.read ptr.[i] :: acc
        if i = start then acc else toListLoop acc (i - 1)

      if start < endIndex then
        toListLoop [] (endIndex - 1)
      else
        []

    xs, Vector(Own.acquire ptr, length, capacity)

  /// Creates a list from a vector.
  let toList (v: Vector<'T>) : 'T list * Vector<'T> =
    let length, v = length v
    sliceToList 0 length v

  /// Creates a vector with the specified length.
  let replicate (len: int) (item: 'T) : Vector<'T> =
    let v = alloc len

    let rec go i v =
      if i < len then go (i + 1) (push item v) else v

    go 0 v

  /// Folds a vector.
  let fold (folder: 'S -> 'T -> 'S) (init: 'S) (v: Vector<'T>) : 'S * Vector<'T> =
    let n, v = length v

    let rec go i state v =
      if i < n then
        let item, v = forceGet i v
        go (i + 1) (folder state item) v
      else
        state, v

    go 0 init v

  /// Sorts a vector in-place.
  ///
  /// (Currently unstable sort algorithm is used.
  ///  This function in future version will use a stable sort algorithm.)
  let sort (itemComparer: 'T -> 'T -> int) (v: Vector<'T>) : Vector<'T> =
    // ptrComparer pl pr = itemComparer (*pl) (*pr)
    let ptrComparer =
      fun (l: VoidInPtr) (r: VoidInPtr) ->
        itemComparer (Ptr.read (Ptr.cast l: InPtr<'T>)) (Ptr.read (Ptr.cast r: InPtr<'T>))

    let (Vector(ptr, length, capacity)) = v
    let ptr = VectorBase.unsafeSortUnstable ptrComparer (Ptr.cast ptr) length sizeof<'T>
    Vector(Ptr.cast ptr, length, capacity)

  /// Creates a vector by copying from a block.
  let ofBlock (block: Block<'T>) : Vector<'T> =
    let srcPtr, len = BlockExt.unsafeToRawParts block

    if len <> 0 then
      let destPtr = Own.release (VectorBase.alloc len sizeof<'T>)

      __nativeStmt ("memcpy({0}, {1}, {2});", destPtr, srcPtr, unativeint len * unativeint sizeof<'T>)
      Vector(Own.acquire (Ptr.cast destPtr), len, len)
    else
      empty ()

  /// Creates a block by copying from a vector.
  let toBlock (v: Vector<'T>) : Block<'T> * Vector<'T> =
    let (Vector(ptr, len, cap)) = v
    let ptr = Own.release ptr
    let block = BlockExt.unsafeOfRawParts (Ptr.asIn ptr) len
    block, Vector(Own.acquire ptr, len, cap)

/// Vector methods that are unsafe or niche.
module VectorExt =
  /// Unsafely restores a vector instance from components.
  let unsafeOfRawParts (ptr: Own<nativeptr<'T>>) (len: int) (capacity: int) : Vector<'T> =
    assert (0 <= len && len <= capacity)

    let ptr = Own.release ptr
    assert (ptr <> Ptr.nullPtr || capacity = 0)
    let ptr = Own.acquire ptr

    Vector(ptr, len, capacity)

  /// Unwraps a vector into components.
  /// Use unsafeOfRawParts to restore.
  let toRawParts (v: Vector<'T>) =
    let (Vector(ptr, len, capacity)) = v
    ptr, len, capacity

  /// Creates a vector with the specified length and capacity. Contents are zero'ed.
  ///
  /// Prefer to use `alloc` to allocate and `push` contents to it.
  ///
  /// Item type must be one of: intNN, uintNN, float, unit, bool, char, obj, voidptr, nativeptr, InPtr.
  let unsafeZeroCreate (length: int) : Vector<'T> =
    assert (length >= 0)
    let ptr = VectorBase.alloc length 1
    Vector((Ptr.cast ptr: Own<nativeptr<'T>>), length, length)

  /// Unsafely exposes the internal buffer of the vector.
  ///
  /// This function exists to build another type on the top of Vector efficiently.
  let unsafeDup (v: Vector<'T>) : nativeptr<'T> * int * int * Vector<'T> =
    let (Vector(ptr, length, capacity)) = v
    let ptr = Own.release ptr
    ptr, length, capacity, Vector(Own.acquire ptr, length, capacity)

  let unsafeSetLength (length: int) (v: Vector<'T>) : Vector<'T> =
    let (Vector(ptr, _, capacity)) = v
    assert (uint length <= uint capacity)
    Vector(ptr, length, capacity)

  /// Gets the capacity.
  let capacity (v: Vector<'T>) : int * Vector<'T> =
    let (Vector(ptr, length, capacity)) = v
    capacity, Vector(ptr, length, capacity)

  /// Reserves more memory for a vector.
  ///
  /// - If the vector already have plenty memory, it's noop.
  /// - Result vector should satisfy `length + addition` is less than or equal to `capacity`.
  ///   Note that new capacity may be greater than required size
  ///   to ensure the number of reallocation is small (O(log N)).
  let reserve (addition: int) (v: Vector<'T>) : Vector<'T> =
    assert (addition >= 0)

    let (Vector(ptr, length, capacity)) = v

    let ptr, capacity =
      VectorBase.unsafeReserve addition (Ptr.cast ptr) length capacity sizeof<'T>

    Vector(Ptr.cast ptr, length, capacity)

  /// Reserves more memory for a vector.
  ///
  /// Prefer to use `reserve` instead of this in most of cases.
  ///
  /// - If the vector already have plenty memory, it's noop.
  /// - Result vector should satisfy `length + addition = capacity`.
  let reserveExact (addition: int) (v: Vector<'T>) : Vector<'T> =
    assert (addition >= 0)

    let (Vector(ptr, length, capacity)) = v

    let ptr, capacity =
      VectorBase.unsafeReserveExact addition (Ptr.cast ptr) length capacity sizeof<'T>

    Vector(Ptr.cast ptr, length, capacity)

// -----------------------------------------------
// VectorWriter
// -----------------------------------------------

/// Wrapper of a vector so that you can write data to it.
///
/// - `destPtr`:
///     Pointer to the **end** of the buffer, where the next data is written.
///     This pointer will advance whenever written.
/// - `destCapacity`:
///     Size of allocation past the `destPtr`.
///     That is, `destPtr.[0..destPtr + destCapacity - 1]` is the range that can be written to.
/// - `offsetLen`:
///     Distance between the start of underlying buffer and the current `destPtr`.
///     In other words, `destPtr.[offsetLen]` is the start of underlying buffer.
type VectorWriter<'T> = private VectorWriter of destPtr: Own<OutPtr<'T>> * destCapacity: int * offsetLen: int

// See StringBuffer for basic usage.

module VectorWriter =
  /// Wraps a vector into a writer.
  ///
  /// - `addition`: A predicted size that would be written in addition. (Can be inexact, or just 0.)
  let start (addition: int) (v: Vector<'T>) : VectorWriter<'T> =
    assert (addition >= 0)

    let v = VectorExt.reserve addition v
    let ptr, len, capacity = VectorExt.toRawParts v

    let destPtr =
      let ptr = Own.release ptr
      let destPtr = Ptr.select ptr.[len]
      Own.acquire (Ptr.cast destPtr: OutPtr<'T>)

    let destCapacity = capacity - len
    VectorWriter(destPtr, destCapacity, len)

  /// Unwraps the underlying vector out of the writer.
  let finish (w: VectorWriter<'T>) : Vector<'T> =
    let (VectorWriter(destPtr, destCapacity, offsetLen)) = w

    let ptr: Own<nativeptr<'T>> =
      let ptr = Own.release destPtr
      let ptr = Ptr.select ptr.[-offsetLen]
      Own.acquire (Ptr.cast ptr)

    let capacity = offsetLen + destCapacity
    VectorExt.unsafeOfRawParts ptr offsetLen capacity

  /// Writes an item.
  let push (value: 'T) (w: VectorWriter<'T>) =
    let (VectorWriter(destPtr, destCapacity, offsetLen)) = w
    assert (destCapacity >= 1)

    let destPtr =
      let destPtr = Own.release destPtr
      __nativeStmt ("memcpy({0}, {1}, 1);", destPtr, &&value)
      Own.acquire (Ptr.select destPtr.[1])

    let destCapacity = destCapacity - 1
    let offsetLen = offsetLen + 1
    VectorWriter(destPtr, destCapacity, offsetLen)

  /// Writes an non-empty block of items.
  let copy (srcPtr: InPtr<'T>) (srcLen: int) (w: VectorWriter<'T>) =
    assert (srcLen >= 1)

    let (VectorWriter(destPtr, destCapacity, offsetLen)) = w
    assert (srcLen <= destCapacity)

    let destPtr =
      let destPtr = Own.release destPtr
      __nativeStmt ("memcpy({0}, {1}, (size_t){2});", destPtr, srcPtr, srcLen)
      Own.acquire (Ptr.select destPtr.[srcLen])

    let destCapacity = destCapacity - srcLen
    let offsetLen = offsetLen + srcLen
    VectorWriter(destPtr, destCapacity, offsetLen)

  /// Performs a write operation.
  ///
  /// writer: bufPtr -> bufOffset -> (writtenLen, bufPtr)
  let writeWith (writer: Own<OutPtr<'T>> -> int -> int * Own<OutPtr<'T>>) (w: VectorWriter<'T>) =
    let (VectorWriter(destPtr, destCapacity, offsetLen)) = w

    let writtenLen, destPtr = writer destPtr destCapacity
    assert (writtenLen <= destCapacity)

    let destPtr =
      let destPtr = Own.release destPtr
      Own.acquire (Ptr.select destPtr.[writtenLen])

    let destCapacity = destCapacity - writtenLen
    let offsetLen = offsetLen + writtenLen
    VectorWriter(destPtr, destCapacity, offsetLen)