tensor/internal/core/narrowdt.go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package core

// The narrow element types: bool and the small integers. They exist so
// byte payloads, masks and the file formats' short integers ride in the
// width they were born in instead of widening to int64 in transit. The
// constructors here mirror the established pair: From…s copies its
// values, …FromArray takes ownership of the caller's slice. The byte
// serialisation stays deliberately Int-only: Bytes answers
// only for an Int array and FromBytes builds Int, so a narrow payload
// converts through Astype first rather than riding Bytes natively.

// FromBools builds a bool array from vals, copying them.
func FromBools(vals []bool, shape ...int) (*Array, error) {
	sh, err := shapeFor(shape, len(vals))
	if err != nil {
		return nil, err
	}
	bools := make([]bool, len(vals))
	copy(bools, vals)
	return &Array{shape: sh, dt: Bool, bools: bools}, nil
}

// FromInt8s builds an int8 array from vals, copying them.
func FromInt8s(vals []int8, shape ...int) (*Array, error) {
	sh, err := shapeFor(shape, len(vals))
	if err != nil {
		return nil, err
	}
	out := make([]int8, len(vals))
	copy(out, vals)
	return &Array{shape: sh, dt: Int8, i8s: out}, nil
}

// FromUint8s builds a uint8 array from vals, copying them.
func FromUint8s(vals []uint8, shape ...int) (*Array, error) {
	sh, err := shapeFor(shape, len(vals))
	if err != nil {
		return nil, err
	}
	out := make([]uint8, len(vals))
	copy(out, vals)
	return &Array{shape: sh, dt: Uint8, u8s: out}, nil
}

// FromInt16s builds an int16 array from vals, copying them.
func FromInt16s(vals []int16, shape ...int) (*Array, error) {
	sh, err := shapeFor(shape, len(vals))
	if err != nil {
		return nil, err
	}
	out := make([]int16, len(vals))
	copy(out, vals)
	return &Array{shape: sh, dt: Int16, i16s: out}, nil
}

// FromUint16s builds a uint16 array from vals, copying them.
func FromUint16s(vals []uint16, shape ...int) (*Array, error) {
	sh, err := shapeFor(shape, len(vals))
	if err != nil {
		return nil, err
	}
	out := make([]uint16, len(vals))
	copy(out, vals)
	return &Array{shape: sh, dt: Uint16, u16s: out}, nil
}

// FromInt32s builds an int32 array from vals, copying them.
func FromInt32s(vals []int32, shape ...int) (*Array, error) {
	sh, err := shapeFor(shape, len(vals))
	if err != nil {
		return nil, err
	}
	out := make([]int32, len(vals))
	copy(out, vals)
	return &Array{shape: sh, dt: Int32, i32s: out}, nil
}

// FromUint32s builds a uint32 array from vals, copying them.
func FromUint32s(vals []uint32, shape ...int) (*Array, error) {
	sh, err := shapeFor(shape, len(vals))
	if err != nil {
		return nil, err
	}
	out := make([]uint32, len(vals))
	copy(out, vals)
	return &Array{shape: sh, dt: Uint32, u32s: out}, nil
}

// BoolsFromArray builds a bool array that takes ownership of vals, with
// the FloatsFromArray contract.
func BoolsFromArray(vals []bool, shape ...int) (*Array, error) {
	sh, err := shapeFor(shape, len(vals))
	if err != nil {
		return nil, err
	}
	return &Array{shape: sh, dt: Bool, bools: vals}, nil
}

// Int8sFromArray builds an int8 array that takes ownership of vals,
// with the FloatsFromArray contract.
func Int8sFromArray(vals []int8, shape ...int) (*Array, error) {
	sh, err := shapeFor(shape, len(vals))
	if err != nil {
		return nil, err
	}
	return &Array{shape: sh, dt: Int8, i8s: vals}, nil
}

