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