swiftlang · rintaro · Apr 10, 2020 · Apr 9, 2020
diff --git a/tools/SourceKit/tools/sourcekitd/include/sourcekitd/CompactArray.h b/tools/SourceKit/tools/sourcekitd/include/sourcekitd/CompactArray.h
@@ -20,7 +20,7 @@ namespace sourcekitd {
 
 class CompactArrayBuilderImpl {
 public:
-  std::unique_ptr<llvm::MemoryBuffer> createBuffer() const;
+  std::unique_ptr<llvm::MemoryBuffer> createBuffer(CustomBufferKind Kind) const;
   void appendTo(llvm::SmallVectorImpl<char> &Buf) const;
   unsigned copyInto(char *BufPtr) const;
   size_t sizeInBytes() const;

diff --git a/tools/SourceKit/tools/sourcekitd/include/sourcekitd/Internal.h b/tools/SourceKit/tools/sourcekitd/include/sourcekitd/Internal.h
@@ -86,7 +86,7 @@ class ResponseBuilder {
     void setBool(SourceKit::UIdent Key, bool val);
     Array setArray(SourceKit::UIdent Key);
     Dictionary setDictionary(SourceKit::UIdent Key);
-    void setCustomBuffer(SourceKit::UIdent Key, CustomBufferKind Kind,
+    void setCustomBuffer(SourceKit::UIdent Key,
                          std::unique_ptr<llvm::MemoryBuffer> MemBuf);
 
   private:

diff --git a/tools/SourceKit/tools/sourcekitd/lib/API/CodeCompletionResultsArray.cpp b/tools/SourceKit/tools/sourcekitd/lib/API/CodeCompletionResultsArray.cpp
@@ -75,7 +75,8 @@ void CodeCompletionResultsArrayBuilder::add(
 
 std::unique_ptr<llvm::MemoryBuffer>
 CodeCompletionResultsArrayBuilder::createBuffer() {
-  return Impl.Builder.createBuffer();
+  return Impl.Builder.createBuffer(
+      CustomBufferKind::CodeCompletionResultsArray);
 }
 
 namespace {

diff --git a/tools/SourceKit/tools/sourcekitd/lib/API/CompactArray.cpp b/tools/SourceKit/tools/sourcekitd/lib/API/CompactArray.cpp
@@ -70,10 +70,13 @@ unsigned CompactArrayBuilderImpl::getOffsetForString(StringRef Str) {
 }
 
 std::unique_ptr<llvm::MemoryBuffer>
-CompactArrayBuilderImpl::createBuffer() const {
+CompactArrayBuilderImpl::createBuffer(CustomBufferKind Kind) const {
+  auto bodySize = sizeInBytes();
   std::unique_ptr<llvm::WritableMemoryBuffer> Buf;
-  Buf = llvm::WritableMemoryBuffer::getNewUninitMemBuffer(sizeInBytes());
-  copyInto(Buf->getBufferStart(), Buf->getBufferSize());
+  Buf = llvm::WritableMemoryBuffer::getNewUninitMemBuffer(
+      sizeof(uint64_t) + bodySize);
+  *reinterpret_cast<uint64_t*>(Buf->getBufferStart()) = (uint64_t)Kind;
+  copyInto(Buf->getBufferStart() + sizeof(uint64_t), bodySize);
   return std::move(Buf);
 }
 

diff --git a/tools/SourceKit/tools/sourcekitd/lib/API/DocStructureArray.cpp b/tools/SourceKit/tools/sourcekitd/lib/API/DocStructureArray.cpp
@@ -244,16 +244,21 @@ std::unique_ptr<llvm::MemoryBuffer> DocStructureArrayBuilder::createBuffer() {
   size_t structureArrayBufferSize = impl.structureArrayBuffer.size();
   size_t structureBufferSize = impl.structureBuilder.sizeInBytes();
 
+  size_t kindSize = sizeof(uint64_t);
+
   // Header:
   // * offset of each section start (5)
   // * offset of top structure array (relative to structure array section) (1)
   size_t headerSize = sizeof(uint64_t) * 6;
 
   auto result = llvm::WritableMemoryBuffer::getNewUninitMemBuffer(
       inheritedTypesBufferSize + attrsBufferSize + elementsBufferSize +
-      structureArrayBufferSize + structureBufferSize + headerSize);
+      structureArrayBufferSize + structureBufferSize + headerSize + kindSize);
+
+  *reinterpret_cast<uint64_t *>(result->getBufferStart()) =
+      (uint64_t)CustomBufferKind::DocStructureArray;
 
-  char *start = result->getBufferStart();
+  char *start = result->getBufferStart() + kindSize;
   char *headerPtr = start;
   char *ptr = start + headerSize;
 

diff --git a/tools/SourceKit/tools/sourcekitd/lib/API/DocSupportAnnotationArray.cpp b/tools/SourceKit/tools/sourcekitd/lib/API/DocSupportAnnotationArray.cpp
@@ -53,7 +53,7 @@ void DocSupportAnnotationArrayBuilder::add(const DocEntityInfo &Info) {
 
 std::unique_ptr<llvm::MemoryBuffer>
 DocSupportAnnotationArrayBuilder::createBuffer() {
-  return Impl.Builder.createBuffer();
+  return Impl.Builder.createBuffer(CustomBufferKind::DocSupportAnnotationArray);
 }
 
 namespace {

diff --git a/tools/SourceKit/tools/sourcekitd/lib/API/ExpressionTypeArray.cpp b/tools/SourceKit/tools/sourcekitd/lib/API/ExpressionTypeArray.cpp
@@ -176,15 +176,18 @@ struct ExpressionTypeArrayBuilder::Implementation {
     return reader.count();
   }
 
-  std::unique_ptr<llvm::MemoryBuffer> createBuffer() {
+  std::unique_ptr<llvm::MemoryBuffer> createBuffer(CustomBufferKind Kind) {
     std::array<CompactArrayBuilderImpl*, 3> builders =
       {&builder, &strBuilder, &protoBuilder};
+    auto kindSize = sizeof(uint64_t);
     size_t headerSize = sizeof(uint64_t) * builders.size();
-    auto allSize = headerSize;
+    auto allSize = kindSize + headerSize;
     for (auto *b: builders)
       allSize += b->sizeInBytes();
     auto result = llvm::WritableMemoryBuffer::getNewUninitMemBuffer(allSize);
-    char *start = result->getBufferStart();
+    *reinterpret_cast<uint64_t*>(result->getBufferStart()) = (uint64_t)Kind;
+
+    char *start = result->getBufferStart() + kindSize;
     char *headerPtr = start;
     char *ptr = start + headerSize;
     auto addBuilder = [&](CompactArrayBuilderImpl& buffer) {
@@ -223,7 +226,7 @@ void ExpressionTypeArrayBuilder::add(const ExpressionType &expType) {
 
 std::unique_ptr<llvm::MemoryBuffer>
 ExpressionTypeArrayBuilder::createBuffer() {
-  return Impl.createBuffer();
+  return Impl.createBuffer(CustomBufferKind::ExpressionTypeArray);
 }
 
 VariantFunctions ExpressionTypeArrayBuilder::Funcs = {

diff --git a/tools/SourceKit/tools/sourcekitd/lib/API/Requests.cpp b/tools/SourceKit/tools/sourcekitd/lib/API/Requests.cpp
@@ -1422,9 +1422,7 @@ class SKDocConsumer : public DocInfoConsumer {
   }
 
   sourcekitd_response_t createResponse() {
-    TopDict.setCustomBuffer(KeyAnnotations,
-        CustomBufferKind::DocSupportAnnotationArray,
-        AnnotationsBuilder.createBuffer());
+    TopDict.setCustomBuffer(KeyAnnotations, AnnotationsBuilder.createBuffer());
     return RespBuilder.createResponse();
   }
 
@@ -1866,8 +1864,7 @@ static void reportExpressionTypeInfo(const RequestResult<ExpressionTypesInFile>
   for (auto &R: Info.Results) {
     ArrBuilder.add(R);
   }
-  Dict.setCustomBuffer(KeyExpressionTypeList, CustomBufferKind::ExpressionTypeArray,
-                       ArrBuilder.createBuffer());
+  Dict.setCustomBuffer(KeyExpressionTypeList, ArrBuilder.createBuffer());
   Rec(Builder.createResponse());
 }
 
@@ -1924,8 +1921,7 @@ class SKCodeCompletionConsumer : public CodeCompletionConsumer {
       return createErrorRequestFailed(ErrorDescription.c_str());
 
     RespBuilder.getDictionary().setCustomBuffer(KeyResults,
-        CustomBufferKind::CodeCompletionResultsArray,
-        ResultsBuilder.createBuffer());
+                                                ResultsBuilder.createBuffer());
     return RespBuilder.createResponse();
   }
 
@@ -2558,18 +2554,13 @@ sourcekitd_response_t SKEditorConsumer::createResponse() {
     return Error;
 
   if (Opts.EnableSyntaxMap) {
-    Dict.setCustomBuffer(KeySyntaxMap,
-        CustomBufferKind::TokenAnnotationsArray,
-        SyntaxMap.createBuffer());
+    Dict.setCustomBuffer(KeySyntaxMap, SyntaxMap.createBuffer());
   }
   if (!SemanticAnnotations.empty()) {
-    Dict.setCustomBuffer(KeyAnnotations,
-        CustomBufferKind::TokenAnnotationsArray,
-        SemanticAnnotations.createBuffer());
+    Dict.setCustomBuffer(KeyAnnotations, SemanticAnnotations.createBuffer());
   }
   if (Opts.EnableStructure) {
-    Dict.setCustomBuffer(KeySubStructure, CustomBufferKind::DocStructureArray,
-                         DocStructure.createBuffer());
+    Dict.setCustomBuffer(KeySubStructure, DocStructure.createBuffer());
   }
 
 
@@ -2756,11 +2747,13 @@ void serializeSyntaxTreeAsByteTree(
                                          *SyntaxTree.getRaw(), UserInfo);
 
   std::unique_ptr<llvm::WritableMemoryBuffer> Buf =
-      llvm::WritableMemoryBuffer::getNewUninitMemBuffer(Stream.data().size());
-  memcpy(Buf->getBufferStart(), Stream.data().data(), Stream.data().size());
+      llvm::WritableMemoryBuffer::getNewUninitMemBuffer(sizeof(uint64_t) + Stream.data().size());
+  *reinterpret_cast<uint64_t*>(Buf->getBufferStart()) =
+      (uint64_t)CustomBufferKind::RawData;
+  memcpy(Buf->getBufferStart() + sizeof(uint64_t),
+         Stream.data().data(), Stream.data().size());
 
-  Dict.setCustomBuffer(KeySerializedSyntaxTree, CustomBufferKind::RawData,
-                       std::move(Buf));
+  Dict.setCustomBuffer(KeySerializedSyntaxTree, std::move(Buf));
 
   auto EndClock = clock();
   LOG_SECTION("incrParse Performance", InfoLowPrio) {

diff --git a/tools/SourceKit/tools/sourcekitd/lib/API/TokenAnnotationsArray.cpp b/tools/SourceKit/tools/sourcekitd/lib/API/TokenAnnotationsArray.cpp
@@ -55,7 +55,7 @@ bool TokenAnnotationsArrayBuilder::empty() const {
 
 std::unique_ptr<llvm::MemoryBuffer>
 TokenAnnotationsArrayBuilder::createBuffer() {
-  return Impl.Builder.createBuffer();
+  return Impl.Builder.createBuffer(CustomBufferKind::TokenAnnotationsArray);
 }
 
 namespace {

diff --git a/tools/SourceKit/tools/sourcekitd/lib/API/sourcekitdAPI-InProc.cpp b/tools/SourceKit/tools/sourcekitd/lib/API/sourcekitdAPI-InProc.cpp
@@ -237,19 +237,15 @@ class SKDBool: public SKDObject {
 
 class SKDCustomData: public SKDObject {
 public:
-  SKDCustomData(CustomBufferKind BufferKind, 
-                std::unique_ptr<llvm::MemoryBuffer>& MemBuf)
-  : SKDObject(ObjectKind::CustomData), BufferKind(BufferKind),
-    BufferPtr(llvm::MemoryBuffer::getMemBufferCopy(
-                                                MemBuf->getBuffer(), 
-                                                MemBuf->getBufferIdentifier())) 
+  SKDCustomData(std::unique_ptr<llvm::MemoryBuffer> MemBuf)
+  : SKDObject(ObjectKind::CustomData), BufferPtr(std::move(MemBuf))
     {}
 
   SKDCustomData(SKDCustomData const&) = delete;
   SKDCustomData &operator=(SKDCustomData const&) = delete;
 
   sourcekitd_variant_type_t getVariantType() const override {
-    switch (BufferKind) {
+    switch (getBufferKind()) {
       case CustomBufferKind::TokenAnnotationsArray:
       case CustomBufferKind::DocSupportAnnotationArray:
       case CustomBufferKind::CodeCompletionResultsArray:
@@ -266,23 +262,24 @@ class SKDCustomData: public SKDObject {
   }
 
   CustomBufferKind getBufferKind() const {
-    return BufferKind;
+    return ((CustomBufferKind)*(const uint64_t*)_getStartPtr());
   }
 
   const void *getDataPtr() const override {
-    return BufferPtr->getBuffer().data();
+    return ((const void*)(((const uint64_t*)_getStartPtr())+1));
   }
 
   size_t getDataSize() const override {
-    return BufferPtr->getBuffer().size();
+    return BufferPtr->getBuffer().size() - sizeof(uint64_t);
   }
 
   static bool classof(const SKDObject *O) {
     return O->getKind() == ObjectKind::CustomData;
   }
 private:
-  CustomBufferKind BufferKind;
   std::unique_ptr<llvm::MemoryBuffer> BufferPtr;
+
+  const void *_getStartPtr() const { return BufferPtr->getBuffer().data(); }
 };
 
 class SKDError: public SKDObject {
@@ -675,10 +672,9 @@ ResponseBuilder::Dictionary::setDictionary(UIdent Key) {
 }
 
 void ResponseBuilder::Dictionary::setCustomBuffer(
-      SourceKit::UIdent Key,
-      CustomBufferKind Kind, std::unique_ptr<llvm::MemoryBuffer> MemBuf) {
+      SourceKit::UIdent Key, std::unique_ptr<llvm::MemoryBuffer> MemBuf) {
   static_cast<SKDObject *>(Impl)->set(SKDUIDFromUIdent(Key), 
-                                      new SKDCustomData(Kind, MemBuf));
+                                      new SKDCustomData(std::move(MemBuf)));
 }
 
 ResponseBuilder::Array

diff --git a/tools/SourceKit/tools/sourcekitd/lib/API/sourcekitdAPI-XPC.cpp b/tools/SourceKit/tools/sourcekitd/lib/API/sourcekitdAPI-XPC.cpp
@@ -258,19 +258,9 @@ ResponseBuilder::Dictionary::setDictionary(UIdent Key) {
 }
 
 void ResponseBuilder::Dictionary::setCustomBuffer(
-      SourceKit::UIdent Key,
-      CustomBufferKind Kind, std::unique_ptr<llvm::MemoryBuffer> MemBuf) {
-
-  std::unique_ptr<llvm::WritableMemoryBuffer> CustomBuf;
-  CustomBuf = llvm::WritableMemoryBuffer::getNewUninitMemBuffer(
-      sizeof(uint64_t) + MemBuf->getBufferSize());
-  char *BufPtr = CustomBuf->getBufferStart();
-  *reinterpret_cast<uint64_t*>(BufPtr) = (uint64_t)Kind;
-  BufPtr += sizeof(uint64_t);
-  memcpy(BufPtr, MemBuf->getBufferStart(), MemBuf->getBufferSize());
-
-  xpc_object_t xdata = xpc_data_create(CustomBuf->getBufferStart(),
-                                       CustomBuf->getBufferSize());
+    SourceKit::UIdent Key, std::unique_ptr<llvm::MemoryBuffer> Buf) {
+  xpc_object_t xdata = xpc_data_create(Buf->getBufferStart(),
+                                       Buf->getBufferSize());
   xpc_dictionary_set_value(Impl, Key.c_str(), xdata);
   xpc_release(xdata);
 }