LowerBitSets: Extend pass to support functions as bitset members.

[pcc]

Summary

This change extends the bitset lowering pass to support bitsets that may
contain either functions or global variables. A function bitset is lowered to
a jump table that is laid out before one of the functions in the bitset.

Also add support for non-string bitset identifier names. This allows for
distinct metadata nodes to stand in for names with internal linkage,
as done in D11857.

[pcc]

• Permit non-string metadata as a bitset identifier

[pcc]

• Update docs to reflect metadata identifier change

[jfb]

lib/Transforms/IPO/LowerBitSets.cpp

227 ↗

(On Diff #31872)

Comment is out of date.

571 ↗

(On Diff #31872)

I'm not sure I understand padding elements. Could you explain further?

579 ↗

(On Diff #31872)

Does this work with thread-local globals?
I'm also a bit worried about visibility / section / linkage type, but I may be misunderstanding the implications of what this patch adds.

589 ↗

(On Diff #31872)

Could this be a virtual function instead? report_fatal_error by default, and only implemented for x86. I'm assuming the creation will also need to be specialized.

651 ↗

(On Diff #31872)

Can Globals ever be empty? It seems a bit silly, but not entirely unlikely if someone just runs the pass.

666 ↗

(On Diff #31872)

Give it some dummy internal name starting with __ so it's easier to figure things out. I'm assuming there's a test that'll check for that function name too.

695 ↗

(On Diff #31872)

clang-format, here and a few other places.

889 ↗

(On Diff #31872)

IIUC duplicate indices can't occur?

906 ↗

(On Diff #31872)

Eek, the above lines make me cry. LLVM doesn't have a way to create ArrayRef taking into account its magical type hierarchy? It would be nice to avoid the reinterpret_cast here, since it relies heavily on the GlobalObject / GlobalVariable / Function class hierarchy layout...

[pcc]

• Address review comments

lib/Transforms/IPO/LowerBitSets.cpp

227 ↗

(On Diff #31872)

Updated.

571 ↗

(On Diff #31872)

I've added a little more explanation to the comments above. Regarding the purpose of the padding elements, this is already documented at http://clang.llvm.org/docs/ControlFlowIntegrityDesign.html#padding-to-powers-of-2

579 ↗

(On Diff #31872)

Does this work with thread-local globals?

It does not. I've added a couple of checks below that ensure that no thread-local globals or globals with an explicit section can be members of bitsets.

I'm also a bit worried about visibility / section / linkage type

Conceptually, linkage and visiblity are both properties of the symbol, rather than of its position in the object file (as thread-localness and section are), so they can be copied from the original entity to the alias.

Linkage is already copied to the alias. I'm also copying visibility now.

I may be misunderstanding the implications of what this patch adds.

This isn't actually new code; see line 579 of the old code below. For some reason the diff doesn't match up correctly here. But thanks all the same for raising these issues :)

589 ↗

(On Diff #31872)

It could be, but I would prefer to do this when we add support for another architecture (which is mostly blocked on an implementation of http://lists.llvm.org/pipermail/llvm-dev/2015-July/088748.html or something similar). Once something like that goes in, we should be able to simplify this code significantly as well.

651 ↗

(On Diff #31872)

It cannot. The buildBitSetsFromGlobalVariables function handles any empty global sets -- see line 894 of the new code. I've added a comment here to make this clear.

666 ↗

(On Diff #31872)

Done, but it isn't too useful as we end up renaming all functions appearing in bitsets anyway. Not sure if there's much we can do about that without making this code even more complicated.

695 ↗

(On Diff #31872)

Done

889 ↗

(On Diff #31872)

Correct, clarified in the comments.

906 ↗

(On Diff #31872)

LLVM doesn't have a way to create ArrayRef taking into account its magical type hierarchy?

I'm not sure if there can be a general way to do this due to covariance/contravariance issues.

It would be nice to avoid the reinterpret_cast here, since it relies heavily on the GlobalObject / GlobalVariable / Function class hierarchy layout...

I decided to just rebuild the array separately on each branch of this conditional. It's a little duplicative, but the code is probably clearest this way.

[jfb]

A few comments, but lgtm overall. It would be nice to get another review on this as it's somewhat tricky stuff :-)

lib/Transforms/IPO/LowerBitSets.cpp

591 ↗

(On Diff #32783)

That sounds fine, though the documentation should be updated the the appropriate x86 limitations, and potentially evolution out of this.

654 ↗

(On Diff #32783)

Could you just assert at the function entry? Code tends to evolve past comments ;-)

[pcc]

• Address comments

[kcc]

docs/BitSets.rst

13 ↗

(On Diff #32789)

maybe a bullet/numbered list?

31 ↗

(On Diff #32789)

You introduce the concept of jump tables w/o explaining them.
Maybe add a paragraph telling what jump table is?

lib/Transforms/IPO/LowerBitSets.cpp

749 ↗

(On Diff #32789)

quite a bit of code in a single function. Maybe split it into several?

902 ↗

(On Diff #32789)

Two pieces of very similar code.
Can we use some C++ magic to not duplicate the code (template? use the parent class?)

[pcc]

• Address kcc comments

docs/BitSets.rst

13 ↗

(On Diff #32789)

Done

31 ↗

(On Diff #32789)

Done

lib/Transforms/IPO/LowerBitSets.cpp

749 ↗

(On Diff #32789)

This and one other part of the code is now split out.

902 ↗

(On Diff #32789)

We could, but I felt that it would be clearest to write the code out on each branch. It's not a lot of code.

[kcc]

I am still trying to wrap around the code, failing so far.
Maybe you could add a high-level algorithm description in the comments?
Also, two more minor request for more comments.

docs/BitSets.rst

33 ↗

(On Diff #33036)

jump table feels like a key concept here and deserves a separate paragraph with some more details,
maybe even an illustration

test/Transforms/LowerBitSets/function-ext.ll

1 ↗

(On Diff #33036)

here and below, please add tests explaining what we are testing for

[pcc]

• Add more docs

[pcc]

Maybe you could add a high-level algorithm description in the comments?

Done

docs/BitSets.rst

33 ↗

(On Diff #33036)

Done,

test/Transforms/LowerBitSets/function-ext.ll

1 ↗

(On Diff #33036)

Done

[kcc]

lib/Transforms/IPO/LowerBitSets.cpp

725 ↗

(On Diff #33176)

Why do we do this?

Yes, this will reduce the size of the jump table by 8 bytes,
but is this optimization worth the extra complexity?
We must try to be as simple as possible

[pcc]

lib/Transforms/IPO/LowerBitSets.cpp

725 ↗

(On Diff #33176)

The motivation wasn't so much optimization but the fact that the jump table is guaranteed to appear in an executable section if we do this. However, if it would work just to set the section name it may be best to do that instead (and fix it if not). Let me take a look.

[krasin]

A few shallow comments...

lib/Transforms/IPO/LowerBitSets.cpp

609 ↗

(On Diff #33176)

What is 0xe9?

781 ↗

(On Diff #33176)

What happens, if the last function isDiscardableIfUnused? Could it be that a consequent pass would drop the whole jump table?

785 ↗

(On Diff #33176)

May be a lame question: what if __function_bitset_dummy is taken?

test/Transforms/LowerBitSets/function-ext.ll

4 ↗

(On Diff #33176)

a minor typo: s/aare/are/

test/Transforms/LowerBitSets/nonstring.ll

21 ↗

(On Diff #33176)

Other newly added tests have CHECK statements inside foo / bar / baz functions, and this test does not. Is it intentional?

[pcc]

• Address comments

lib/Transforms/IPO/LowerBitSets.cpp

609 ↗

(On Diff #33176)

This would be the opcode for the PC-relative jmp instruction. I thought this was clear enough from the variable name and how it is used, but maybe not?

725 ↗

(On Diff #33176)

Okay, done. The code is much simpler now.

781 ↗

(On Diff #33176)

This comment is obsoleted by the move to a global variable for the jump table. (But any uses of the jump table would also count as uses of the function, so this shouldn't matter.)

785 ↗

(On Diff #33176)

Also obsolete

test/Transforms/LowerBitSets/function-ext.ll

4 ↗

(On Diff #33176)

Fixed and a few other typos here

test/Transforms/LowerBitSets/nonstring.ll

21 ↗

(On Diff #33176)

I suppose we could do that here as well; done.

[kcc]

I think I will soon run out of comments. :)

lib/Transforms/IPO/LowerBitSets.cpp

672 ↗

(On Diff #33286)

instead of 0xe9, 8, 5, and 0xcc in this function please define something like

static const int kIntCode = 0xcc;
static const int kJmpCode = 0xe9;

static const int kOffsetDispSize = 4;
...

use getJumpTableEntrySize instead of '8'

703 ↗

(On Diff #33286)

Should this be named "Functions"?

739 ↗

(On Diff #33286)

Is this part of comment still relevant?

[pcc]

• Update documentation
• Address review comments

[pcc]

lib/Transforms/IPO/LowerBitSets.cpp

620 ↗

(On Diff #34100)

Done

651 ↗

(On Diff #34100)

Done

687 ↗

(On Diff #34100)

Done

[kcc]

LGTM

I am out of useful comments (except one minor one below)

lib/Transforms/IPO/LowerBitSets.cpp

656 ↗

(On Diff #34100)

static, please (or anon namespace, or a class member )

llvm/trunk/docs/BitSets.rst

 =======
 Bitsets
 =======
 
 This is a mechanism that allows IR modules to co-operatively build pointer
 sets corresponding to addresses within a given set of globals. One example
 of a use case for this is to allow a C++ program to efficiently verify (at
 each call site) that a vtable pointer is in the set of valid vtable pointers
 for the type of the class or its derived classes.
 
 To use the mechanism, a client creates a global metadata node named
 ``llvm.bitsets``.  Each element is a metadata node with three elements:
-the first is a metadata string containing an identifier for the bitset,
-the second is a global variable and the third is a byte offset into the
-global variable.
+
+1. a metadata object representing an identifier for the bitset
+2. either a global variable or a function
+3. a byte offset into the global (generally zero for functions)
+
+Each bitset must exclusively contain either global variables or functions.
+
+.. admonition:: Limitation
+
+  The current implementation only supports functions as members of bitsets on
+  the x86-32 and x86-64 architectures.
 
 This will cause a link-time optimization pass to generate bitsets from the
-memory addresses referenced from the elements of the bitset metadata. The pass
-will lay out the referenced globals consecutively, so their definitions must
-be available at LTO time. The `GlobalLayoutBuilder`_ class is responsible for
-laying out the globals efficiently to minimize the sizes of the underlying
-bitsets. An intrinsic, :ref:`llvm.bitset.test <bitset.test>`, generates code
-to test whether a given pointer is a member of a bitset.
+memory addresses referenced from the elements of the bitset metadata. The
+pass will lay out referenced global variables consecutively, so their
+definitions must be available at LTO time.
+
+A bit set containing functions is transformed into a jump table, which
+is a block of code consisting of one branch instruction for each of the
+functions in the bit set that branches to the target function, and redirect
+any taken function addresses to the corresponding jump table entry. In the
+object file's symbol table, the jump table entries take the identities of
+the original functions, so that addresses taken outside the module will pass
+any verification done inside the module.
+
+Jump tables may call external functions, so their definitions need not
+be available at LTO time. Note that if an externally defined function is a
+member of a bitset, there is no guarantee that its identity within the module
+will be the same as its identity outside of the module, as the former will
+be the jump table entry if a jump table is necessary.
+
+The `GlobalLayoutBuilder`_ class is responsible for laying out the globals
+efficiently to minimize the sizes of the underlying bitsets. An intrinsic,
+:ref:`llvm.bitset.test <bitset.test>`, generates code to test whether a
+given pointer is a member of a bitset.
 
 :Example:
 
 ::
 
     target datalayout = "e-p:32:32"
 
     @a = internal global i32 0
     @b = internal global i32 0
     @c = internal global i32 0
     @d = internal global [2 x i32] [i32 0, i32 0]
 
-    !llvm.bitsets = !{!0, !1, !2, !3, !4}
+    define void @e() {
+      ret void
+    }
+
+    define void @f() {
+      ret void
+    }
+
+    declare void @g()
+
+    !llvm.bitsets = !{!0, !1, !2, !3, !4, !5, !6}
 
     !0 = !{!"bitset1", i32* @a, i32 0}
     !1 = !{!"bitset1", i32* @b, i32 0}
     !2 = !{!"bitset2", i32* @b, i32 0}
     !3 = !{!"bitset2", i32* @c, i32 0}
     !4 = !{!"bitset2", i32* @d, i32 4}
+    !5 = !{!"bitset3", void ()* @e, i32 0}
+    !6 = !{!"bitset3", void ()* @g, i32 0}
 
     declare i1 @llvm.bitset.test(i8* %ptr, metadata %bitset) nounwind readnone
 
     define i1 @foo(i32* %p) {
       %pi8 = bitcast i32* %p to i8*
       %x = call i1 @llvm.bitset.test(i8* %pi8, metadata !"bitset1")
       ret i1 %x
     }
 
     define i1 @bar(i32* %p) {
       %pi8 = bitcast i32* %p to i8*
       %x = call i1 @llvm.bitset.test(i8* %pi8, metadata !"bitset2")
       ret i1 %x
     }
 
+    define i1 @baz(void ()* %p) {
+      %pi8 = bitcast void ()* %p to i8*
+      %x = call i1 @llvm.bitset.test(i8* %pi8, metadata !"bitset3")
+      ret i1 %x
+    }
+
     define void @main() {
       %a1 = call i1 @foo(i32* @a) ; returns 1
       %b1 = call i1 @foo(i32* @b) ; returns 1
       %c1 = call i1 @foo(i32* @c) ; returns 0
       %a2 = call i1 @bar(i32* @a) ; returns 0
       %b2 = call i1 @bar(i32* @b) ; returns 1
       %c2 = call i1 @bar(i32* @c) ; returns 1
       %d02 = call i1 @bar(i32* getelementptr ([2 x i32]* @d, i32 0, i32 0)) ; returns 0
       %d12 = call i1 @bar(i32* getelementptr ([2 x i32]* @d, i32 0, i32 1)) ; returns 1
+      %e = call i1 @baz(void ()* @e) ; returns 1
+      %f = call i1 @baz(void ()* @f) ; returns 0
+      %g = call i1 @baz(void ()* @g) ; returns 1
       ret void
     }
 
 .. _GlobalLayoutBuilder: http://llvm.org/klaus/llvm/blob/master/include/llvm/Transforms/IPO/LowerBitSets.h

llvm/trunk/docs/LangRef.rst

::
 
       declare i1 @llvm.bitset.test(i8* %ptr, metadata %bitset) nounwind readnone
 
 
 Arguments:
 """"""""""
 
 The first argument is a pointer to be tested. The second argument is a
-metadata string containing the name of a :doc:`bitset <BitSets>`.
+metadata object representing an identifier for a :doc:`bitset <BitSets>`.
 
 Overview:
 """""""""
 
 The ``llvm.bitset.test`` intrinsic tests whether the given pointer is a
 member of the given bitset.
 
 '``llvm.donothing``' Intrinsic

llvm/trunk/include/llvm/Transforms/IPO/LowerBitSets.h

 #include <stdint.h>
 #include <limits>
 #include <set>
 #include <vector>
 
 namespace llvm {
 
 class DataLayout;
-class GlobalVariable;
+class GlobalObject;
 class Value;
 class raw_ostream;
 
 struct BitSetInfo {
   // The indices of the set bits in the bitset.
   std::set<uint64_t> Bits;
 
   // The byte offset into the combined global represented by the bitset.
struct BitSetInfo {
 
   bool isAllOnes() const {
     return Bits.size() == BitSize;
   }
 
   bool containsGlobalOffset(uint64_t Offset) const;
 
   bool containsValue(const DataLayout &DL,
-                     const DenseMap<GlobalVariable *, uint64_t> &GlobalLayout,
+                     const DenseMap<GlobalObject *, uint64_t> &GlobalLayout,
                      Value *V, uint64_t COffset = 0) const;
 
   void print(raw_ostream &OS) const;
 };
 
 struct BitSetBuilder {
   SmallVector<uint64_t, 16> Offsets;
   uint64_t Min, Max;

llvm/trunk/lib/Transforms/IPO/LowerBitSets.cpp

 
 #include "llvm/Transforms/IPO/LowerBitSets.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/ADT/EquivalenceClasses.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/GlobalObject.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Debug.h"
bool BitSetInfo::containsGlobalOffset(uint64_t Offset) const {
   if (BitOffset >= BitSize)
     return false;
 
   return Bits.count(BitOffset);
 }
 
 bool BitSetInfo::containsValue(
     const DataLayout &DL,
-    const DenseMap<GlobalVariable *, uint64_t> &GlobalLayout, Value *V,
+    const DenseMap<GlobalObject *, uint64_t> &GlobalLayout, Value *V,
     uint64_t COffset) const {
-  if (auto GV = dyn_cast<GlobalVariable>(V)) {
+  if (auto GV = dyn_cast<GlobalObject>(V)) {
     auto I = GlobalLayout.find(GV);
     if (I == GlobalLayout.end())
       return false;
     return containsGlobalOffset(I->second + COffset);
   }
 
   if (auto GEP = dyn_cast<GEPOperator>(V)) {
     APInt APOffset(DL.getPointerSizeInBits(0), 0);
struct LowerBitSets : public ModulePass {
   static char ID;
   LowerBitSets() : ModulePass(ID) {
     initializeLowerBitSetsPass(*PassRegistry::getPassRegistry());
   }
 
   Module *M;
 
   bool LinkerSubsectionsViaSymbols;
+  Triple::ArchType Arch;
+  Triple::ObjectFormatType ObjectFormat;
   IntegerType *Int1Ty;
   IntegerType *Int8Ty;
   IntegerType *Int32Ty;
   Type *Int32PtrTy;
   IntegerType *Int64Ty;
-  Type *IntPtrTy;
+  IntegerType *IntPtrTy;
 
   // The llvm.bitsets named metadata.
   NamedMDNode *BitSetNM;
 
-  // Mapping from bitset mdstrings to the call sites that test them.
-  DenseMap<MDString *, std::vector<CallInst *>> BitSetTestCallSites;
+  // Mapping from bitset identifiers to the call sites that test them.
+  DenseMap<Metadata *, std::vector<CallInst *>> BitSetTestCallSites;
 
   std::vector<ByteArrayInfo> ByteArrayInfos;
 
   BitSetInfo
-  buildBitSet(MDString *BitSet,
-              const DenseMap<GlobalVariable *, uint64_t> &GlobalLayout);
+  buildBitSet(Metadata *BitSet,
+              const DenseMap<GlobalObject *, uint64_t> &GlobalLayout);
   ByteArrayInfo *createByteArray(BitSetInfo &BSI);
   void allocateByteArrays();
   Value *createBitSetTest(IRBuilder<> &B, BitSetInfo &BSI, ByteArrayInfo *&BAI,
                           Value *BitOffset);
+  void lowerBitSetCalls(ArrayRef<Metadata *> BitSets,
+                        Constant *CombinedGlobalAddr,
+                        const DenseMap<GlobalObject *, uint64_t> &GlobalLayout);
   Value *
   lowerBitSetCall(CallInst *CI, BitSetInfo &BSI, ByteArrayInfo *&BAI,
-                  GlobalVariable *CombinedGlobal,
-                  const DenseMap<GlobalVariable *, uint64_t> &GlobalLayout);
-  void buildBitSetsFromGlobals(const std::vector<MDString *> &BitSets,
-                               const std::vector<GlobalVariable *> &Globals);
+                  Constant *CombinedGlobal,
+                  const DenseMap<GlobalObject *, uint64_t> &GlobalLayout);
+  void buildBitSetsFromGlobalVariables(ArrayRef<Metadata *> BitSets,
+                                       ArrayRef<GlobalVariable *> Globals);
+  unsigned getJumpTableEntrySize();
+  Type *getJumpTableEntryType();
+  Constant *createJumpTableEntry(GlobalObject *Src, Function *Dest,
+                                 unsigned Distance);
+  void verifyBitSetMDNode(MDNode *Op);
+  void buildBitSetsFromFunctions(ArrayRef<Metadata *> BitSets,
+                                 ArrayRef<Function *> Functions);
+  void buildBitSetsFromDisjointSet(ArrayRef<Metadata *> BitSets,
+                                   ArrayRef<GlobalObject *> Globals);
   bool buildBitSets();
   bool eraseBitSetMetadata();
 
   bool doInitialization(Module &M) override;
   bool runOnModule(Module &M) override;
 };
 
 } // namespace
 
 INITIALIZE_PASS_BEGIN(LowerBitSets, "lowerbitsets",
                 "Lower bitset metadata", false, false)
 INITIALIZE_PASS_END(LowerBitSets, "lowerbitsets",
                 "Lower bitset metadata", false, false)
 char LowerBitSets::ID = 0;
 
 ModulePass *llvm::createLowerBitSetsPass() { return new LowerBitSets; }
 
 bool LowerBitSets::doInitialization(Module &Mod) {
   M = &Mod;
   const DataLayout &DL = Mod.getDataLayout();
 
   Triple TargetTriple(M->getTargetTriple());
   LinkerSubsectionsViaSymbols = TargetTriple.isMacOSX();
+  Arch = TargetTriple.getArch();
+  ObjectFormat = TargetTriple.getObjectFormat();
 
   Int1Ty = Type::getInt1Ty(M->getContext());
   Int8Ty = Type::getInt8Ty(M->getContext());
   Int32Ty = Type::getInt32Ty(M->getContext());
   Int32PtrTy = PointerType::getUnqual(Int32Ty);
   Int64Ty = Type::getInt64Ty(M->getContext());
   IntPtrTy = DL.getIntPtrType(M->getContext(), 0);
 
   BitSetNM = M->getNamedMetadata("llvm.bitsets");
 
   BitSetTestCallSites.clear();
 
   return false;
 }
 
 /// Build a bit set for BitSet using the object layouts in
 /// GlobalLayout.
 BitSetInfo LowerBitSets::buildBitSet(
-    MDString *BitSet,
-    const DenseMap<GlobalVariable *, uint64_t> &GlobalLayout) {
+    Metadata *BitSet,
+    const DenseMap<GlobalObject *, uint64_t> &GlobalLayout) {
   BitSetBuilder BSB;
 
   // Compute the byte offset of each element of this bitset.
   if (BitSetNM) {
     for (MDNode *Op : BitSetNM->operands()) {
       if (Op->getOperand(0) != BitSet || !Op->getOperand(1))
         continue;
-      auto OpGlobal = dyn_cast<GlobalVariable>(
-          cast<ConstantAsMetadata>(Op->getOperand(1))->getValue());
+      Constant *OpConst =
+          cast<ConstantAsMetadata>(Op->getOperand(1))->getValue();
+      if (auto GA = dyn_cast<GlobalAlias>(OpConst))
+        OpConst = GA->getAliasee();
+      auto OpGlobal = dyn_cast<GlobalObject>(OpConst);
       if (!OpGlobal)
         continue;
       uint64_t Offset =
           cast<ConstantInt>(cast<ConstantAsMetadata>(Op->getOperand(2))
                                 ->getValue())->getZExtValue();
 
       Offset += GlobalLayout.find(OpGlobal)->second;
 
if (BSI.BitSize <= 64) {
     return B.CreateICmpNE(ByteAndMask, ConstantInt::get(Int8Ty, 0));
   }
 }
 
 /// Lower a llvm.bitset.test call to its implementation. Returns the value to
 /// replace the call with.
 Value *LowerBitSets::lowerBitSetCall(
     CallInst *CI, BitSetInfo &BSI, ByteArrayInfo *&BAI,
-    GlobalVariable *CombinedGlobal,
-    const DenseMap<GlobalVariable *, uint64_t> &GlobalLayout) {
+    Constant *CombinedGlobalIntAddr,
+    const DenseMap<GlobalObject *, uint64_t> &GlobalLayout) {
   Value *Ptr = CI->getArgOperand(0);
   const DataLayout &DL = M->getDataLayout();
 
   if (BSI.containsValue(DL, GlobalLayout, Ptr))
-    return ConstantInt::getTrue(CombinedGlobal->getParent()->getContext());
+    return ConstantInt::getTrue(M->getContext());
 
-  Constant *GlobalAsInt = ConstantExpr::getPtrToInt(CombinedGlobal, IntPtrTy);
   Constant *OffsetedGlobalAsInt = ConstantExpr::getAdd(
-      GlobalAsInt, ConstantInt::get(IntPtrTy, BSI.ByteOffset));
+      CombinedGlobalIntAddr, ConstantInt::get(IntPtrTy, BSI.ByteOffset));
 
   BasicBlock *InitialBB = CI->getParent();
 
   IRBuilder<> B(CI);
 
   Value *PtrAsInt = B.CreatePtrToInt(Ptr, IntPtrTy);
 
   if (BSI.isSingleOffset())
Value *LowerBitSets::lowerBitSetCall(
   PHINode *P = B.CreatePHI(Int1Ty, 2);
   P->addIncoming(ConstantInt::get(Int1Ty, 0), InitialBB);
   P->addIncoming(Bit, ThenB.GetInsertBlock());
   return P;
 }
 
 /// Given a disjoint set of bitsets and globals, layout the globals, build the
 /// bit sets and lower the llvm.bitset.test calls.
-void LowerBitSets::buildBitSetsFromGlobals(
-    const std::vector<MDString *> &BitSets,
-    const std::vector<GlobalVariable *> &Globals) {
+void LowerBitSets::buildBitSetsFromGlobalVariables(
+    ArrayRef<Metadata *> BitSets, ArrayRef<GlobalVariable *> Globals) {
   // Build a new global with the combined contents of the referenced globals.
+  // This global is a struct whose even-indexed elements contain the original
+  // contents of the referenced globals and whose odd-indexed elements contain
+  // any padding required to align the next element to the next power of 2.
   std::vector<Constant *> GlobalInits;
   const DataLayout &DL = M->getDataLayout();
   for (GlobalVariable *G : Globals) {
     GlobalInits.push_back(G->getInitializer());
     uint64_t InitSize = DL.getTypeAllocSize(G->getInitializer()->getType());
 
-    // Compute the amount of padding required to align the next element to the
-    // next power of 2.
+    // Compute the amount of padding required.
     uint64_t Padding = NextPowerOf2(InitSize - 1) - InitSize;
 
     // Cap at 128 was found experimentally to have a good data/instruction
     // overhead tradeoff.
     if (Padding > 128)
       Padding = RoundUpToAlignment(InitSize, 128) - InitSize;
 
     GlobalInits.push_back(
         ConstantAggregateZero::get(ArrayType::get(Int8Ty, Padding)));
   }
   if (!GlobalInits.empty())
     GlobalInits.pop_back();
   Constant *NewInit = ConstantStruct::getAnon(M->getContext(), GlobalInits);
   auto CombinedGlobal =
       new GlobalVariable(*M, NewInit->getType(), /*isConstant=*/true,
                          GlobalValue::PrivateLinkage, NewInit);
 
   const StructLayout *CombinedGlobalLayout =
       DL.getStructLayout(cast<StructType>(NewInit->getType()));
 
   // Compute the offsets of the original globals within the new global.
-  DenseMap<GlobalVariable *, uint64_t> GlobalLayout;
+  DenseMap<GlobalObject *, uint64_t> GlobalLayout;
   for (unsigned I = 0; I != Globals.size(); ++I)
     // Multiply by 2 to account for padding elements.
     GlobalLayout[Globals[I]] = CombinedGlobalLayout->getElementOffset(I * 2);
 
-  // For each bitset in this disjoint set...
-  for (MDString *BS : BitSets) {
-    // Build the bitset.
-    BitSetInfo BSI = buildBitSet(BS, GlobalLayout);
-    DEBUG({
-      dbgs() << BS->getString() << ": ";
-      BSI.print(dbgs());
-    });
-
-    ByteArrayInfo *BAI = 0;
-
-    // Lower each call to llvm.bitset.test for this bitset.
-    for (CallInst *CI : BitSetTestCallSites[BS]) {
-      ++NumBitSetCallsLowered;
-      Value *Lowered = lowerBitSetCall(CI, BSI, BAI, CombinedGlobal, GlobalLayout);
-      CI->replaceAllUsesWith(Lowered);
-      CI->eraseFromParent();
-    }
-  }
+  lowerBitSetCalls(BitSets, CombinedGlobal, GlobalLayout);
 
   // Build aliases pointing to offsets into the combined global for each
   // global from which we built the combined global, and replace references
   // to the original globals with references to the aliases.
   for (unsigned I = 0; I != Globals.size(); ++I) {
     // Multiply by 2 to account for padding elements.
     Constant *CombinedGlobalIdxs[] = {ConstantInt::get(Int32Ty, 0),
                                       ConstantInt::get(Int32Ty, I * 2)};
     Constant *CombinedGlobalElemPtr = ConstantExpr::getGetElementPtr(
         NewInit->getType(), CombinedGlobal, CombinedGlobalIdxs);
     if (LinkerSubsectionsViaSymbols) {
       Globals[I]->replaceAllUsesWith(CombinedGlobalElemPtr);
     } else {
       GlobalAlias *GAlias =
           GlobalAlias::create(Globals[I]->getType(), Globals[I]->getLinkage(),
                               "", CombinedGlobalElemPtr, M);
+      GAlias->setVisibility(Globals[I]->getVisibility());
       GAlias->takeName(Globals[I]);
       Globals[I]->replaceAllUsesWith(GAlias);
     }
     Globals[I]->eraseFromParent();
   }
 }
 
+void LowerBitSets::lowerBitSetCalls(
+    ArrayRef<Metadata *> BitSets, Constant *CombinedGlobalAddr,
+    const DenseMap<GlobalObject *, uint64_t> &GlobalLayout) {
+  Constant *CombinedGlobalIntAddr =
+      ConstantExpr::getPtrToInt(CombinedGlobalAddr, IntPtrTy);
+
+  // For each bitset in this disjoint set...
+  for (Metadata *BS : BitSets) {
+    // Build the bitset.
+    BitSetInfo BSI = buildBitSet(BS, GlobalLayout);
+    DEBUG({
+      if (auto BSS = dyn_cast<MDString>(BS))
+        dbgs() << BSS->getString() << ": ";
+      else
+        dbgs() << "<unnamed>: ";
+      BSI.print(dbgs());
+    });
+
+    ByteArrayInfo *BAI = 0;
+
+    // Lower each call to llvm.bitset.test for this bitset.
+    for (CallInst *CI : BitSetTestCallSites[BS]) {
+      ++NumBitSetCallsLowered;
+      Value *Lowered =
+          lowerBitSetCall(CI, BSI, BAI, CombinedGlobalIntAddr, GlobalLayout);
+      CI->replaceAllUsesWith(Lowered);
+      CI->eraseFromParent();
+    }
+  }
+}
+
+void LowerBitSets::verifyBitSetMDNode(MDNode *Op) {
+  if (Op->getNumOperands() != 3)
+    report_fatal_error(
+        "All operands of llvm.bitsets metadata must have 3 elements");
+  if (!Op->getOperand(1))
+    return;
+
+  auto OpConstMD = dyn_cast<ConstantAsMetadata>(Op->getOperand(1));
+  if (!OpConstMD)
+    report_fatal_error("Bit set element must be a constant");
+  auto OpGlobal = dyn_cast<GlobalObject>(OpConstMD->getValue());
+  if (!OpGlobal)
+    return;
+
+  if (OpGlobal->isThreadLocal())
+    report_fatal_error("Bit set element may not be thread-local");
+  if (OpGlobal->hasSection())
+    report_fatal_error("Bit set element may not have an explicit section");
+
+  if (isa<GlobalVariable>(OpGlobal) && OpGlobal->isDeclarationForLinker())
+    report_fatal_error("Bit set global var element must be a definition");
+
+  auto OffsetConstMD = dyn_cast<ConstantAsMetadata>(Op->getOperand(2));
+  if (!OffsetConstMD)
+    report_fatal_error("Bit set element offset must be a constant");
+  auto OffsetInt = dyn_cast<ConstantInt>(OffsetConstMD->getValue());
+  if (!OffsetInt)
+    report_fatal_error("Bit set element offset must be an integer constant");
+}
+
+static const unsigned kX86JumpTableEntrySize = 8;
+
+unsigned LowerBitSets::getJumpTableEntrySize() {
+  if (Arch != Triple::x86 && Arch != Triple::x86_64)
+    report_fatal_error("Unsupported architecture for jump tables");
+
+  return kX86JumpTableEntrySize;
+}
+
+// Create a constant representing a jump table entry for the target. This
+// consists of an instruction sequence containing a relative branch to Dest. The
+// constant will be laid out at address Src+(Len*Distance) where Len is the
+// target-specific jump table entry size.
+Constant *LowerBitSets::createJumpTableEntry(GlobalObject *Src, Function *Dest,
+                                             unsigned Distance) {
+  if (Arch != Triple::x86 && Arch != Triple::x86_64)
+    report_fatal_error("Unsupported architecture for jump tables");
+
+  const unsigned kJmpPCRel32Code = 0xe9;
+  const unsigned kInt3Code = 0xcc;
+
+  ConstantInt *Jmp = ConstantInt::get(Int8Ty, kJmpPCRel32Code);
+
+  // Build a constant representing the displacement between the constant's
+  // address and Dest. This will resolve to a PC32 relocation referring to Dest.
+  Constant *DestInt = ConstantExpr::getPtrToInt(Dest, IntPtrTy);
+  Constant *SrcInt = ConstantExpr::getPtrToInt(Src, IntPtrTy);
+  Constant *Disp = ConstantExpr::getSub(DestInt, SrcInt);
+  ConstantInt *DispOffset =
+      ConstantInt::get(IntPtrTy, Distance * kX86JumpTableEntrySize + 5);
+  Constant *OffsetedDisp = ConstantExpr::getSub(Disp, DispOffset);
+  OffsetedDisp = ConstantExpr::getTrunc(OffsetedDisp, Int32Ty);
+
+  ConstantInt *Int3 = ConstantInt::get(Int8Ty, kInt3Code);
+
+  Constant *Fields[] = {
+      Jmp, OffsetedDisp, Int3, Int3, Int3,
+  };
+  return ConstantStruct::getAnon(Fields, /*Packed=*/true);
+}
+
+Type *LowerBitSets::getJumpTableEntryType() {
+  if (Arch != Triple::x86 && Arch != Triple::x86_64)
+    report_fatal_error("Unsupported architecture for jump tables");
+
+  return StructType::get(M->getContext(),
+                         {Int8Ty, Int32Ty, Int8Ty, Int8Ty, Int8Ty},
+                         /*Packed=*/true);
+}
+
+/// Given a disjoint set of bitsets and functions, build a jump table for the
+/// functions, build the bit sets and lower the llvm.bitset.test calls.
+void LowerBitSets::buildBitSetsFromFunctions(ArrayRef<Metadata *> BitSets,
+                                             ArrayRef<Function *> Functions) {
+  // Unlike the global bitset builder, the function bitset builder cannot
+  // re-arrange functions in a particular order and base its calculations on the
+  // layout of the functions' entry points, as we have no idea how large a
+  // particular function will end up being (the size could even depend on what
+  // this pass does!) Instead, we build a jump table, which is a block of code
+  // consisting of one branch instruction for each of the functions in the bit
+  // set that branches to the target function, and redirect any taken function
+  // addresses to the corresponding jump table entry. In the object file's
+  // symbol table, the symbols for the target functions also refer to the jump
+  // table entries, so that addresses taken outside the module will pass any
+  // verification done inside the module.
+  //
+  // In more concrete terms, suppose we have three functions f, g, h which are
+  // members of a single bitset, and a function foo that returns their
+  // addresses:
+  //
+  // f:
+  // mov 0, %eax
+  // ret
+  //
+  // g:
+  // mov 1, %eax
+  // ret
+  //
+  // h:
+  // mov 2, %eax
+  // ret
+  //
+  // foo:
+  // mov f, %eax
+  // mov g, %edx
+  // mov h, %ecx
+  // ret
+  //
+  // To create a jump table for these functions, we instruct the LLVM code
+  // generator to output a jump table in the .text section. This is done by
+  // representing the instructions in the jump table as an LLVM constant and
+  // placing them in a global variable in the .text section. The end result will
+  // (conceptually) look like this:
+  //
+  // f:
+  // jmp .Ltmp0 ; 5 bytes
+  // int3       ; 1 byte
+  // int3       ; 1 byte
+  // int3       ; 1 byte
+  //
+  // g:
+  // jmp .Ltmp1 ; 5 bytes
+  // int3       ; 1 byte
+  // int3       ; 1 byte
+  // int3       ; 1 byte
+  //
+  // h:
+  // jmp .Ltmp2 ; 5 bytes
+  // int3       ; 1 byte
+  // int3       ; 1 byte
+  // int3       ; 1 byte
+  //
+  // .Ltmp0:
+  // mov 0, %eax
+  // ret
+  //
+  // .Ltmp1:
+  // mov 1, %eax
+  // ret
+  //
+  // .Ltmp2:
+  // mov 2, %eax
+  // ret
+  //
+  // foo:
+  // mov f, %eax
+  // mov g, %edx
+  // mov h, %ecx
+  // ret
+  //
+  // Because the addresses of f, g, h are evenly spaced at a power of 2, in the
+  // normal case the check can be carried out using the same kind of simple
+  // arithmetic that we normally use for globals.
+
+  assert(!Functions.empty());
+
+  // Build a simple layout based on the regular layout of jump tables.
+  DenseMap<GlobalObject *, uint64_t> GlobalLayout;
+  unsigned EntrySize = getJumpTableEntrySize();
+  for (unsigned I = 0; I != Functions.size(); ++I)
+    GlobalLayout[Functions[I]] = I * EntrySize;
+
+  // Create a constant to hold the jump table.
+  ArrayType *JumpTableType =
+      ArrayType::get(getJumpTableEntryType(), Functions.size());
+  auto JumpTable = new GlobalVariable(*M, JumpTableType,
+                                      /*isConstant=*/true,
+                                      GlobalValue::PrivateLinkage, nullptr);
+  JumpTable->setSection(ObjectFormat == Triple::MachO
+                            ? "__TEXT,__text,regular,pure_instructions"
+                            : ".text");
+  lowerBitSetCalls(BitSets, JumpTable, GlobalLayout);
+
+  // Build aliases pointing to offsets into the jump table, and replace
+  // references to the original functions with references to the aliases.
+  for (unsigned I = 0; I != Functions.size(); ++I) {
+    Constant *CombinedGlobalElemPtr = ConstantExpr::getBitCast(
+        ConstantExpr::getGetElementPtr(
+            JumpTableType, JumpTable,
+            ArrayRef<Constant *>{ConstantInt::get(IntPtrTy, 0),
+                                 ConstantInt::get(IntPtrTy, I)}),
+        Functions[I]->getType());
+    if (LinkerSubsectionsViaSymbols || Functions[I]->isDeclarationForLinker()) {
+      Functions[I]->replaceAllUsesWith(CombinedGlobalElemPtr);
+    } else {
+      GlobalAlias *GAlias = GlobalAlias::create(Functions[I]->getType(),
+                                                Functions[I]->getLinkage(), "",
+                                                CombinedGlobalElemPtr, M);
+      GAlias->setVisibility(Functions[I]->getVisibility());
+      GAlias->takeName(Functions[I]);
+      Functions[I]->replaceAllUsesWith(GAlias);
+    }
+    if (!Functions[I]->isDeclarationForLinker())
+      Functions[I]->setLinkage(GlobalValue::PrivateLinkage);
+  }
+
+  // Build and set the jump table's initializer.
+  std::vector<Constant *> JumpTableEntries;
+  for (unsigned I = 0; I != Functions.size(); ++I)
+    JumpTableEntries.push_back(
+        createJumpTableEntry(JumpTable, Functions[I], I));
+  JumpTable->setInitializer(
+      ConstantArray::get(JumpTableType, JumpTableEntries));
+}
+
+void LowerBitSets::buildBitSetsFromDisjointSet(
+    ArrayRef<Metadata *> BitSets, ArrayRef<GlobalObject *> Globals) {
+  llvm::DenseMap<Metadata *, uint64_t> BitSetIndices;
+  llvm::DenseMap<GlobalObject *, uint64_t> GlobalIndices;
+  for (auto B : BitSets)
+    BitSetIndices[B] = BitSetIndices.size();
+  for (auto G : Globals)
+    GlobalIndices[G] = GlobalIndices.size();
+
+  // For each bitset, build a set of indices that refer to globals referenced by
+  // the bitset.
+  std::vector<std::set<uint64_t>> BitSetMembers(BitSets.size());
+  if (BitSetNM) {
+    for (MDNode *Op : BitSetNM->operands()) {
+      // Op = { bitset name, global, offset }
+      if (!Op->getOperand(1))
+        continue;
+      auto I = BitSetIndices.find(Op->getOperand(0));
+      if (I == BitSetIndices.end())
+        continue;
+
+      auto OpGlobal = dyn_cast<GlobalObject>(
+          cast<ConstantAsMetadata>(Op->getOperand(1))->getValue());
+      if (!OpGlobal)
+        continue;
+      BitSetMembers[I->second].insert(GlobalIndices[OpGlobal]);
+    }
+  }
+
+  // Order the sets of indices by size. The GlobalLayoutBuilder works best
+  // when given small index sets first.
+  std::stable_sort(
+      BitSetMembers.begin(), BitSetMembers.end(),
+      [](const std::set<uint64_t> &O1, const std::set<uint64_t> &O2) {
+        return O1.size() < O2.size();
+      });
+
+  // Create a GlobalLayoutBuilder and provide it with index sets as layout
+  // fragments. The GlobalLayoutBuilder tries to lay out members of fragments as
+  // close together as possible.
+  GlobalLayoutBuilder GLB(Globals.size());
+  for (auto &&MemSet : BitSetMembers)
+    GLB.addFragment(MemSet);
+
+  // Build the bitsets from this disjoint set.
+  if (Globals.empty() || isa<GlobalVariable>(Globals[0])) {
+    // Build a vector of global variables with the computed layout.
+    std::vector<GlobalVariable *> OrderedGVs(Globals.size());
+    auto OGI = OrderedGVs.begin();
+    for (auto &&F : GLB.Fragments) {
+      for (auto &&Offset : F) {
+        auto GV = dyn_cast<GlobalVariable>(Globals[Offset]);
+        if (!GV)
+          report_fatal_error(
+              "Bit set may not contain both global variables and functions");
+        *OGI++ = GV;
+      }
+    }
+
+    buildBitSetsFromGlobalVariables(BitSets, OrderedGVs);
+  } else {
+    // Build a vector of functions with the computed layout.
+    std::vector<Function *> OrderedFns(Globals.size());
+    auto OFI = OrderedFns.begin();
+    for (auto &&F : GLB.Fragments) {
+      for (auto &&Offset : F) {
+        auto Fn = dyn_cast<Function>(Globals[Offset]);
+        if (!Fn)
+          report_fatal_error(
+              "Bit set may not contain both global variables and functions");
+        *OFI++ = Fn;
+      }
+    }
+
+    buildBitSetsFromFunctions(BitSets, OrderedFns);
+  }
+}
+
 /// Lower all bit sets in this module.
 bool LowerBitSets::buildBitSets() {
   Function *BitSetTestFunc =
       M->getFunction(Intrinsic::getName(Intrinsic::bitset_test));
   if (!BitSetTestFunc)
     return false;
 
   // Equivalence class set containing bitsets and the globals they reference.
   // This is used to partition the set of bitsets in the module into disjoint
   // sets.
-  typedef EquivalenceClasses<PointerUnion<GlobalVariable *, MDString *>>
+  typedef EquivalenceClasses<PointerUnion<GlobalObject *, Metadata *>>
       GlobalClassesTy;
   GlobalClassesTy GlobalClasses;
 
+  // Verify the bitset metadata and build a mapping from bitset identifiers to
+  // their last observed index in BitSetNM. This will used later to
+  // deterministically order the list of bitset identifiers.
+  llvm::DenseMap<Metadata *, unsigned> BitSetIdIndices;
+  if (BitSetNM) {
+    for (unsigned I = 0, E = BitSetNM->getNumOperands(); I != E; ++I) {
+      MDNode *Op = BitSetNM->getOperand(I);
+      verifyBitSetMDNode(Op);
+      BitSetIdIndices[Op] = I;
+    }
+  }
+
   for (const Use &U : BitSetTestFunc->uses()) {
     auto CI = cast<CallInst>(U.getUser());
 
     auto BitSetMDVal = dyn_cast<MetadataAsValue>(CI->getArgOperand(1));
-    if (!BitSetMDVal || !isa<MDString>(BitSetMDVal->getMetadata()))
+    if (!BitSetMDVal)
       report_fatal_error(
-          "Second argument of llvm.bitset.test must be metadata string");
-    auto BitSet = cast<MDString>(BitSetMDVal->getMetadata());
+          "Second argument of llvm.bitset.test must be metadata");
+    auto BitSet = BitSetMDVal->getMetadata();
 
     // Add the call site to the list of call sites for this bit set. We also use
     // BitSetTestCallSites to keep track of whether we have seen this bit set
     // before. If we have, we don't need to re-add the referenced globals to the
     // equivalence class.
-    std::pair<DenseMap<MDString *, std::vector<CallInst *>>::iterator,
+    std::pair<DenseMap<Metadata *, std::vector<CallInst *>>::iterator,
               bool> Ins =
         BitSetTestCallSites.insert(
             std::make_pair(BitSet, std::vector<CallInst *>()));
     Ins.first->second.push_back(CI);
     if (!Ins.second)
       continue;
 
     // Add the bitset to the equivalence class.
     GlobalClassesTy::iterator GCI = GlobalClasses.insert(BitSet);
     GlobalClassesTy::member_iterator CurSet = GlobalClasses.findLeader(GCI);
 
     if (!BitSetNM)
       continue;
 
-    // Verify the bitset metadata and add the referenced globals to the bitset's
-    // equivalence class.
+    // Add the referenced globals to the bitset's equivalence class.
     for (MDNode *Op : BitSetNM->operands()) {
-      if (Op->getNumOperands() != 3)
-        report_fatal_error(
-            "All operands of llvm.bitsets metadata must have 3 elements");
-
       if (Op->getOperand(0) != BitSet || !Op->getOperand(1))
         continue;
 
-      auto OpConstMD = dyn_cast<ConstantAsMetadata>(Op->getOperand(1));
-      if (!OpConstMD)
-        report_fatal_error("Bit set element must be a constant");
-      auto OpGlobal = dyn_cast<GlobalVariable>(OpConstMD->getValue());
+      auto OpGlobal = dyn_cast<GlobalObject>(
+          cast<ConstantAsMetadata>(Op->getOperand(1))->getValue());
       if (!OpGlobal)
         continue;
 
-      auto OffsetConstMD = dyn_cast<ConstantAsMetadata>(Op->getOperand(2));
-      if (!OffsetConstMD)
-        report_fatal_error("Bit set element offset must be a constant");
-      auto OffsetInt = dyn_cast<ConstantInt>(OffsetConstMD->getValue());
-      if (!OffsetInt)
-        report_fatal_error(
-            "Bit set element offset must be an integer constant");
-
       CurSet = GlobalClasses.unionSets(
           CurSet, GlobalClasses.findLeader(GlobalClasses.insert(OpGlobal)));
     }
   }
 
   if (GlobalClasses.empty())
     return false;
 
   // For each disjoint set we found...
   for (GlobalClassesTy::iterator I = GlobalClasses.begin(),
                                  E = GlobalClasses.end();
        I != E; ++I) {
     if (!I->isLeader()) continue;
 
     ++NumBitSetDisjointSets;
 
-    // Build the list of bitsets and referenced globals in this disjoint set.
-    std::vector<MDString *> BitSets;
-    std::vector<GlobalVariable *> Globals;
-    llvm::DenseMap<MDString *, uint64_t> BitSetIndices;
-    llvm::DenseMap<GlobalVariable *, uint64_t> GlobalIndices;
+    // Build the list of bitsets in this disjoint set.
+    std::vector<Metadata *> BitSets;
+    std::vector<GlobalObject *> Globals;
     for (GlobalClassesTy::member_iterator MI = GlobalClasses.member_begin(I);
          MI != GlobalClasses.member_end(); ++MI) {
-      if ((*MI).is<MDString *>()) {
-        BitSetIndices[MI->get<MDString *>()] = BitSets.size();
-        BitSets.push_back(MI->get<MDString *>());
-      } else {
-        GlobalIndices[MI->get<GlobalVariable *>()] = Globals.size();
-        Globals.push_back(MI->get<GlobalVariable *>());
-      }
-    }
-
-    // For each bitset, build a set of indices that refer to globals referenced
-    // by the bitset.
-    std::vector<std::set<uint64_t>> BitSetMembers(BitSets.size());
-    if (BitSetNM) {
-      for (MDNode *Op : BitSetNM->operands()) {
-        // Op = { bitset name, global, offset }
-        if (!Op->getOperand(1))
-          continue;
-        auto I = BitSetIndices.find(cast<MDString>(Op->getOperand(0)));
-        if (I == BitSetIndices.end())
-          continue;
-
-        auto OpGlobal = dyn_cast<GlobalVariable>(
-            cast<ConstantAsMetadata>(Op->getOperand(1))->getValue());
-        if (!OpGlobal)
-          continue;
-        BitSetMembers[I->second].insert(GlobalIndices[OpGlobal]);
-      }
+      if ((*MI).is<Metadata *>())
+        BitSets.push_back(MI->get<Metadata *>());
+      else
+        Globals.push_back(MI->get<GlobalObject *>());
     }
 
-    // Order the sets of indices by size. The GlobalLayoutBuilder works best
-    // when given small index sets first.
-    std::stable_sort(
-        BitSetMembers.begin(), BitSetMembers.end(),
-        [](const std::set<uint64_t> &O1, const std::set<uint64_t> &O2) {
-          return O1.size() < O2.size();
-        });
-
-    // Create a GlobalLayoutBuilder and provide it with index sets as layout
-    // fragments. The GlobalLayoutBuilder tries to lay out members of fragments
-    // as close together as possible.
-    GlobalLayoutBuilder GLB(Globals.size());
-    for (auto &&MemSet : BitSetMembers)
-      GLB.addFragment(MemSet);
-
-    // Build a vector of globals with the computed layout.
-    std::vector<GlobalVariable *> OrderedGlobals(Globals.size());
-    auto OGI = OrderedGlobals.begin();
-    for (auto &&F : GLB.Fragments)
-      for (auto &&Offset : F)
-        *OGI++ = Globals[Offset];
-
-    // Order bitsets by name for determinism.
-    std::sort(BitSets.begin(), BitSets.end(), [](MDString *S1, MDString *S2) {
-      return S1->getString() < S2->getString();
+    // Order bitsets by BitSetNM index for determinism. This ordering is stable
+    // as there is a one-to-one mapping between metadata and indices.
+    std::sort(BitSets.begin(), BitSets.end(), [&](Metadata *M1, Metadata *M2) {
+      return BitSetIdIndices[M1] < BitSetIdIndices[M2];
     });
 
-    // Build the bitsets from this disjoint set.
-    buildBitSetsFromGlobals(BitSets, OrderedGlobals);
+    // Lower the bitsets in this disjoint set.
+    buildBitSetsFromDisjointSet(BitSets, Globals);
   }
 
   allocateByteArrays();
 
   return true;
 }
 
 bool LowerBitSets::eraseBitSetMetadata() {

llvm/trunk/test/Transforms/LowerBitSets/function-ext.ll

+; RUN: opt -S -lowerbitsets < %s | FileCheck %s
+
+; Tests that we correctly handle external references, including the case where
+; all functions in a bitset are external references.
+
+target triple = "x86_64-unknown-linux-gnu"
+
+declare void @foo()
+
+; CHECK: @[[JT:.*]] = private constant [1 x <{ i8, i32, i8, i8, i8 }>] [<{ i8, i32, i8, i8, i8 }> <{ i8 -23, i32 trunc (i64 sub (i64 sub (i64 ptrtoint (void ()* @foo to i64), i64 ptrtoint ([1 x <{ i8, i32, i8, i8, i8 }>]* @[[JT]] to i64)), i64 5) to i32), i8 -52, i8 -52, i8 -52 }>], section ".text"
+
+define i1 @bar(i8* %ptr) {
+  ; CHECK: icmp eq i64 {{.*}}, ptrtoint ([1 x <{ i8, i32, i8, i8, i8 }>]* @[[JT]] to i64)
+  %p = call i1 @llvm.bitset.test(i8* %ptr, metadata !"void")
+  ret i1 %p
+}
+
+declare i1 @llvm.bitset.test(i8* %ptr, metadata %bitset) nounwind readnone
+
+!0 = !{!"void", void ()* @foo, i64 0}
+
+!llvm.bitsets = !{!0}

llvm/trunk/test/Transforms/LowerBitSets/function.ll

+; RUN: opt -S -lowerbitsets < %s | FileCheck %s
+
+; Tests that we correctly create a jump table for bitsets containing 2 or more
+; functions.
+
+target triple = "x86_64-unknown-linux-gnu"
+target datalayout = "e-p:64:64"
+
+; CHECK: @[[JT:.*]] = private constant [2 x <{ i8, i32, i8, i8, i8 }>] [<{ i8, i32, i8, i8, i8 }> <{ i8 -23, i32 trunc (i64 sub (i64 sub (i64 ptrtoint (void ()* @[[FNAME:.*]] to i64), i64 ptrtoint ([2 x <{ i8, i32, i8, i8, i8 }>]* @[[JT]] to i64)), i64 5) to i32), i8 -52, i8 -52, i8 -52 }>, <{ i8, i32, i8, i8, i8 }> <{ i8 -23, i32 trunc (i64 sub (i64 sub (i64 ptrtoint (void ()* @[[GNAME:.*]] to i64), i64 ptrtoint ([2 x <{ i8, i32, i8, i8, i8 }>]* @[[JT]] to i64)), i64 13) to i32), i8 -52, i8 -52, i8 -52 }>], section ".text"
+
+; CHECK: @f = alias bitcast ([2 x <{ i8, i32, i8, i8, i8 }>]* @[[JT]] to void ()*)
+; CHECK: @g = alias bitcast (<{ i8, i32, i8, i8, i8 }>* getelementptr inbounds ([2 x <{ i8, i32, i8, i8, i8 }>], [2 x <{ i8, i32, i8, i8, i8 }>]* @[[JT]], i64 0, i64 1) to void ()*)
+
+; CHECK: define private void @[[FNAME]]() {
+define void @f() {
+  ret void
+}
+
+; CHECK: define private void @[[GNAME]]() {
+define void @g() {
+  ret void
+}
+
+!0 = !{!"bitset1", void ()* @f, i32 0}
+!1 = !{!"bitset1", void ()* @g, i32 0}
+
+!llvm.bitsets = !{ !0, !1 }
+
+declare i1 @llvm.bitset.test(i8* %ptr, metadata %bitset) nounwind readnone
+
+define i1 @foo(i8* %p) {
+  ; CHECK: sub i64 {{.*}}, ptrtoint ([2 x <{ i8, i32, i8, i8, i8 }>]* @[[JT]] to i64)
+  %x = call i1 @llvm.bitset.test(i8* %p, metadata !"bitset1")
+  ret i1 %x
+}

llvm/trunk/test/Transforms/LowerBitSets/nonstring.ll

+; RUN: opt -S -lowerbitsets < %s | FileCheck %s
+
+; Tests that non-string metadata nodes may be used as bitset identifiers.
+
+target datalayout = "e-p:32:32"
+
+; CHECK: @[[BNAME:.*]] = private constant { [2 x i32] }
+; CHECK: @[[ANAME:.*]] = private constant { i32 }
+
+@a = constant i32 1
+@b = constant [2 x i32] [i32 2, i32 3]
+
+!0 = !{!2, i32* @a, i32 0}
+!1 = !{!3, [2 x i32]* @b, i32 0}
+!2 = distinct !{}
+!3 = distinct !{}
+
+!llvm.bitsets = !{ !0, !1 }
+
+declare i1 @llvm.bitset.test(i8* %ptr, metadata %bitset) nounwind readnone
+
+; CHECK-LABEL: @foo
+define i1 @foo(i8* %p) {
+  ; CHECK: icmp eq i32 {{.*}}, ptrtoint ({ i32 }* @[[ANAME]] to i32)
+  %x = call i1 @llvm.bitset.test(i8* %p, metadata !2)
+  ret i1 %x
+}
+
+; CHECK-LABEL: @bar
+define i1 @bar(i8* %p) {
+  ; CHECK: icmp eq i32 {{.*}}, ptrtoint ({ [2 x i32] }* @[[BNAME]] to i32)
+  %x = call i1 @llvm.bitset.test(i8* %p, metadata !3)
+  ret i1 %x
+}

llvm/trunk/test/Transforms/LowerBitSets/simple.ll

 ; RUN: opt -S -lowerbitsets < %s | FileCheck %s
 ; RUN: opt -S -lowerbitsets -mtriple=x86_64-apple-macosx10.8.0 < %s | FileCheck -check-prefix=CHECK-DARWIN %s
 ; RUN: opt -S -O3 < %s | FileCheck -check-prefix=CHECK-NODISCARD %s
 
 target datalayout = "e-p:32:32"
 
 ; CHECK: [[G:@[^ ]*]] = private constant { i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] } { i32 1, [0 x i8] zeroinitializer, [63 x i32] zeroinitializer, [4 x i8] zeroinitializer, i32 3, [0 x i8] zeroinitializer, [2 x i32] [i32 4, i32 5] }
 @a = constant i32 1
-@b = constant [63 x i32] zeroinitializer
-@c = constant i32 3
+@b = hidden constant [63 x i32] zeroinitializer
+@c = protected constant i32 3
 @d = constant [2 x i32] [i32 4, i32 5]
 
 ; CHECK: [[BA:@[^ ]*]] = private constant [68 x i8] c"\03\01\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\02\00\01"
 
 ; Offset 0, 4 byte alignment
 !0 = !{!"bitset1", i32* @a, i32 0}
 ; CHECK-NODISCARD-DAG: !{!"bitset1", i32* @a, i32 0}
 !1 = !{!"bitset1", [63 x i32]* @b, i32 0}

 
 !llvm.bitsets = !{ !0, !1, !2, !3, !4, !5, !6, !7 }
 
 ; CHECK: @bits_use{{[0-9]*}} = private alias i8* @bits{{[0-9]*}}
 ; CHECK: @bits_use.{{[0-9]*}} = private alias i8* @bits{{[0-9]*}}
 ; CHECK: @bits_use.{{[0-9]*}} = private alias i8* @bits{{[0-9]*}}
 
 ; CHECK: @a = alias getelementptr inbounds ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }, { i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G]], i32 0, i32 0)
-; CHECK: @b = alias getelementptr inbounds ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }, { i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G]], i32 0, i32 2)
-; CHECK: @c = alias getelementptr inbounds ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }, { i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G]], i32 0, i32 4)
+; CHECK: @b = hidden alias getelementptr inbounds ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }, { i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G]], i32 0, i32 2)
+; CHECK: @c = protected alias getelementptr inbounds ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }, { i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G]], i32 0, i32 4)
 ; CHECK: @d = alias getelementptr inbounds ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }, { i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G]], i32 0, i32 6)
 
 ; CHECK-DARWIN: @aptr = constant i32* getelementptr inbounds ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }, { i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G:@[^ ]*]], i32 0, i32 0)
 @aptr = constant i32* @a
 
 ; CHECK-DARWIN: @bptr = constant [63 x i32]* getelementptr inbounds ({ i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }, { i32, [0 x i8], [63 x i32], [4 x i8], i32, [0 x i8], [2 x i32] }* [[G]], i32 0, i32 2)
 @bptr = constant [63 x i32]* @b