Line data Source code
1 : // Copyright 2022 The LevelDB-Go and Pebble Authors. All rights reserved. Use
2 : // of this source code is governed by a BSD-style license that can be found in
3 : // the LICENSE file.
4 :
5 : package keyspan
6 :
7 : import "github.com/cockroachdb/pebble/internal/base"
8 :
9 : // FilterFunc defines a transform from the input Span into the output Span. The
10 : // function returns true if the Span should be returned by the iterator, and
11 : // false if the Span should be skipped. The FilterFunc is permitted to mutate
12 : // the output Span, for example, to elice certain keys, or update the Span's
13 : // bounds if so desired. The output Span's Keys slice may be reused to reduce
14 : // allocations.
15 : type FilterFunc func(in *Span, out *Span) (keep bool)
16 :
17 : // filteringIter is a FragmentIterator that uses a FilterFunc to select which
18 : // Spans from the input iterator are returned in the output.
19 : //
20 : // A note on Span lifetimes: as the FilterFunc reuses a Span with a mutable
21 : // slice of Keys to reduce allocations, Spans returned by this iterator are only
22 : // valid until the next relative or absolute positioning method is called.
23 : type filteringIter struct {
24 : iter FragmentIterator
25 : filterFn FilterFunc
26 : cmp base.Compare
27 :
28 : // span is a mutable Span passed to the filterFn. The filterFn is free to
29 : // mutate this Span. The slice of Keys in the Span is reused with every call
30 : // to the filterFn.
31 : span Span
32 : }
33 :
34 : var _ FragmentIterator = (*filteringIter)(nil)
35 :
36 : // Filter returns a new filteringIter that will filter the Spans from the
37 : // provided child iterator using the provided FilterFunc.
38 1 : func Filter(iter FragmentIterator, filter FilterFunc, cmp base.Compare) FragmentIterator {
39 1 : return MaybeAssert(&filteringIter{iter: iter, filterFn: filter, cmp: cmp}, cmp)
40 1 : }
41 :
42 : // SeekGE implements FragmentIterator.
43 1 : func (i *filteringIter) SeekGE(key []byte) (*Span, error) {
44 1 : s, err := i.iter.SeekGE(key)
45 1 : if err != nil {
46 0 : return nil, err
47 0 : }
48 1 : s, err = i.filter(s, +1)
49 1 : if err != nil {
50 0 : return nil, err
51 0 : }
52 : // i.filter could return a span that's less than key, _if_ the filterFunc
53 : // (which has no knowledge of the seek key) mutated the span to end at a key
54 : // less than or equal to `key`. Detect this case and next/invalidate the iter.
55 1 : if s != nil && i.cmp(s.End, key) <= 0 {
56 1 : return i.Next()
57 1 : }
58 1 : return s, nil
59 : }
60 :
61 : // SeekLT implements FragmentIterator.
62 1 : func (i *filteringIter) SeekLT(key []byte) (*Span, error) {
63 1 : span, err := i.iter.SeekLT(key)
64 1 : if err != nil {
65 0 : return nil, err
66 0 : }
67 1 : span, err = i.filter(span, -1)
68 1 : if err != nil {
69 0 : return nil, err
70 0 : }
71 : // i.filter could return a span that's >= key, _if_ the filterFunc (which has
72 : // no knowledge of the seek key) mutated the span to start at a key greater
73 : // than or equal to `key`. Detect this case and prev/invalidate the iter.
74 1 : if span != nil && i.cmp(span.Start, key) >= 0 {
75 1 : return i.Prev()
76 1 : }
77 1 : return span, nil
78 : }
79 :
80 : // First implements FragmentIterator.
81 1 : func (i *filteringIter) First() (*Span, error) {
82 1 : s, err := i.iter.First()
83 1 : if err != nil {
84 0 : return nil, err
85 0 : }
86 1 : return i.filter(s, +1)
87 : }
88 :
89 : // Last implements FragmentIterator.
90 1 : func (i *filteringIter) Last() (*Span, error) {
91 1 : s, err := i.iter.Last()
92 1 : if err != nil {
93 0 : return nil, err
94 0 : }
95 1 : return i.filter(s, -1)
96 : }
97 :
98 : // Next implements FragmentIterator.
99 1 : func (i *filteringIter) Next() (*Span, error) {
100 1 : s, err := i.iter.Next()
101 1 : if err != nil {
102 0 : return nil, err
103 0 : }
104 1 : return i.filter(s, +1)
105 : }
106 :
107 : // Prev implements FragmentIterator.
108 1 : func (i *filteringIter) Prev() (*Span, error) {
109 1 : s, err := i.iter.Prev()
110 1 : if err != nil {
111 0 : return nil, err
112 0 : }
113 1 : return i.filter(s, -1)
114 : }
115 :
116 : // Close implements FragmentIterator.
117 1 : func (i *filteringIter) Close() error {
118 1 : return i.iter.Close()
119 1 : }
120 :
121 : // filter uses the filterFn (if configured) to filter and possibly mutate the
122 : // given Span. If the current Span is to be skipped, the iterator continues
123 : // iterating in the given direction until it lands on a Span that should be
124 : // returned, or the iterator becomes invalid.
125 1 : func (i *filteringIter) filter(span *Span, dir int8) (*Span, error) {
126 1 : if i.filterFn == nil {
127 0 : return span, nil
128 0 : }
129 1 : var err error
130 1 : for span != nil {
131 1 : if keep := i.filterFn(span, &i.span); keep {
132 1 : return &i.span, nil
133 1 : }
134 1 : if dir == +1 {
135 1 : span, err = i.iter.Next()
136 1 : } else {
137 1 : span, err = i.iter.Prev()
138 1 : }
139 : }
140 : // NB: err may be nil or non-nil.
141 1 : return span, err
142 : }
143 :
144 : // WrapChildren implements FragmentIterator.
145 0 : func (i *filteringIter) WrapChildren(wrap WrapFn) {
146 0 : i.iter = wrap(i.iter)
147 0 : }
|
