<commit_before>// Copyright (C) 2014 <NAME>

package hfile

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"os"

	"github.com/edsrzf/mmap-go"
	"github.com/golang/snappy"
)

type Reader struct {
	mmap mmap.MMap

	majorVersion uint32
	minorVersion uint32

	header Header
	index  []Block
}

type Block struct {
	buf           *bytes.Reader
	offset        uint64
	size          uint32
	firstKeyBytes []byte
}

func NewReader(file *os.File) (Reader, error) {
	hfile := Reader{}
	var err error
	hfile.mmap, err = mmap.Map(file, mmap.RDONLY, 0)
	if err != nil {
		return hfile, err
	}

	v := binary.BigEndian.Uint32(hfile.mmap[len(hfile.mmap)-4:])
	hfile.majorVersion = v & 0x00ffffff
	hfile.minorVersion = v >> 24

	hfile.header, err = hfile.newHeader(hfile.mmap)
	if err != nil {
		return hfile, err
	}
	err = hfile.loadIndex(hfile.mmap)
	if err != nil {
		return hfile, err
	}

	return hfile, nil
}

func getDataBlock(key []byte, blocks *[]Block) (*Block, bool) {
	// TODO(dan): Binary search instead.
	for i := len(*blocks) - 1; i >= 0; i-- {
		block := (*blocks)[i]
		if cmp := bytes.Compare(key, block.firstKeyBytes); cmp == 0 || cmp == 1 {
			return &block, true
		}
	}
	return nil, false
}
func (hfile *Reader) Get(key []byte) ([]byte, bool) {
	dataBlock, found := getDataBlock(key, &hfile.index)
	if !found {
		return nil, false
	}
	v, _, found := dataBlock.get(key, true)
	return v, found
}

func (r *Reader) PrintDebugInfo(out io.Writer) {
	fmt.Fprintln(out, "entries: ", r.header.entryCount)
	fmt.Fprintln(out, "blocks: ", len(r.index))
	for i, blk := range r.index {
		fmt.Fprintf(out, "\t#%d: %s (%v)\n", i, blk.firstKeyBytes, blk.firstKeyBytes)
	}
}

func (r *Reader) newHeader(mmap mmap.MMap) (Header, error) {
	header := Header{}

	if r.majorVersion != 1 || r.minorVersion != 0 {
		return header, errors.New("wrong version")
	}

	header.index = len(mmap) - 60
	buf := bytes.NewReader(mmap[header.index:])

	headerMagic := make([]byte, 8)
	buf.Read(headerMagic)
	if bytes.Compare(headerMagic, []byte("TRABLK\"$")) != 0 {
		return header, errors.New("bad header magic")
	}

	binary.Read(buf, binary.BigEndian, &header.fileInfoOffset)
	binary.Read(buf, binary.BigEndian, &header.dataIndexOffset)
	binary.Read(buf, binary.BigEndian, &header.dataIndexCount)
	binary.Read(buf, binary.BigEndian, &header.metaIndexOffset)
	binary.Read(buf, binary.BigEndian, &header.metaIndexCount)
	binary.Read(buf, binary.BigEndian, &header.totalUncompressedDataBytes)
	binary.Read(buf, binary.BigEndian, &header.entryCount)
	binary.Read(buf, binary.BigEndian, &header.compressionCodec)
	return header, nil
}

type Header struct {
	index int

	fileInfoOffset             uint64
	dataIndexOffset            uint64
	dataIndexCount             uint32
	metaIndexOffset            uint64
	metaIndexCount             uint32
	totalUncompressedDataBytes uint64
	entryCount                 uint32
	compressionCodec           uint32
}

func (r *Reader) loadIndex(mmap mmap.MMap) error {

	dataIndexEnd := r.header.metaIndexOffset
	if r.header.metaIndexOffset == 0 {
		dataIndexEnd = uint64(r.header.index)
	}
	buf := bytes.NewReader(mmap[r.header.dataIndexOffset:dataIndexEnd])<commit_msg>move structs to top<commit_after>// Copyright (C) 2014 <NAME>

package hfile

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"os"

	"github.com/edsrzf/mmap-go"
	"github.com/golang/snappy"
)

