[llvm-ar] Support N [count] modifier

[rupprecht]

Summary

GNU ar supports the 'N' count modifier for the extract (x) and delete (d) operations. When an archive contains multiple members with the same name, this can be used to extract (or delete) them individually. For example:

$ llvm-ar t archive.a
foo
foo
$ llvm-ar x archive.a
-> Writes foo twice, overwriting it the second time :( :(
$ llvm-ar xN 1 archive.a foo && mv foo foo.1
$ llvm-ar xN 2 archive.a foo && mv foo foo.2
-> Write foo twice, renaming it in between invocations to preserve all versions

[ruiu]

Wow, this option is too good... I'm using Unix more than 20 years now but didn't know this option existed. In rare occasions where I had to extract a file with a duplicate name, I opened an archive file using a binary editor and rename a file name before extracting.

llvm/tools/llvm-ar/llvm-ar.cpp

194

If this is 1-indexed, you could use 0 to indicate that the option was not given, which allows you to remove UseCountParam. (But this is perhaps more like a personal preference, so you can ignore this comment.)

[rupprecht]

• Remove UseCountParam, as CountParam != 0 implies the same thing

[ruiu]

LGTM

llvm/test/tools/llvm-ar/count.test

+# Test the 'N' count parameter.
+
+# Get a temp clean cwd to extract into.
+RUN: rm -rf %t && mkdir -p %t && cd %t
+
+RUN: mkdir -p %t/x %t/y %t/z
+RUN: echo hello > %t/x/foo.txt
+RUN: echo cool  > %t/y/foo.txt
+RUN: echo world > %t/z/foo.txt
+RUN: echo fizz   > %t/x/bar.txt
+RUN: echo buzz   > %t/y/bar.txt
+RUN: echo fizbuz > %t/z/bar.txt
+RUN: llvm-ar rc %t/archive.a %t/x/foo.txt %t/y/foo.txt %t/z/foo.txt \
+RUN:     %t/x/bar.txt %t/y/bar.txt %t/z/bar.txt
+RUN: llvm-ar t %t/archive.a | FileCheck %s --check-prefix=LIST-MEMBERS
+
+# Make sure we set it up correctly.
+LIST-MEMBERS:      foo.txt
+LIST-MEMBERS-NEXT: foo.txt
+LIST-MEMBERS-NEXT: foo.txt
+LIST-MEMBERS-NEXT: bar.txt
+LIST-MEMBERS-NEXT: bar.txt
+LIST-MEMBERS-NEXT: bar.txt
+
+# Must be a number.
+RUN: not llvm-ar xN abc %t/archive.a foo.txt 2>&1 | FileCheck %s --check-prefix=ERR-NOT-NUM
+RUN: not llvm-ar xN 0x1 %t/archive.a foo.txt 2>&1 | FileCheck %s --check-prefix=ERR-NOT-NUM
+# Only three members named foo, so 1 <= N <= 3.
+RUN: not llvm-ar xN 0 %t/archive.a foo.txt 2>&1 | FileCheck %s --check-prefix=ERR-NOT-POS
+RUN: not llvm-ar xN 4 %t/archive.a foo.txt 2>&1 | FileCheck %s --check-prefix=ERR-NOT-FOUND
+# N only applies to x/d.
+RUN: not llvm-ar rN 1 %t/archive.a foo.txt 2>&1 | FileCheck %s --check-prefix=ERR-BAD-OP
+
+ERR-NOT-NUM:   error: Value for [count] must be numeric
+ERR-NOT-POS:   error: Value for [count] must be positive
+ERR-BAD-OP:    error: The 'N' modifier can only be specified with the 'x' or 'd' operations
+ERR-NOT-FOUND: error: 'foo.txt' was not found
+
+# Extract individual items.
+
+RUN: rm -f foo.txt bar.txt
+RUN: llvm-ar xN 1 %t/archive.a foo.txt bar.txt
+RUN: cat %t/foo.txt | FileCheck %s --check-prefix=FOO-1
+RUN: cat %t/bar.txt | FileCheck %s --check-prefix=BAR-1
+
+RUN: rm -f foo.txt bar.txt
+RUN: llvm-ar xN 2 %t/archive.a foo.txt bar.txt
+RUN: cat %t/foo.txt | FileCheck %s --check-prefix=FOO-2
+RUN: cat %t/bar.txt | FileCheck %s --check-prefix=BAR-2
+
+RUN: rm -f foo.txt bar.txt
+RUN: llvm-ar xN 3 %t/archive.a foo.txt bar.txt
+RUN: cat %t/foo.txt | FileCheck %s --check-prefix=FOO-3
+RUN: cat %t/bar.txt | FileCheck %s --check-prefix=BAR-3
+
+# Delete individual items.
+
+# Deleting the second member named foo means the new second member of the
+# archive is what used to be the third element.
+RUN: rm -f foo.txt bar.txt
+RUN: llvm-ar dN 2 %t/archive.a foo.txt
+RUN: llvm-ar xN 2 %t/archive.a foo.txt bar.txt
+RUN: cat %t/foo.txt | FileCheck %s --check-prefix=FOO-3
+RUN: cat %t/bar.txt | FileCheck %s --check-prefix=BAR-2
+
+# Deleting the first member from *both* archives means the new first member
+# named foo is the what used to be the third member, and the new first member
+# named bar is what used to be the second member.
+RUN: rm -f foo.txt bar.txt
+RUN: llvm-ar dN 1 %t/archive.a foo.txt bar.txt
+RUN: llvm-ar xN 1 %t/archive.a foo.txt bar.txt
+RUN: cat %t/foo.txt | FileCheck %s --check-prefix=FOO-3
+RUN: cat %t/bar.txt | FileCheck %s --check-prefix=BAR-2
+
+FOO-1: hello
+FOO-2: cool
+FOO-3: world
+BAR-1: fizz
+BAR-2: buzz
+BAR-3: fizbuz

llvm/tools/llvm-ar/llvm-ar.cpp

 OPTIONS:
   -help                             - Display available options
   -version                          - Display the version of this program
 )";
 
 const char ArHelp[] = R"(
 OVERVIEW: LLVM Archiver
 
-USAGE: llvm-ar [options] [-]<operation>[modifiers] [relpos] <archive> [files]
+USAGE: llvm-ar [options] [-]<operation>[modifiers] [relpos] [count] <archive> [files]
        llvm-ar -M [<mri-script]
 
 OPTIONS:
   --format              - Archive format to create
     =default            -   default
     =gnu                -   gnu
     =darwin             -   darwin
     =bsd                -   bsd

 MODIFIERS:
   [a] - put [files] after [relpos]
   [b] - put [files] before [relpos] (same as [i])
   [c] - do not warn if archive had to be created
   [D] - use zero for timestamps and uids/gids (default)
   [i] - put [files] before [relpos] (same as [b])
   [l] - ignored for compatibility
   [L] - add archive's contents
+  [N] - use instance [count] of name
   [o] - preserve original dates
   [P] - use full names when matching (implied for thin archives)
   [s] - create an archive index (cf. ranlib)
   [S] - do not build a symbol table
   [T] - create a thin archive
   [u] - update only [files] newer than archive contents
   [U] - use actual timestamps and uids/gids
   [v] - be verbose about actions taken

 static bool AddBefore = false;       ///< 'b' modifier
 static bool Create = false;          ///< 'c' modifier
 static bool OriginalDates = false;   ///< 'o' modifier
 static bool CompareFullPath = false; ///< 'P' modifier
 static bool OnlyUpdate = false;      ///< 'u' modifier
 static bool Verbose = false;         ///< 'v' modifier
 static bool Symtab = true;           ///< 's' modifier
 static bool Deterministic = true;    ///< 'D' and 'U' modifiers
+static bool UseCountParam = false;   ///< 'N' modifier
 static bool Thin = false;            ///< 'T' modifier
 static bool AddLibrary = false;      ///< 'L' modifier
 
 // Relative Positional Argument (for insert/move). This variable holds
 // the name of the archive member to which the 'a', 'b' or 'i' modifier
 // refers. Only one of 'a', 'b' or 'i' can be specified so we only need
 // one variable.
 static std::string RelPos;
 
+// Count parameter for 'N' modifier. This variable specifies which file should
+// match for extract/delete operations when there are multiple matches. This is
+// 1-indexed.
+static int CountParam;
+
 // This variable holds the name of the archive file as given on the
 // command line.
 static std::string ArchiveName;
 
 // This variable holds the list of member files to proecess, as given
 // on the command line.
 static std::vector<StringRef> Members;
 
 // Static buffer to hold StringRefs.
 static BumpPtrAllocator Alloc;
 
 // Extract the member filename from the command line for the [relpos] argument
 // associated with a, b, and i modifiers
 static void getRelPos() {
   if (PositionalArgs.empty())
     fail("Expected [relpos] for a, b, or i modifier");
   RelPos = PositionalArgs[0];
   PositionalArgs.erase(PositionalArgs.begin());
 }
 
+// Extract the parameter from the command line for the [count] argument
+// associated with the N modifier
+static void getCountParam() {
+  if (PositionalArgs.empty())
+    fail("Expected [count] for N modifier");
+  auto CountParamArg = StringRef(PositionalArgs[0]);
+  if (CountParamArg.getAsInteger(10, CountParam))
+    fail("Value for [count] must be numeric, got: " + CountParamArg);
+  if (CountParam < 1)
+    fail("Value for [count] must be positive, got: " + CountParamArg);
+  PositionalArgs.erase(PositionalArgs.begin());
+}
+
 // Get the archive file name from the command line
 static void getArchive() {
   if (PositionalArgs.empty())
     fail("An archive name must be specified");
   ArchiveName = PositionalArgs[0];
   PositionalArgs.erase(PositionalArgs.begin());
 }
 
case 'i':
       NumPositional++;
       break;
     case 'D':
       Deterministic = true;
       break;
     case 'U':
       Deterministic = false;
       break;
+    case 'N':
+      getCountParam();
+      UseCountParam = true;
+      break;
     case 'T':
       Thin = true;
       // Thin archives store path names, so P should be forced.
       CompareFullPath = true;
       break;
     case 'L':
       AddLibrary = true;
       break;
static ArchiveOperation parseCommandLine() {
   // Perform various checks on the operation/modifier specification
   // to make sure we are dealing with a legal request.
   if (NumOperations == 0)
     fail("You must specify at least one of the operations");
   if (NumOperations > 1)
     fail("Only one operation may be specified");
   if (NumPositional > 1)
     fail("You may only specify one of a, b, and i modifiers");
-  if (AddAfter || AddBefore) {
+  if (AddAfter || AddBefore)
     if (Operation != Move && Operation != ReplaceOrInsert)
       fail("The 'a', 'b' and 'i' modifiers can only be specified with "
            "the 'm' or 'r' operations");
-  }
+  if (UseCountParam)
+    if (Operation != Extract && Operation != Delete)
+      fail("The 'N' modifier can only be specified with the 'x' or 'd' "
+           "operations");
   if (OriginalDates && Operation != Extract)
     fail("The 'o' modifier is only applicable to the 'x' operation");
   if (OnlyUpdate && Operation != ReplaceOrInsert)
     fail("The 'u' modifier is only applicable to the 'r' operation");
   if (AddLibrary && Operation != QuickAppend)
     fail("The 'L' modifier is only applicable to the 'q' operation");
 
   // Return the parsed operation to the caller

 }
 
 static void performReadOperation(ArchiveOperation Operation,
                                  object::Archive *OldArchive) {
   if (Operation == Extract && OldArchive->isThin())
     fail("extracting from a thin archive is not supported");
 
   bool Filter = !Members.empty();
+  StringMap<int> MemberCount;
   {
     Error Err = Error::success();
     for (auto &C : OldArchive->children(Err)) {
       Expected<StringRef> NameOrErr = C.getName();
       failIfError(NameOrErr.takeError());
       StringRef Name = NameOrErr.get();
 
       if (Filter) {
         auto I = find_if(Members, [Name](StringRef Path) {
           return Name == normalizePath(Path);
         });
         if (I == Members.end())
           continue;
+        if (UseCountParam && ++MemberCount[Name] != CountParam)
+          continue;
         Members.erase(I);
       }
 
       switch (Operation) {
       default:
         llvm_unreachable("Not a read operation");
       case Print:
         doPrint(Name, C);
enum InsertAction {
   IA_Delete,
   IA_MoveOldMember,
   IA_MoveNewMember
 };
 
 static InsertAction computeInsertAction(ArchiveOperation Operation,
                                         const object::Archive::Child &Member,
                                         StringRef Name,
-                                        std::vector<StringRef>::iterator &Pos) {
+                                        std::vector<StringRef>::iterator &Pos,
+                                        StringMap<int> &MemberCount) {
   if (Operation == QuickAppend || Members.empty())
     return IA_AddOldMember;
-
   auto MI = find_if(
       Members, [Name](StringRef Path) { return Name == normalizePath(Path); });
 
   if (MI == Members.end())
     return IA_AddOldMember;
 
   Pos = MI;
 
-  if (Operation == Delete)
+  if (Operation == Delete) {
+    if (UseCountParam && ++MemberCount[Name] != CountParam)
+      return IA_AddOldMember;
     return IA_Delete;
+  }
 
   if (Operation == Move)
     return IA_MoveOldMember;
 
   if (Operation == ReplaceOrInsert) {
     StringRef PosName = normalizePath(RelPos);
     if (!OnlyUpdate) {
       if (PosName.empty())

 computeNewArchiveMembers(ArchiveOperation Operation,
                          object::Archive *OldArchive) {
   std::vector<NewArchiveMember> Ret;
   std::vector<NewArchiveMember> Moved;
   int InsertPos = -1;
   StringRef PosName = normalizePath(RelPos);
   if (OldArchive) {
     Error Err = Error::success();
+    StringMap<int> MemberCount;
     for (auto &Child : OldArchive->children(Err)) {
       int Pos = Ret.size();
       Expected<StringRef> NameOrErr = Child.getName();
       failIfError(NameOrErr.takeError());
       StringRef Name = NameOrErr.get();
       if (Name == PosName) {
         assert(AddAfter || AddBefore);
         if (AddBefore)
           InsertPos = Pos;
         else
           InsertPos = Pos + 1;
       }
 
       std::vector<StringRef>::iterator MemberI = Members.end();
       InsertAction Action =
-          computeInsertAction(Operation, Child, Name, MemberI);
+          computeInsertAction(Operation, Child, Name, MemberI, MemberCount);
       switch (Action) {
       case IA_AddOldMember:
         addChildMember(Ret, Child, /*FlattenArchive=*/Thin);
         break;
       case IA_AddNewMember:
         addMember(Ret, *MemberI);
         break;
       case IA_Delete:
         break;
       case IA_MoveOldMember:
         addChildMember(Moved, Child, /*FlattenArchive=*/Thin);
         break;
       case IA_MoveNewMember:
         addMember(Moved, *MemberI);
         break;
       }
-      if (MemberI != Members.end())
+      // When processing elements with the count param, we need to preserve the
+      // full members list when iterating over all archive members. For
+      // instance, "llvm-ar dN 2 archive.a member.o" should delete the second
+      // file named member.o it sees; we are not done with member.o the first
+      // time we see it in the archive.
+      if (MemberI != Members.end() && !UseCountParam)
         Members.erase(MemberI);
     }
     failIfError(std::move(Err));
   }
 
   if (Operation == Delete)
     return Ret;