[modules] Refactor handling of cases where we write an offset to a prior record into the bitstream and simplify a little, in preparation for doing this in more cases.

llvm-svn: 266160

3 files changed, 45 insertions, 29 deletions

diff --git a/clang/include/clang/Serialization/ASTWriter.h b/clang/include/clang/Serialization/ASTWriter.h
index 900806676f9..094af456d8d 100644
--- a/clang/include/clang/Serialization/ASTWriter.h
+++ b/clang/include/clang/Serialization/ASTWriter.h
@@ -787,11 +787,29 @@ class ASTRecordWriter {
   /// declaration or type.
   SmallVector<Stmt *, 16> StmtsToEmit;
 
+  static constexpr int MaxOffsetIndices = 4;
+  /// \brief Indices of record elements that describe offsets within the
+  /// bitcode. These will be converted to offsets relative to the current
+  /// record when emitted.
+  unsigned OffsetIndices[MaxOffsetIndices];
+  unsigned NumOffsetIndices = 0;
+
   /// \brief Flush all of the statements and expressions that have
   /// been added to the queue via AddStmt().
   void FlushStmts();
   void FlushSubStmts();
 
+  void PrepareToEmit(uint64_t MyOffset) {
+    // Convert offsets into relative form.
+    for (unsigned I = 0; I != NumOffsetIndices; ++I) {
+      auto &StoredOffset = (*Record)[OffsetIndices[I]];
+      assert(StoredOffset < MyOffset && "invalid offset");
+      if (StoredOffset)
+        StoredOffset = MyOffset - StoredOffset;
+    }
+    NumOffsetIndices = 0;
+  }
+
 public:
   /// Construct a ASTRecordWriter that uses the default encoding scheme.
   ASTRecordWriter(ASTWriter &Writer, ASTWriter::RecordDataImpl &Record)
@@ -802,6 +820,10 @@ public:
   ASTRecordWriter(ASTRecordWriter &Parent, ASTWriter::RecordDataImpl &Record)
       : Writer(Parent.Writer), Record(&Record) {}
 
+  /// Copying an ASTRecordWriter is almost certainly a bug.
+  ASTRecordWriter(const ASTRecordWriter&) = delete;
+  void operator=(const ASTRecordWriter&) = delete;
+
   /// \brief Extract the underlying record storage.
   ASTWriter::RecordDataImpl &getRecordData() const { return *Record; }
 
@@ -822,6 +844,7 @@ public:
   // FIXME: Allow record producers to suggest Abbrevs.
   uint64_t Emit(unsigned Code, unsigned Abbrev = 0) {
     uint64_t Offset = Writer->Stream.GetCurrentBitNo();
+    PrepareToEmit(Offset);
     Writer->Stream.EmitRecord(Code, *Record, Abbrev);
     FlushStmts();
     return Offset;
@@ -830,10 +853,19 @@ public:
   /// \brief Emit the record to the stream, preceded by its substatements.
   uint64_t EmitStmt(unsigned Code, unsigned Abbrev = 0) {
     FlushSubStmts();
+    PrepareToEmit(Writer->Stream.GetCurrentBitNo());
     Writer->Stream.EmitRecord(Code, *Record, Abbrev);
     return Writer->Stream.GetCurrentBitNo();
   }
 
+  /// \brief Add a bit offset into the record. This will be converted into an
+  /// offset relative to the current record when emitted.
+  void AddOffset(uint64_t BitOffset) {
+    assert(NumOffsetIndices != MaxOffsetIndices && "too many offset indices");
+    OffsetIndices[NumOffsetIndices++] = Record->size();
+    Record->push_back(BitOffset);
+  }
+
   /// \brief Add the given statement or expression to the queue of
   /// statements to emit.
   ///
diff --git a/clang/lib/Serialization/ASTWriter.cpp b/clang/lib/Serialization/ASTWriter.cpp
index 06e9558536e..cf4872f2efb 100644
--- a/clang/lib/Serialization/ASTWriter.cpp
+++ b/clang/lib/Serialization/ASTWriter.cpp
@@ -4694,6 +4694,7 @@ void ASTWriter::WriteDeclUpdatesBlocks(RecordDataImpl &OffsetsRecord) {
         auto *RD = cast<CXXRecordDecl>(D);
         UpdatedDeclContexts.insert(RD->getPrimaryContext());
         Record.AddCXXDefinitionData(RD);
+        // FIXME: Use AddOffset here.
         Record.push_back(WriteDeclContextLexicalBlock(
             *Context, const_cast<CXXRecordDecl *>(RD)));
 
diff --git a/clang/lib/Serialization/ASTWriterDecl.cpp b/clang/lib/Serialization/ASTWriterDecl.cpp
index 21fca17c248..c06e272e2b7 100644
--- a/clang/lib/Serialization/ASTWriterDecl.cpp
+++ b/clang/lib/Serialization/ASTWriterDecl.cpp
@@ -125,8 +125,7 @@ namespace clang {
     void VisitCapturedDecl(CapturedDecl *D);
     void VisitEmptyDecl(EmptyDecl *D);
 
-    void VisitDeclContext(DeclContext *DC, uint64_t LexicalOffset,
-                          uint64_t VisibleOffset);
+    void VisitDeclContext(DeclContext *DC);
     template <typename T> void VisitRedeclarable(Redeclarable<T> *D);
 
 
@@ -149,12 +148,6 @@ namespace clang {
     void VisitOMPDeclareReductionDecl(OMPDeclareReductionDecl *D);
     void VisitOMPCapturedExprDecl(OMPCapturedExprDecl *D);
 
-    void AddLocalOffset(uint64_t LocalOffset) {
-      uint64_t Offset = Writer.Stream.GetCurrentBitNo();
-      assert(LocalOffset < Offset && "invalid offset");
-      Record.push_back(LocalOffset ? Offset - LocalOffset : 0);
-    }
-
     /// Add an Objective-C type parameter list to the given record.
     void AddObjCTypeParamList(ObjCTypeParamList *typeParams) {
       // Empty type parameter list.
@@ -284,6 +277,12 @@ void ASTDeclWriter::Visit(Decl *D) {
     if (FD->doesThisDeclarationHaveABody())
       Record.AddFunctionDefinition(FD);
   }
+
+  // If this declaration is also a DeclContext, write blocks for the
+  // declarations that lexically stored inside its context and those
+  // declarations that are visible from its context.
+  if (DeclContext *DC = dyn_cast<DeclContext>(D))
+    VisitDeclContext(DC);
 }
 
 void ASTDeclWriter::VisitDecl(Decl *D) {
@@ -1553,10 +1552,9 @@ void ASTDeclWriter::VisitStaticAssertDecl(StaticAssertDecl *D) {
 /// that there are no declarations visible from this context. Note
 /// that this value will not be emitted for non-primary declaration
 /// contexts.
-void ASTDeclWriter::VisitDeclContext(DeclContext *DC, uint64_t LexicalOffset,
-                                     uint64_t VisibleOffset) {
-  AddLocalOffset(LexicalOffset);
-  AddLocalOffset(VisibleOffset);
+void ASTDeclWriter::VisitDeclContext(DeclContext *DC) {
+  Record.AddOffset(Writer.WriteDeclContextLexicalBlock(Context, DC));
+  Record.AddOffset(Writer.WriteDeclContextVisibleBlock(Context, DC));
 }
 
 const Decl *ASTWriter::getFirstLocalDecl(const Decl *D) {
@@ -1624,9 +1622,8 @@ void ASTDeclWriter::VisitRedeclarable(Redeclarable<T> *D) {
       // the declaration itself.
       if (LocalRedecls.empty())
         Record.push_back(0);
-      else {
-        AddLocalOffset(LocalRedeclWriter.Emit(LOCAL_REDECLARATIONS));
-      }
+      else
+        Record.AddOffset(LocalRedeclWriter.Emit(LOCAL_REDECLARATIONS));
     } else {
       Record.push_back(0);
       Record.AddDeclRef(FirstLocal);
@@ -2148,26 +2145,12 @@ void ASTWriter::WriteDecl(ASTContext &Context, Decl *D) {
   ID = IDR;
 
   assert(ID >= FirstDeclID && "invalid decl ID");
-
-  // If this declaration is also a DeclContext, write blocks for the
-  // declarations that lexically stored inside its context and those
-  // declarations that are visible from its context. These blocks
-  // are written before the declaration itself so that we can put
-  // their offsets into the record for the declaration.
-  uint64_t LexicalOffset = 0;
-  uint64_t VisibleOffset = 0;
-  DeclContext *DC = dyn_cast<DeclContext>(D);
-  if (DC) {
-    LexicalOffset = WriteDeclContextLexicalBlock(Context, DC);
-    VisibleOffset = WriteDeclContextVisibleBlock(Context, DC);
-  }
   
   RecordData Record;
   ASTDeclWriter W(*this, Context, Record);
 
   // Build a record for this declaration
   W.Visit(D);
-  if (DC) W.VisitDeclContext(DC, LexicalOffset, VisibleOffset);
 
   // Emit this declaration to the bitstream.
   uint64_t Offset = W.Emit(D);