<commit_before>package cete

import (
	"bytes"
	"log"
	"os"

	"github.com/dgraph-io/badger/badger"
	"gopkg.in/vmihailenco/msgpack.v2"
)

// Bounds is the type for variables which represent a bound for Between.
type Bounds int

// Valid bounds.
var (
	MinBounds Bounds = 1
	MaxBounds Bounds = 2
)

// NewIndex creates a new index on the table, using the name as the Query.
// The index name must not be empty, and must be no more than 125 bytes
// long. ErrAlreadyExists will be returned if the index already exists.
//
// NewIndex may take a while if there are already values in the
// table, as it needs to index all the existing values in the table.
func (t *Table) NewIndex(name string) error {
	if name == "" || len(name) > 125 {
		return ErrBadIdentifier
	}

	t.db.configMutex.Lock()

	tableName := t.name()
	tableConfigKey := -1

	for key, table := range t.db.config.Tables {
		if table.TableName == tableName {
			tableConfigKey = key
			for _, index := range table.Indexes {
				if index.IndexName == name {
					t.db.configMutex.Unlock()
					return ErrAlreadyExists
				}
			}
		}
	}

	if tableConfigKey < 0 {
		log.Println("cete: attempt to call new index on a non-existent table")
		t.db.configMutex.Unlock()
		return ErrNotFound
	}

	kv, err := t.db.newKV(Name(tableName), Name(name))
	if err != nil {
		t.db.configMutex.Unlock()
		return err
	}

	indexes := t.db.config.Tables[tableConfigKey].Indexes
	indexes = append(indexes, indexConfig{IndexName: name})
	t.db.config.Tables[tableConfigKey].Indexes = indexes
	if err = t.db.writeConfig(); err != nil {
		t.db.configMutex.Unlock()
		return err
	}

	t.db.configMutex.Unlock()

	idx := &Index{
		index: kv,
		table: t,
	}

	if err = idx.fill(name); err != nil {
		return err
	}

	t.indexes[Name(name)] = idx

	return nil
}

func (i *Index) fill(name string) error {
	r := i.table.Between(MinBounds, MaxBounds)

	var entry bufferEntry
	var results []interface{}
	var err error

	for {
		entry = <-r.buffer
		if entry.err == ErrEndOfRange {
			break
		} else if entry.err != nil {
			return entry.err
		}

		results, err = msgpack.NewDecoder(bytes.NewReader(entry.data)).Query(name)
		if err != nil {
			continue
		}

		for _, result := range results {
			err = i.addToIndex(valueToBytes(result), entry.key)
			if err != nil {
				log.Println("cete: index error for index \""+name+"\":", err)
			}
		}
	}

	return nil
}

// One puts the first matching value with the index's key into dst. dst
// must either be a pointer or nil if you would like to only get the key/counter
// and check for existence. Note that indexes are non-unique, a single index key
// can map to multiple values. Use GetAll to get all such matching values.
func (i *Index) One(key interface{}, dst interface{}) (string, int, error) {
	r, err := i.GetAll(key)
	if err != nil {
		return "", 0, err
	}

	defer r.Close()

	tableKey, counter, err := r.Next(dst)
	if err == ErrEndOfRange {
		log.Println("cete: warning: corrupt index detected:", i.name())
		return tableKey, counter, ErrNotFound
	}

	return tableKey, counter, err
}

// GetAll returns all the matching values as a range for the provided index key.
func (i *Index) GetAll(key interface{}) (*Range, error) {
	var item badger.KVItem
	err := i.index.Get(valueToBytes(key), &item)
	if err != nil {
		return nil, err
	}

	if item.Value() == nil {
		return nil, ErrNotFound
	}

	var keys []string
	err = msgpack.Unmarshal(item.Value(), &keys)
	if err != nil {
		log.Println("cete: corrupt index \""+i.name()+"\":", err)
		return nil, ErrIndexError
	}

	if len(keys) == 0 {
		log.Println("cete: corrupt index \""+i.name()+"\":", err)
		return nil, ErrNotFound
	}

	c := 0
	var value []byte

	return newRange(func() (string, []byte, int, error) {
		for {
			if c >= len(keys) {
				return "", nil, 0, ErrEndOfRange
			}

			err = i.table.data.Get([]byte(keys[c]), &item)
			if err != nil {
				return "", nil, 0, err
			}

			if item.Value() == nil {
				c++
				continue
			}

			value = make([]byte, len(item.Value()))
			copy(value, item.Value())

			c++
			return keys[c-1], value, int(item.Counter()), nil
		}
	}, func() {}), nil
}

