[Object] Create MutableELFObject Class for Doing Mutations on ELFObjectFiles [Part 4]

[abrachet]

Summary

Added addSection and addSymbol methods

[jhenderson]

I think some other interesting test cases are a) adding more than one symbol/section, b) adding and then removing a symbol/section, c) removing stuff and then adding stuff, showing that the iteration still works correctly.

llvm/include/llvm/Object/MutableELFObject.h

172

My instinct is that this should be a void return type.

340

I think this should just be a void return type.

345

I think this should just be a void return type.

393–398

This seems like an unrelated change?

llvm/unittests/Object/MutableELFObjectTest.cpp

568–570

Comment why there's 4 sections already.

589–591

Comment why there's already a symbol here.

593

It's not a header. Symbols are just that - symbols (unlike sections which have a header and actual contents). So this should just be called "Sym" or something like that.

[abrachet]

Addressed review comments

[abrachet]

llvm/include/llvm/Object/MutableELFObject.h

172

Hmm, I'm not sure about this and the other comments. I think it is more convenient this way, but also I wonder when you would add something and then immediately change something. The cost of this return is very small though I would say.

I don't have a strong opinion either way.

[jhenderson]

llvm/include/llvm/Object/MutableELFObject.h

172

My thinking is that add is akin to vector's push_back method, which doesn't return anything.

llvm/unittests/Object/MutableELFObjectTest.cpp

568

"yaml2obj creates the .symtab ..."

590

"... creates the index 0 symbol even when ..."

636

for (size_t I = 0; I < 10000; ++I) would seem more natural.

I'd also reduce it to 1000 or 100 for speed.

647

AddRemoveSymbol

682

AddRemoveSection

711–728

Could you add some new lines in this block between the groups of setup followed by checks, to make it easier to see where the interesting tests actually are, please.

[abrachet]

Addressed review comments

[abrachet]

Fixed failing test for std::distance call.

[jhenderson]

LGTM, though @rupprecht should confirm too.

[abrachet]

rebase

[MaskRay]

llvm/include/llvm/Object/MutableELFObject.h

340

void addSymbol(MutableELFSymbol<ELFT> Sym)

If you don't have a lvalue overload, taking an rvalue is usually inferior to pass-by-value because you'll not able to call addSymbol with a lvalue.

343

void addSection(MutableELFSection<ELFT> Sym)

I think the method is self-explanatory. No need for a comment.

[MaskRay]

llvm/include/llvm/Object/MutableELFObject.h

42

Omit Other, but why do you need the defaulted move ctor?

174

Swap the two statements to omit - 1.

[abrachet]

Addressed comments. Made unit tests better.

[jhenderson]

Latest changes LGTM.

llvm/include/llvm/Object/MutableELFObject.h

template <typename ELFT> struct MutableELFSection {
   MutableELFSection(const Elf_Shdr &Header, StringRef Name, void *Base)
       : Header(Header), Name(Name),
         Data(OwningArrayRef<uint8_t>(Header.sh_size)) {
     std::memcpy(Data.data(),
                 reinterpret_cast<const uint8_t *>(Base) + Header.sh_offset,
                 Header.sh_size);
   }
 
+  MutableELFSection(const Elf_Shdr &Header, StringRef Name,
+                    OwningArrayRef<uint8_t> &&Data)
+      : Header(Header), Name(Name), Data(std::move(Data)) {}
+
   void setData(ArrayRef<uint8_t> Ref) {
     Data = OwningArrayRef<uint8_t>(Ref);
     Header.sh_size = Data.size();
   }
 
   operator const Elf_Shdr &() const { return Header; }
 };
 
 /// This class is used for doing mutations on an ELFObjectFile. The
public:
     }
 
     /// Returns the index of the last element. This is different to size()
     /// because size() returns the number of entries that haven't been removed.
     size_t getEndIndex() const { return Mappings.size(); }
 
     /// Removes the entry at \param Index.
     void remove(size_t Index) {
-      assert(Index < OriginalValues.size() && "Out of bounds");
+      assert(Index < Mappings.size() && "Out of bounds");
       assert(Mappings[Index].Type != MappingType::Removed &&
              "Entry already removed");
       Mappings[Index].Type = MappingType::Removed;
     }
 
