forgejo/vendor/github.com/pingcap/tidb/kv/btree_buffer.go

143 lines
3 KiB
Go
Raw Normal View History

// Copyright 2015 PingCAP, Inc.
//
// Copyright 2015 Wenbin Xiao
//
// 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 kv
import (
"io"
"github.com/juju/errors"
"github.com/pingcap/tidb/kv/memkv"
"github.com/pingcap/tidb/terror"
"github.com/pingcap/tidb/util/types"
)
type btreeBuffer struct {
tree *memkv.Tree
}
// NewBTreeBuffer returns a breeBuffer.
func NewBTreeBuffer() MemBuffer {
return &btreeBuffer{
tree: memkv.NewTree(types.Collators[true]),
}
}
// Get returns the value associated with the key; ErrNotExist error if the key does not exist.
func (b *btreeBuffer) Get(k Key) ([]byte, error) {
v, ok := b.tree.Get(toIfaces(k))
if !ok {
return nil, ErrNotExist
}
return fromIfaces(v), nil
}
// Set associates the key with the value.
func (b *btreeBuffer) Set(k Key, v []byte) error {
if len(v) == 0 {
return errors.Trace(ErrCannotSetNilValue)
}
b.tree.Set(toIfaces(k), toIfaces(v))
return nil
}
// Delete removes the entry from buffer with provided key.
func (b *btreeBuffer) Delete(k Key) error {
b.tree.Set(toIfaces(k), nil)
return nil
}
// Release clear the whole buffer.
func (b *btreeBuffer) Release() {
b.tree.Clear()
}
type btreeIter struct {
e *memkv.Enumerator
k Key
v []byte
ok bool
}
// Seek creates a new Iterator based on the provided key.
func (b *btreeBuffer) Seek(k Key) (Iterator, error) {
var e *memkv.Enumerator
var err error
if k == nil {
e, err = b.tree.SeekFirst()
if err != nil {
if terror.ErrorEqual(err, io.EOF) {
return &btreeIter{ok: false}, nil
}
return &btreeIter{ok: false}, errors.Trace(err)
}
} else {
key := toIfaces([]byte(k))
e, _ = b.tree.Seek(key)
}
iter := &btreeIter{e: e}
// the initial push...
err = iter.Next()
if err != nil {
return &btreeIter{ok: false}, errors.Trace(err)
}
return iter, nil
}
// Close implements Iterator Close.
func (i *btreeIter) Close() {
//noop
}
// Key implements Iterator Key.
func (i *btreeIter) Key() Key {
return i.k
}
// Value implements Iterator Value.
func (i *btreeIter) Value() []byte {
return i.v
}
// Next implements Iterator Next.
func (i *btreeIter) Next() error {
k, v, err := i.e.Next()
if err != nil {
i.ok = false
if terror.ErrorEqual(err, io.EOF) {
return nil
}
return errors.Trace(err)
}
i.k, i.v, i.ok = fromIfaces(k), fromIfaces(v), true
return nil
}
// Valid implements Iterator Valid.
func (i *btreeIter) Valid() bool {
return i.ok
}
func toIfaces(v []byte) []interface{} {
return []interface{}{v}
}
func fromIfaces(v []interface{}) []byte {
if v == nil {
return nil
}
return v[0].([]byte)
}