github.com/nspcc-dev/neo-go@v0.105.2-0.20240517133400-6be757af3eba/pkg/core/interop/storage/find.go (about)

     1  package storage
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"fmt"
     7  
     8  	"github.com/nspcc-dev/neo-go/pkg/core/interop"
     9  	"github.com/nspcc-dev/neo-go/pkg/core/storage"
    10  	"github.com/nspcc-dev/neo-go/pkg/vm/stackitem"
    11  )
    12  
    13  // Storage iterator options.
    14  const (
    15  	FindDefault      = 0
    16  	FindKeysOnly     = 1 << 0
    17  	FindRemovePrefix = 1 << 1
    18  	FindValuesOnly   = 1 << 2
    19  	FindDeserialize  = 1 << 3
    20  	FindPick0        = 1 << 4
    21  	FindPick1        = 1 << 5
    22  	FindBackwards    = 1 << 7
    23  
    24  	FindAll = FindDefault | FindKeysOnly | FindRemovePrefix | FindValuesOnly |
    25  		FindDeserialize | FindPick0 | FindPick1 | FindBackwards
    26  )
    27  
    28  // Iterator is an iterator state representation.
    29  type Iterator struct {
    30  	seekCh chan storage.KeyValue
    31  	curr   storage.KeyValue
    32  	next   bool
    33  	opts   int64
    34  	// prefix is the storage item key prefix Find is performed with. It must be
    35  	// copied if no FindRemovePrefix option specified since it's shared between all
    36  	// iterator items.
    37  	prefix []byte
    38  }
    39  
    40  // NewIterator creates a new Iterator with the given options for the given channel of store.Seek results.
    41  func NewIterator(seekCh chan storage.KeyValue, prefix []byte, opts int64) *Iterator {
    42  	return &Iterator{
    43  		seekCh: seekCh,
    44  		opts:   opts,
    45  		prefix: bytes.Clone(prefix),
    46  	}
    47  }
    48  
    49  // Next advances the iterator and returns true if Value can be called at the
    50  // current position.
    51  func (s *Iterator) Next() bool {
    52  	s.curr, s.next = <-s.seekCh
    53  	return s.next
    54  }
    55  
    56  // Value returns current iterators value (exact type depends on options this
    57  // iterator was created with).
    58  func (s *Iterator) Value() stackitem.Item {
    59  	if !s.next {
    60  		panic("iterator index out of range")
    61  	}
    62  	key := s.curr.Key
    63  	if s.opts&FindRemovePrefix == 0 {
    64  		key = append(bytes.Clone(s.prefix), key...)
    65  	}
    66  	if s.opts&FindKeysOnly != 0 {
    67  		return stackitem.NewByteArray(key)
    68  	}
    69  	value := stackitem.Item(stackitem.NewByteArray(s.curr.Value))
    70  	if s.opts&FindDeserialize != 0 {
    71  		bs := s.curr.Value
    72  		var err error
    73  		value, err = stackitem.Deserialize(bs)
    74  		if err != nil {
    75  			panic(err)
    76  		}
    77  	}
    78  	if s.opts&FindPick0 != 0 {
    79  		value = value.Value().([]stackitem.Item)[0]
    80  	} else if s.opts&FindPick1 != 0 {
    81  		value = value.Value().([]stackitem.Item)[1]
    82  	}
    83  	if s.opts&FindValuesOnly != 0 {
    84  		return value
    85  	}
    86  	return stackitem.NewStruct([]stackitem.Item{
    87  		stackitem.NewByteArray(key),
    88  		value,
    89  	})
    90  }
    91  
    92  // Find finds stored key-value pair.
    93  func Find(ic *interop.Context) error {
    94  	stcInterface := ic.VM.Estack().Pop().Value()
    95  	stc, ok := stcInterface.(*Context)
    96  	if !ok {
    97  		return fmt.Errorf("%T is not a storage,Context", stcInterface)
    98  	}
    99  	prefix := ic.VM.Estack().Pop().Bytes()
   100  	opts := ic.VM.Estack().Pop().BigInt().Int64()
   101  	if opts&^FindAll != 0 {
   102  		return fmt.Errorf("%w: unknown flag", errFindInvalidOptions)
   103  	}
   104  	if opts&FindKeysOnly != 0 &&
   105  		opts&(FindDeserialize|FindPick0|FindPick1) != 0 {
   106  		return fmt.Errorf("%w KeysOnly conflicts with other options", errFindInvalidOptions)
   107  	}
   108  	if opts&FindValuesOnly != 0 &&
   109  		opts&(FindKeysOnly|FindRemovePrefix) != 0 {
   110  		return fmt.Errorf("%w: KeysOnly conflicts with ValuesOnly", errFindInvalidOptions)
   111  	}
   112  	if opts&FindPick0 != 0 && opts&FindPick1 != 0 {
   113  		return fmt.Errorf("%w: Pick0 conflicts with Pick1", errFindInvalidOptions)
   114  	}
   115  	if opts&FindDeserialize == 0 && (opts&FindPick0 != 0 || opts&FindPick1 != 0) {
   116  		return fmt.Errorf("%w: PickN is specified without Deserialize", errFindInvalidOptions)
   117  	}
   118  	bkwrds := opts&FindBackwards != 0
   119  	ctx, cancel := context.WithCancel(context.Background())
   120  	seekres := ic.DAO.SeekAsync(ctx, stc.ID, storage.SeekRange{Prefix: prefix, Backwards: bkwrds})
   121  	item := NewIterator(seekres, prefix, opts)
   122  	ic.VM.Estack().PushItem(stackitem.NewInterop(item))
   123  	ic.RegisterCancelFunc(func() {
   124  		cancel()
   125  		// Underlying persistent store is likely to be a private MemCachedStore. Thus,
   126  		// to avoid concurrent map iteration and map write we need to wait until internal
   127  		// seek goroutine is finished, because it can access underlying persistent store.
   128  		for range seekres { //nolint:revive //empty-block
   129  		}
   130  	})
   131  
   132  	return nil
   133  }