Address review feedback.

Author: dan-zheng

include/swift/AST/DiagnosticsParse.def

@@ -687,12 +687,12 @@ ERROR(sil_witness_protocol_conformance_not_found,none,
       "sil protocol conformance not found", ())
 
 // SIL differentiability witnesses
-ERROR(sil_diff_witness_expected_keyword,PointsToFirstBadToken,
+ERROR(sil_diff_witness_expected_token,PointsToFirstBadToken,
       "expected '%0' in differentiability witness", (StringRef))
 ERROR(sil_diff_witness_expected_index_list,PointsToFirstBadToken,
       "expected a space-separated list of indices, e.g. '0 1'", ())
 ERROR(sil_diff_witness_expected_parameter_index,PointsToFirstBadToken,
-      "expected a parameter index to differentiate with respect to.", ())
+      "expected a parameter index to differentiate with respect to", ())
 ERROR(sil_diff_witness_expected_result_index,PointsToFirstBadToken,
       "expected a result index to differentiate with respect to", ())
 

include/swift/SIL/SILDifferentiabilityWitness.h

@@ -67,11 +67,6 @@ class SILDifferentiabilityWitness
   /// deserialized.
   DeclAttribute *attribute = nullptr;
 
-  static AutoDiffConfig *
-  getAutoDiffConfig(SILModule &module, IndexSubset *parameterIndices,
-                    IndexSubset *resultIndices,
-                    GenericSignature *derivativeGenSig);
-
   SILDifferentiabilityWitness(SILModule &module, SILLinkage linkage,
                               SILFunction *originalFunction,
                               IndexSubset *parameterIndices,
@@ -112,6 +107,8 @@ class SILDifferentiabilityWitness
     case AutoDiffDerivativeFunctionKind::VJP: return vjp;
     }
   }