// Between returns a Range of documents between the lower and upper index values
// provided. The range will be sorted in ascending order by index value. You can
// reverse the sorting by specifying true to the optional reverse parameter.
// The bounds are inclusive on both ends. It is possible to have
// duplicate documents if the same document has multiple unique index values.
//
// You can use cete.MinBounds and cete.MaxBounds to specify minimum and maximum
// bound values.
func (i *Index) Between(lower, upper interface{},
	reverse ...bool) *Range {
	shouldReverse := (len(reverse) > 0) && reverse[0]

	itOpts := badger.DefaultIteratorOptions
	itOpts.PrefetchSize = 5
	itOpts.Reverse = shouldReverse
	it := i.index.NewIterator(itOpts)

	upperBytes := valueToBytes(upper)
	lowerBytes := valueToBytes(lower)

	if !shouldReverse {
		if lower == MinBounds {
			it.Rewind()
		} else {
			it.Seek(lowerBytes)
		}
	} else {
		if upper == MaxBounds {
			it.Rewind()
		} else {
			it.Seek(upperBytes)
		}
	}

	var lastRange *Range

	return newRange(i.betweenNext(it, lastRange, shouldReverse, lower, upper),
		func() {
			if lastRange != nil {
				lastRange.Close()
			}
			it.Close()
		})
}

func (i *Index) betweenNext(it *badger.Iterator, lastRange *Range,
	shouldReverse bool, lower,
	upper interface{}) func() (string, []byte, int, error) {
	upperBytes := valueToBytes(upper)
	lowerBytes := valueToBytes(lower)

	var entry bufferEntry

	return func() (string, []byte, int, error) {
		if lastRange != nil {
			entry = <-lastRange.buffer
			if entry.err != ErrEndOfRange {
				return entry.key, entry.data, entry.counter, entry.err
			}

			lastRange.Close()
		}

		for it.Valid() {
			if !shouldReverse && upper != MaxBounds &&
				bytes.Compare(it.Item().Key(), upperBytes) > 0 {
				return "", nil, 0, ErrEndOfRange
			} else if shouldReverse && lower != MinBounds &&
				bytes.Compare(it.Item().Key(), lowerBytes) < 0 {
				return "", nil, 0, ErrEndOfRange
			}

			r, err := i.GetAll(it.Item().Key())
			it.Next()
			if err != nil {
				continue
			}

			lastRange = r

			entry = <-lastRange.buffer
			if entry.err != ErrEndOfRange {
				return entry.key, entry.data, entry.counter, entry.err
			}

			lastRange.Close()
		}

		return "", nil, 0, ErrEndOfRange
	}
}

// All returns all the documents which have an index value. It is shorthand
// for Between(MinBounds, MaxBounds, reverse...)
func (i *Index) All(reverse ...bool) *Range {
	return i.Between(MinBounds, MaxBounds, reverse...)
}

// Drop drops the index from the table, deleting its folder from the disk.
// All further calls to the index will result in undefined behaviour.
// Note that table.Index("deleted index") will be nil.
func (i *Index) Drop() error {
	i.table.db.configMutex.Lock()
	defer i.table.db.configMutex.Unlock()

	tableName := i.table.name()

	var indexName string

	for idxName, index := range i.table.indexes {
		if index == i {
			indexName = string(idxName)
		}
	}

	if indexName == "" {
		return ErrNotFound
	}

tableLoop:
	for key, table := range i.table.db.config.Tables {
		if table.TableName == tableName {
			for indexKey, index := range table.Indexes {
				if index.IndexName == indexName {
					indexes := i.table.db.config.Tables[key].Indexes
					indexes = append(indexes[:indexKey], indexes[indexKey+1:]...)
					i.table.db.config.Tables[key].Indexes = indexes
					break tableLoop
				}
			}
		}
	}

	if err := i.table.db.writeConfig(); err != nil {
		return err
	}

	i.index.Close()

	delete(i.table.indexes, Name(indexName))

	return os.RemoveAll(i.table.db.path + "/" + Name(tableName).Hex() + "/" +
		Name(indexName).Hex())
}
<commit_msg>Prevent race condition with index creation<commit_after>package cete

import (
	"bytes"
	"log"
	"os"

	"github.com/dgraph-io/badger/badger"
	"gopkg.in/vmihailenco/msgpack.v2"
)

