[flang] Improve API for runtime allocator (I/O runtime work part 3)

[klausler]

Summary

New<A> used to return an A&; now it returns an OwningPtr<A>
to force better ownership tracking of allocations. Its API
has also been split into New<A> and SizedNew<A> to allow
allocations with a size override.

flang/runtime/file.cpp

if (at >= 0) {
     return true;
   } else {
     return false;
   }
 }
 
 int OpenFile::PendingResult(const Terminator &terminator, int iostat) {
   int id{nextId_++};
-  pending_.reset(&New<Pending>{}(terminator, id, iostat, std::move(pending_)));
+  pending_ = New<Pending>{terminator}(id, iostat, std::move(pending_));
   return id;
 }
 } // namespace Fortran::runtime::io

flang/runtime/io-api.cpp

 
 namespace Fortran::runtime::io {
 
 template <Direction DIR>
 Cookie BeginInternalArrayListIO(const Descriptor &descriptor,
     void ** /*scratchArea*/, std::size_t /*scratchBytes*/,
     const char *sourceFile, int sourceLine) {
   Terminator oom{sourceFile, sourceLine};
-  return &New<InternalListIoStatementState<DIR>>{}(
-      oom, descriptor, sourceFile, sourceLine)
-              .ioStatementState();
+  return &New<InternalListIoStatementState<DIR>>{oom}(
+      descriptor, sourceFile, sourceLine)
+              .release()
+              ->ioStatementState();
 }
 
 Cookie IONAME(BeginInternalArrayListOutput)(const Descriptor &descriptor,
     void **scratchArea, std::size_t scratchBytes, const char *sourceFile,
     int sourceLine) {
   return BeginInternalArrayListIO<Direction::Output>(
       descriptor, scratchArea, scratchBytes, sourceFile, sourceLine);
 }
 
 Cookie IONAME(BeginInternalArrayListInput)(const Descriptor &descriptor,
     void **scratchArea, std::size_t scratchBytes, const char *sourceFile,
     int sourceLine) {
   return BeginInternalArrayListIO<Direction::Input>(
       descriptor, scratchArea, scratchBytes, sourceFile, sourceLine);
 }
 
 template <Direction DIR>
 Cookie BeginInternalArrayFormattedIO(const Descriptor &descriptor,
     const char *format, std::size_t formatLength, void ** /*scratchArea*/,
     std::size_t /*scratchBytes*/, const char *sourceFile, int sourceLine) {
   Terminator oom{sourceFile, sourceLine};
-  return &New<InternalFormattedIoStatementState<DIR>>{}(
-      oom, descriptor, format, formatLength, sourceFile, sourceLine)
-              .ioStatementState();
+  return &New<InternalFormattedIoStatementState<DIR>>{oom}(
+      descriptor, format, formatLength, sourceFile, sourceLine)
+              .release()
+              ->ioStatementState();
 }
 
 Cookie IONAME(BeginInternalArrayFormattedOutput)(const Descriptor &descriptor,
     const char *format, std::size_t formatLength, void **scratchArea,
     std::size_t scratchBytes, const char *sourceFile, int sourceLine) {
   return BeginInternalArrayFormattedIO<Direction::Output>(descriptor, format,
       formatLength, scratchArea, scratchBytes, sourceFile, sourceLine);
 }
 
 Cookie IONAME(BeginInternalArrayFormattedInput)(const Descriptor &descriptor,
     const char *format, std::size_t formatLength, void **scratchArea,
     std::size_t scratchBytes, const char *sourceFile, int sourceLine) {
   return BeginInternalArrayFormattedIO<Direction::Input>(descriptor, format,
       formatLength, scratchArea, scratchBytes, sourceFile, sourceLine);
 }
 
 template <Direction DIR>
 Cookie BeginInternalFormattedIO(
     std::conditional_t<DIR == Direction::Input, const char, char> *internal,
     std::size_t internalLength, const char *format, std::size_t formatLength,
     void ** /*scratchArea*/, std::size_t /*scratchBytes*/,
     const char *sourceFile, int sourceLine) {
   Terminator oom{sourceFile, sourceLine};
-  return &New<InternalFormattedIoStatementState<DIR>>{}(oom, internal,
-      internalLength, format, formatLength, sourceFile, sourceLine)
-              .ioStatementState();
+  return &New<InternalFormattedIoStatementState<DIR>>{oom}(
+      internal, internalLength, format, formatLength, sourceFile, sourceLine)
+              .release()
+              ->ioStatementState();
 }
 
 Cookie IONAME(BeginInternalFormattedOutput)(char *internal,
     std::size_t internalLength, const char *format, std::size_t formatLength,
     void **scratchArea, std::size_t scratchBytes, const char *sourceFile,
     int sourceLine) {
   Terminator oom{sourceFile, sourceLine};
   return BeginInternalFormattedIO<Direction::Output>(internal, internalLength,

 Cookie IONAME(BeginClose)(
     ExternalUnit unitNumber, const char *sourceFile, int sourceLine) {
   if (ExternalFileUnit * unit{ExternalFileUnit::LookUpForClose(unitNumber)}) {
     return &unit->BeginIoStatement<CloseStatementState>(
         *unit, sourceFile, sourceLine);
   } else {
     // CLOSE(UNIT=bad unit) is just a no-op
     Terminator oom{sourceFile, sourceLine};
-    return &New<NoopCloseStatementState>{}(oom, sourceFile, sourceLine)
-                .ioStatementState();
+    return &New<NoopCloseStatementState>{oom}(sourceFile, sourceLine)
+                .release()
+                ->ioStatementState();
   }
 }
 
 // Control list items
 
 void IONAME(EnableHandlers)(Cookie cookie, bool hasIoStat, bool hasErr,
     bool hasEnd, bool hasEor, bool hasIoMsg) {
   IoErrorHandler &handler{cookie->GetIoErrorHandler()};

flang/runtime/memory.h

 template <typename A> void FreeMemory(A *p) {
   FreeMemory(reinterpret_cast<void *>(p));
 }
 template <typename A> void FreeMemoryAndNullify(A *&p) {
   FreeMemory(p);
   p = nullptr;
 }
 
-template <typename A> struct New {
-  template <typename... X>
-  [[nodiscard]] A &operator()(const Terminator &terminator, X &&... x) {
-    return *new (AllocateMemoryOrCrash(terminator, sizeof(A)))
-        A{std::forward<X>(x)...};
-  }
-};
-
 template <typename A> struct OwningPtrDeleter {
   void operator()(A *p) { FreeMemory(p); }
 };
 
 template <typename A> using OwningPtr = std::unique_ptr<A, OwningPtrDeleter<A>>;
 
+template <typename A> class SizedNew {
+public:
+  explicit SizedNew(const Terminator &terminator) : terminator_{terminator} {}
+  template <typename... X>
+  [[nodiscard]] OwningPtr<A> operator()(std::size_t bytes, X &&... x) {
+    return OwningPtr<A>{new (AllocateMemoryOrCrash(terminator_, bytes))
+            A{std::forward<X>(x)...}};
+  }
+
+private:
+  const Terminator &terminator_;
+};
+
+template <typename A> struct New : public SizedNew<A> {
+  using SizedNew<A>::SizedNew;
+  template <typename... X> [[nodiscard]] OwningPtr<A> operator()(X &&... x) {
+    return SizedNew<A>::operator()(sizeof(A), std::forward<X>(x)...);
+  }
+};
+
 template <typename A> struct Allocator {
   using value_type = A;
   explicit Allocator(const Terminator &t) : terminator{t} {}
   template <typename B>
   explicit constexpr Allocator(const Allocator<B> &that) noexcept
       : terminator{that.terminator} {}
   Allocator(const Allocator &) = default;
   Allocator(Allocator &&) = default;

flang/runtime/unit-map.cpp

while (Chain * p{bucket_[j].get()}) {
       p->unit.CloseUnit(CloseStatus::Keep, handler);
       p->unit.~ExternalFileUnit();
       FreeMemory(p);
     }
   }
 }
 
 ExternalFileUnit &UnitMap::Create(int n, const Terminator &terminator) {
-  Chain &chain{New<Chain>{}(terminator, n)};
+  Chain &chain{*New<Chain>{terminator}(n).release()};
   chain.next.reset(&chain);
   bucket_[Hash(n)].swap(chain.next); // pushes new node as list head
   return chain.unit;
 }
 } // namespace Fortran::runtime::io

flang/runtime/unit.cpp

UnitMap &ExternalFileUnit::GetUnitMap() {
   if (unitMap) {
     return *unitMap;
   }
   CriticalSection critical{unitMapLock};
   if (unitMap) {
     return *unitMap;
   }
   Terminator terminator{__FILE__, __LINE__};
-  unitMap = &New<UnitMap>{}(terminator);
+  unitMap = New<UnitMap>{terminator}().release();
   ExternalFileUnit &out{ExternalFileUnit::LookUpOrCreate(6, terminator)};
   out.Predefine(1);
   out.set_mayRead(false);
   out.set_mayWrite(true);
   out.set_mayPosition(false);
   defaultOutput = &out;
   ExternalFileUnit &in{ExternalFileUnit::LookUpOrCreate(5, terminator)};
   in.Predefine(0);