sink.go

package storm

import (
	"reflect"
	"sort"

	"github.com/asdine/storm/index"
	"github.com/asdine/storm/q"
	"github.com/coreos/bbolt"
)

type item struct {
	value  *reflect.Value
	bucket *bolt.Bucket
	k      []byte
	v      []byte
}

func newSorter(n Node, snk sink) *sorter {
	return &sorter{
		node:  n,
		sink:  snk,
		skip:  0,
		limit: -1,
		list:  make([]*item, 0),
		err:   make(chan error),
		done:  make(chan struct{}),
	}
}

type sorter struct {
	node    Node
	sink    sink
	list    []*item
	skip    int
	limit   int
	orderBy []string
	reverse bool
	err     chan error
	done    chan struct{}
}

func (s *sorter) filter(tree q.Matcher, bucket *bolt.Bucket, k, v []byte) (bool, error) {
	itm := &item{
		bucket: bucket,
		k:      k,
		v:      v,
	}
	rsink, ok := s.sink.(reflectSink)
	if !ok {
		return s.add(itm)
	}

	newElem := rsink.elem()
	if err := s.node.Codec().Unmarshal(v, newElem.Interface()); err != nil {
		return false, err
	}
	itm.value = &newElem

	if tree != nil {
		ok, err := tree.Match(newElem.Interface())
		if err != nil {
			return false, err
		}
		if !ok {
			return false, nil
		}
	}

	if len(s.orderBy) == 0 {
		return s.add(itm)
	}

	if _, ok := s.sink.(sliceSink); ok {
		// add directly to sink, we'll apply skip/limits after sorting
		return false, s.sink.add(itm)
	}

	s.list = append(s.list, itm)

	return false, nil
}

func (s *sorter) add(itm *item) (stop bool, err error) {
	if s.limit == 0 {
		return true, nil
	}

	if s.skip > 0 {
		s.skip--
		return false, nil
	}

	if s.limit > 0 {
		s.limit--
	}

	err = s.sink.add(itm)

	return s.limit == 0, err
}

func (s *sorter) compareValue(left reflect.Value, right reflect.Value) int {
	if !left.IsValid() || !right.IsValid() {
		if left.IsValid() {
			return 1
		}
		return -1
	}

	switch left.Kind() {
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		l, r := left.Int(), right.Int()
		if l < r {
			return -1
		}
		if l > r {
			return 1
		}
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		l, r := left.Uint(), right.Uint()
		if l < r {
			return -1
		}
		if l > r {
			return 1
		}
	case reflect.Float32, reflect.Float64:
		l, r := left.Float(), right.Float()
		if l < r {
			return -1
		}
		if l > r {
			return 1
		}
	case reflect.String:
		l, r := left.String(), right.String()
		if l < r {
			return -1
		}
		if l > r {
			return 1
		}
	default:
		rawLeft, err := toBytes(left.Interface(), s.node.Codec())
		if err != nil {
			return -1
		}
		rawRight, err := toBytes(right.Interface(), s.node.Codec())
		if err != nil {
			return 1
		}

		l, r := string(rawLeft), string(rawRight)
		if l < r {
			return -1
		}
		if l > r {
			return 1
		}
	}

	return 0
}

func (s *sorter) less(leftElem reflect.Value, rightElem reflect.Value) bool {
	for _, orderBy := range s.orderBy {
		leftField := reflect.Indirect(leftElem).FieldByName(orderBy)
		if !leftField.IsValid() {
			s.err <- ErrNotFound
			return false
		}
		rightField := reflect.Indirect(rightElem).FieldByName(orderBy)
		if !rightField.IsValid() {
			s.err <- ErrNotFound
			return false
		}

		direction := 1
		if s.reverse {
			direction = -1
		}

		switch s.compareValue(leftField, rightField) * direction {
		case -1:
			return true
		case 1:
			return false
		default:
			continue
		}
	}

	return false
}