// Bounds is the type for variables which represent a bound for Between.
type Bounds int

// Valid bounds.
var (
	MinBounds Bounds = 1
	MaxBounds Bounds = 2
)

// NewIndex creates a new index on the table, using the name as the Query.
// The index name must not be empty, and must be no more than 125 bytes
// long. ErrAlreadyExists will be returned if the index already exists.
//
// NewIndex may take a while if there are already values in the
// table, as it needs to index all the existing values in the table.
func (t *Table) NewIndex(name string) error {
	if name == "" || len(name) > 125 {
		return ErrBadIdentifier
	}

	t.db.configMutex.Lock()

	tableName := t.name()
	tableConfigKey := -1

	for key, table := range t.db.config.Tables {
		if table.TableName == tableName {
			tableConfigKey = key
			for _, index := range table.Indexes {
				if index.IndexName == name {
					t.db.configMutex.Unlock()
					return ErrAlreadyExists
				}
			}
		}
	}

	if tableConfigKey < 0 {
		log.Println("cete: attempt to call new index on a non-existent table")
		t.db.configMutex.Unlock()
		return ErrNotFound
	}

	kv, err := t.db.newKV(Name(tableName), Name(name))
	if err != nil {
		t.db.configMutex.Unlock()
		return err
	}

	indexes := t.db.config.Tables[tableConfigKey].Indexes
	indexes = append(indexes, indexConfig{IndexName: name})
	t.db.config.Tables[tableConfigKey].Indexes = indexes
	if err = t.db.writeConfig(); err != nil {
		t.db.configMutex.Unlock()
		return err
	}

	t.db.configMutex.Unlock()

	idx := &Index{
		index: kv,
		table: t,
	}

	t.indexes[Name(name)] = idx

	if err = idx.indexValues(name); err != nil {
		log.Println("cete: error while indexing \""+
			idx.name()+"\", index likely corrupt:",
			err)
		return nil
	}

	return nil
}

func (i *Index) indexValues(name string) error {
	r := i.table.Between(MinBounds, MaxBounds)

	var entry bufferEntry
	var results []interface{}
	var err error

	for {
		entry = <-r.buffer
		if entry.err == ErrEndOfRange {
			break
		} else if entry.err != nil {
			return entry.err
		}

		results, err = msgpack.NewDecoder(bytes.NewReader(entry.data)).Query(name)
		if err != nil {
			continue
		}

		for _, result := range results {
			err = i.addToIndex(valueToBytes(result), entry.key)
			if err != nil {
				log.Println("cete: index error for index \""+name+"\":", err)
			}
		}
	}

	return nil
}

// One puts the first matching value with the index's key into dst. dst
// must either be a pointer or nil if you would like to only get the key/counter
// and check for existence. Note that indexes are non-unique, a single index key
// can map to multiple values. Use GetAll to get all such matching values.
func (i *Index) One(key interface{}, dst interface{}) (string, int, error) {
	r, err := i.GetAll(key)
	if err != nil {
		return "", 0, err
	}

	defer r.Close()

	tableKey, counter, err := r.Next(dst)
	if err == ErrEndOfRange {
		log.Println("cete: warning: corrupt index detected:", i.name())
		return tableKey, counter, ErrNotFound
	}

	return tableKey, counter, err
}

// GetAll returns all the matching values as a range for the provided index key.
func (i *Index) GetAll(key interface{}) (*Range, error) {
	var item badger.KVItem
	err := i.index.Get(valueToBytes(key), &item)
	if err != nil {
		return nil, err
	}

	if item.Value() == nil {
		return nil, ErrNotFound
	}

	var keys []string
	err = msgpack.Unmarshal(item.Value(), &keys)
	if err != nil {
		log.Println("cete: corrupt index \""+i.name()+"\":", err)
		return nil, ErrIndexError
	}

	if len(keys) == 0 {
		log.Println("cete: corrupt index \""+i.name()+"\":", err)
		return nil, ErrNotFound
	}

	c := 0
	var value []byte

	return newRange(func() (string, []byte, int, error) {
		for {
			if c >= len(keys) {
				return "", nil, 0, ErrEndOfRange
			}

			err = i.table.data.Get([]byte(keys[c]), &item)
			if err != nil {
				return "", nil, 0, err
			}

			if item.Value() == nil {
				c++
				continue
			}

			value = make([]byte, len(item.Value()))
			copy(value, item.Value())

			c++
			return keys[c-1], value, int(item.Counter()), nil
		}
	}, func() {}), nil
}

