package storm import ( "encoding/binary" "reflect" "github.com/asdine/storm/index" "github.com/asdine/storm/q" "github.com/boltdb/bolt" rbt "github.com/emirpasic/gods/trees/redblacktree" ) type item struct { value *reflect.Value bucket *bolt.Bucket k []byte v []byte } func newSorter(node Node) *sorter { return &sorter{ node: node, rbTree: rbt.NewWithStringComparator(), } } type sorter struct { node Node rbTree *rbt.Tree orderBy string reverse bool counter int64 } func (s *sorter) filter(snk sink, tree q.Matcher, bucket *bolt.Bucket, k, v []byte) (bool, error) { s.counter++ rsnk, ok := snk.(reflectSink) if !ok { return snk.add(&item{ bucket: bucket, k: k, v: v, }) } newElem := rsnk.elem() err := s.node.Codec().Unmarshal(v, newElem.Interface()) if err != nil { return false, err } ok = tree == nil if !ok { ok, err = tree.Match(newElem.Interface()) if err != nil { return false, err } } if ok { it := item{ bucket: bucket, value: &newElem, k: k, v: v, } if s.orderBy != "" { elm := reflect.Indirect(newElem).FieldByName(s.orderBy) if !elm.IsValid() { return false, ErrNotFound } raw, err := toBytes(elm.Interface(), s.node.Codec()) if err != nil { return false, err } key := make([]byte, len(raw)+8) for i := 0; i < len(raw); i++ { key[i] = raw[i] } binary.PutVarint(key[len(raw):], s.counter) s.rbTree.Put(string(key), &it) return false, nil } return snk.add(&it) } return false, nil } func (s *sorter) flush(snk sink) error { if s.orderBy == "" { return snk.flush() } s.orderBy = "" var err error var stop bool it := s.rbTree.Iterator() if s.reverse { it.End() } else { it.Begin() } for (s.reverse && it.Prev()) || (!s.reverse && it.Next()) { item := it.Value().(*item) stop, err = snk.add(item) if err != nil { return err } if stop { break } } return snk.flush() } type sink interface { bucketName() string flush() error add(*item) (bool, error) readOnly() bool } type reflectSink interface { elem() reflect.Value } 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(), limit: -1, }, nil } type listSink struct { node Node ref reflect.Value results reflect.Value elemType reflect.Type name string isPtr bool skip int limit int idx int } 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) (bool, error) { if l.limit == 0 { return true, nil } if l.skip > 0 { l.skip-- return false, nil } if !l.results.IsValid() { l.results = reflect.MakeSlice(reflect.Indirect(l.ref).Type(), 0, 0) } if l.limit > 0 { l.limit-- } 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 l.limit == 0, 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 skip int 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) (bool, error) { if f.skip > 0 { f.skip-- return false, nil } reflect.Indirect(f.ref).Set(i.value.Elem()) f.found = true return true, 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 skip int limit int 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) (bool, error) { if d.skip > 0 { d.skip-- return false, nil } if d.limit > 0 { d.limit-- } info, err := extract(&d.ref) if err != nil { return false, err } for fieldName, fieldCfg := range info.Fields { if fieldCfg.Index == "" { continue } idx, err := getIndex(i.bucket, fieldCfg.Index, fieldName) if err != nil { return false, err } err = idx.RemoveID(i.k) if err != nil { if err == index.ErrNotFound { return false, ErrNotFound } return false, err } } d.removed++ return d.limit == 0, 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 skip int limit int 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) (bool, error) { if c.skip > 0 { c.skip-- return false, nil } if c.limit > 0 { c.limit-- } c.counter++ return c.limit == 0, nil } func (c *countSink) flush() error { return nil } func (c *countSink) readOnly() bool { return true } func newRawSink() *rawSink { return &rawSink{ limit: -1, } } type rawSink struct { results [][]byte skip int limit int execFn func([]byte, []byte) error } func (r *rawSink) add(i *item) (bool, error) { if r.limit == 0 { return true, nil } if r.skip > 0 { r.skip-- return false, nil } if r.limit > 0 { r.limit-- } if r.execFn != nil { err := r.execFn(i.k, i.v) if err != nil { return false, err } } else { r.results = append(r.results, i.v) } return r.limit == 0, 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 { skip int limit int 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) (bool, error) { if e.limit == 0 { return true, nil } if e.skip > 0 { e.skip-- return false, nil } if e.limit > 0 { e.limit-- } err := e.execFn(i.value.Interface()) if err != nil { return false, err } return e.limit == 0, nil } func (e *eachSink) flush() error { return nil } func (e *eachSink) readOnly() bool { return true }