<commit_before>package db

import (
	"fmt"

	"bytes"

	"github.com/boltdb/bolt"
	"github.com/pkg/errors"
)

const indiceBucket string = "indices"

// getLeagueIndexBucketRO returns the bucket corresponding
// to a specific league's index. This will never write
// and can be used safely with a readonly transaction.
//
// Will either panic or return a valid bucket.
func getLeagueIndexBucket(league LeagueHeapID, tx *bolt.Tx) *bolt.Bucket {
	// Grab league bucket
	leagueBucket := getLeagueBucket(league, tx)

	// This can never fail, its a guarantee that the itemStoreBucket was registered
	// and will always appear on a valid leagueBucket
	indices := leagueBucket.Bucket([]byte(indiceBucket))
	if indices == nil {
		panic(fmt.Sprintf("%s bucket not found when expected", itemStoreBucket))
	}

	return indices
}

// getItemModIndexBucket returns a bucket which a given mod can be put
// when considering the item containing it.
//
// This WILL write if a bucket is not found. Hence, readonly tx unsafe.
func getItemModIndexBucket(rootType, rootFlavor, mod StringHeapID,
	league LeagueHeapID, tx *bolt.Tx) (*bolt.Bucket, error) {
	// Keys towards the bucket we want to return, they may or may not exist
	keys := []StringHeapID{rootType, rootFlavor, mod}

	// Start at the index bucket
	currentBucket := getLeagueIndexBucket(league, tx)

	// Create all of the intervening keys
	for _, key := range keys {
		keyBytes := key.ToBytes()
		prevBucket := currentBucket.Bucket(keyBytes)
		if prevBucket == nil {
			// Create the bucket
			var err error
			prevBucket, err = currentBucket.CreateBucket(keyBytes)
			if err != nil {
				return nil,
					errors.Wrapf(err,
						"failed to add index intermediary bucket bucket, bucket=%s, chain=%v",
						key, keys)
			}
		}
		currentBucket = prevBucket
	}

	// If we made it through, our currentBucket should be the one we want
	return currentBucket, nil
}

// getItemModIndexBucketRO returns a bucket which a given mod can be put
// when considering the item containing it.
//
// This WILL NOT write if a bucket is not found. Hence, readonly tx unsafe.
func getItemModIndexBucketRO(rootType, rootFlavor, mod StringHeapID,
	league LeagueHeapID, tx *bolt.Tx) (*bolt.Bucket, error) {
	// Keys towards the bucket we want to return, they may or may not exist
	keys := []StringHeapID{rootType, rootFlavor, mod}

	// Start at the index bucket
	currentBucket := getLeagueIndexBucket(league, tx)

	// Traverse all intervening buckets
	for _, key := range keys {
		keyBytes := key.ToBytes()
		prevBucket := currentBucket.Bucket(keyBytes)
		if prevBucket == nil {
			return nil, errors.Errorf("invalid bucket, key=%d, chain=%v", key, keys)
		}
		currentBucket = prevBucket
	}

	// If we made it through, our currentBucket should be the one we want
	return currentBucket, nil
}<commit_msg>db getIemModIndexBucket reduce allocations by unrolling<commit_after>package db

import (
	"fmt"

	"bytes"

	"github.com/boltdb/bolt"
	"github.com/pkg/errors"
)

const indiceBucket string = "indices"

// getLeagueIndexBucketRO returns the bucket corresponding
// to a specific league's index. This will never write
// and can be used safely with a readonly transaction.
//
// Will either panic or return a valid bucket.
func getLeagueIndexBucket(league LeagueHeapID, tx *bolt.Tx) *bolt.Bucket {
	// Grab league bucket
	leagueBucket := getLeagueBucket(league, tx)

	// This can never fail, its a guarantee that the itemStoreBucket was registered
	// and will always appear on a valid leagueBucket
	indices := leagueBucket.Bucket([]byte(indiceBucket))
	if indices == nil {
		panic(fmt.Sprintf("%s bucket not found when expected", indiceBucket))
	}

	return indices
}

// getItemModIndexBucket returns a bucket which a given mod can be put
// when considering the item containing it.
//
// This WILL write if a bucket is not found. Hence, readonly tx unsafe.
func getItemModIndexBucket(rootType, rootFlavor, mod StringHeapID,
	league LeagueHeapID, tx *bolt.Tx) (*bolt.Bucket, error) {

	var err error

	// Start at the index bucket
	indexBucket := getLeagueIndexBucket(league, tx)
	rootTypeBucket := indexBucket.Bucket(rootType.ToBytes())
	if rootTypeBucket == nil {
		rootTypeBucket, err = indexBucket.CreateBucket(rootType.ToBytes())
		if err != nil {
			return nil, errors.Wrapf(err, "failed to add index rootType bucket")
		}
	}

	rootFlavorBucket := rootTypeBucket.Bucket(rootFlavor.ToBytes())
	if rootFlavorBucket == nil {
		rootFlavorBucket, err = rootTypeBucket.CreateBucket(rootFlavor.ToBytes())
		if err != nil {
			return nil, errors.Wrapf(err, "failed to add index rootFlavor bucket")
		}
	}

	modBucket := rootFlavorBucket.Bucket(mod.ToBytes())
	if modBucket == nil {
		modBucket, err = rootFlavorBucket.CreateBucket(mod.ToBytes())
		if err != nil {
			return nil, errors.Wrapf(err, "failed to add index mod bucket")
		}
	}

	// If we made it through, our currentBucket should be the one we want
	return modBucket, nil
}

// getItemModIndexBucketRO returns a bucket which a given mod can be put
// when considering the item containing it.
//
// This WILL NOT write if a bucket is not found. Hence, readonly tx unsafe.
func getItemModIndexBucketRO(rootType, rootFlavor, mod StringHeapID,
	league LeagueHeapID, tx *bolt.Tx) (*bolt.Bucket, error) {

	// Start at the index bucket
	indexBucket := getLeagueIndexBucket(league, tx)

	rootTypeBucket := indexBucket.Bucket(rootType.ToBytes())
	if rootTypeBucket == nil {
		return nil, errors.New("invalid rootType, no bucket found")
	}

	rootFlavorBucket := rootTypeBucket.Bucket(rootFlavor.ToBytes())
	if rootFlavorBucket == nil {
		return nil, errors.New("invalid rootFlavor, no bucket found")
	}

	modBucket := rootFlavorBucket.Bucket(mod.ToBytes())
	if modBucket == nil {
		return nil, errors.New("invalid mod, no bucket found")
	}

	// If we made it through, our currentBucket should be the one we want
	return modBucket, nil
}
{'fraction_non_alphanumeric': 0.09588344909368118, 'fraction_numerical': 0.0, 'mean_word_length': 4.107266435986159, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '1005886', 'n_tokens_mistral': 1854, 'n_tokens_neox': 1731, 'n_words': 648}