feat(kzip_merge): allow applying vname rules during merge (#4366)

* feat(kzip_merge): allow applying vname rules during merge

* fix(kzip_merge): address comments

kythe/go/platform/tools/kzip/mergecmd/BUILD

@@ -4,12 +4,17 @@ package(default_visibility = ["//kythe:default_visibility"])
 
 go_library(
     name = "mergecmd",
-    srcs = ["mergecmd.go"],
+    srcs = [
+        "flag.go",
+        "mergecmd.go",
+    ],
     deps = [
         "//kythe/go/platform/kzip",
         "//kythe/go/platform/tools/kzip/flags",
         "//kythe/go/platform/vfs",
         "//kythe/go/util/cmdutil",
+        "//kythe/go/util/vnameutil",
+        "@com_github_golang_protobuf//proto:go_default_library",
         "@com_github_google_subcommands//:go_default_library",
         "@org_bitbucket_creachadair_stringset//:go_default_library",
     ],

kythe/go/platform/tools/kzip/mergecmd/flag.go

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2020 The Kythe Authors. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package mergecmd
+
+import (
+	"fmt"
+	"io/ioutil"
+
+	"kythe.io/kythe/go/util/vnameutil"
+)
+
+// vnameRules is a path-valued flag used for loading VName mapping rules from a vnames.json file.
+type vnameRules struct {
+	filename string
+	vnameutil.Rules
+}
+
+// Set implements part of the flag.Value interface and will set a new filename for the flag.
+func (f *vnameRules) Set(s string) error {
+	f.filename = s
+	data, err := ioutil.ReadFile(f.filename)
+	if err != nil {
+		return fmt.Errorf("reading vname rules: %v", err)
+	}
+	f.Rules, err = vnameutil.ParseRules(data)
+	if err != nil {
+		return fmt.Errorf("reading vname rules: %v", err)
+	}
+	return nil
+}
+
+// String implements part of the flag.Value interface and returns a string-ish value for the flag.
+func (f *vnameRules) String() string {
+	return f.filename
+}
+
+// Get implements part of the flag.Getter interface.
+func (f *vnameRules) Get() interface{} {
+	return f
+}

kythe/go/platform/tools/kzip/mergecmd/mergecmd.go

@@ -43,6 +43,7 @@ type mergeCommand struct {
 	append    bool
 	encoding  flags.EncodingFlag
 	recursive bool
+	rules     vnameRules
 }
 
 // New creates a new subcommand for merging kzip files.
@@ -60,6 +61,7 @@ func (c *mergeCommand) SetFlags(fs *flag.FlagSet) {
 	fs.BoolVar(&c.append, "append", false, "Whether to additionally merge the contents of the existing output file, if it exists")
 	fs.Var(&c.encoding, "encoding", "Encoding to use on output, one of JSON, PROTO, or ALL")
 	fs.BoolVar(&c.recursive, "recursive", false, "Recurisvely merge .kzip files from directories")
+	fs.Var(&c.rules, "rules", "VName rules to apply while merging (optional)")
 }
 
 // Execute implements the subcommands interface and merges the provided files.
@@ -99,7 +101,7 @@ func (c *mergeCommand) Execute(ctx context.Context, fs *flag.FlagSet, _ ...inter
 			}
 		}
 	}
-	if err := mergeArchives(ctx, tmpOut, archives, opt); err != nil {
+	if err := c.mergeArchives(ctx, tmpOut, archives, opt); err != nil {
 		return c.Fail("Error merging archives: %v", err)
 	}
 	if err := vfs.Rename(ctx, tmpName, c.output); err != nil {
@@ -108,7 +110,7 @@ func (c *mergeCommand) Execute(ctx context.Context, fs *flag.FlagSet, _ ...inter
 	return subcommands.ExitSuccess
 }
 
-func mergeArchives(ctx context.Context, out io.WriteCloser, archives []string, opts ...kzip.WriterOption) error {
+func (c *mergeCommand) mergeArchives(ctx context.Context, out io.WriteCloser, archives []string, opts ...kzip.WriterOption) error {
 	wr, err := kzip.NewWriteCloser(out, opts...)
 	if err != nil {
 		out.Close()
@@ -117,7 +119,7 @@ func mergeArchives(ctx context.Context, out io.WriteCloser, archives []string, o
 
 	filesAdded := stringset.New()
 	for _, path := range archives {
-		if err := mergeInto(ctx, wr, path, filesAdded); err != nil {
+		if err := c.mergeInto(ctx, wr, path, filesAdded); err != nil {
 			wr.Close()
 			return err
 		}
@@ -129,7 +131,7 @@ func mergeArchives(ctx context.Context, out io.WriteCloser, archives []string, o
 	return nil
 }
 
-func mergeInto(ctx context.Context, wr *kzip.Writer, path string, filesAdded stringset.Set) error {
+func (c *mergeCommand) mergeInto(ctx context.Context, wr *kzip.Writer, path string, filesAdded stringset.Set) error {
 	f, err := vfs.Open(ctx, path)
 	if err != nil {
 		return fmt.Errorf("error opening archive: %v", err)
@@ -165,6 +167,9 @@ func mergeInto(ctx context.Context, wr *kzip.Writer, path string, filesAdded str
 					return fmt.Errorf("error closing file: %v", err)
 				}
 			}
+			if vname, match := c.rules.Apply(ri.Info.Path); match {
+				ri.VName = vname
+			}
 		}
 		// TODO(schroederc): duplicate compilations with different revisions
 		_, err = wr.AddUnit(u.Proto, u.Index)