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 {
      44           1 :         span := i.filter(i.iter.SeekGE(key), +1)
      45           1 :         // i.filter could return a span that's less than key, _if_ the filterFunc
      46           1 :         // (which has no knowledge of the seek key) mutated the span to end at a key
      47           1 :         // less than or equal to `key`. Detect this case and next/invalidate the iter.
      48           1 :         if span != nil && i.cmp(span.End, key) <= 0 {
      49           1 :                 return i.Next()
      50           1 :         }
      51           1 :         return span
      52             : }
      53             : 
      54             : // SeekLT implements FragmentIterator.
      55           1 : func (i *filteringIter) SeekLT(key []byte) *Span {
      56           1 :         span := i.filter(i.iter.SeekLT(key), -1)
      57           1 :         // i.filter could return a span that's >= key, _if_ the filterFunc (which has
      58           1 :         // no knowledge of the seek key) mutated the span to start at a key greater
      59           1 :         // than or equal to `key`. Detect this case and prev/invalidate the iter.
      60           1 :         if span != nil && i.cmp(span.Start, key) >= 0 {
      61           1 :                 return i.Prev()
      62           1 :         }
      63           1 :         return span
      64             : }
      65             : 
      66             : // First implements FragmentIterator.
      67           1 : func (i *filteringIter) First() *Span {
      68           1 :         return i.filter(i.iter.First(), +1)
      69           1 : }
      70             : 
      71             : // Last implements FragmentIterator.
      72           1 : func (i *filteringIter) Last() *Span {
      73           1 :         return i.filter(i.iter.Last(), -1)
      74           1 : }
      75             : 
      76             : // Next implements FragmentIterator.
      77           1 : func (i *filteringIter) Next() *Span {
      78           1 :         return i.filter(i.iter.Next(), +1)
      79           1 : }
      80             : 
      81             : // Prev implements FragmentIterator.
      82           1 : func (i *filteringIter) Prev() *Span {
      83           1 :         return i.filter(i.iter.Prev(), -1)
      84           1 : }
      85             : 
      86             : // Error implements FragmentIterator.
      87           1 : func (i *filteringIter) Error() error {
      88           1 :         return i.iter.Error()
      89           1 : }
      90             : 
      91             : // Close implements FragmentIterator.
      92           1 : func (i *filteringIter) Close() error {
      93           1 :         return i.iter.Close()
      94           1 : }
      95             : 
      96             : // filter uses the filterFn (if configured) to filter and possibly mutate the
      97             : // given Span. If the current Span is to be skipped, the iterator continues
      98             : // iterating in the given direction until it lands on a Span that should be
      99             : // returned, or the iterator becomes invalid.
     100           1 : func (i *filteringIter) filter(span *Span, dir int8) *Span {
     101           1 :         if i.filterFn == nil {
     102           1 :                 return span
     103           1 :         }
     104           1 :         for i.Error() == nil && span != nil {
     105           1 :                 if keep := i.filterFn(span, &i.span); keep {
     106           1 :                         return &i.span
     107           1 :                 }
     108           1 :                 if dir == +1 {
     109           1 :                         span = i.iter.Next()
     110           1 :                 } else {
     111           1 :                         span = i.iter.Prev()
     112           1 :                 }
     113             :         }
     114           1 :         return span
     115             : }
     116             : 
     117             : // WrapChildren implements FragmentIterator.
     118           0 : func (i *filteringIter) WrapChildren(wrap WrapFn) {
     119           0 :         i.iter = wrap(i.iter)
     120           0 : }