[llvm-objcopy] Add objcopy::elf::Object::allocSections to simplify loops on SHF_ALLOC sections

[jhenderson]

LGTM! Thanks for doing this.

llvm/trunk/tools/llvm-objcopy/ELF/Object.h

 
 public:
   using iterator = pointee_iterator<std::unique_ptr<SectionBase> *>;
 
   explicit SectionTableRef(MutableArrayRef<std::unique_ptr<SectionBase>> Secs)
       : Sections(Secs) {}
   SectionTableRef(const SectionTableRef &) = default;
 
-  iterator begin() { return iterator(Sections.data()); }
-  iterator end() { return iterator(Sections.data() + Sections.size()); }
+  iterator begin() const { return iterator(Sections.data()); }
+  iterator end() const { return iterator(Sections.data() + Sections.size()); }
   size_t size() const { return Sections.size(); }
 
   SectionBase *getSection(uint32_t Index, Twine ErrMsg);
 
   template <class T>
   T *getSectionOfType(uint32_t Index, Twine IndexErrMsg, Twine TypeErrMsg);
 };
 

 private:
   using SecPtr = std::unique_ptr<SectionBase>;
   using SegPtr = std::unique_ptr<Segment>;
 
   std::vector<SecPtr> Sections;
   std::vector<SegPtr> Segments;
   std::vector<SecPtr> RemovedSections;
 
+  static bool sectionIsAlloc(const SectionBase &Sec) {
+    return Sec.Flags & ELF::SHF_ALLOC;
+  };
+
 public:
   template <class T>
   using Range = iterator_range<
       pointee_iterator<typename std::vector<std::unique_ptr<T>>::iterator>>;
 
   template <class T>
   using ConstRange = iterator_range<pointee_iterator<
       typename std::vector<std::unique_ptr<T>>::const_iterator>>;
public:
   SymbolTableSection *SymbolTable = nullptr;
   SectionIndexSection *SectionIndexTable = nullptr;
 
   void sortSections();
   SectionTableRef sections() { return SectionTableRef(Sections); }
   ConstRange<SectionBase> sections() const {
     return make_pointee_range(Sections);
   }
+  iterator_range<
+      filter_iterator<pointee_iterator<std::vector<SecPtr>::const_iterator>,
+                      decltype(&sectionIsAlloc)>>
+  allocSections() const {
+    return make_filter_range(make_pointee_range(Sections), sectionIsAlloc);
+  }
+
   SectionBase *findSection(StringRef Name) {
     auto SecIt =
         find_if(Sections, [&](const SecPtr &Sec) { return Sec->Name == Name; });
     return SecIt == Sections.end() ? nullptr : SecIt->get();
   }
   SectionTableRef removedSections() { return SectionTableRef(RemovedSections); }
 
   Range<Segment> segments() { return make_pointee_range(Segments); }

llvm/trunk/tools/llvm-objcopy/ELF/Object.cpp

template <class ELFT> Error ELFWriter<ELFT>::finalize() {
 
   if (Error E = Buf.allocate(totalSize()))
     return E;
   SecWriter = std::make_unique<ELFSectionWriter<ELFT>>(Buf);
   return Error::success();
 }
 
 Error BinaryWriter::write() {
-  for (const SectionBase &Sec : Obj.sections())
-    if (Sec.Flags & SHF_ALLOC)
-      Sec.accept(*SecWriter);
+  for (const SectionBase &Sec : Obj.allocSections())
   return Buf.commit();
 }
 
 Error BinaryWriter::finalize() {
-  // TODO: Create a filter range to construct OrderedSegments from so that this
-  // code can be deduped with assignOffsets above. This should also solve the
-  // todo below for LayoutSections.
   // We need a temporary list of segments that has a special order to it
   // so that we know that anytime ->ParentSegment is set that segment has
   // already had it's offset properly set. We only want to consider the segments
   // that will affect layout of allocated sections so we only add those.
   std::vector<Segment *> OrderedSegments;
-  for (const SectionBase &Sec : Obj.sections())
-    if ((Sec.Flags & SHF_ALLOC) != 0 && Sec.ParentSegment != nullptr)
+  for (const SectionBase &Sec : Obj.allocSections())
+    if (Sec.ParentSegment != nullptr)
       OrderedSegments.push_back(Sec.ParentSegment);
 
   // For binary output, we're going to use physical addresses instead of
   // virtual addresses, since a binary output is used for cases like ROM
   // loading and physical addresses are intended for ROM loading.
   // However, if no segment has a physical address, we'll fallback to using
   // virtual addresses for all.
   if (all_of(OrderedSegments,
if (!OrderedSegments.empty()) {
     Seg->PAddr += Diff;
     uint64_t LowestPAddr = Seg->PAddr;
     for (Segment *Segment : OrderedSegments) {
       Segment->Offset = Segment->PAddr - LowestPAddr;
       Offset = std::max(Offset, Segment->Offset + Segment->FileSize);
     }
   }
 
-  // TODO: generalize layoutSections to take a range. Pass a special range
-  // constructed from an iterator that skips values for which a predicate does
-  // not hold. Then pass such a range to layoutSections instead of constructing
-  // AllocatedSections here.
-  std::vector<SectionBase *> AllocatedSections;
-  for (SectionBase &Sec : Obj.sections())
-    if (Sec.Flags & SHF_ALLOC)
-      AllocatedSections.push_back(&Sec);
-  layoutSections(make_pointee_range(AllocatedSections), Offset);
+  layoutSections(Obj.allocSections(), Offset);
 
   // Now that every section has been laid out we just need to compute the total
   // file size. This might not be the same as the offset returned by
   // layoutSections, because we want to truncate the last segment to the end of
   // its last section, to match GNU objcopy's behaviour.
   TotalSize = 0;
-  for (SectionBase *Sec : AllocatedSections)
-    if (Sec->Type != SHT_NOBITS)
-      TotalSize = std::max(TotalSize, Sec->Offset + Sec->Size);
+  for (const SectionBase &Sec : Obj.allocSections())
+    if (Sec.Type != SHT_NOBITS)
+      TotalSize = std::max(TotalSize, Sec.Offset + Sec.Size);
 
   if (Error E = Buf.allocate(TotalSize))
     return E;
   SecWriter = std::make_unique<BinarySectionWriter>(Buf);
   return Error::success();
 }
 
 bool IHexWriter::SectionCompare::operator()(const SectionBase *Lhs,