// Between returns a Range of documents between the lower and upper index values
// provided. The range will be sorted in ascending order by index value. You can
// reverse the sorting by specifying true to the optional reverse parameter.
// The bounds are inclusive on both ends. It is possible to have
// duplicate documents if the same document has multiple unique index values.
//
// You can use cete.MinBounds and cete.MaxBounds to specify minimum and maximum
// bound values.
func (i *Index) Between(lower, upper interface{},
	reverse ...bool) *Range {
	shouldReverse := (len(reverse) > 0) && reverse[0]

	itOpts := badger.DefaultIteratorOptions
	itOpts.PrefetchSize = 5
	itOpts.Reverse = shouldReverse
	it := i.index.NewIterator(itOpts)

	upperBytes := valueToBytes(upper)
	lowerBytes := valueToBytes(lower)

	if !shouldReverse {
		if lower == MinBounds {
			it.Rewind()
		} else {
			it.Seek(lowerBytes)
		}
	} else {
		if upper == MaxBounds {
			it.Rewind()
		} else {
			it.Seek(upperBytes)
		}
	}

	var lastRange *Range

	return newRange(i.betweenNext(it, lastRange, shouldReverse, lower, upper),
		func() {
			if lastRange != nil {
				lastRange.Close()
			}
			it.Close()
		})
}

func (i *Index) betweenNext(it *badger.Iterator, lastRange *Range,
	shouldReverse bool, lower,
	upper interface{}) func() (string, []byte, int, error) {
	upperBytes := valueToBytes(upper)
	lowerBytes := valueToBytes(lower)

	var entry bufferEntry

	return func() (string, []byte, int, error) {
		if lastRange != nil {
			entry = <-lastRange.buffer
			if entry.err != ErrEndOfRange {
				return entry.key, entry.data, entry.counter, entry.err
			}

			lastRange.Close()
		}

		for it.Valid() {
			if !shouldReverse && upper != MaxBounds &&
				bytes.Compare(it.Item().Key(), upperBytes) > 0 {
				return "", nil, 0, ErrEndOfRange
			} else if shouldReverse && lower != MinBounds &&
				bytes.Compare(it.Item().Key(), lowerBytes) < 0 {
				return "", nil, 0, ErrEndOfRange
			}

			r, err := i.GetAll(it.Item().Key())
			it.Next()
			if err != nil {
				continue
			}

			lastRange = r

			entry = <-lastRange.buffer
			if entry.err != ErrEndOfRange {
				return entry.key, entry.data, entry.counter, entry.err
			}

			lastRange.Close()
		}

		return "", nil, 0, ErrEndOfRange
	}
}

// All returns all the documents which have an index value. It is shorthand
// for Between(MinBounds, MaxBounds, reverse...)
func (i *Index) All(reverse ...bool) *Range {
	return i.Between(MinBounds, MaxBounds, reverse...)
}

// Drop drops the index from the table, deleting its folder from the disk.
// All further calls to the index will result in undefined behaviour.
// Note that table.Index("deleted index") will be nil.
func (i *Index) Drop() error {
	i.table.db.configMutex.Lock()
	defer i.table.db.configMutex.Unlock()

	tableName := i.table.name()

	var indexName string

	for idxName, index := range i.table.indexes {
		if index == i {
			indexName = string(idxName)
		}
	}

	if indexName == "" {
		return ErrNotFound
	}

tableLoop:
	for key, table := range i.table.db.config.Tables {
		if table.TableName == tableName {
			for indexKey, index := range table.Indexes {
				if index.IndexName == indexName {
					indexes := i.table.db.config.Tables[key].Indexes
					indexes = append(indexes[:indexKey], indexes[indexKey+1:]...)
					i.table.db.config.Tables[key].Indexes = indexes
					break tableLoop
				}
			}
		}
	}

	if err := i.table.db.writeConfig(); err != nil {
		return err
	}

	i.index.Close()

	delete(i.table.indexes, Name(indexName))

	return os.RemoveAll(i.table.db.path + "/" + Name(tableName).Hex() + "/" +
		Name(indexName).Hex())
}

{'fraction_non_alphanumeric': 0.13667367884235354, 'fraction_numerical': 0.003314846688340664, 'mean_word_length': 3.2285714285714286, 'pattern_counts': {'":': 6, '<': 13, '<?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': '5459700', 'n_tokens_mistral': 5877, 'n_tokens_neox': 5250, 'n_words': 1632}