<commit_before>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package mocktikv

import (
	"bytes"
	"sync"

	"github.com/petar/GoLLRB/llrb"
	"github.com/pingcap/kvproto/pkg/kvrpcpb"
)

type mvccValue struct {
	startTS  uint64
	commitTS uint64
	value    []byte
}

type mvccLock struct {
	startTS uint64
	primary []byte
	value   []byte
	op      kvrpcpb.Op
}

type mvccEntry struct {
	key    []byte
	values []mvccValue
	lock   *mvccLock
}

func newEntry(key []byte) *mvccEntry {
	return &mvccEntry{
		key: key,
	}
}

func (e *mvccEntry) Less(than llrb.Item) bool {
	return bytes.Compare(e.key, than.(*mvccEntry).key) < 0
}

func (e *mvccEntry) lockErr() error {
	return &ErrLocked{
		Key:     e.key,
		Primary: e.lock.primary,
		StartTS: e.lock.startTS,
	}
}

func (e *mvccEntry) Get(ts uint64) ([]byte, error) {
	if e.lock != nil {
		if e.lock.startTS <= ts {
			return nil, e.lockErr()
		}
	}
	for _, v := range e.values {
		if v.commitTS <= ts {
			return v.value, nil
		}
	}
	return nil, nil
}

func (e *mvccEntry) Prewrite(mutation *kvrpcpb.Mutation, startTS uint64, primary []byte) error {
	if len(e.values) > 0 {
		if e.values[0].commitTS >= startTS {
			return ErrRetryable("write conflict")
		}
	}
	if e.lock != nil {
		if e.lock.startTS != startTS {
			return e.lockErr()
		}
		return nil
	}
	e.lock = &mvccLock{
		startTS: startTS,
		primary: primary,
		value:   mutation.Value,
		op:      mutation.GetOp(),
	}
	return nil
}

func (e *mvccEntry) checkTxnCommitted(startTS uint64) (uint64, bool) {
	for _, v := range e.values {
		if v.startTS == startTS {
			return v.commitTS, true
		}
	}
	return 0, false
}

func (e *mvccEntry) Commit(startTS, commitTS uint64) error {
	if e.lock == nil || e.lock.startTS != startTS {
		if _, ok := e.checkTxnCommitted(startTS); ok {
			return nil
		}
		return ErrRetryable("txn not found")
	}
	if e.lock.op != kvrpcpb.Op_Lock {
		e.values = append([]mvccValue{{
			startTS:  startTS,
			commitTS: commitTS,
			value:    e.lock.value,
		}}, e.values...)
	}
	e.lock = nil
	return nil
}

func (e *mvccEntry) Rollback(startTS uint64) error {
	if e.lock == nil || e.lock.startTS != startTS {
		if commitTS, ok := e.checkTxnCommitted(startTS); ok {
			return ErrAlreadyCommitted(commitTS)
		}
		return nil
	}
	e.lock = nil
	return nil
}

// MvccStore is an in-memory, multi-versioned, transaction-supported kv storage.
type MvccStore struct {
	mu   sync.RWMutex
	tree *llrb.LLRB
}

// NewMvccStore creates a MvccStore.
func NewMvccStore() *MvccStore {
	return &MvccStore{
		tree: llrb.New(),
	}
}

// Get reads a key by ts.
func (s *MvccStore) Get(key []byte, startTS uint64) ([]byte, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.get(key, startTS)
}

func (s *MvccStore) get(key []byte, startTS uint64) ([]byte, error) {
	entry := s.tree.Get(newEntry(key))
	if entry == nil {
		return nil, nil
	}
	return entry.(*mvccEntry).Get(startTS)
}

// A Pair is a KV pair read from MvccStore or an error if any occurs.
type Pair struct {
	Key   []byte
	Value []byte
	Err   error
}

// BatchGet gets values with keys and ts.
func (s *MvccStore) BatchGet(ks [][]byte, startTS uint64) []Pair {
	s.mu.RLock()
	defer s.mu.RUnlock()

	var pairs []Pair
	for _, k := range ks {
		val, err := s.get(k, startTS)
		if val == nil && err == nil {
			continue
		}
		pairs = append(pairs, Pair{
			Key:   k,
			Value: val,
			Err:   err,
		})
	}
	return pairs
}