// Uint8sFromArray builds a uint8 array that takes ownership of vals,
// with the FloatsFromArray contract.
func Uint8sFromArray(vals []uint8, shape ...int) (*Array, error) {
	sh, err := shapeFor(shape, len(vals))
	if err != nil {
		return nil, err
	}
	return &Array{shape: sh, dt: Uint8, u8s: vals}, nil
}

// Int16sFromArray builds an int16 array that takes ownership of vals,
// with the FloatsFromArray contract.
func Int16sFromArray(vals []int16, shape ...int) (*Array, error) {
	sh, err := shapeFor(shape, len(vals))
	if err != nil {
		return nil, err
	}
	return &Array{shape: sh, dt: Int16, i16s: vals}, nil
}

// Uint16sFromArray builds a uint16 array that takes ownership of vals,
// with the FloatsFromArray contract.
func Uint16sFromArray(vals []uint16, shape ...int) (*Array, error) {
	sh, err := shapeFor(shape, len(vals))
	if err != nil {
		return nil, err
	}
	return &Array{shape: sh, dt: Uint16, u16s: vals}, nil
}

// Int32sFromArray builds an int32 array that takes ownership of vals,
// with the FloatsFromArray contract.
func Int32sFromArray(vals []int32, shape ...int) (*Array, error) {
	sh, err := shapeFor(shape, len(vals))
	if err != nil {
		return nil, err
	}
	return &Array{shape: sh, dt: Int32, i32s: vals}, nil
}

// Uint32sFromArray builds a uint32 array that takes ownership of vals,
// with the FloatsFromArray contract.
func Uint32sFromArray(vals []uint32, shape ...int) (*Array, error) {
	sh, err := shapeFor(shape, len(vals))
	if err != nil {
		return nil, err
	}
	return &Array{shape: sh, dt: Uint32, u32s: vals}, nil
}

// RawBools returns the bool payload with the RawFloats contract.
func (a *Array) RawBools() []bool { return a.bools }

// RawInt8s returns the int8 payload with the RawFloats contract.
func (a *Array) RawInt8s() []int8 { return a.i8s }

// RawUint8s returns the uint8 payload with the RawFloats contract.
func (a *Array) RawUint8s() []uint8 { return a.u8s }

// RawInt16s returns the int16 payload with the RawFloats contract.
func (a *Array) RawInt16s() []int16 { return a.i16s }

// RawUint16s returns the uint16 payload with the RawFloats contract.
func (a *Array) RawUint16s() []uint16 { return a.u16s }

// RawInt32s returns the int32 payload with the RawFloats contract.
func (a *Array) RawInt32s() []int32 { return a.i32s }

// RawUint32s returns the uint32 payload with the RawFloats contract.
func (a *Array) RawUint32s() []uint32 { return a.u32s }

// boolAt returns element i as a bool: the boolean payload's own value,
// and everywhere else the test against zero. A NaN compares unequal to
// zero, so a NaN element reads true, the same masked-read decision the
// int mask path makes on it.
func (a *Array) boolAt(i int) bool {
	if a.strides != nil {
		i = a.physIndex(i)
	}
	switch a.dt {
	case Bool:
		return a.bools[i]
	case Int:
		return a.ints[i] != 0
	case Float16:
		return a.halves[i]&0x7FFF != 0
	case Float32:
		return a.floats32[i] != 0
	case Float:
		return a.floats[i] != 0
	case Complex:
		return a.complexes[i] != 0
	case Int8:
		return a.i8s[i] != 0
	case Uint8:
		return a.u8s[i] != 0
	case Int16:
		return a.i16s[i] != 0
	case Uint16:
		return a.u16s[i] != 0
	case Int32:
		return a.i32s[i] != 0
	default:
		return a.u32s[i] != 0
	}
}

// BoolAt returns the element at the given index as a bool: the boolean
// payload's own value, and every other dtype read against zero. It is
// the widening reader the small integer types share with the internal
// accessors; unlike IntAt and FloatAt it takes any dtype, because the
// mask question makes sense over all of them.
func BoolAt(a *Array, index ...int) (bool, error) {
	off, err := flatIndex(a.shape, index)
	if err != nil {
		return false, err
	}
	return a.boolAt(off), nil
}

