Line data Source code
1 : // Copyright 2019 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 (
8 : "context"
9 :
10 : "github.com/cockroachdb/pebble/internal/base"
11 : "github.com/cockroachdb/pebble/internal/invariants"
12 : "github.com/cockroachdb/pebble/internal/treeprinter"
13 : )
14 :
15 : // Truncate creates a new iterator where every span in the supplied iterator is
16 : // truncated to be contained within the given user key bounds.
17 : //
18 : // Note that fragment iterator Spans always have exclusive end-keys; if the
19 : // given bounds have an inclusive end key, then the input iterator must not
20 : // produce a span that contains that key. The only difference between bounds.End
21 : // being inclusive vs exclusive is this extra check.
22 1 : func Truncate(cmp base.Compare, iter FragmentIterator, bounds base.UserKeyBounds) FragmentIterator {
23 1 : return &truncatingIter{
24 1 : iter: iter,
25 1 : cmp: cmp,
26 1 : bounds: bounds,
27 1 : }
28 1 : }
29 :
30 : type truncatingIter struct {
31 : iter FragmentIterator
32 : cmp base.Compare
33 :
34 : bounds base.UserKeyBounds
35 :
36 : span Span
37 : }
38 :
39 : // SeekGE implements FragmentIterator.
40 1 : func (i *truncatingIter) SeekGE(key []byte) (*Span, error) {
41 1 : span, err := i.iter.SeekGE(key)
42 1 : if err != nil {
43 0 : return nil, err
44 0 : }
45 1 : span, spanBoundsChanged, err := i.nextSpanWithinBounds(span, +1)
46 1 : if err != nil {
47 0 : return nil, err
48 0 : }
49 : // nextSpanWithinBounds could return a span that's less than key, if the end
50 : // bound was truncated to end at a key less than or equal to `key`. Detect
51 : // this case and next/invalidate the iter.
52 1 : if spanBoundsChanged && i.cmp(span.End, key) <= 0 {
53 1 : return i.Next()
54 1 : }
55 1 : return span, nil
56 : }
57 :
58 : // SeekLT implements FragmentIterator.
59 1 : func (i *truncatingIter) SeekLT(key []byte) (*Span, error) {
60 1 : span, err := i.iter.SeekLT(key)
61 1 : if err != nil {
62 0 : return nil, err
63 0 : }
64 1 : span, spanBoundsChanged, err := i.nextSpanWithinBounds(span, -1)
65 1 : if err != nil {
66 0 : return nil, err
67 0 : }
68 : // nextSpanWithinBounds could return a span that's >= key, if the start bound
69 : // was truncated to start at a key greater than or equal to `key`. Detect this
70 : // case and prev/invalidate the iter.
71 1 : if spanBoundsChanged && i.cmp(span.Start, key) >= 0 {
72 1 : return i.Prev()
73 1 : }
74 1 : return span, nil
75 : }
76 :
77 : // First implements FragmentIterator.
78 1 : func (i *truncatingIter) First() (*Span, error) {
79 1 : span, err := i.iter.First()
80 1 : if err != nil {
81 0 : return nil, err
82 0 : }
83 1 : span, _, err = i.nextSpanWithinBounds(span, +1)
84 1 : return span, err
85 : }
86 :
87 : // Last implements FragmentIterator.
88 1 : func (i *truncatingIter) Last() (*Span, error) {
89 1 : span, err := i.iter.Last()
90 1 : if err != nil {
91 0 : return nil, err
92 0 : }
93 1 : span, _, err = i.nextSpanWithinBounds(span, -1)
94 1 : return span, err
95 : }
96 :
97 : // Next implements FragmentIterator.
98 1 : func (i *truncatingIter) Next() (*Span, error) {
99 1 : span, err := i.iter.Next()
100 1 : if err != nil {
101 0 : return nil, err
102 0 : }
103 1 : span, _, err = i.nextSpanWithinBounds(span, +1)
104 1 : return span, err
105 : }
106 :
107 : // Prev implements FragmentIterator.
108 1 : func (i *truncatingIter) Prev() (*Span, error) {
109 1 : span, err := i.iter.Prev()
110 1 : if err != nil {
111 0 : return nil, err
112 0 : }
113 1 : span, _, err = i.nextSpanWithinBounds(span, -1)
114 1 : return span, err
115 : }
116 :
117 : // SetContext is part of the FragmentIterator interface.
118 0 : func (i *truncatingIter) SetContext(ctx context.Context) {
119 0 : i.iter.SetContext(ctx)
120 0 : }
121 :
122 : // Close implements FragmentIterator.
123 1 : func (i *truncatingIter) Close() {
124 1 : i.iter.Close()
125 1 : }
126 :
127 : // nextSpanWithinBounds returns the first span (starting with the given span and
128 : // advancing in the given direction) that intersects the bounds. It returns a
129 : // span that is entirely within the bounds; spanBoundsChanged indicates if the span
130 : // bounds had to be truncated.
131 : func (i *truncatingIter) nextSpanWithinBounds(
132 : span *Span, dir int8,
133 1 : ) (_ *Span, spanBoundsChanged bool, _ error) {
134 1 : var err error
135 1 : for span != nil {
136 1 : if i.bounds.End.Kind == base.Inclusive && span.Contains(i.cmp, i.bounds.End.Key) {
137 0 : err := base.AssertionFailedf("inclusive upper bound %q inside span %s", i.bounds.End.Key, span)
138 0 : if invariants.Enabled {
139 0 : panic(err)
140 : }
141 0 : return nil, false, err
142 : }
143 : // Intersect [span.Start, span.End) with [i.bounds.Start, i.bounds.End.Key).
144 1 : spanBoundsChanged = false
145 1 : start := span.Start
146 1 : if i.cmp(start, i.bounds.Start) < 0 {
147 1 : spanBoundsChanged = true
148 1 : start = i.bounds.Start
149 1 : }
150 1 : end := span.End
151 1 : if i.cmp(end, i.bounds.End.Key) > 0 {
152 1 : spanBoundsChanged = true
153 1 : end = i.bounds.End.Key
154 1 : }
155 1 : if !spanBoundsChanged {
156 1 : return span, false, nil
157 1 : }
158 1 : if i.cmp(start, end) < 0 {
159 1 : i.span = Span{
160 1 : Start: start,
161 1 : End: end,
162 1 : Keys: span.Keys,
163 1 : KeysOrder: span.KeysOrder,
164 1 : }
165 1 : return &i.span, true, nil
166 1 : }
167 : // Span is outside of bounds, find the next one.
168 1 : if dir == +1 {
169 1 : span, err = i.iter.Next()
170 1 : } else {
171 1 : span, err = i.iter.Prev()
172 1 : }
173 : }
174 : // NB: err may be nil or non-nil.
175 1 : return nil, false, err
176 : }
177 :
178 : // WrapChildren implements FragmentIterator.
179 0 : func (i *truncatingIter) WrapChildren(wrap WrapFn) {
180 0 : i.iter = wrap(i.iter)
181 0 : }
182 :
183 : // DebugTree is part of the FragmentIterator interface.
184 0 : func (i *truncatingIter) DebugTree(tp treeprinter.Node) {
185 0 : n := tp.Childf("%T(%p)", i, i)
186 0 : if i.iter != nil {
187 0 : i.iter.DebugTree(n)
188 0 : }
189 : }
|