shapearray.nim

########## Shape Array
const MAXRANK* = 8

type
  ShapeArray* = object
    data: array[MAXRANK, int]
    len: int

proc newShapeArray*(len: int): ShapeArray {.inline.} =
  result.len = len

converter toShapeArray*(s: varargs[int]): ShapeArray {.inline.} =
  assert s.len <= MAXRANK
  result.len = s.len
  for i in 0..<s.len:
    result.data[i] = s[i]

proc copyFrom*(a: var ShapeArray, s: openarray[int]) {.inline.} =
  assert s.len <= MAXRANK
  a.len = s.len
  for i in 0..<s.len:
    a.data[i] = s[i]

proc copyFrom*(a: var ShapeArray, s: ShapeArray) {.inline.} =
  assert s.len <= MAXRANK
  a.len = s.len
  for i in 0..<s.len:
    a.data[i] = s.data[i]

proc setLen*(a: var ShapeArray, len: int) {.inline.} =
  assert len <= MAXRANK
  a.len = len

proc len*(a: ShapeArray): int =
  a.len

proc low*(a: ShapeArray): int {.inline.} =
  0

proc high*(a: ShapeArray): int {.inline.} =
  a.len-1

proc `[]`*(a: ShapeArray, idx: int): int {.inline.} =
  assert idx >= 0 and idx < MAXRANK
  a.data[idx]

proc `[]`*(a: var ShapeArray, idx: int): var int {.inline.} =
  assert idx >= 0 and idx < MAXRANK
  a.data[idx]

proc `[]=`*(a: var ShapeArray, idx: int, v: int) {.inline.} =
  assert idx >= 0 and idx < MAXRANK
  a.data[idx] = v

proc `[]`*(a: ShapeArray, slice: Slice[int]): ShapeArray {.inline.} =
  if slice.b >= slice.a:
    assert slice.a >= 0 and slice.b < a.len
    result.len = (slice.b - slice.a + 1)
    for i in 0..result.len:
      result[i] = a[slice.a+i]

iterator items*(a: ShapeArray): int {.inline.} =
  for i in 0..<a.len:
    yield a.data[i]

iterator mitems*(a: var ShapeArray): var int {.inline.} =
  for i in 0..<a.len:
    yield a.data[i]

iterator pairs*(a: ShapeArray): (int, int) {.inline.} =
  for i in 0..<a.len:
    yield (i,a.data[i])

proc `@`*(a: ShapeArray): seq[int] {.inline.} =
  result = newSeq[int](a.len)
  for i in 0..<a.len:
    result[i] = a.data[i]

proc `$`*(a: ShapeArray): string {.inline.} =
  result = "["
  var firstElement = true
  for value in items(a):
    if not firstElement: result.add(", ")
    result.add($value)
    firstElement = false
  result.add("]")

proc product*(a: ShapeArray): int {.inline.} =
  result = 1
  for value in items(a):
    result *= value

proc insert*(a: var ShapeArray, value: int, index: int = 0) {.inline.} =
  assert a.len+1 < MAXRANK
  assert index >= 0 and index <= a.len
  for i in countdown(a.len, index+1):
    a[i] = a[i-1]
  a[index] = value
  inc a.len

proc delete*(a: var ShapeArray, index: int) {.inline.} =
  assert a.len > 0
  assert index >= 0 and index < a.len
  for i in countdown(a.len-1, index+1):
    a[i-1] = a[i]
  a[a.len] = 0
  dec a.len

proc add*(a: var ShapeArray, value: int) {.inline.} =
  assert a.len+1 < MAXRANK
  a[a.len] = value
  inc a.len

proc reversed*(a: ShapeArray): ShapeArray {.inline.} =
  for i in 0..<a.len:
    result[a.len-i-1] = a[i]
  result.len = a.len

proc reversed*(a: ShapeArray, result: var ShapeArray) {.inline.} =
  for i in 0..<a.len:
    result[a.len-i-1] = a[i]
  for i in a.len..<result.len:
    result[i] = 0
  result.len = a.len

proc `==`*(a: ShapeArray, s: openarray[int]): bool {.inline.} =
  if a.len != s.len:
    return false
  for i in 0..<s.len:
    if a[i] != s[i]:
      return false
  return true

proc `==`*(a: ShapeArray, s: ShapeArray): bool {.inline.} =
  if a.len != s.len:
    return false
  for i in 0..<s.len:
    if a[i] != s[i]:
      return false
  return true