+  void setJVP(SILFunction *jvp) { this->jvp = jvp; }
+  void setVJP(SILFunction *vjp) { this->vjp = vjp; }
   void setDerivative(AutoDiffDerivativeFunctionKind kind,
                      SILFunction *derivative) {
     switch (kind) {
@@ -123,9 +120,9 @@ class SILDifferentiabilityWitness
   DeclAttribute *getAttribute() const { return attribute; }
 
   /// Verify that the differentiability witness is well-formed.
-  void verify(const SILModule &M) const;
+  void verify(const SILModule &module) const;
 
-  void print(llvm::raw_ostream &OS, bool verbose = false) const;
+  void print(llvm::raw_ostream &os, bool verbose = false) const;
   void dump() const;
 };
 

lib/ParseSIL/ParseSIL.cpp

@@ -6747,7 +6747,10 @@ bool SILParserTUState::parseSILDefaultWitnessTable(Parser &P) {
 
 // SWIFT_ENABLE_TENSORFLOW
 // TODO(TF-893): Dedupe with `SILParser::convertRequirements` upstream.
-// Consider defining this as `Parser::convertRequirements`.
+// Currently, this utility is defined on `SILParser`, but SIL differentiability
+// witness is defined on `SILParserTUState` and only has access to `Parser`.
+// Consider redefining `SILParser::convertRequirements`as
+// `Parser::convertRequirements`.
 static void convertRequirements(Parser &P, SILFunction *F,
                                 ArrayRef<RequirementRepr> From,
                                 SmallVectorImpl<Requirement> &To) {
@@ -6827,8 +6830,8 @@ bool SILParserTUState::parseSILDifferentiabilityWitness(Parser &P) {
   if (!linkage)
     linkage = SILLinkage::PublicExternal;
 
-  Scope S(&P, ScopeKind::TopLevel);
-  Scope Body(&P, ScopeKind::FunctionBody);
+  Scope scope(&P, ScopeKind::TopLevel);
+  Scope body(&P, ScopeKind::FunctionBody);
 
   // Parse a SIL function name.
   auto parseFunctionName = [&](SILFunction *&fn) -> bool {
@@ -6858,11 +6861,10 @@ bool SILParserTUState::parseSILDifferentiabilityWitness(Parser &P) {
   // Parse an index subset, prefaced with the given label.
   auto parseIndexSubset =
       [&](StringRef label, IndexSubset *& indexSubset) -> bool {
-    if (P.parseToken(tok::l_square, diag::sil_diff_witness_expected_keyword,
-                     "["))
+    if (P.parseToken(tok::l_square, diag::sil_diff_witness_expected_token, "["))
       return true;
     if (P.parseSpecificIdentifier(
-          label, diag::sil_diff_witness_expected_keyword, label))
+          label, diag::sil_diff_witness_expected_token, label))
       return true;
     // Parse parameter index list.
     SmallVector<unsigned, 8> paramIndices;
@@ -6883,8 +6885,7 @@ bool SILParserTUState::parseSILDifferentiabilityWitness(Parser &P) {
     while (P.Tok.isNot(tok::r_square))
       if (parseParam())
         return true;
-    if (P.parseToken(tok::r_square, diag::sil_diff_witness_expected_keyword,
-                     "]"))
+    if (P.parseToken(tok::r_square, diag::sil_diff_witness_expected_token, "]"))
       return true;
     auto maxIndexRef =
         std::max_element(paramIndices.begin(), paramIndices.end());
@@ -6911,8 +6912,7 @@ bool SILParserTUState::parseSILDifferentiabilityWitness(Parser &P) {
     P.parseGenericWhereClause(whereLoc, derivativeRequirementReprs,
                               firstTypeInComplete,
                               /*AllowLayoutConstraints*/ false);
-    if (P.parseToken(tok::r_square, diag::sil_diff_witness_expected_keyword,
-                     "]"))
+    if (P.parseToken(tok::r_square, diag::sil_diff_witness_expected_token, "]"))
       return true;
   }
 
@@ -6944,25 +6944,23 @@ bool SILParserTUState::parseSILDifferentiabilityWitness(Parser &P) {
   SILFunction *vjp = nullptr;
   if (P.Tok.is(tok::l_brace)) {
     // Parse '{'.
-    SourceLoc lBraceLoc = P.Tok.getLoc();
-    P.consumeToken(tok::l_brace);
+    SourceLoc lBraceLoc;
+    P.consumeIf(tok::l_brace, lBraceLoc);
     // Parse JVP (optional).
     if (P.Tok.is(tok::identifier) && P.Tok.getText() == "jvp") {
       P.consumeToken(tok::identifier);
-      if (P.parseToken(tok::colon, diag::sil_diff_witness_expected_keyword,
-                       ":"))
+      if (P.parseToken(tok::colon, diag::sil_diff_witness_expected_token, ":"))
         return true;
-      Scope Body(&P, ScopeKind::FunctionBody);
+      Scope body(&P, ScopeKind::FunctionBody);
       if (parseFunctionName(jvp))
         return true;
     }
     // Parse VJP (optional).
     if (P.Tok.is(tok::identifier) && P.Tok.getText() == "vjp") {
       P.consumeToken(tok::identifier);
-      if (P.parseToken(tok::colon, diag::sil_diff_witness_expected_keyword,
-                       ":"))
+      if (P.parseToken(tok::colon, diag::sil_diff_witness_expected_token, ":"))
         return true;
-      Scope Body(&P, ScopeKind::FunctionBody);
+      Scope body(&P, ScopeKind::FunctionBody);
       if (parseFunctionName(vjp))
         return true;
     }

lib/SIL/SILDifferentiabilityWitness.cpp

@@ -22,9 +22,7 @@ SILDifferentiabilityWitness *SILDifferentiabilityWitness::create(
     IndexSubset *parameterIndices, IndexSubset *resultIndices,
     GenericSignature *derivativeGenSig, SILFunction *jvp, SILFunction *vjp,
     bool isSerialized, DeclAttribute *attribute) {
-  void *buf = module.allocate(sizeof(SILDifferentiabilityWitness),
-                              alignof(SILDifferentiabilityWitness));
-  auto *diffWitness = ::new (buf) SILDifferentiabilityWitness(
+  auto *diffWitness = new (module) SILDifferentiabilityWitness(
       module, linkage, originalFunction, parameterIndices, resultIndices,
       derivativeGenSig, jvp, vjp, isSerialized, attribute);
   // Register the differentiability witness in the module.

lib/SIL/SILPrinter.cpp

@@ -3061,23 +3061,23 @@ void SILDefaultWitnessTable::dump() const {
 void SILDifferentiabilityWitness::print(
     llvm::raw_ostream &OS, bool verbose) const {
   OS << "// differentiability witness for "
-     << demangleSymbol(originalFunction->getName()) << "\n";
-  // sil_differentiability_witness @original-function-name : $original-sil-type
+     << demangleSymbol(originalFunction->getName()) << '\n';
   PrintOptions qualifiedSILTypeOptions = PrintOptions::printQualifiedSILType();
+  // sil_differentiability_witness (linkage)?
   OS << "sil_differentiability_witness ";
   printLinkage(OS, linkage, ForDefinition);
-  // [parameters 0 1 ...]
+  // [parameters ...]
   OS << "[parameters ";
   interleave(getParameterIndices()->getIndices(),
              [&](unsigned index) { OS << index; },
-             [&] { OS << " "; });
-  // [results 0 1 ...]
+             [&] { OS << ' '; });
+  // [results ...]
   OS << "] [results ";
   interleave(getResultIndices()->getIndices(),
              [&](unsigned index) { OS << index; },
-             [&] { OS << " "; });
+             [&] { OS << ' '; });
   OS << ']';
-  // [where ...]
+  // ([where ...])?
   if (auto *derivativeGenSig = getDerivativeGenericSignature()) {
     ArrayRef<Requirement> requirements;
     SmallVector<Requirement, 4> requirementsScratch;
@@ -3093,17 +3093,17 @@ void SILDifferentiabilityWitness::print(
     }
     if (!requirements.empty()) {
       OS << " [where ";
-      auto SubPrinter = PrintOptions::printSIL();
+      auto subPrinter = PrintOptions::printSIL();
       interleave(requirements,
                  [&](Requirement req) {
-                   req.print(OS, SubPrinter);
+                   req.print(OS, subPrinter);
                    return;
                  },
                  [&] { OS << ", "; });
       OS << ']';
     }
   }
-  // original: @original-function-name : $original-sil-type
+  // @original-function-name : $original-sil-type
   OS << " @" << originalFunction->getName() << " : "
      << originalFunction->getLoweredType();
   // {
@@ -3112,9 +3112,9 @@ void SILDifferentiabilityWitness::print(
   // }
   OS << " {\n";
   if (jvp)
-    OS << "  jvp: @" << jvp->getName() << " : " << jvp->getLoweredType() << "\n";
+    OS << "  jvp: @" << jvp->getName() << " : " << jvp->getLoweredType() << '\n';
   if (vjp)
-    OS << "  vjp: @" << vjp->getName() << " : " << vjp->getLoweredType() << "\n";
+    OS << "  vjp: @" << vjp->getName() << " : " << vjp->getLoweredType() << '\n';
   OS << "}\n\n";
 }
 

lib/Serialization/DeserializeSIL.cpp

@@ -3369,9 +3369,7 @@ SILDifferentiabilityWitness *SILDeserializer::lookupDifferentiabilityWitness(
   auto iter = DifferentiabilityWitnessList->find(mangledDiffWitnessKey);
   if (iter == DifferentiabilityWitnessList->end())
     return nullptr;
-
-  auto *diffWitness = readDifferentiabilityWitness(*iter);
-  return diffWitness;
+  return readDifferentiabilityWitness(*iter);
 }
 
 void SILDeserializer::getAllDifferentiabilityWitnesses() {

lib/Serialization/SerializeSIL.cpp

@@ -2300,6 +2300,7 @@ void SILSerializer::writeIndexTables() {
     Offset.emit(ScratchRecord, sil_index_block::SIL_PROPERTY_OFFSETS,
                 PropertyOffset);
   }
+
 }
 
 void SILSerializer::writeSILGlobalVar(const SILGlobalVariable &g) {
@@ -2517,7 +2518,7 @@ writeSILDifferentiabilityWitness(const SILDifferentiabilityWitness &dw) {
     vjpID = S.addUniquedStringRef(vjp->getName());
   }
   SmallVector<unsigned, 8> parameterAndResultIndices(
-       dw.getParameterIndices()->begin(), dw.getParameterIndices()->end());
+      dw.getParameterIndices()->begin(), dw.getParameterIndices()->end());
   parameterAndResultIndices.append(dw.getResultIndices()->begin(),
                                    dw.getResultIndices()->end());
   auto originalFnType = original->getLoweredFunctionType();