// intClassOrder lists the integer-class dtypes in containment order,
// the axis the promotion table below is indexed by.
var intClassOrder = [...]Dtype{Bool, Int8, Uint8, Int16, Uint16, Int32, Uint32, Int}

// intClassIndex maps an integer-class dtype to its row in the promotion
// table; an unlisted ordinal lands on Int, the recorded default.
func intClassIndex(d Dtype) int {
	for i, c := range intClassOrder {
		if c == d {
			return i
		}
	}
	return len(intClassOrder) - 1
}

// intClass reports whether d belongs to the integer class: Bool or one
// of the integer dtypes, everything promote resolves by containment.
func intClass(d Dtype) bool {
	switch d {
	case Bool, Int8, Uint8, Int16, Uint16, Int32, Uint32, Int:
		return true
	}
	return false
}

// cloneArray returns a deep copy of the array: same shape, same dtype,
// own payload. The narrow element types clone here, where their payload
// fields live; cloneData keeps serving the five legacy slices to the
// callers that consume them individually.
func (a *Array) cloneArray() *Array {
	out := &Array{shape: a.Shape(), dt: a.dt}
	n := a.Len()
	out.alloc(n)
	if a.strides != nil {
		for i := range n {
			out.setFrom(i, a, i)
		}
		return out
	}
	switch a.dt {
	case Int:
		copy(out.ints, a.ints[:n])
	case Float16:
		copy(out.halves, a.halves[:n])
	case Float32:
		copy(out.floats32, a.floats32[:n])
	case Float:
		copy(out.floats, a.floats[:n])
	case Complex:
		copy(out.complexes, a.complexes[:n])
	case Bool:
		copy(out.bools, a.bools[:n])
	case Int8:
		copy(out.i8s, a.i8s[:n])
	case Uint8:
		copy(out.u8s, a.u8s[:n])
	case Int16:
		copy(out.i16s, a.i16s[:n])
	case Uint16:
		copy(out.u16s, a.u16s[:n])
	case Int32:
		copy(out.i32s, a.i32s[:n])
	default:
		copy(out.u32s, a.u32s[:n])
	}
	return out
}

// intPromote answers the smallest dtype of the integer class whose
// value range contains both operands' ranges. Mixed signedness pairs
// therefore widen instead of losing negative values: int8 with uint8
// answers int16, int16 with uint16 answers int32, int32 with uint32
// answers int. promote answers a same-dtype pair before the table is
// ever consulted, and the signed rows widen their own kind to the next
// width as the containment rule dictates; the table itself is
// symmetric.
var intPromote = [len(intClassOrder)][len(intClassOrder)]Dtype{
	// Bool row: every integer dtype contains {0, 1}.
	{Bool, Int8, Uint8, Int16, Uint16, Int32, Uint32, Int},
	// Int8 [-128, 127]: with Uint8 needs the 16-bit signed width.
	{Int8, Int16, Int16, Int16, Int32, Int32, Int, Int},
	// Uint8 [0, 255]: contained by every wider dtype of either sign.
	{Uint8, Int16, Uint8, Int16, Uint16, Int32, Uint32, Int},
	// Int16: with Uint16 needs 32-bit signed; with Uint32 needs Int.
	{Int16, Int16, Int16, Int32, Int32, Int32, Int, Int},
	// Uint16 [0, 65535]: contained by Int32 and everything wider.
	{Uint16, Int32, Uint16, Int32, Uint16, Int32, Uint32, Int},
	// Int32: with Uint32 needs the 64-bit signed width.
	{Int32, Int32, Int32, Int32, Int32, Int, Int, Int},
	// Uint32 [0, 2^32-1]: contained by Int.
	{Uint32, Int, Uint32, Int, Uint32, Int, Uint32, Int},
	// Int contains the whole class.
	{Int, Int, Int, Int, Int, Int, Int, Int},
}