github.com/dolthub/go-mysql-server@v0.18.0/sql/in_mem_table/multimap.go (about)

     1  // Copyright 2023 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 in_mem_table
    16  
    17  type MultiMap[V any] struct {
    18  	Equals  func(v1, v2 V) bool
    19  	entries map[any][]V
    20  }
    21  
    22  func NewMultiMap[V any](f func(v1, v2 V) bool) MultiMap[V] {
    23  	return MultiMap[V]{
    24  		f,
    25  		make(map[any][]V),
    26  	}
    27  }
    28  
    29  func (m MultiMap[V]) GetMany(k any) []V {
    30  	if vs, ok := m.entries[k]; ok {
    31  		vscopy := make([]V, len(vs))
    32  		copy(vscopy, vs)
    33  		return vscopy
    34  	}
    35  	return nil
    36  }
    37  
    38  func (m MultiMap[V]) Get(k any, v V) (res V, found bool) {
    39  	if vs, ok := m.entries[k]; ok {
    40  		for _, vp := range vs {
    41  			if m.Equals(v, vp) {
    42  				return vp, true
    43  			}
    44  		}
    45  	}
    46  	return
    47  }
    48  
    49  func (m MultiMap[V]) Put(k any, v V) {
    50  	m.entries[k] = append(m.entries[k], v)
    51  }
    52  
    53  func (m MultiMap[V]) Clear() {
    54  	for k := range m.entries {
    55  		delete(m.entries, k)
    56  	}
    57  }
    58  
    59  func (m MultiMap[V]) Remove(k any, v V) (res V, found bool) {
    60  	if vs, ok := m.entries[k]; ok {
    61  		var newvs []V
    62  		for _, vp := range vs {
    63  			if !m.Equals(v, vp) {
    64  				newvs = append(newvs, vp)
    65  			} else {
    66  				res = v
    67  				found = true
    68  			}
    69  		}
    70  		if len(newvs) > 0 {
    71  			m.entries[k] = newvs
    72  		} else {
    73  			delete(m.entries, k)
    74  		}
    75  	}
    76  	return
    77  }
    78  
    79  func (m MultiMap[V]) VisitEntries(f func(v V)) {
    80  	for _, es := range m.entries {
    81  		for _, e := range es {
    82  			f(e)
    83  		}
    84  	}
    85  }
    86  
    87  type Keyer[V any] interface {
    88  	GetKey(V) any
    89  }
    90  
    91  type IndexedSet[V any] struct {
    92  	Keyers  []Keyer[V]
    93  	Indexes []MultiMap[V]
    94  }
    95  
    96  func NewIndexedSet[V any](eqf func(V, V) bool, keyers []Keyer[V]) IndexedSet[V] {
    97  	cp := make([]Keyer[V], len(keyers))
    98  	copy(cp, keyers)
    99  	idxs := make([]MultiMap[V], len(keyers))
   100  	for i := range idxs {
   101  		idxs[i] = NewMultiMap[V](eqf)
   102  	}
   103  	return IndexedSet[V]{
   104  		cp,
   105  		idxs,
   106  	}
   107  }
   108  
   109  func (is IndexedSet[V]) Put(v V) {
   110  	for i, keyer := range is.Keyers {
   111  		k := keyer.GetKey(v)
   112  		is.Indexes[i].Put(k, v)
   113  	}
   114  }
   115  
   116  func (is IndexedSet[V]) GetMany(keyer Keyer[V], k any) []V {
   117  	for i, x := range is.Keyers {
   118  		if x == keyer {
   119  			return is.Indexes[i].GetMany(k)
   120  		}
   121  	}
   122  	return nil
   123  }
   124  
   125  func (is IndexedSet[V]) Get(v V) (V, bool) {
   126  	k := is.Keyers[0].GetKey(v)
   127  	return is.Indexes[0].Get(k, v)
   128  }
   129  
   130  func (is IndexedSet[V]) Remove(v V) (res V, found bool) {
   131  	for i, keyer := range is.Keyers {
   132  		k := keyer.GetKey(v)
   133  		if fv, ok := is.Indexes[i].Remove(k, v); ok {
   134  			res = fv
   135  			found = true
   136  		}
   137  	}
   138  	return
   139  }
   140  
   141  func (is IndexedSet[V]) RemoveMany(keyer Keyer[V], k any) {
   142  	for i, x := range is.Keyers {
   143  		if x == keyer {
   144  			vs := is.Indexes[i].GetMany(k)
   145  			for _, v := range vs {
   146  				is.Remove(v)
   147  			}
   148  		}
   149  	}
   150  }
   151  
   152  func (is IndexedSet[V]) Count() int {
   153  	var c int
   154  	is.VisitEntries(func(V) {
   155  		c += 1
   156  	})
   157  	return c
   158  }
   159  
   160  func (is IndexedSet[V]) Clear() {
   161  	for _, i := range is.Indexes {
   162  		i.Clear()
   163  	}
   164  }
   165  
   166  func (is IndexedSet[V]) VisitEntries(f func(v V)) {
   167  	// Every multimap has every entry, so you we just iterate over the entries of the first one.
   168  	if len(is.Indexes) > 0 {
   169  		is.Indexes[0].VisitEntries(f)
   170  	}
   171  }