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 is a callback that allows filtering keys from a Span. The result
      10             : // is the set of keys that should be retained (using buf as a buffer). If the
      11             : // result has no keys, the span is skipped altogether.
      12             : type FilterFunc func(span *Span, buf []Key) []Key
      13             : 
      14             : // filteringIter is a FragmentIterator that uses a FilterFunc to select which
      15             : // Spans from the input iterator are returned in the output.
      16             : //
      17             : // A note on Span lifetimes: as the FilterFunc reuses a Span with a mutable
      18             : // slice of Keys to reduce allocations, Spans returned by this iterator are only
      19             : // valid until the next relative or absolute positioning method is called.
      20             : type filteringIter struct {
      21             :         iter     FragmentIterator
      22             :         filterFn FilterFunc
      23             :         cmp      base.Compare
      24             : 
      25             :         // span is a mutable Span passed to the filterFn. The filterFn is free to
      26             :         // mutate this Span. The slice of Keys in the Span is reused with every call
      27             :         // to the filterFn.
      28             :         span Span
      29             : }
      30             : 
      31             : var _ FragmentIterator = (*filteringIter)(nil)
      32             : 
      33             : // Filter returns a new filteringIter that will filter the Spans from the
      34             : // provided child iterator using the provided FilterFunc.
      35           1 : func Filter(iter FragmentIterator, filter FilterFunc, cmp base.Compare) FragmentIterator {
      36           1 :         return MaybeAssert(&filteringIter{iter: iter, filterFn: filter, cmp: cmp}, cmp)
      37           1 : }
      38             : 
      39             : // SeekGE implements FragmentIterator.
      40           0 : func (i *filteringIter) SeekGE(key []byte) (*Span, error) {
      41           0 :         s, err := i.iter.SeekGE(key)
      42           0 :         if err != nil {
      43           0 :                 return nil, err
      44           0 :         }
      45           0 :         return i.filter(s, +1)
      46             : }
      47             : 
      48             : // SeekLT implements FragmentIterator.
      49           0 : func (i *filteringIter) SeekLT(key []byte) (*Span, error) {
      50           0 :         span, err := i.iter.SeekLT(key)
      51           0 :         if err != nil {
      52           0 :                 return nil, err
      53           0 :         }
      54           0 :         return i.filter(span, -1)
      55             : }
      56             : 
      57             : // First implements FragmentIterator.
      58           1 : func (i *filteringIter) First() (*Span, error) {
      59           1 :         s, err := i.iter.First()
      60           1 :         if err != nil {
      61           0 :                 return nil, err
      62           0 :         }
      63           1 :         return i.filter(s, +1)
      64             : }
      65             : 
      66             : // Last implements FragmentIterator.
      67           0 : func (i *filteringIter) Last() (*Span, error) {
      68           0 :         s, err := i.iter.Last()
      69           0 :         if err != nil {
      70           0 :                 return nil, err
      71           0 :         }
      72           0 :         return i.filter(s, -1)
      73             : }
      74             : 
      75             : // Next implements FragmentIterator.
      76           1 : func (i *filteringIter) Next() (*Span, error) {
      77           1 :         s, err := i.iter.Next()
      78           1 :         if err != nil {
      79           0 :                 return nil, err
      80           0 :         }
      81           1 :         return i.filter(s, +1)
      82             : }
      83             : 
      84             : // Prev implements FragmentIterator.
      85           0 : func (i *filteringIter) Prev() (*Span, error) {
      86           0 :         s, err := i.iter.Prev()
      87           0 :         if err != nil {
      88           0 :                 return nil, err
      89           0 :         }
      90           0 :         return i.filter(s, -1)
      91             : }
      92             : 
      93             : // Close implements FragmentIterator.
      94           1 : func (i *filteringIter) Close() {
      95           1 :         i.iter.Close()
      96           1 : }
      97             : 
      98             : // filter uses the filterFn (if configured) to filter and possibly mutate the
      99             : // given Span. If the current Span is to be skipped, the iterator continues
     100             : // iterating in the given direction until it lands on a Span that should be
     101             : // returned, or the iterator becomes invalid.
     102           1 : func (i *filteringIter) filter(span *Span, dir int8) (*Span, error) {
     103           1 :         if i.filterFn == nil {
     104           0 :                 return span, nil
     105           0 :         }
     106           1 :         var err error
     107           1 :         for span != nil {
     108           1 :                 keys := i.filterFn(span, i.span.Keys[:0])
     109           1 :                 if len(keys) > 0 {
     110           1 :                         i.span = Span{
     111           1 :                                 Start:     span.Start,
     112           1 :                                 End:       span.End,
     113           1 :                                 Keys:      keys,
     114           1 :                                 KeysOrder: span.KeysOrder,
     115           1 :                         }
     116           1 :                         return &i.span, nil
     117           1 :                 }
     118             : 
     119           1 :                 if dir == +1 {
     120           1 :                         span, err = i.iter.Next()
     121           1 :                 } else {
     122           0 :                         span, err = i.iter.Prev()
     123           0 :                 }
     124             :         }
     125             :         // NB: err may be nil or non-nil.
     126           1 :         return span, err
     127             : }
     128             : 
     129             : // WrapChildren implements FragmentIterator.
     130           0 : func (i *filteringIter) WrapChildren(wrap WrapFn) {
     131           0 :         i.iter = wrap(i.iter)
     132           0 : }