SIL: add a utility which can manage per-block bitfields and flags efficiently.

It is very efficient: no memory allocation is needed an initialization is at zero cost.

Author: eeckstein

include/swift/SIL/SILBasicBlock.h

@@ -34,6 +34,7 @@ public llvm::ilist_node<SILBasicBlock>, public SILAllocated<SILBasicBlock> {
   friend class SILFunction;
   friend class SILGlobalVariable;
   template <typename Data, typename Vector> friend class BasicBlockData;
+  friend class BasicBlockBitfield;
 
 public:
   using InstListType = llvm::iplist<SILInstruction>;
@@ -57,6 +58,25 @@ public llvm::ilist_node<SILBasicBlock>, public SILAllocated<SILBasicBlock> {
   /// A value of -1 means that the index is not initialized yet.
   int index = -1;
 
+  /// Custom bits managed by BasicBlockBitfield.
+  uint32_t customBits = 0;
+  
+  /// The BasicBlockBitfield ID of the last initialized bitfield in customBits.
+  /// Example:
+  ///
+  ///                   Last initialized field:
+  ///           lastInitializedBitfieldID == C.bitfieldID
+  ///                              |
+  ///                              V
+  /// customBits:  <unused> EE DDD C BB AAA
+  ///              31         ...         0
+  ///
+  /// -> AAA, BB and C are initialized,
+  ///    DD and EEE are uninitialized
+  ///
+  /// See also: BasicBlockBitfield::bitfieldID, SILFunction::currentBitfieldID.
+  uint64_t lastInitializedBitfieldID = 0;
+
   friend struct llvm::ilist_traits<SILBasicBlock>;
   SILBasicBlock() : Parent(nullptr) {}
   void operator=(const SILBasicBlock &) = delete;

include/swift/SIL/SILBitfield.h

@@ -0,0 +1,169 @@
+//===--- SILBitField.h - Defines the bitfield utilities ---------*- 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 BasicBlockBitfield and BasicBlockFlag utilities.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_SIL_SILBITFIELD_H
+#define SWIFT_SIL_SILBITFIELD_H
+
+#include "swift/SIL/SILFunction.h"
+
+namespace swift {
+
+/// Utility to add a custom bitfield to a function's basic blocks.
+///
+/// This can be used by transforms to store temporary flags or tiny values per
+/// basic block.
+/// It is very efficient: no memory allocation is needed, no hash set or map is
+/// needed for lookup and there is no initialization cost (in contrast to
+/// BasicBlockData which needs to iterate over all blocks at initialization).
+///
+/// Restrictions:
+/// * BasicBlockBitfield instances must be allocated and deallocated
+///   following a strict stack discipline. This means, it's fine to use them as
+///   (or in) local variables in transformations. But it's e.g. not possible to
+///   store a BasicBlockBitfield in an Analysis.
+/// * The total number of bits which are alive at the same time must not exceed
+///   32.
+class BasicBlockBitfield {
+  /// The bitfield is "added" to the blocks of this function.
+  SILFunction *function;
+
+  /// A single linked list of currently alive BasicBlockBitfields (oldest is
+  /// last, newest is first).
+  /// The head of the list is function->lastAllocatedBitfield.
+  BasicBlockBitfield *parent;
+  
+  /// Initialized with the monotonically increasing currentBitfieldID of the
+  /// function.
+  /// Used to check if the bitfield in a block is initialized.
+  /// If a block's lastInitializedBitfieldID is less than this ID, it means
+  /// that the bits of that block are not initialized yet.
+  /// See also: SILBasicBlock::lastInitializedBitfieldID,
+  ///           SILFunction::currentBitfieldID
+  unsigned bitfieldID;
+
+  short startBit;
+  short endBit;
+  uint32_t mask;
+
+public:
+  BasicBlockBitfield(SILFunction *function, int size) :
+      function(function),
+      parent(function->newestAliveBitfield),
+      bitfieldID(function->currentBitfieldID),
+      startBit(parent ? parent->endBit : 0),
+      endBit(startBit + size),
+      mask(0xffffffffu >> (32 - size) << startBit) {
+    assert(size > 0 && "bit field size must be > 0");
+    assert(endBit <= 32 && "too many/large bit fields allocated in function");
+    assert((!parent || bitfieldID > parent->bitfieldID) &&
+           "BasicBlockBitfield indices are not in order");
+    function->newestAliveBitfield = this;
+    ++function->currentBitfieldID;
+    assert(function->currentBitfieldID != 0 && "currentBitfieldID overflow");
+  }
+
+  ~BasicBlockBitfield() {
+    assert(function->newestAliveBitfield == this &&
+           "BasicBlockBitfield destructed too early");
+    function->newestAliveBitfield = parent;
+  }
+
+  BasicBlockBitfield(const BasicBlockBitfield &) = delete;
+  BasicBlockBitfield(BasicBlockBitfield &&) = delete;
+  BasicBlockBitfield &operator=(const BasicBlockBitfield &) = delete;
+  BasicBlockBitfield &operator=(BasicBlockBitfield &&) = delete;
+
+  SILFunction *getFunction() const { return function; }
+
+  unsigned get(SILBasicBlock *block) const {
+    if (bitfieldID > block->lastInitializedBitfieldID) {
+      // The bitfield is not initialized yet in this block.
+      return 0;
+    }
+    return (block->customBits & mask) >> startBit;
+  }
+
+  void set(SILBasicBlock *block, unsigned value) {
+    assert(((value << startBit) & ~mask) == 0 &&
+           "value too large for BasicBlockBitfield");
+    unsigned clearMask = mask;
+    if (bitfieldID > block->lastInitializedBitfieldID) {
+
+      // The bitfield is not initialized yet in this block.
+      // Initialize the bitfield, and also initialize all parent bitfields,
+      // which are not initialized, yet. Example:
+      //
+      //  This field   Last initialized field
+      //    |            |
+      //    V            V
+      //   EE DDD C BB AAA
+      //
+      // block->lastInitializedBitfieldID == AAA.bitfieldID
+      // -> we have to initialize the fields: BB, C, DDD and EE
+      //
+      BasicBlockBitfield *bf = parent;
+      while (bf && bf->bitfieldID > block->lastInitializedBitfieldID) {
+        clearMask |= bf->mask;
+        bf = bf->parent;
+      }
+      block->lastInitializedBitfieldID = bitfieldID;
+    }
+    block->customBits = (block->customBits & ~clearMask) | (value << startBit);
+  }
+};
+
+/// A BasicBlockBitfield containing a single bit - a flag.
+class BasicBlockFlag {
+  BasicBlockBitfield bit;
+
+public:
+  BasicBlockFlag(SILFunction *function) : bit(function, 1) {}
+
+  SILFunction *getFunction() const { return bit.getFunction(); }
+  
+  bool get(SILBasicBlock *block) const { return (bool)bit.get(block); }
+
+  void set(SILBasicBlock *block) { bit.set(block, 1); }
+  void reset(SILBasicBlock *block) { bit.set(block, 0); }
+
+  /// Sets the flag and returns the old value.
+  bool testAndSet(SILBasicBlock *block) {
+    bool oldValue = get(block);
+    set(block);
+    return oldValue;
+  }
+};
+
+/// A BasicBlockFlag with a set-like API.
+class BasicBlockSet {
+  BasicBlockFlag flag;
+
+public:
+  BasicBlockSet(SILFunction *function) : flag(function) {}
+
+  SILFunction *getFunction() const { return flag.getFunction(); }
+
+  bool contains(SILBasicBlock *block) const { return flag.get(block); }
+
+  /// Returns true if \p block was not contained in the set before inserting.
+  bool insert(SILBasicBlock *block) { return !flag.testAndSet(block); }
+
+  void remove(SILBasicBlock *block) { flag.reset(block); }
+};
+
+} // namespace swift
+
+#endif

include/swift/SIL/SILFunction.h

@@ -37,6 +37,7 @@ class SILInstruction;
 class SILModule;
 class SILFunctionBuilder;
 class SILProfiler;
+class BasicBlockBitfield;
 
 namespace Lowering {
 class TypeLowering;
@@ -148,6 +149,7 @@ class SILFunction
   friend class SILModule;
   friend class SILFunctionBuilder;
   template <typename Data, typename Vector> friend class BasicBlockData;
+  friend class BasicBlockBitfield;
 
   /// Module - The SIL module that the function belongs to.
   SILModule &Module;
@@ -192,6 +194,16 @@ class SILFunction
 
   Identifier ObjCReplacementFor;
 
+  /// The head of a single-linked list of currently alive BasicBlockBitfield.
+  BasicBlockBitfield *newestAliveBitfield = nullptr;
+
+  /// A monotonically increasing ID which is incremented whenever a
+  /// BasicBlockBitfield is constructed.
+  /// Usually this stays below 1000, so a 32-bit unsigned is more than
+  /// sufficient.
+  /// For details see BasicBlockBitfield::bitfieldID;
+  unsigned currentBitfieldID = 1;
+
   /// The function's set of semantics attributes.
   ///
   /// TODO: Why is this using a std::string? Why don't we use uniqued
@@ -1012,9 +1024,6 @@ 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);
   }
 
@@ -1026,7 +1035,6 @@ 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.

lib/SIL/IR/SILBasicBlock.cpp

@@ -309,6 +309,8 @@ transferNodesFromList(llvm::ilist_traits<SILBasicBlock> &SrcTraits,
   // If splicing blocks not in the same function, update the parent pointers.
   for (; First != Last; ++First) {
     First->Parent = Parent;
+    First->index = -1;
+    First->lastInitializedBitfieldID = 0;
     for (auto &II : *First)
       II.setDebugScope(ScopeCloner.getOrCreateClonedScope(II.getDebugScope()));
   }

lib/SIL/IR/SILFunction.cpp

@@ -210,6 +210,8 @@ SILFunction::~SILFunction() {
 
   assert(RefCount == 0 &&
          "Function cannot be deleted while function_ref's still exist");
+  assert(!newestAliveBitfield &&
+         "Not all BasicBlockBitfields deleted at function destruction");
 }
 
 void SILFunction::createProfiler(ASTNode Root, SILDeclRef forDecl,

unittests/CMakeLists.txt

@@ -14,6 +14,7 @@ if(SWIFT_INCLUDE_TOOLS)
   add_subdirectory(IDE)
   add_subdirectory(Parse)
   add_subdirectory(Sema)
+  add_subdirectory(SIL)
   add_subdirectory(SwiftDemangle)
   add_subdirectory(Syntax)
   if(SWIFT_BUILD_SYNTAXPARSERLIB)

unittests/SIL/CMakeLists.txt

@@ -0,0 +1,8 @@
+add_swift_unittest(SwiftSILTests
+  SILBitfieldTest.cpp
+)
+
+target_link_libraries(SwiftSILTests
+   PRIVATE
+   swiftSIL
+)

unittests/SIL/SILBitfieldTest.cpp

@@ -0,0 +1,109 @@
+//===--- SILBitfieldTest.cpp ----------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2014 - 2017 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
+//
+//===----------------------------------------------------------------------===//
+
+#include "gtest/gtest.h"
+
+
+namespace swift {
+
+class BasicBlockBitfield;
+
+// Fake SILFunction and SILBasicBlock.
+
+struct SILFunction {
+  BasicBlockBitfield *newestAliveBitfield = nullptr;
+  uint64_t currentBitfieldID = 1;
+};
+
+struct SILBasicBlock {
+  uint32_t customBits = 0;
+  uint64_t lastInitializedBitfieldID = 0;
+};
+
+}
+
+#define SWIFT_SIL_SILFUNCTION_H
+
+#include "swift/SIL/SILBitfield.h"
+
+using namespace swift;
+
+namespace {
+
+TEST(SILBitfieldTest, Basic) {
+  SILFunction f;
+  SILBasicBlock b;
+
+  {
+    BasicBlockFlag A(&f);
+    EXPECT_FALSE(A.get(&b));
+
+    EXPECT_FALSE(A.testAndSet(&b));
+    EXPECT_TRUE(A.get(&b));
+
+    EXPECT_TRUE(A.testAndSet(&b));
+    EXPECT_TRUE(A.get(&b));
+
+    A.reset(&b);
+    EXPECT_FALSE(A.get(&b));
+
+    {
+      BasicBlockBitfield B(&f, 5);
+      EXPECT_EQ(B.get(&b), 0u);
+
+      B.set(&b, 27);
+      EXPECT_FALSE(A.get(&b));
+      EXPECT_EQ(B.get(&b), 27u);
+
+      A.set(&b);
+      EXPECT_TRUE(A.get(&b));
+      EXPECT_EQ(B.get(&b), 27u);
+
+      B.set(&b, 2);
+      EXPECT_TRUE(A.get(&b));
+      EXPECT_EQ(B.get(&b), 2u);
+
+      B.set(&b, 31);
+      EXPECT_TRUE(A.get(&b));
+      EXPECT_EQ(B.get(&b), 31u);
+    }
+    {
+      BasicBlockBitfield C(&f, 2);
+      EXPECT_EQ(C.get(&b), 0u);
+
+      BasicBlockFlag D(&f);
+      EXPECT_FALSE(D.get(&b));
+
+      BasicBlockBitfield E(&f, 3);
+
+      E.set(&b, 5);
+      EXPECT_TRUE(A.get(&b));
+      EXPECT_EQ(C.get(&b), 0u);
+      EXPECT_FALSE(D.get(&b));
+      EXPECT_EQ(E.get(&b), 5u);
+
+      C.set(&b, 1);
+      EXPECT_TRUE(A.get(&b));
+      EXPECT_EQ(C.get(&b), 1u);
+      EXPECT_FALSE(D.get(&b));
+      EXPECT_EQ(E.get(&b), 5u);
+    }
+  }
+  {
+    BasicBlockBitfield F(&f, 32);
+    EXPECT_EQ(F.get(&b), 0u);
+    F.set(&b, 0xdeadbeaf);
+    EXPECT_EQ(F.get(&b), 0xdeadbeaf);
+  }
+}
+
+}