[Bitcode] Encode alloca address space

Since D101045, allocas are no longer required to be part of the
default alloca address space. There may be allocas in multiple
different address spaces. However, the bitcode reader would
simply assume the default alloca address space, resulting in
either an error or incorrect IR.

Add an optional record for allocas which encodes the address
space.

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

@@ -5294,7 +5294,7 @@ Error BitcodeReader::parseFunctionBody(Function *F) {
     }
 
     case bitc::FUNC_CODE_INST_ALLOCA: { // ALLOCA: [instty, opty, op, align]
-      if (Record.size() != 4)
+      if (Record.size() != 4 && Record.size() != 5)
         return error("Invalid record");
       using APV = AllocaPackedValues;
       const uint64_t Rec = Record[3];
@@ -5321,9 +5321,8 @@ Error BitcodeReader::parseFunctionBody(Function *F) {
       if (!Ty || !Size)
         return error("Invalid record");
 
-      // FIXME: Make this an optional field.
       const DataLayout &DL = TheModule->getDataLayout();
-      unsigned AS = DL.getAllocaAddrSpace();
+      unsigned AS = Record.size() == 5 ? Record[4] : DL.getAllocaAddrSpace();
 
       SmallPtrSet<Type *, 4> Visited;
       if (!Align && !Ty->isSized(&Visited))

llvm/lib/Bitcode/Writer/BitcodeWriter.cpp

@@ -3072,6 +3072,10 @@ void ModuleBitcodeWriter::writeInstruction(const Instruction &I,
     Bitfield::set<APV::ExplicitType>(Record, true);
     Bitfield::set<APV::SwiftError>(Record, AI.isSwiftError());
     Vals.push_back(Record);
+
+    unsigned AS = AI.getAddressSpace();
+    if (AS != M.getDataLayout().getAllocaAddrSpace())
+      Vals.push_back(AS);
     break;
   }
 

llvm/test/Bitcode/alloca-addrspace.ll

@@ -0,0 +1,18 @@
+; RUN: llvm-as < %s | llvm-dis | FileCheck %s
+
+target datalayout = "A2"
+
+; CHECK-LABEL: define i8 addrspace(2)* @alloca_addrspace_2() {
+; CHECK: %alloca = alloca i8, align 1, addrspace(2)
+define i8 addrspace(2)* @alloca_addrspace_2() {
+  %alloca = alloca i8, addrspace(2)
+  ret i8 addrspace(2)* %alloca
+}
+
+; CHECK-LABEL: define i8 addrspace(5)* @alloca_addrspace_5() {
+; CHECK: %alloca = alloca i8, align 1, addrspace(5)
+define i8 addrspace(5)* @alloca_addrspace_5() {
+  %alloca = alloca i8, addrspace(5)
+  ret i8 addrspace(5)* %alloca
+}
+