func (s *sorter) flush() error {
	if len(s.orderBy) == 0 {
		return s.sink.flush()
	}

	go func() {
		sort.Sort(s)
		close(s.err)
	}()
	err := <-s.err
	close(s.done)

	if err != nil {
		return err
	}

	if ssink, ok := s.sink.(sliceSink); ok {
		if !ssink.slice().IsValid() {
			return s.sink.flush()
		}
		if s.skip >= ssink.slice().Len() {
			ssink.reset()
			return s.sink.flush()
		}
		leftBound := s.skip
		if leftBound < 0 {
			leftBound = 0
		}
		limit := s.limit
		if s.limit < 0 {
			limit = 0
		}

		rightBound := leftBound + limit
		if rightBound > ssink.slice().Len() || rightBound == leftBound {
			rightBound = ssink.slice().Len()
		}
		ssink.setSlice(ssink.slice().Slice(leftBound, rightBound))
		return s.sink.flush()
	}

	for _, itm := range s.list {
		if itm == nil {
			break
		}
		stop, err := s.add(itm)
		if err != nil {
			return err
		}
		if stop {
			break
		}
	}

	return s.sink.flush()
}

func (s *sorter) Len() int {
	// skip if we encountered an earlier error
	select {
	case <-s.done:
		return 0
	default:
	}
	if ssink, ok := s.sink.(sliceSink); ok {
		return ssink.slice().Len()
	}
	return len(s.list)

}

func (s *sorter) Less(i, j int) bool {
	// skip if we encountered an earlier error
	select {
	case <-s.done:
		return false
	default:
	}

	if ssink, ok := s.sink.(sliceSink); ok {
		return s.less(ssink.slice().Index(i), ssink.slice().Index(j))
	}
	return s.less(*s.list[i].value, *s.list[j].value)
}

type sink interface {
	bucketName() string
	flush() error
	add(*item) error
	readOnly() bool
}

type reflectSink interface {
	elem() reflect.Value
}

type sliceSink interface {
	slice() reflect.Value
	setSlice(reflect.Value)
	reset()
}

func newListSink(node Node, to interface{}) (*listSink, error) {
	ref := reflect.ValueOf(to)

	if ref.Kind() != reflect.Ptr || reflect.Indirect(ref).Kind() != reflect.Slice {
		return nil, ErrSlicePtrNeeded
	}

	sliceType := reflect.Indirect(ref).Type()
	elemType := sliceType.Elem()

	if elemType.Kind() == reflect.Ptr {
		elemType = elemType.Elem()
	}

	if elemType.Name() == "" {
		return nil, ErrNoName
	}

	return &listSink{
		node:     node,
		ref:      ref,
		isPtr:    sliceType.Elem().Kind() == reflect.Ptr,
		elemType: elemType,
		name:     elemType.Name(),
		results:  reflect.MakeSlice(reflect.Indirect(ref).Type(), 0, 0),
	}, nil
}

type listSink struct {
	node     Node
	ref      reflect.Value
	results  reflect.Value
	elemType reflect.Type
	name     string
	isPtr    bool
	idx      int
}

func (l *listSink) slice() reflect.Value {
	return l.results
}

func (l *listSink) setSlice(s reflect.Value) {
	l.results = s
}

func (l *listSink) reset() {
	l.results = reflect.MakeSlice(reflect.Indirect(l.ref).Type(), 0, 0)
}

func (l *listSink) elem() reflect.Value {
	if l.results.IsValid() && l.idx < l.results.Len() {
		return l.results.Index(l.idx).Addr()
	}
	return reflect.New(l.elemType)
}

func (l *listSink) bucketName() string {
	return l.name
}

func (l *listSink) add(i *item) error {
	if l.idx == l.results.Len() {
		if l.isPtr {
			l.results = reflect.Append(l.results, *i.value)
		} else {
			l.results = reflect.Append(l.results, reflect.Indirect(*i.value))
		}
	}

	l.idx++

	return nil
}

func (l *listSink) flush() error {
	if l.results.IsValid() && l.results.Len() > 0 {
		reflect.Indirect(l.ref).Set(l.results)
		return nil
	}

	return ErrNotFound
}

func (l *listSink) readOnly() bool {
	return true
}

func newFirstSink(node Node, to interface{}) (*firstSink, error) {
	ref := reflect.ValueOf(to)

	if !ref.IsValid() || ref.Kind() != reflect.Ptr || ref.Elem().Kind() != reflect.Struct {
		return nil, ErrStructPtrNeeded
	}

	return &firstSink{
		node: node,
		ref:  ref,
	}, nil
}

type firstSink struct {
	node  Node
	ref   reflect.Value
	found bool
}

