diff --git a/kythe/go/platform/tools/kzip/mergecmd/mergecmd.go b/kythe/go/platform/tools/kzip/mergecmd/mergecmd.go index 9f359efdc7..9184a7ced2 100644 --- a/kythe/go/platform/tools/kzip/mergecmd/mergecmd.go +++ b/kythe/go/platform/tools/kzip/mergecmd/mergecmd.go @@ -45,6 +45,8 @@ type mergeCommand struct { recursive bool ignoreDuplicateCUs bool rules vnameRules + + unitsBeforeFiles bool } // New creates a new subcommand for merging kzip files. @@ -64,6 +66,7 @@ func (c *mergeCommand) SetFlags(fs *flag.FlagSet) { fs.BoolVar(&c.recursive, "recursive", false, "Recurisvely merge .kzip files from directories") fs.Var(&c.rules, "rules", "VName rules to apply while merging (optional)") fs.BoolVar(&c.ignoreDuplicateCUs, "ignore_duplicate_cus", false, "Do not fail if we try to add the same CU twice") + fs.BoolVar(&c.unitsBeforeFiles, "experimental_write_units_first", false, "When writing the kzip file, puts CU entries before files") } // Execute implements the subcommands interface and merges the provided files. @@ -155,26 +158,38 @@ func (c *mergeCommand) mergeInto(ctx context.Context, wr *kzip.Writer, path stri return fmt.Errorf("error creating reader: %v", err) } + if c.unitsBeforeFiles { + var requiredDigests []string + if err := c.mergeUnitsInto(ctx, wr, rd, func(digest string) error { + requiredDigests = append(requiredDigests, digest) + return nil + }); err != nil { + return err + } + for _, digest := range requiredDigests { + if err := copyFileInto(wr, rd, digest, filesAdded); err != nil { + return err + } + } + return nil + } + return c.mergeUnitsInto(ctx, wr, rd, func(digest string) error { + return copyFileInto(wr, rd, digest, filesAdded) + }) +} + +func (c *mergeCommand) mergeUnitsInto(ctx context.Context, wr *kzip.Writer, rd *kzip.Reader, f func(digest string) error) error { return rd.Scan(func(u *kzip.Unit) error { for _, ri := range u.Proto.RequiredInput { - if filesAdded.Add(ri.Info.Digest) { - r, err := rd.Open(ri.Info.Digest) - if err != nil { - return fmt.Errorf("error opening file: %v", err) - } - if _, err := wr.AddFile(r); err != nil { - r.Close() - return fmt.Errorf("error adding file: %v", err) - } else if err := r.Close(); err != nil { - return fmt.Errorf("error closing file: %v", err) - } + if err := f(ri.Info.Digest); err != nil { + return 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) + _, err := wr.AddUnit(u.Proto, u.Index) if c.ignoreDuplicateCUs && err == kzip.ErrUnitExists { log.Printf("Found duplicate CU: %v", u.Proto.GetDetails()) return nil @@ -183,6 +198,22 @@ func (c *mergeCommand) mergeInto(ctx context.Context, wr *kzip.Writer, path stri }) } +func copyFileInto(wr *kzip.Writer, rd *kzip.Reader, digest string, filesAdded stringset.Set) error { + if filesAdded.Add(digest) { + r, err := rd.Open(digest) + if err != nil { + return fmt.Errorf("error opening file: %v", err) + } + if _, err := wr.AddFile(r); err != nil { + r.Close() + return fmt.Errorf("error adding file: %v", err) + } else if err := r.Close(); err != nil { + return fmt.Errorf("error closing file: %v", err) + } + } + return nil +} + func recurseDirectories(ctx context.Context, archives []string) ([]string, error) { var files []string for _, path := range archives {