Index: llvm/include/llvm/Object/MutableELFObject.h =================================================================== --- llvm/include/llvm/Object/MutableELFObject.h +++ llvm/include/llvm/Object/MutableELFObject.h @@ -35,6 +35,12 @@ Header.sh_size); } + MutableELFSection(const Elf_Shdr &Header, StringRef Name, + OwningArrayRef &&Data) + : Header(Header), Name(Name), Data(std::move(Data)) {} + + MutableELFSection(MutableELFSection &&Other) = default; + void setData(ArrayRef Ref) { Data = OwningArrayRef(Ref); Header.sh_size = Data.size(); @@ -155,12 +161,19 @@ /// 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. @@ -329,6 +342,16 @@ return Sections.makeMutable(Sec.getRawDataRefImpl().p, Header, *Name, this); } + /// Add symbol \param Sym. + MutableELFSymbol &addSymbol(MutableELFSymbol &&Sym) { + return Symbols.add(std::move(Sym)); + } + + /// Add section \param Sec. + MutableELFSection &addSection(MutableELFSection &&Sec) { + return Sections.add(std::move(Sec)); + } + /// Removes a symbol. void removeSymbol(SymbolRef Sym) { assert(Sections.getOriginal(Sym.getRawDataRefImpl().d.a).sh_type == Index: llvm/unittests/Object/MutableELFObjectTest.cpp =================================================================== --- llvm/unittests/Object/MutableELFObjectTest.cpp +++ llvm/unittests/Object/MutableELFObjectTest.cpp @@ -526,3 +526,181 @@ ++Iter; EXPECT_EQ(MutableObject.section_end(), Iter); } + +TEST(MutableELFObject, AddSectionAndSymbols) { + SmallString<0> Storage; + Expected> 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>(ErrOrObj->get()); + ASSERT_TRUE(ELFObjFile); + MutableELFObject 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 Header; + Header.sh_size = 64; + + MutableObject.addSection({Header, ".newsec", OwningArrayRef(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 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 NameOrErr = NewSym.getName(); + ASSERT_THAT_EXPECTED(NameOrErr, Succeeded()); + EXPECT_EQ(*NameOrErr, "newsym"); +} + +// Add many symbols. +TEST(MutableELFObject, AddMany) { + SmallString<0> Storage; + Expected> 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>(ErrOrObj->get()); + ASSERT_TRUE(ELFObjFile); + MutableELFObject MutableObject(std::move(*ELFObjFile)); + + EXPECT_EQ( + std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()), + 1); + + Elf_Sym_Impl 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> 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>(ErrOrObj->get()); + ASSERT_TRUE(ELFObjFile); + MutableELFObject MutableObject(std::move(*ELFObjFile)); + + EXPECT_EQ( + std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()), + 1); + + Elf_Sym_Impl 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> 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>(ErrOrObj->get()); + ASSERT_TRUE(ELFObjFile); + MutableELFObject 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 Header; + Header.sh_size = 2; + MutableObject.addSection({Header, "", OwningArrayRef(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(0)}); + EXPECT_EQ(Distance(), 5); + EXPECT_EQ(IterAtIndex(4)->getSize(), 4u); + MutableObject.addSection({Header, "", OwningArrayRef(0)}); + MutableObject.addSection({Header, "", OwningArrayRef(0)}); + Header.sh_size = 100; + MutableObject.addSection({Header, "", OwningArrayRef(0)}); + EXPECT_EQ(Distance(), 8); + EXPECT_EQ(IterAtIndex(7)->getSize(), 100u); +}