func (f *firstSink) elem() reflect.Value {
	return reflect.New(reflect.Indirect(f.ref).Type())
}

func (f *firstSink) bucketName() string {
	return reflect.Indirect(f.ref).Type().Name()
}

func (f *firstSink) add(i *item) error {
	reflect.Indirect(f.ref).Set(i.value.Elem())
	f.found = true
	return nil
}

func (f *firstSink) flush() error {
	if !f.found {
		return ErrNotFound
	}

	return nil
}

func (f *firstSink) readOnly() bool {
	return true
}

func newDeleteSink(node Node, kind interface{}) (*deleteSink, error) {
	ref := reflect.ValueOf(kind)

	if !ref.IsValid() || ref.Kind() != reflect.Ptr || ref.Elem().Kind() != reflect.Struct {
		return nil, ErrStructPtrNeeded
	}

	return &deleteSink{
		node: node,
		ref:  ref,
	}, nil
}

type deleteSink struct {
	node    Node
	ref     reflect.Value
	removed int
}

func (d *deleteSink) elem() reflect.Value {
	return reflect.New(reflect.Indirect(d.ref).Type())
}

func (d *deleteSink) bucketName() string {
	return reflect.Indirect(d.ref).Type().Name()
}

func (d *deleteSink) add(i *item) error {
	info, err := extract(&d.ref)
	if err != nil {
		return err
	}

	for fieldName, fieldCfg := range info.Fields {
		if fieldCfg.Index == "" {
			continue
		}
		idx, err := getIndex(i.bucket, fieldCfg.Index, fieldName)
		if err != nil {
			return err
		}

		err = idx.RemoveID(i.k)
		if err != nil {
			if err == index.ErrNotFound {
				return ErrNotFound
			}
			return err
		}
	}

	d.removed++
	return i.bucket.Delete(i.k)
}

func (d *deleteSink) flush() error {
	if d.removed == 0 {
		return ErrNotFound
	}

	return nil
}

func (d *deleteSink) readOnly() bool {
	return false
}

func newCountSink(node Node, kind interface{}) (*countSink, error) {
	ref := reflect.ValueOf(kind)

	if !ref.IsValid() || ref.Kind() != reflect.Ptr || ref.Elem().Kind() != reflect.Struct {
		return nil, ErrStructPtrNeeded
	}

	return &countSink{
		node: node,
		ref:  ref,
	}, nil
}

type countSink struct {
	node    Node
	ref     reflect.Value
	counter int
}

func (c *countSink) elem() reflect.Value {
	return reflect.New(reflect.Indirect(c.ref).Type())
}

func (c *countSink) bucketName() string {
	return reflect.Indirect(c.ref).Type().Name()
}

func (c *countSink) add(i *item) error {
	c.counter++
	return nil
}

func (c *countSink) flush() error {
	return nil
}

func (c *countSink) readOnly() bool {
	return true
}

func newRawSink() *rawSink {
	return &rawSink{}
}

type rawSink struct {
	results [][]byte
	execFn  func([]byte, []byte) error
}

func (r *rawSink) add(i *item) error {
	if r.execFn != nil {
		err := r.execFn(i.k, i.v)
		if err != nil {
			return err
		}
	} else {
		r.results = append(r.results, i.v)
	}

	return nil
}

func (r *rawSink) bucketName() string {
	return ""
}

func (r *rawSink) flush() error {
	return nil
}

func (r *rawSink) readOnly() bool {
	return true
}

func newEachSink(to interface{}) (*eachSink, error) {
	ref := reflect.ValueOf(to)

	if !ref.IsValid() || ref.Kind() != reflect.Ptr || ref.Elem().Kind() != reflect.Struct {
		return nil, ErrStructPtrNeeded
	}

	return &eachSink{
		ref: ref,
	}, nil
}

type eachSink struct {
	ref    reflect.Value
	execFn func(interface{}) error
}

func (e *eachSink) elem() reflect.Value {
	return reflect.New(reflect.Indirect(e.ref).Type())
}

func (e *eachSink) bucketName() string {
	return reflect.Indirect(e.ref).Type().Name()
}

func (e *eachSink) add(i *item) error {
	return e.execFn(i.value.Interface())
}

func (e *eachSink) flush() error {
	return nil
}

func (e *eachSink) readOnly() bool {
	return true
}