mmapbinfuse-read.go

package mmapbinfuse

import (
	"bytes"
	"errors"
	"fmt"
	"math/bits"
	"os"
	"runtime"

	"github.com/kelindar/binary"
)

type BinFuse[T Unsigned] struct {
	Seed               uint64
	SegmentLength      uint32
	SegmentLengthMask  uint32
	SegmentCount       uint32
	SegmentCountLength uint32

	Fingerprints *GenericMMap[T]
}

func Read[T Unsigned](src string) (BinFuse[T], error) {
	srcFile, err := os.Open(src)
	if err != nil {
		return BinFuse[T]{}, fmt.Errorf("failed to open file: %w", err)
	}
	defer srcFile.Close()

	headerBytes, err := readHeaderBytesFromFile(srcFile)
	if err != nil {
		return BinFuse[T]{}, fmt.Errorf("failed to read header: %w", err)
	}
	header, err := unmarshalFilterHeader(headerBytes)
	if err != nil {
		return BinFuse[T]{}, fmt.Errorf("failed to unmarshal header: %w", err)
	}

	switch any(*new(T)).(type) {
	case uint8:
		if header.TypeOfT != TypeOfTUint8 {
			return BinFuse[T]{}, fmt.Errorf("type mismatch; the file contains fingerprints of type %s, but you're trying to open it as %s", typeNumberToString(header.TypeOfT), typeNumberToString(TypeOfTUint8))
		}
	case uint16:
		if header.TypeOfT != TypeOfTUint16 {
			return BinFuse[T]{}, fmt.Errorf("type mismatch; the file contains fingerprints of type %s, but you're trying to open it as %s", typeNumberToString(header.TypeOfT), typeNumberToString(TypeOfTUint16))
		}
	case uint32:
		if header.TypeOfT != TypeOfTUint32 {
			return BinFuse[T]{}, fmt.Errorf("type mismatch; the file contains fingerprints of type %s, but you're trying to open it as %s", typeNumberToString(header.TypeOfT), typeNumberToString(TypeOfTUint32))
		}
	default:
		return BinFuse[T]{}, errors.New("unsupported type")
	}

	// read from file to memory
	mm, err := NewGenericMMapWithOffset[T](src, int64(len(headerBytes)))
	if err != nil {
		return BinFuse[T]{}, fmt.Errorf("failed to create mmap: %w", err)
	}

	bf := BinFuse[T]{
		Seed:               header.Seed,
		SegmentLength:      header.SegmentLength,
		SegmentLengthMask:  header.SegmentLengthMask,
		SegmentCount:       header.SegmentCount,
		SegmentCountLength: header.SegmentCountLength,
		Fingerprints:       mm,
	}

	// set finalizer to unmap the memory-mapped file
	runtime.SetFinalizer(&bf, func(bf *BinFuse[T]) {
		bf.Fingerprints.Close()
	})

	return bf, nil
}

// Close unmaps the memory-mapped file.
func (filter *BinFuse[T]) Close() error {
	return filter.Fingerprints.Close()
}

func readHeaderBytesFromFile(file *os.File) ([]byte, error) {
	// read the header
	headerLen := len(MagicNumber) +
		8 + // Seed
		4 + // SegmentLength
		4 + // SegmentLengthMask
		4 + // SegmentCount
		4 + // SegmentCountLength
		8 + // NumFingerprints
		1 // type of T
	header := make([]byte, headerLen)
	_, err := file.Read(header)
	if err != nil {
		return nil, err
	}
	return header, nil
}

type Header struct {
	Seed               uint64
	SegmentLength      uint32
	SegmentLengthMask  uint32
	SegmentCount       uint32
	SegmentCountLength uint32

	NumFingerprints uint64
	TypeOfT         byte
}

func unmarshalFilterHeader(data []byte) (*Header, error) {
	header := new(Header)
	decoder := binary.NewDecoder(bytes.NewReader(data))
	var err error
	magic := make([]byte, len(MagicNumber))
	_, err = decoder.Read(magic)
	if err != nil {
		return nil, fmt.Errorf("failed to read magic number: %w", err)
	}
	if !bytes.Equal(magic, MagicNumber) {
		return nil, fmt.Errorf("invalid magic number; got %q, want %q", magic, MagicNumber)
	}
	header.Seed, err = decoder.ReadUint64()
	if err != nil {
		return nil, fmt.Errorf("failed to read seed: %w", err)
	}
	header.SegmentLength, err = decoder.ReadUint32()
	if err != nil {
		return nil, fmt.Errorf("failed to read segment length: %w", err)
	}
	header.SegmentLengthMask, err = decoder.ReadUint32()
	if err != nil {
		return nil, fmt.Errorf("failed to read segment length mask: %w", err)
	}
	header.SegmentCount, err = decoder.ReadUint32()
	if err != nil {
		return nil, fmt.Errorf("failed to read segment count: %w", err)
	}
	header.SegmentCountLength, err = decoder.ReadUint32()
	if err != nil {
		return nil, fmt.Errorf("failed to read segment count length: %w", err)
	}
	len, err := decoder.ReadUint64()
	if err != nil {
		return nil, fmt.Errorf("failed to read fingerprints length: %w", err)
	}
	header.NumFingerprints = len

	t := make([]byte, 1)
	_, err = decoder.Read(t)
	if err != nil {
		return nil, fmt.Errorf("failed to read type of T: %w", err)
	}
	header.TypeOfT = t[0]
	if header.TypeOfT != TypeOfTUint8 && header.TypeOfT != TypeOfTUint16 && header.TypeOfT != TypeOfTUint32 {
		return nil, fmt.Errorf("unsupported type: %d", header.TypeOfT)
	}

	return header, nil
}

func (filter *BinFuse[T]) getHashFromHash(hash uint64) (uint32, uint32, uint32) {
	hi, _ := bits.Mul64(hash, uint64(filter.SegmentCountLength))
	h0 := uint32(hi)
	h1 := h0 + filter.SegmentLength
	h2 := h1 + filter.SegmentLength
	h1 ^= uint32(hash>>18) & filter.SegmentLengthMask
	h2 ^= uint32(hash) & filter.SegmentLengthMask
	return h0, h1, h2
}

// Contains returns `true` if key is part of the set with a false positive probability.
func (filter *BinFuse[T]) Contains(key uint64) bool {
	hash := mixsplit(key, filter.Seed)
	f := T(fingerprint(hash))
	h0, h1, h2 := filter.getHashFromHash(hash)
	f ^= filter.Fingerprints.Get(h0) ^ filter.Fingerprints.Get(h1) ^ filter.Fingerprints.Get(h2)
	return f == 0
}