// Scan reads up to limit numbers of Pairs from a key.
func (s *MvccStore) Scan(startKey, endKey []byte, limit int, startTS uint64) []Pair {
	s.mu.RLock()
	defer s.mu.RUnlock()

	var pairs []Pair
	iterator := func(item llrb.Item) bool {
		if len(pairs) >= limit {
			return false
		}
		k := item.(*mvccEntry).key
		if !regionContains(startKey, endKey, k) {
			return false
		}
		val, err := s.Get(k, startTS)
		if val != nil || err != nil {
			pairs = append(pairs, Pair{
				Key:   k,
				Value: val,<commit_msg>store/tikv: add missed function to pass test. (#1185)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package mocktikv

import (
	"bytes"
	"sync"

	"github.com/petar/GoLLRB/llrb"
	"github.com/pingcap/kvproto/pkg/kvrpcpb"
)

type mvccValue struct {
	startTS  uint64
	commitTS uint64
	value    []byte
}

type mvccLock struct {
	startTS uint64
	primary []byte
	value   []byte
	op      kvrpcpb.Op
}

type mvccEntry struct {
	key    []byte
	values []mvccValue
	lock   *mvccLock
}

func newEntry(key []byte) *mvccEntry {
	return &mvccEntry{
		key: key,
	}
}

func (e *mvccEntry) Less(than llrb.Item) bool {
	return bytes.Compare(e.key, than.(*mvccEntry).key) < 0
}

func (e *mvccEntry) lockErr() error {
	return &ErrLocked{
		Key:     e.key,
		Primary: e.lock.primary,
		StartTS: e.lock.startTS,
	}
}

func (e *mvccEntry) Get(ts uint64) ([]byte, error) {
	if e.lock != nil {
		if e.lock.startTS <= ts {
			return nil, e.lockErr()
		}
	}
	for _, v := range e.values {
		if v.commitTS <= ts {
			return v.value, nil
		}
	}
	return nil, nil
}

func (e *mvccEntry) Prewrite(mutation *kvrpcpb.Mutation, startTS uint64, primary []byte) error {
	if len(e.values) > 0 {
		if e.values[0].commitTS >= startTS {
			return ErrRetryable("write conflict")
		}
	}
	if e.lock != nil {
		if e.lock.startTS != startTS {
			return e.lockErr()
		}
		return nil
	}
	e.lock = &mvccLock{
		startTS: startTS,
		primary: primary,
		value:   mutation.Value,
		op:      mutation.GetOp(),
	}
	return nil
}

func (e *mvccEntry) checkTxnCommitted(startTS uint64) (uint64, bool) {
	for _, v := range e.values {
		if v.startTS == startTS {
			return v.commitTS, true
		}
	}
	return 0, false
}

func (e *mvccEntry) Commit(startTS, commitTS uint64) error {
	if e.lock == nil || e.lock.startTS != startTS {
		if _, ok := e.checkTxnCommitted(startTS); ok {
			return nil
		}
		return ErrRetryable("txn not found")
	}
	if e.lock.op != kvrpcpb.Op_Lock {
		e.values = append([]mvccValue{{
			startTS:  startTS,
			commitTS: commitTS,
			value:    e.lock.value,
		}}, e.values...)
	}
	e.lock = nil
	return nil
}

func (e *mvccEntry) Rollback(startTS uint64) error {
	if e.lock == nil || e.lock.startTS != startTS {
		if commitTS, ok := e.checkTxnCommitted(startTS); ok {
			return ErrAlreadyCommitted(commitTS)
		}
		return nil
	}
	e.lock = nil
	return nil
}

// MvccStore is an in-memory, multi-versioned, transaction-supported kv storage.
type MvccStore struct {
	mu   sync.RWMutex
	tree *llrb.LLRB
}

// NewMvccStore creates a MvccStore.
func NewMvccStore() *MvccStore {
	return &MvccStore{
		tree: llrb.New(),
	}
}

// Get reads a key by ts.
func (s *MvccStore) Get(key []byte, startTS uint64) ([]byte, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.get(key, startTS)
}

func (s *MvccStore) get(key []byte, startTS uint64) ([]byte, error) {
	entry := s.tree.Get(newEntry(key))
	if entry == nil {
		return nil, nil
	}
	return entry.(*mvccEntry).Get(startTS)
}

// A Pair is a KV pair read from MvccStore or an error if any occurs.
type Pair struct {
	Key   []byte
	Value []byte
	Err   error
}

// BatchGet gets values with keys and ts.
func (s *MvccStore) BatchGet(ks [][]byte, startTS uint64) []Pair {
	s.mu.RLock()
	defer s.mu.RUnlock()

	var pairs []Pair
	for _, k := range ks {
		val, err := s.get(k, startTS)
		if val == nil && err == nil {
			continue
		}
		pairs = append(pairs, Pair{
			Key:   k,
			Value: val,
			Err:   err,
		})
	}
	return pairs
}

func regionContains(startKey []byte, endKey []byte, key []byte) bool {
	return bytes.Compare(startKey, key) <= 0 &&
		(bytes.Compare(key, endKey) < 0 || len(endKey) == 0)
}

// Scan reads up to limit numbers of Pairs from a key.
func (s *MvccStore) Scan(startKey, endKey []byte, limit int, startTS uint64) []Pair {
	s.mu.RLock()
	defer s.mu.RUnlock()

	var pairs []Pair
	iterator := func(item llrb.Item) bool {
		if len(pairs) >= limit {
			return false
		}
		k := item.(*mvccEntry).key
		if !regionContains(startKey, endKey, k) {
			return false
		}
		val, err := s.Get(k, startTS)
		if val != nil || err != nil {
			pairs = append(pairs, Pair{
				Key:   k,
				Value: val,
{'fraction_non_alphanumeric': 0.1380196691998212, 'fraction_numerical': 0.009275815824765311, 'mean_word_length': 3.158457249070632, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, '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': '7293171', 'n_tokens_mistral': 3657, 'n_tokens_neox': 3432, 'n_words': 999}