Line data    Source code

       1             : // Copyright 2024 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 metamorphic
       6             : 
       7             : import (
       8             :         "sort"
       9             : 
      10             :         "github.com/cockroachdb/pebble/internal/base"
      11             : )
      12             : 
      13             : // TryToSimplifyKeys parses the operations data and tries to reassign keys to
      14             : // single lowercase characters. Note that the result is not necessarily
      15             : // semantically equivalent.
      16             : //
      17             : // On success it returns the new operations data.
      18             : //
      19             : // If there are too many distinct keys, returns nil.
      20           0 : func TryToSimplifyKeys(keyFormat KeyFormat, opsData []byte, retainSuffixes bool) []byte {
      21           0 :         ops, err := parse(opsData, parserOpts{
      22           0 :                 parseFormattedUserKey:       keyFormat.ParseFormattedKey,
      23           0 :                 parseFormattedUserKeySuffix: keyFormat.ParseFormattedKeySuffix,
      24           0 :         })
      25           0 :         if err != nil {
      26           0 :                 panic(err)
      27             :         }
      28           0 :         keys := make(map[string]struct{})
      29           0 :         for i := range ops {
      30           0 :                 ops[i].rewriteKeys(func(k UserKey) UserKey {
      31           0 :                         if retainSuffixes {
      32           0 :                                 keys[string(k[:keyFormat.Comparer.Split(k)])] = struct{}{}
      33           0 :                         } else {
      34           0 :                                 keys[string(k)] = struct{}{}
      35           0 :                         }
      36           0 :                         return k
      37             :                 })
      38             :         }
      39           0 :         if len(keys) > ('z' - 'a' + 1) {
      40           0 :                 return nil
      41           0 :         }
      42           0 :         sorted := sortedKeys(keyFormat.Comparer.Compare, keys)
      43           0 :         ordinals := make(map[string]int, len(sorted))
      44           0 :         for i, k := range sorted {
      45           0 :                 ordinals[k] = i
      46           0 :         }
      47           0 :         for i := range ops {
      48           0 :                 ops[i].rewriteKeys(func(k UserKey) UserKey {
      49           0 :                         var suffix []byte
      50           0 :                         if retainSuffixes {
      51           0 :                                 n := keyFormat.Comparer.Split(k)
      52           0 :                                 suffix = k[n:]
      53           0 :                                 k = k[:n]
      54           0 :                         }
      55           0 :                         idx := ordinals[string(k)]
      56           0 :                         newKey := []byte{'a' + byte(idx)}
      57           0 :                         return append(newKey, suffix...)
      58             :                 })
      59             :         }
      60           0 :         return []byte(formatOps(keyFormat, ops))
      61             : }
      62             : 
      63           0 : func sortedKeys(cmp base.Compare, in map[string]struct{}) []string {
      64           0 :         var sorted []string
      65           0 :         for k := range in {
      66           0 :                 sorted = append(sorted, k)
      67           0 :         }
      68           0 :         sort.Slice(sorted, func(i, j int) bool {
      69           0 :                 return cmp([]byte(sorted[i]), []byte(sorted[j])) < 0
      70           0 :         })
      71           0 :         return sorted
      72             : }