SIL: add a utility which let's manage per-block data efficiently.

It can be used by transforms to store temporary data per basic block.
It is very efficient: only a single memory allocation is needed and no maps are used to lookup data.

Author: eeckstein

include/swift/SIL/BasicBlockData.h

@@ -0,0 +1,207 @@
+//===--- BasicBlockData.h - Defines the BasicBlockData utility --*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2021 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines the BasicBlockData utility
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_SIL_BASICBLOCKDATA_H
+#define SWIFT_SIL_BASICBLOCKDATA_H
+
+#include "swift/SIL/SILFunction.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/STLExtras.h"
+
+namespace swift {
+
+/// Manage additional data for a function's basic blocks.
+///
+/// This can be used by transforms to store temporary data per basic block.
+/// It is very efficient: only a single memory allocation is needed and no maps
+/// are used to lookup data.
+///
+/// It is still permitted to use a BasicBlockData after the block list of a
+/// function has changed. With two exceptions:
+///  * If new blocks are added during the lifetime of a BasicBlockData, no data
+///    can be queried for a new block.
+///  * If a new BasicBlockData is created after the block list changed, then any
+///    older BasicBlockData (which was created before the change) cannot be
+///    used anymore.
+///
+/// This means, it's totally fine to store a BasicBlockData in an Analysis, as
+/// long as the invalidation mechanism ensures that the analysis data is not
+/// retrieved after the function's block list changed.
+///
+/// The Vector template argument specifies how the data is stored. By default
+/// it's a SmallVector with an inline-size of 4. This avoids memory allocation
+/// for about 70% of all functions.
+template <typename Data, typename Vector = llvm::SmallVector<Data, 4>>
+class BasicBlockData {
+  SILFunction *function;
+  Vector data;
+
+  /// The data is valid if validForBlockOrder matches the function's
+  /// BlockListChangeIdx.
+  unsigned validForBlockOrder = 0;
+
+  int getIndex(SILBasicBlock *block) const {
+    assert(block->index >= 0 && "Cannot get data for a new block");
+    assert(validForBlockOrder == function->BlockListChangeIdx &&
+           "BasicBlockData invalid because the function's block list changed");
+    return block->index;
+  }
+
+public:
+
+  /// For iterating over the function's blocks, yielding the block and its data.
+  template <typename BlockIterTy, bool IsConst> class IteratorTy {
+    using DataTy = typename std::conditional<IsConst, const Data, Data>::type;
+    using BasicBlockDataTy = typename std::conditional<IsConst,
+      const BasicBlockData<Data>, BasicBlockData<Data>>::type;
+
+    /// The underlying SILBasicBlock iterator.
+    BlockIterTy blockIter;
+
+    /// Reference to the constructing BasicBlockData. Used to retrieve the data.
+    BasicBlockDataTy &blockData;
+
+    friend class BasicBlockData<Data>;
+
+    IteratorTy(BlockIterTy blockIter, BasicBlockDataTy &blockData) :
+      blockIter(blockIter), blockData(blockData) {}
+
+  public:
+    /// Result of re-referencing the iterator.
+    struct BlockAndData {
+      SILBasicBlock &block;
+      DataTy &data;
+
+      BlockAndData(SILBasicBlock &block, DataTy &data) :
+        block(block), data(data) {}
+    };
+
+    IteratorTy & operator++() {  // Preincrement
+      blockIter++; return *this;
+    }
+    IteratorTy & operator++(int) {  // Postincrement
+      IteratorTy tmp = *this; ++*this; return tmp;
+    }
+
+    SILBasicBlock &block() const { return *blockIter; }
+    DataTy &data() const { return blockData[&block()]; }
+
+    BlockAndData operator*() const { return BlockAndData(block(), data()); }
+
+    bool operator==(const IteratorTy &rhs) const {
+      return blockIter == rhs.blockIter;
+    }
+    bool operator!=(const IteratorTy &rhs) const {
+      return blockIter != rhs.blockIter;
+    }
+  };
+
+  using iterator = IteratorTy<SILFunction::iterator, false>;
+  using const_iterator = IteratorTy<SILFunction::iterator, true>;
+  using reverse_iterator = IteratorTy<SILFunction::reverse_iterator, false>;
+  using const_reverse_iterator = IteratorTy<SILFunction::reverse_iterator, true>;
+
+  static Data constructDefaultData(SILBasicBlock *) { return Data(); }
+
+  /// Initialize Data for all basic blocks.
+  ///
+  /// The \p init function must return the initial data for a block. If not
+  /// provided, all data is constructed with the Data() default constructor.
+  BasicBlockData(SILFunction *function,
+                 llvm::function_ref<Data(SILBasicBlock *block)> init =
+                   // Would be nice to simply write constructDefaultData in
+                   // closure syntax, but I didn't get it compile under Windows.
+                   constructDefaultData) :
+      function(function) {
+    // Reserve enough space. Though SILFunction::size() iterates over all
+    // blocks it is still better than to risk multiple mallocs.
+    data.reserve(function->size());
+    bool blockListChanged = false;
+    for (auto blockAndIdx : llvm::enumerate(*function)) {
+      SILBasicBlock &block = blockAndIdx.value();
+      int idx = blockAndIdx.index();
+      if (block.index != idx) {
+        blockListChanged = true;
+        block.index = idx;
+      }
+      data.push_back(init(&block));
+    }
+    function->incrementRefCount();
+
+    // If we assigned new block indices, it invalidates all older BasicBlockData
+    // instances.
+    if (blockListChanged)
+      ++function->BlockListChangeIdx;
+    validForBlockOrder = function->BlockListChangeIdx;
+  }
+
+  ~BasicBlockData() {
+    function->decrementRefCount();
+  }
+
+  // For iterating over the function's basic blocks, yielding the block and its
+  // data. For example:
+  //
+  //    BasicBlockData<Mydata> blockData;
+  //    for (auto bd : blockData) {
+  //      SILBasicBlock &block = bd.block;
+  //      Mydata &dataForBlock = bd.data;
+  //    }
+
+  iterator begin() { return iterator(function->begin(), *this); }
+  iterator end() { return iterator(function->end(), *this); }
+  const_iterator begin() const {
+    return const_iterator(function->begin(), *this);
+  }
+  const_iterator end() const {
+    return const_iterator(function->end(), *this);
+  }
+  reverse_iterator rbegin() {
+    return reverse_iterator(function->rbegin(), *this);
+  }
+  reverse_iterator rend() {
+    return reverse_iterator(function->rend(), *this);
+  }
+  const_reverse_iterator rbegin() const {
+    return const_reverse_iterator(function->rbegin(), *this);
+  }
+  const_reverse_iterator rend() const {
+    return const_reverse_iterator(function->rend(), *this);
+  }
+
+  /// Returns the function's entry block and its data.
+  typename iterator::BlockAndData entry() { return *begin(); }
+
+  /// Returns the function's entry block and its data (const-version).
+  typename const_iterator::BlockAndData entry() const { return *begin(); }
+
+  /// Returns the underlying function.
+  SILFunction *getFunction() const { return function; }
+
+  /// Gets a mutable reference to a block's data.
+  Data &operator[] (SILBasicBlock *block) {
+    return data[getIndex(block)];
+  }
+
+  /// Gets a constant reference to a block's data.
+  const Data &operator[] (SILBasicBlock *block) const {
+    return data[getIndex(block)];
+  }
+};
+
+} // end swift namespace
+
+#endif