proc `^`*(x: int; a: ShapeArray): int {.inline.} =
  a.len - x


#### TensorBuffer


const TENSOR_DATA_ALIGN* = 8

type
  TensorBuffer*[T]  = object
    dataRef: ref[ptr T]
    data*: ptr UncheckedArray[T]
    len*: int

proc deallocTensorBuffer[T](dataRef: ref[ptr T]) =
  if not dataRef[].isNil:
    deallocShared(dataRef[])
    dataRef[] = nil

proc alignedAddress(p: pointer, align: csize): pointer =
  result = cast[pointer]((cast[csize](p) + TENSOR_DATA_ALIGN.csize - 1.csize) and not(align - 1.csize))

proc gcc_assume_aligned*[T](data: ptr T, n: csize): ptr T {.importc: "__builtin_assume_aligned",noDecl.}

proc init*[T](result: var TensorBuffer[T], len: int) {.inline.} =
  new(result.dataRef, deallocTensorBuffer)
  let address = allocShared(T.sizeof * len + TENSOR_DATA_ALIGN - 1)
  result.dataRef[] = cast[ptr T](address)
  result.data = cast[ptr UncheckedArray[T]](alignedAddress(address, TENSOR_DATA_ALIGN))
  result.len = len

proc init0*[T](result: var TensorBuffer[T], len: int) {.inline.} =
  result.init(len)
  zeroMem(result.data, T.sizeof * len)

proc initWith*[T](result: var TensorBuffer[T], len: int, value: T) {.inline.} =
  result.init(len)
  for i in 0..<len:
    result.data[i] = value

proc initFrom*[T](result: var TensorBuffer[T], s: ptr T, len: int) {.inline.} =
  result.init(len)
  copyMem(result.data, cast[pointer](s), T.sizeof * len)

proc cloneFrom*[T](result: var TensorBuffer[T], a: TensorBuffer[T]) {.inline.} =
  result.init(a.len)
  copyMem(result.data, a.data, T.sizeof * a.len)

proc copyFrom*[T](a: var TensorBuffer[T], s: TensorBuffer[T]) {.inline.} =
  assert a.len == s.len
  copyMem(a.data, s.data, T.sizeof * s.len)

proc setLen*[T](a: var TensorBuffer[T], len: int) {.inline.} =
  assert a.dataRef[] == data.pointer
  a.dataRef[] = reallocShared(a.dataRef[], T.sizeof * len)
  a.len = len

proc dataPtr*[T](a: TensorBuffer[T]): ptr T {.inline.} =
  a.data[0].unsafeAddr

proc len*(a: TensorBuffer): int {.inline.} =
  a.len

proc low*(a: TensorBuffer): int {.inline.} =
  0

proc high*(a: TensorBuffer): int {.inline.} =
  a.len-1

proc `[]`*[T](a: TensorBuffer[T], idx: int): T {.inline.} =
  assert idx >= 0 and idx < a.len
  a.data[idx]

proc `[]`*[T](a: var TensorBuffer[T], idx: int): var T {.inline.} =
  assert idx >= 0 and idx < a.len
  a.data[idx]

proc `[]=`*[T](a: var TensorBuffer[T], idx: int, v: T) {.inline.} =
  assert idx >= 0 and idx < a.len
  a.data[idx] = v

iterator items*[T](a: TensorBuffer[T]): T {.inline.} =
  for i in 0..<a.len:
    yield a.data[i]

iterator mitems*[T](a: var TensorBuffer[T]): var T {.inline.} =
  for i in 0..<a.len:
    yield a.data[i]

proc `==`*[T](a, b: TensorBuffer[T]): bool {.inline.} =
  if a.len != b.len:
    return false
  return equalMem(a.data[0].unsafeAddr, b.data[0].unsafeAddr, T.sizeof * a.len)

#[
proc `=`*[T](dest: var TensorBuffer[T], src: TensorBuffer[T]) {.inline.} =
  dest.cloneFrom(src)

proc `=`*[T](dest: var TensorBuffer[T], src: TensorBuffer[T]{call, `let`}) {.inline.} =
  shallowCopy(dest, src)
]#

#### Tensor

type
  Tensor*[T] = object
    shape*: ShapeArray
    strides*: ShapeArray
    data*: TensorBuffer[T]