+    /// Add's \param Value to the end of the table.
+    void add(NewType &&Value) {
+      Mappings.emplace_back(NewValues.size(), MappingType::New);
+      NewValues.emplace_back(std::move(Value));
+    }
+
     /// If the entry at index Index has already been made mutable, this returns
     /// a reference to that. Otherwise, this replaces the current entry at the
     /// specified index with a NewType constructued with Arguments.
     template <typename... Args>
     NewType &makeMutable(size_t Index, Args &&... Arguments) {
       assert(Index < Mappings.size() && "Out of bounds");
       assert(Mappings[Index].Type != MappingType::Removed);
       if (Mappings[Index].Type == MappingType::New)
public:
   Expected<MutableELFSection<ELFT> &> getMutableSection(SectionRef Sec) {
     const Elf_Shdr_Impl<ELFT> &Header = Sections[Sec.getRawDataRefImpl().p];
     Expected<StringRef> Name = getSectionName(Sec.getRawDataRefImpl());
     if (!Name)
       return Name.takeError();
     return Sections.makeMutable(Sec.getRawDataRefImpl().p, Header, *Name, this);
   }
 
+  void addSymbol(MutableELFSymbol<ELFT> Sym) { Symbols.add(std::move(Sym)); }
+
+  void addSection(MutableELFSection<ELFT> Sec) { Sections.add(std::move(Sec)); }
+
   /// Removes a symbol.
   void removeSymbol(SymbolRef Sym) {
     assert(Sections.getOriginal(Sym.getRawDataRefImpl().d.a).sh_type ==
                ELF::SHT_SYMTAB &&
            "Not pointing to symbol table");
     assert(Sym.getRawDataRefImpl().d.b && "Cannot remove index 0 symbol");
     Symbols.remove(Sym.getRawDataRefImpl().d.b);
   }
 
   /// Removes a section.
   void removeSection(SectionRef Sec) {

 template <typename ELFT>
 void MutableELFObject<ELFT>::moveSymbolNext(DataRefImpl &Sym) const {
   const MutableSymbolTable &SymbolTable = getSymbolTable(Sym);
   Sym.d.b = SymbolTable.getNextIndex(Sym.d.b);
 }
 
 template <typename ELFT>
 Expected<StringRef>
 MutableELFObject<ELFT>::getSymbolName(DataRefImpl Sym) const {
   const MutableSymbolTable &SymbolTable = getSymbolTable(Sym);
   if (const MutableELFSymbol<ELFT> *SymOrNull =
           SymbolTable.getConstIfNew(Sym.d.b))
     return SymOrNull->Name;
   return ELFObjectFile<ELFT>::getSymbolName(Sym);
 }
 
 template <typename ELFT>
 ArrayRef<typename ELFT::Sym>
 MutableELFObject<ELFT>::findSymbolTable(uint64_t ShType) const {
   assert(ShType == ELF::SHT_SYMTAB ||
          ShType == ELF::SHT_DYNSYM && "Not a symbol table type");
   for (const auto &Sec : this->sections())

llvm/unittests/Object/MutableELFObjectTest.cpp

std::vector<T> collect(
     Range TheRange,
     function_ref<T(const decltype(*TheRange.begin()) &)> Extractor =
         [](auto It) { return *It; }) {
   std::vector<T> Ret;
   llvm::transform(TheRange, std::back_inserter(Ret), Extractor);
   return Ret;
 }
 
-StringRef getSectionName(const SectionRef &Sec) {
-  Expected<StringRef> NameOrErr = Sec.getName();
+template <typename T> StringRef getName(const T &Named) {
+  Expected<StringRef> NameOrErr = Named.getName();
   if (!NameOrErr) {
     consumeError(NameOrErr.takeError());
-    return "Couldn't get section name";
+    return "Couldn't get name";
   }
   return *NameOrErr;
 }
 
+StringRef getSectionName(const SectionRef &Sec) { return getName(Sec); }
+
+StringRef getSymbolName(const SymbolRef &Sym) { return getName(Sym); }
+
 // Test that when no modifications have been made SectionRef's methods work
 // the same in both ELFObjectFile and MutableELFObject.
 TEST(MutableELFObject, NoChange) {
   StringRef Yaml = R"(
 --- !ELF
 FileHeader:
   Class:     ELFCLASS64
   Data:      ELFDATA2LSB
Sections:
   EXPECT_EQ(Iter->getIndex(), 1u);
   ++Iter;
   EXPECT_EQ(Iter->getIndex(), 2u);
   ++Iter;
   EXPECT_EQ(Iter->getIndex(), 3u);
   ++Iter;
   EXPECT_EQ(MutableObject.section_end(), Iter);
 }
+
+TEST(MutableELFObject, AddSectionAndSymbols) {
+  SmallString<0> Storage;
+  Expected<std::unique_ptr<ObjectFile>> ErrOrObj = yaml2ObjectFile(Storage, R"(
+--- !ELF
+FileHeader:
+  Class:     ELFCLASS64
+  Data:      ELFDATA2LSB
+  Type:      ET_REL
+  Machine:   EM_X86_64)");
+
+  ASSERT_THAT_EXPECTED(ErrOrObj, Succeeded());
+  auto *ELFObjFile = dyn_cast<ELFObjectFile<ELF64LE>>(ErrOrObj->get());
+  ASSERT_TRUE(ELFObjFile);
+  MutableELFObject<ELF64LE> MutableObject(std::move(*ELFObjFile));
+
+  std::vector<StringRef> SectionNames =
+      collect<StringRef>(MutableObject.sections(), getSectionName);
+  EXPECT_THAT(SectionNames,
+              ::testing::ElementsAre("", ".symtab", ".strtab", ".shstrtab"));
+
+  Elf_Shdr_Impl<ELF64LE> Header;
+  Header.sh_size = 64;
+
+  MutableObject.addSection({Header, ".newsec", OwningArrayRef<uint8_t>(64)});
+
+  SectionNames = collect<StringRef>(MutableObject.sections(), getSectionName);
+  EXPECT_THAT(SectionNames, ::testing::ElementsAre("", ".symtab", ".strtab",
+                                                   ".shstrtab", ".newsec"));
+
+  auto Iter = MutableObject.section_begin();
+  std::advance(Iter, 4);
+
+  EXPECT_EQ(Iter->getSize(), 64u);
+  EXPECT_EQ(getSectionName(*Iter), ".newsec");
+
+  // yaml2obj creates the index 0 symbol even when no other symbols exist.
+  EXPECT_EQ(
+      std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()),
+      1);
+  std::vector<StringRef> SymbolNames =
+      collect<StringRef>(MutableObject.symbols(), getSymbolName);
+  EXPECT_THAT(SymbolNames, ::testing::ElementsAre(""));
+
+  Elf_Sym_Impl<ELF64LE> Sym;
+  Sym.st_shndx = ELF::SHN_ABS;
+  Sym.st_value = 1234;
+  Sym.st_size = 2;
+  MutableObject.addSymbol({Sym, "newsym"});
+
+  SymbolNames = collect<StringRef>(MutableObject.symbols(), getSymbolName);
+  EXPECT_THAT(SymbolNames, ::testing::ElementsAre("", "newsym"));
+
+  ELFSymbolRef NewSym = *symbol_iterator(++MutableObject.symbol_begin());
+  EXPECT_EQ(NewSym.getValue(), 1234u);
+  EXPECT_EQ(NewSym.getSize(), 2u);
+  Expected<StringRef> NameOrErr = NewSym.getName();
+  ASSERT_THAT_EXPECTED(NameOrErr, Succeeded());
+  EXPECT_EQ(*NameOrErr, "newsym");
+}
+
+// Add many symbols.
+TEST(MutableELFObject, AddMany) {
+  SmallString<0> Storage;
+  Expected<std::unique_ptr<ObjectFile>> ErrOrObj = yaml2ObjectFile(Storage, R"(
+--- !ELF
+FileHeader:
+  Class:     ELFCLASS64
+  Data:      ELFDATA2LSB
+  Type:      ET_REL
+  Machine:   EM_X86_64)");
+
+  ASSERT_THAT_EXPECTED(ErrOrObj, Succeeded());
+  auto *ELFObjFile = dyn_cast<ELFObjectFile<ELF64LE>>(ErrOrObj->get());
+  ASSERT_TRUE(ELFObjFile);
+  MutableELFObject<ELF64LE> MutableObject(std::move(*ELFObjFile));
+
+  EXPECT_EQ(
+      std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()),
+      1);
+
+  Elf_Sym_Impl<ELF64LE> Sym;
+  for (size_t I = 0; I < 100; ++I)
+    MutableObject.addSymbol({Sym, ""});
+
+  EXPECT_EQ(
+      std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()),
+      101);
+}
+
+// Add and then remove a symbol.
+TEST(MutableELFObject, AddRemoveSymbol) {
+  SmallString<0> Storage;
+  Expected<std::unique_ptr<ObjectFile>> ErrOrObj = yaml2ObjectFile(Storage, R"(
+--- !ELF
+FileHeader:
+  Class:     ELFCLASS64
+  Data:      ELFDATA2LSB
+  Type:      ET_REL
+  Machine:   EM_X86_64)");
+
+  ASSERT_THAT_EXPECTED(ErrOrObj, Succeeded());
+  auto *ELFObjFile = dyn_cast<ELFObjectFile<ELF64LE>>(ErrOrObj->get());
+  ASSERT_TRUE(ELFObjFile);
+  MutableELFObject<ELF64LE> MutableObject(std::move(*ELFObjFile));
+
+  EXPECT_EQ(
+      std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()),
+      1);
+
+  Elf_Sym_Impl<ELF64LE> Sym;
+  MutableObject.addSymbol({Sym, ""});
+
+  EXPECT_EQ(
+      std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()),
+      2);
+
+  basic_symbol_iterator NewSym = ++MutableObject.symbol_begin();
+  MutableObject.removeSymbol(*NewSym);
+
+  EXPECT_EQ(
+      std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()),
+      1);
+}
+
+// Add and then remove sections.
+TEST(MutableELFObject, AddRemoveSection) {
+  SmallString<0> Storage;
+  Expected<std::unique_ptr<ObjectFile>> ErrOrObj = yaml2ObjectFile(Storage, R"(
+--- !ELF
+FileHeader:
+  Class:     ELFCLASS64
+  Data:      ELFDATA2LSB
+  Type:      ET_REL
+  Machine:   EM_X86_64)");
+
+  ASSERT_THAT_EXPECTED(ErrOrObj, Succeeded());
+  auto *ELFObjFile = dyn_cast<ELFObjectFile<ELF64LE>>(ErrOrObj->get());
+  ASSERT_TRUE(ELFObjFile);
+  MutableELFObject<ELF64LE> MutableObject(std::move(*ELFObjFile));
+
+  auto Distance = [&MutableObject] {
+    return std::distance(MutableObject.section_begin(),
+                         MutableObject.section_end());
+  };
+
+  auto IterAtIndex = [&MutableObject, Distance](unsigned Index) {
+    assert(Index < Distance());
+    auto Iter = MutableObject.section_begin();
+    std::advance(Iter, Index);
+    return Iter;
+  };
+
+  EXPECT_EQ(Distance(), 4);
+
+  Elf_Shdr_Impl<ELF64LE> Header;
+  Header.sh_size = 2;
+  MutableObject.addSection({Header, "", OwningArrayRef<uint8_t>(0)});
+  EXPECT_EQ(Distance(), 5);
+  EXPECT_EQ(IterAtIndex(4)->getSize(), 2u);
+
+  MutableObject.removeSection(*(IterAtIndex(1)));
+  EXPECT_EQ(Distance(), 4);
+  EXPECT_EQ(IterAtIndex(3)->getSize(), 2u);
+
+  Header.sh_size = 4;
+  MutableObject.addSection({Header, "", OwningArrayRef<uint8_t>(0)});
+  EXPECT_EQ(Distance(), 5);
+  EXPECT_EQ(IterAtIndex(4)->getSize(), 4u);
+
+  MutableObject.addSection({Header, "", OwningArrayRef<uint8_t>(0)});
+  MutableObject.addSection({Header, "", OwningArrayRef<uint8_t>(0)});
+  Header.sh_size = 100;
+  MutableObject.addSection({Header, "", OwningArrayRef<uint8_t>(0)});
+  EXPECT_EQ(Distance(), 8);
+  EXPECT_EQ(IterAtIndex(7)->getSize(), 100u);
+}