[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

171

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

346

I think this should just be a void return type.

351

I think this should just be a void return type.

403–408

This seems like an unrelated change?

llvm/unittests/Object/MutableELFObjectTest.cpp

545–547

Comment why there's 4 sections already.

566–568

Comment why there's already a symbol here.

570

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

171

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

171

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

llvm/unittests/Object/MutableELFObjectTest.cpp

545

"yaml2obj creates the .symtab ..."

567

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

613

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

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

624

AddRemoveSymbol

659

AddRemoveSection

688–705

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

346

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.

349

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?

173

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)) {}
+
+  MutableELFSection(MutableELFSection &&Other) = default;
+
   void setData(ArrayRef<uint8_t> Ref) {
     Data = OwningArrayRef<uint8_t>(Ref);
     Header.sh_size = Data.size();
   }
 
   operator const Elf_Shdr &() const { return Header; }
 };
 
public:
     }
 
     /// 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(uint64_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.
+    NewType &add(NewType &&Value) {
+      NewValues.emplace_back(std::move(Value));
+      Mappings.emplace_back(NewValues.size() - 1, MappingType::New);
+      return NewValues.back();
+    }
+
     /// 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(uint64_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);
   }
 
+  /// Add symbol \param Sym.
+  MutableELFSymbol<ELFT> &addSymbol(MutableELFSymbol<ELFT> &&Sym) {
+    return Symbols.add(std::move(Sym));
+  }
+
+  /// Add section \param Sec.
+  MutableELFSection<ELFT> &addSection(MutableELFSection<ELFT> &&Sec) {
+    return 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");
     Symbols.remove(Sym.getRawDataRefImpl().d.b);
   }
 

 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

Sections:
   EXPECT_EQ(Name, ".strtab");
   ++Iter;
   EXPECT_EQ(Iter->getIndex(), 3u);
   ASSERT_FALSE(Iter->getName(Name));
   EXPECT_EQ(Name, ".shstrtab");
   ++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));
+
+  // yaml2obj puts in .symtab, .strtab and .shstrtab sections by default.
+  EXPECT_EQ(
+      std::distance(MutableObject.section_begin(), MutableObject.section_end()),
+      4);
+
+  Elf_Shdr_Impl<ELF64LE> Header;
+  Header.sh_size = 64;
+
+  MutableObject.addSection({Header, ".newsec", OwningArrayRef<uint8_t>(64)});
+
+  EXPECT_EQ(
+      std::distance(MutableObject.section_begin(), MutableObject.section_end()),
+      5);
+
+  auto Iter = MutableObject.section_begin();
+  std::advance(Iter, 4);
+
+  EXPECT_EQ(Iter->getSize(), 64u);
+  StringRef Name;
+  EXPECT_FALSE(Iter->getName(Name));
+  EXPECT_EQ(Name, ".newsec");
+
+  // yaml2obj puts in the index 0 symbol in the symbol table even when no
+  // other symbols exist.
+  EXPECT_EQ(
+      std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()),
+      1);
+
+  Elf_Sym_Impl<ELF64LE> Sym;
+  Sym.st_shndx = ELF::SHN_ABS;
+  Sym.st_value = 1234;
+  Sym.st_size = 2;
+  auto ReturnedSym = MutableObject.addSymbol({Sym, "newsym"});
+  EXPECT_EQ(ReturnedSym.Header.st_value, 1234u);
+
+  EXPECT_EQ(
+      std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()),
+      2);
+
+  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 (int _ : seq(0, 10000)) {
+    (void)_;
+    MutableObject.addSymbol({Sym, ""});
+  }
+
+  EXPECT_EQ(
+      std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()),
+      10001);
+}
+
+// Add and then remove a symbol.
+TEST(MutableELFObject, AddRemove) {
+  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, AddThenRemove) {
+  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);
+}