tensor/io/mmap_test.go

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

package io

import (
	"encoding/binary"
	"math"
	"os"
	"path/filepath"
	"sourcedock.dev/petrbalvin/tensor/internal/core"
	"testing"
)

// TestMapFloatsRoundTrip writes floats natively, maps them back, and
// checks spot values far apart, which is exactly the sparse-touch
// pattern big-file mapping exists for.
func TestMapFloatsRoundTrip(t *testing.T) {
	const n = 4096
	path := filepath.Join(t.TempDir(), "values.bin")
	values := make([]float64, n)
	for i := range n {
		values[i] = math.Sin(float64(i)) * 1e6
	}
	if err := SaveNativeFloats(path, values); err != nil {
		t.Fatalf("SaveNativeFloats: %v", err)
	}
	a, release, err := MapFloats(path, 0, n)
	if err != nil {
		t.Fatalf("MapFloats: %v", err)
	}
	if a.Len() != n || a.Dtype() != core.Float {
		t.Fatalf("mapped array %s of length %d", a.Dtype(), a.Len())
	}
	for _, i := range []int{0, 1, 999, 2048, n - 1} {
		if a.FloatAt(i) != values[i] {
			t.Fatalf("value %d = %.14g, want %.14g", i, a.FloatAt(i), values[i])
		}
	}
	// A reshape keeps the view: strides read straight from the mapping.
	square, err := core.Reshape(a, 64, 64)
	if err != nil {
		t.Fatalf("Reshape: %v", err)
	}
	if square.FloatAt(3*64+7) != values[3*64+7] {
		t.Fatal("the reshaped view lost the mapping")
	}
	if err := release(); err != nil {
		t.Fatalf("release: %v", err)
	}
}

// TestMapFloat32sAndInts covers the other two element types.
func TestMapFloat32sAndInts(t *testing.T) {
	dir := t.TempDir()
	const n = 100
	f32path := filepath.Join(dir, "f32.bin")
	var buf []byte
	var word [4]byte
	want32 := make([]float32, n)
	for i := range n {
		want32[i] = float32(i) / 7
		binary.NativeEndian.PutUint32(word[:], math.Float32bits(want32[i]))
		buf = append(buf, word[:]...)
	}
	if err := os.WriteFile(f32path, buf, 0o644); err != nil {
		t.Fatal(err)
	}
	a, release, err := MapFloat32s(f32path, 0, n)
	if err != nil {
		t.Fatalf("MapFloat32s: %v", err)
	}
	if a.RawFloat32s()[42] != want32[42] {
		t.Fatalf("f32[42] = %v, want %v", a.RawFloat32s()[42], want32[42])
	}
	if err := release(); err != nil {
		t.Fatal(err)
	}

	ipath := filepath.Join(dir, "i64.bin")
	var ibuf []byte
	var iword [8]byte
	wantInt := make([]int64, n)
	for i := range n {
		wantInt[i] = int64(i) * 1_000_000
		binary.NativeEndian.PutUint64(iword[:], uint64(wantInt[i]))
		ibuf = append(ibuf, iword[:]...)
	}
	if err := os.WriteFile(ipath, ibuf, 0o644); err != nil {
		t.Fatal(err)
	}
	b, releaseInt, err := MapInts(ipath, 0, n)
	if err != nil {
		t.Fatalf("MapInts: %v", err)
	}
	if b.RawInts()[13] != wantInt[13] {
		t.Fatalf("i64[13] = %d, want %d", b.RawInts()[13], wantInt[13])
	}
	if err := releaseInt(); err != nil {
		t.Fatal(err)
	}
}

// TestMapFloatsOffset maps from a non-zero offset inside a file that
// carries a small header first. The header is a multiple of the
// element size: the typed view the reader builds is built with unsafe
// and must be aligned.
func TestMapFloatsOffset(t *testing.T) {
	path := filepath.Join(t.TempDir(), "headered.bin")
	header := []byte("TENSOR HEADER 24 BYTES!!")
	if len(header)%8 != 0 {
		t.Fatalf("the test header is %d bytes, want a multiple of 8", len(header))
	}
	values := []float64{3.5, 1.25, -9}
	var buf []byte
	buf = append(buf, header...)
	var word [8]byte
	for _, v := range values {
		binary.NativeEndian.PutUint64(word[:], math.Float64bits(v))
		buf = append(buf, word[:]...)
	}
	if err := os.WriteFile(path, buf, 0o644); err != nil {
		t.Fatal(err)
	}
	a, release, err := MapFloats(path, int64(len(header)), 3)
	if err != nil {
		t.Fatalf("MapFloats: %v", err)
	}
	for i, want := range values {
		if a.FloatAt(i) != want {
			t.Fatalf("value %d = %.14g, want %.14g", i, a.FloatAt(i), want)
		}
	}
	if err := release(); err != nil {
		t.Fatal(err)
	}
}

// TestMapErrors pins the validation contract: short files, negative
// offsets, zero counts and directories are errors, never mappings.
func TestMapErrors(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "small.bin")
	if err := SaveNativeFloats(path, []float64{1, 2, 3}); err != nil {
		t.Fatal(err)
	}
	if _, _, err := MapFloats(path, 4, 1); err == nil {
		t.Fatal("expected an error for an offset that misaligns the elements")
	}
	if _, _, err := MapFloats(path, 0, 4); err == nil {
		t.Fatal("expected an error when the file is too short")
	}
	if _, _, err := MapFloats(path, -8, 2); err == nil {
		t.Fatal("expected an error for a negative offset")
	}
	if _, _, err := MapFloats(path, 0, 0); err == nil {
		t.Fatal("expected an error for a zero count")
	}
	if _, _, err := MapFloats(dir, 0, 1); err == nil {
		t.Fatal("expected an error for a directory")
	}
	if _, _, err := MapFloats(filepath.Join(dir, "missing.bin"), 0, 1); err == nil {
		t.Fatal("expected an error for a missing file")
	}
}