metamorphic: support retaining prefixes when simplifying keys

internal/metamorphic/reduce_test.go

@@ -206,10 +206,14 @@ func (r *reducer) Run(t *testing.T) {
 		}
 	}
 	// Try to simplify the keys.
-	newOpsData, err := metamorphic.TryToSimplifyKeys([]byte(strings.Join(ops, "\n")))
-	require.NoError(t, err)
-	o := strings.Split(strings.TrimSpace(string(newOpsData)), "\n")
-	r.try(t, o)
+	opsData := []byte(strings.Join(ops, "\n"))
+	for _, retainSuffixes := range []bool{false, true} {
+		newOpsData := metamorphic.TryToSimplifyKeys(opsData, retainSuffixes)
+		o := strings.Split(strings.TrimSpace(string(newOpsData)), "\n")
+		if r.try(t, o) {
+			return
+		}
+	}
 }
 
 func randomSubset(t *testing.T, ops []string, removeProbability float64) []string {

metamorphic/simplify.go

@@ -16,36 +16,46 @@ import (
 //
 // On success it returns the new operations data.
 //
-// If there are too many distinct keys, returns nil (and no error).
-func TryToSimplifyKeys(opsData []byte) ([]byte, error) {
+// If there are too many distinct keys, returns nil.
+func TryToSimplifyKeys(opsData []byte, retainSuffixes bool) []byte {
 	ops, err := parse(opsData, parserOpts{})
 	if err != nil {
-		return nil, err
+		panic(err)
 	}
 	keys := make(map[string]struct{})
 	for i := range ops {
 		for _, k := range ops[i].keys() {
-			keys[string(*k)] = struct{}{}
+			key := *k
+			if retainSuffixes {
+				key = key[:testkeys.Comparer.Split(key)]
+			}
+			keys[string(key)] = struct{}{}
 		}
 	}
 	if len(keys) > ('z' - 'a' + 1) {
-		return nil, nil
+		return nil
 	}
 	sorted := sortedKeys(keys)
 	ordinals := make(map[string]int, len(sorted))
 	for i, k := range sorted {
 		ordinals[k] = i
 	}
-	// TODO(radu): We reassign the user keys, ignoring the prefix and suffix
-	// composition. This is sufficient to reproduce a class of problems, but if it
-	// fails, we should try to simplify just the prefix and retain the suffix.
 	for i := range ops {
 		for _, k := range ops[i].keys() {
-			idx := ordinals[string(*k)]
-			*k = []byte{'a' + byte(idx)}
+			key := *k
+			var suffix []byte
+			if retainSuffixes {
+				n := testkeys.Comparer.Split(key)
+				suffix = key[n:]
+				key = key[:n]
+			}
+			idx := ordinals[string(key)]
+			newKey := []byte{'a' + byte(idx)}
+			newKey = append(newKey, suffix...)
+			*k = newKey
 		}
 	}
-	return []byte(formatOps(ops)), nil
+	return []byte(formatOps(ops))
 }
 
 func sortedKeys(in map[string]struct{}) []string {

metamorphic/simplify_test.go

@@ -15,10 +15,8 @@ func TestSimplifyKeys(t *testing.T) {
 	datadriven.RunTest(t, "testdata/simplify", func(t *testing.T, d *datadriven.TestData) string {
 		switch d.Cmd {
 		case "simplify-keys":
-			res, err := TryToSimplifyKeys([]byte(d.Input))
-			if err != nil {
-				return err.Error()
-			}
+			retainSuffixes := d.HasArg("retain-suffixes")
+			res := TryToSimplifyKeys([]byte(d.Input), retainSuffixes)
 			return string(res)
 
 		default:

metamorphic/testdata/simplify

@@ -15,3 +15,21 @@ db2.Compact("apple", "raspberry", true /* parallelize */)
 db2.RangeKeySet("a", "b", "", "")
 snap9 = db2.NewSnapshot("a", "b")
 db2.Compact("a", "b", true /* parallelize */)
+
+simplify-keys
+db2.RangeKeySet("apple@1", "raspberry", "", "")
+snap9 = db2.NewSnapshot("apple@1", "raspberry@4")
+db2.Compact("apple@1", "raspberry@4", true /* parallelize */)
+----
+db2.RangeKeySet("a", "b", "", "")
+snap9 = db2.NewSnapshot("a", "c")
+db2.Compact("a", "c", true /* parallelize */)
+
+simplify-keys retain-suffixes
+db2.RangeKeySet("apple@1", "raspberry", "", "")
+snap9 = db2.NewSnapshot("apple@1", "raspberry@4")
+db2.Compact("apple@1", "raspberry@4", true /* parallelize */)
+----
+db2.RangeKeySet("a@1", "b", "", "")
+snap9 = db2.NewSnapshot("a@1", "b@4")
+db2.Compact("a@1", "b@4", true /* parallelize */)