[AffineMap] Move result exprs into trailing storage. NFCI.

mlir/lib/IR/AffineMap.cpp

@@ -299,22 +299,17 @@ unsigned AffineMap::getNumSymbols() const {
   assert(map && "uninitialized map storage");
   return map->numSymbols;
 }
-unsigned AffineMap::getNumResults() const {
-  assert(map && "uninitialized map storage");
-  return map->results.size();
-}
+unsigned AffineMap::getNumResults() const { return getResults().size(); }
 unsigned AffineMap::getNumInputs() const {
   assert(map && "uninitialized map storage");
   return map->numDims + map->numSymbols;
 }
-
 ArrayRef<AffineExpr> AffineMap::getResults() const {
   assert(map && "uninitialized map storage");
-  return map->results;
+  return map->results();
 }
 AffineExpr AffineMap::getResult(unsigned idx) const {
-  assert(map && "uninitialized map storage");
-  return map->results[idx];
+  return getResults()[idx];
 }
 
 unsigned AffineMap::getDimPosition(unsigned idx) const {

mlir/lib/IR/AffineMapDetail.h

@@ -17,36 +17,47 @@
 #include "mlir/IR/AffineMap.h"
 #include "mlir/Support/StorageUniquer.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/Support/TrailingObjects.h"
 
 namespace mlir {
 namespace detail {
 
-struct AffineMapStorage : public StorageUniquer::BaseStorage {
+struct AffineMapStorage final
+    : public StorageUniquer::BaseStorage,
+      public llvm::TrailingObjects<AffineMapStorage, AffineExpr> {
   /// The hash key used for uniquing.
   using KeyTy = std::tuple<unsigned, unsigned, ArrayRef<AffineExpr>>;
 
   unsigned numDims;
   unsigned numSymbols;
-
-  /// The affine expressions for this (multi-dimensional) map.
-  /// TODO: use trailing objects for this.
-  ArrayRef<AffineExpr> results;
+  unsigned numResults;
 
   MLIRContext *context;
 
+  /// The affine expressions for this (multi-dimensional) map.
+  ArrayRef<AffineExpr> results() const {
+    return {getTrailingObjects<AffineExpr>(), numResults};
+  }
+
   bool operator==(const KeyTy &key) const {
     return std::get<0>(key) == numDims && std::get<1>(key) == numSymbols &&
-           std::get<2>(key) == results;
+           std::get<2>(key) == results();
   }
 
   // Constructs an AffineMapStorage from a key. The context must be set by the
   // caller.
   static AffineMapStorage *
   construct(StorageUniquer::StorageAllocator &allocator, const KeyTy &key) {
-    auto *res = new (allocator.allocate<AffineMapStorage>()) AffineMapStorage();
+    auto results = std::get<2>(key);
+    auto byteSize =
+        AffineMapStorage::totalSizeToAlloc<AffineExpr>(results.size());
+    auto *rawMem = allocator.allocate(byteSize, alignof(AffineMapStorage));
+    auto *res = new (rawMem) AffineMapStorage();
     res->numDims = std::get<0>(key);
     res->numSymbols = std::get<1>(key);
-    res->results = allocator.copyInto(std::get<2>(key));
+    res->numResults = results.size();
+    std::uninitialized_copy(results.begin(), results.end(),
+                            res->getTrailingObjects<AffineExpr>());
     return res;
   }
 };