github.com/dolthub/dolt/go@v0.40.5-0.20240520175717-68db7794bea6/store/prolly/tree/mutable_map.go (about)

     1  // Copyright 2021 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package tree
    16  
    17  import (
    18  	"context"
    19  
    20  	"github.com/dolthub/dolt/go/store/skip"
    21  )
    22  
    23  // MutableMap is a mutable prolly Static with ordered elements.
    24  type MutableMap[K, V ~[]byte, O Ordering[K]] struct {
    25  	Edits  *skip.List
    26  	Static StaticMap[K, V, O]
    27  }
    28  
    29  func (m MutableMap[K, V, O]) Put(_ context.Context, key K, value V) error {
    30  	m.Edits.Put(key, value)
    31  	return nil
    32  }
    33  
    34  func (m MutableMap[K, V, O]) Delete(_ context.Context, key K) error {
    35  	m.Edits.Put(key, nil)
    36  	return nil
    37  }
    38  
    39  func (m MutableMap[K, V, O]) Get(ctx context.Context, key K, cb KeyValueFn[K, V]) (err error) {
    40  	value, ok := m.Edits.Get(key)
    41  	if ok {
    42  		if value == nil {
    43  			key = nil // there is a pending delete of |key| in |m.Edits|.
    44  		}
    45  		return cb(key, value)
    46  	}
    47  
    48  	return m.Static.Get(ctx, key, cb)
    49  }
    50  
    51  func (m MutableMap[K, V, O]) GetPrefix(ctx context.Context, key K, prefixOrder O, cb KeyValueFn[K, V]) (err error) {
    52  	iter := m.Edits.GetIterFromSeekFn(func(k []byte) (advance bool) {
    53  		if k != nil { // seek until |k| >= |key|
    54  			advance = prefixOrder.Compare(k, key) < 0
    55  		}
    56  		return
    57  	})
    58  	k, v := iter.Current()
    59  	if k != nil && prefixOrder.Compare(k, key) == 0 {
    60  		if v == nil {
    61  			k = nil // there is a pending delete of |key| in |m.Edits|.
    62  		}
    63  		return cb(k, v)
    64  	}
    65  	return m.Static.GetPrefix(ctx, key, prefixOrder, cb)
    66  }
    67  
    68  func (m MutableMap[K, V, O]) Has(ctx context.Context, key K) (present bool, err error) {
    69  	value, ok := m.Edits.Get(key)
    70  	if ok {
    71  		present = value != nil
    72  		return
    73  	}
    74  	return m.Static.Has(ctx, key)
    75  }
    76  
    77  func (m MutableMap[K, V, O]) HasPrefix(ctx context.Context, key K, prefixOrder O) (present bool, err error) {
    78  	iter := m.Edits.GetIterFromSeekFn(func(k []byte) (advance bool) {
    79  		if k != nil { // seek until |k| >= |key|
    80  			advance = prefixOrder.Compare(k, key) < 0
    81  		}
    82  		return
    83  	})
    84  	k, v := iter.Current()
    85  	if k != nil && prefixOrder.Compare(k, key) == 0 {
    86  		present = v != nil
    87  		return
    88  	}
    89  	return m.Static.HasPrefix(ctx, key, prefixOrder)
    90  }
    91  
    92  func (m MutableMap[K, V, O]) Copy() MutableMap[K, V, O] {
    93  	return MutableMap[K, V, O]{
    94  		Edits:  m.Edits.Copy(),
    95  		Static: m.Static,
    96  	}
    97  }
    98  
    99  func (m MutableMap[K, V, O]) Mutations() MutationIter {
   100  	return orderedListIter[K, V]{iter: m.Edits.IterAtStart()}
   101  }
   102  
   103  type orderedListIter[K, V ~[]byte] struct {
   104  	iter *skip.ListIter
   105  }
   106  
   107  var _ MutationIter = &orderedListIter[Item, Item]{}
   108  
   109  func (it orderedListIter[K, V]) NextMutation(context.Context) (Item, Item) {
   110  	k, v := it.iter.Current()
   111  	if k == nil {
   112  		return nil, nil
   113  	}
   114  	it.iter.Advance()
   115  	return k, v
   116  }
   117  
   118  func (it orderedListIter[K, V]) Close() error {
   119  	return nil
   120  }