include/swift/SIL/SILBasicBlock.h

@@ -33,6 +33,8 @@ public llvm::ilist_node<SILBasicBlock>, public SILAllocated<SILBasicBlock> {
   friend class SILSuccessor;
   friend class SILFunction;
   friend class SILGlobalVariable;
+  template <typename Data, typename Vector> friend class BasicBlockData;
+
 public:
   using InstListType = llvm::iplist<SILInstruction>;
 private:
@@ -50,6 +52,11 @@ public llvm::ilist_node<SILBasicBlock>, public SILAllocated<SILBasicBlock> {
   /// The ordered set of instructions in the SILBasicBlock.
   InstListType InstList;
 
+  /// Used by BasicBlockData to index the Data vector.
+  ///
+  /// A value of -1 means that the index is not initialized yet.
+  int index = -1;
+
   friend struct llvm::ilist_traits<SILBasicBlock>;
   SILBasicBlock() : Parent(nullptr) {}
   void operator=(const SILBasicBlock &) = delete;

include/swift/SIL/SILFunction.h

@@ -147,6 +147,7 @@ class SILFunction
   friend class SILBasicBlock;
   friend class SILModule;
   friend class SILFunctionBuilder;
+  template <typename Data, typename Vector> friend class BasicBlockData;
 
   /// Module - The SIL module that the function belongs to.
   SILModule &Module;
@@ -214,6 +215,11 @@ class SILFunction
   /// It does not include references from debug scopes.
   unsigned RefCount = 0;
 
+  /// Used to verify if a BasicBlockData is not valid anymore.
+  /// This counter is incremented every time a BasicBlockData re-assigns new
+  /// block indices.
+  unsigned BlockListChangeIdx = 0;
+
   /// The function's bare attribute. Bare means that the function is SIL-only
   /// and does not require debug info.
   unsigned Bare : 1;
@@ -1006,6 +1012,9 @@ class SILFunction
   /// Transfer all blocks of \p F into this function, at the begin of the block
   /// list.
   void moveAllBlocksFromOtherFunction(SILFunction *F) {
+    for (SILBasicBlock &block : *F) {
+      block.index = -1;
+    }
     BlockList.splice(begin(), F->BlockList);
   }
 
@@ -1017,6 +1026,7 @@ class SILFunction
     assert(otherFunc != this);
     BlockList.splice(insertPointInThisFunction, otherFunc->BlockList,
                      blockInOtherFunction);
+    blockInOtherFunction->index = -1;
   }
 
   /// Move block \p BB to immediately before the iterator \p IP.