type Reader struct {
	mmap mmap.MMap

	majorVersion uint32
	minorVersion uint32

	header Header
	index  []Block
}

type Header struct {
	index int

	fileInfoOffset             uint64
	dataIndexOffset            uint64
	dataIndexCount             uint32
	metaIndexOffset            uint64
	metaIndexCount             uint32
	totalUncompressedDataBytes uint64
	entryCount                 uint32
	compressionCodec           uint32
}

type Block struct {
	buf           *bytes.Reader
	offset        uint64
	size          uint32
	firstKeyBytes []byte
}

func NewReader(file *os.File) (Reader, error) {
	hfile := Reader{}
	var err error
	hfile.mmap, err = mmap.Map(file, mmap.RDONLY, 0)
	if err != nil {
		return hfile, err
	}

	v := binary.BigEndian.Uint32(hfile.mmap[len(hfile.mmap)-4:])
	hfile.majorVersion = v & 0x00ffffff
	hfile.minorVersion = v >> 24

	hfile.header, err = hfile.newHeader(hfile.mmap)
	if err != nil {
		return hfile, err
	}
	err = hfile.loadIndex(hfile.mmap)
	if err != nil {
		return hfile, err
	}

	return hfile, nil
}

func getDataBlock(key []byte, blocks *[]Block) (*Block, bool) {
	// TODO(dan): Binary search instead.
	for i := len(*blocks) - 1; i >= 0; i-- {
		block := (*blocks)[i]
		if cmp := bytes.Compare(key, block.firstKeyBytes); cmp == 0 || cmp == 1 {
			return &block, true
		}
	}
	return nil, false
}
func (hfile *Reader) Get(key []byte) ([]byte, bool) {
	dataBlock, found := getDataBlock(key, &hfile.index)
	if !found {
		return nil, false
	}
	v, _, found := dataBlock.get(key, true)
	return v, found
}

func (r *Reader) PrintDebugInfo(out io.Writer) {
	fmt.Fprintln(out, "entries: ", r.header.entryCount)
	fmt.Fprintln(out, "blocks: ", len(r.index))
	for i, blk := range r.index {
		fmt.Fprintf(out, "\t#%d: %s (%v)\n", i, blk.firstKeyBytes, blk.firstKeyBytes)
	}
}

func (r *Reader) newHeader(mmap mmap.MMap) (Header, error) {
	header := Header{}

	if r.majorVersion != 1 || r.minorVersion != 0 {
		return header, errors.New("wrong version")
	}

	header.index = len(mmap) - 60
	buf := bytes.NewReader(mmap[header.index:])

	headerMagic := make([]byte, 8)
	buf.Read(headerMagic)
	if bytes.Compare(headerMagic, []byte("TRABLK\"$")) != 0 {
		return header, errors.New("bad header magic")
	}

	binary.Read(buf, binary.BigEndian, &header.fileInfoOffset)
	binary.Read(buf, binary.BigEndian, &header.dataIndexOffset)
	binary.Read(buf, binary.BigEndian, &header.dataIndexCount)
	binary.Read(buf, binary.BigEndian, &header.metaIndexOffset)
	binary.Read(buf, binary.BigEndian, &header.metaIndexCount)
	binary.Read(buf, binary.BigEndian, &header.totalUncompressedDataBytes)
	binary.Read(buf, binary.BigEndian, &header.entryCount)
	binary.Read(buf, binary.BigEndian, &header.compressionCodec)
	return header, nil
}

func (r *Reader) loadIndex(mmap mmap.MMap) error {

	dataIndexEnd := r.header.metaIndexOffset
	if r.header.metaIndexOffset == 0 {
		dataIndexEnd = uint64(r.header.index)
	}
	buf := bytes.NewReader(mmap[r.header.dataIndexOffset:dataIndexEnd])
{'fraction_non_alphanumeric': 0.13949104618284638, 'fraction_numerical': 0.015708451146716932, 'mean_word_length': 3.6954277286135695, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6017831', 'n_tokens_mistral': 2483, 'n_tokens_neox': 2342, 'n_words': 538}