Add Printer::Sub for finer configuration of how a variable substitution is processed.

In particular, this allows for configuring what the "consume after" characters that follow
a $...$ are, which previously were ";," for callback subs and "" for everything else.

In a followup the three-argument constructor for Sub will be removed and replaced with
calls to Annotate().

Also, calls to WithVars still reject closure arguments, but this check happens at runtime
now. This is a minor papercut for users that simplifies the implementation of Sub and
related private types immensely.

PiperOrigin-RevId: 494254545
diff --git a/src/google/protobuf/io/printer.cc b/src/google/protobuf/io/printer.cc
index 5802908..d46c009 100644
--- a/src/google/protobuf/io/printer.cc
+++ b/src/google/protobuf/io/printer.cc
@@ -261,9 +261,10 @@
 }
 
 absl::string_view Printer::LookupVar(absl::string_view var) {
-  LookupResult result = LookupInFrameStack(var, absl::MakeSpan(var_lookups_));
+  auto result = LookupInFrameStack(var, absl::MakeSpan(var_lookups_));
   GOOGLE_ABSL_CHECK(result.has_value()) << "could not find " << var;
-  auto* view = absl::get_if<absl::string_view>(&*result);
+
+  auto* view = result->AsString();
   GOOGLE_ABSL_CHECK(view != nullptr)
       << "could not find " << var << "; found callback instead";
 
@@ -299,16 +300,13 @@
   indent_ -= options_.spaces_per_indent;
 }
 
-void Printer::Emit(
-    std::initializer_list<
-        VarDefinition<absl::string_view, /*allow_callbacks=*/true>>
-        vars,
-    absl::string_view format, SourceLocation loc) {
+void Printer::Emit(std::initializer_list<Sub> vars, absl::string_view format,
+                   SourceLocation loc) {
   PrintOptions opts;
   opts.strip_raw_string_indentation = true;
   opts.loc = loc;
 
-  auto defs = WithDefs(vars);
+  auto defs = WithDefs(vars, /*allow_callbacks=*/true);
 
   PrintImpl(format, {}, opts);
 }
@@ -323,9 +321,10 @@
   }
 
   std::pair<size_t, size_t> range = it->second;
-  if (!Validate(range.first <= range.second, opts, [varname] {
-        return absl::StrCat(
-            "variable used for annotation used multiple times: ", varname);
+  if (!Validate(range.first <= range.second, opts, [range, varname] {
+        return absl::StrFormat(
+            "variable used for annotation used multiple times: %s (%d..%d)",
+            varname, range.first, range.second);
       })) {
     return absl::nullopt;
   }
@@ -614,7 +613,7 @@
         continue;
       }
 
-      LookupResult sub;
+      absl::optional<ValueView> sub;
       absl::optional<AnnotationRecord> same_name_record;
       if (opts.allow_digit_substitutions && absl::ascii_isdigit(var[0])) {
         if (!Validate(var.size() == 1u, opts,
@@ -652,7 +651,7 @@
       size_t range_start = sink_.bytes_written();
       size_t range_end = sink_.bytes_written();
 
-      if (auto* str = absl::get_if<absl::string_view>(&*sub)) {
+      if (const absl::string_view* str = sub->AsString()) {
         if (at_start_of_line_ && str->empty()) {
           line_start_variables_.emplace_back(var);
         }
@@ -666,7 +665,7 @@
           PrintRaw(suffix);
         }
       } else {
-        auto* fnc = absl::get_if<std::function<void()>>(&*sub);
+        const ValueView::Callback* fnc = sub->AsCallback();
         GOOGLE_ABSL_CHECK(fnc != nullptr);
 
         Validate(
@@ -676,34 +675,42 @@
         range_start = sink_.bytes_written();
         (*fnc)();
         range_end = sink_.bytes_written();
+      }
 
-        // If we just evaluated a closure, and we are at the start of a line,
-        // that means it finished with a newline. If a newline follows
-        // immediately after, we drop it. This helps callback formatting "work
-        // as expected" with respect to forms like
-        //
-        //   class Foo {
-        //     $methods$;
-        //   };
-        //
-        // Without this line, this would turn into something like
-        //
-        //   class Foo {
-        //     void Bar() {}
-        //
-        //   };
-        //
-        // in many cases. We *also* do this if a ; or , follows the
-        // substitution, because this helps clang-format keep its head on in
-        // many cases. Users that need to keep the semi can write $foo$/**/;
-        ++chunk_idx;
-        if (chunk_idx < line.chunks.size()) {
-          absl::string_view text = line.chunks[chunk_idx].text;
-          if (!absl::ConsumePrefix(&text, ";")) {
-            absl::ConsumePrefix(&text, ",");
+      // If we just evaluated a value which specifies end-of-line consume-after
+      // characters, and we're at the start of a line, that means we finished
+      // with a newline.
+      //
+      // We trim a single end-of-line `consume_after` character in this case.
+      //
+      // This helps callback formatting "work as expected" with respect to forms
+      // like
+      //
+      //   class Foo {
+      //     $methods$;
+      //   };
+      //
+      // Without this post-processing, it would turn into
+      //
+      //   class Foo {
+      //     void Bar() {};
+      //   };
+      //
+      // in many cases. Without the `;`, clang-format may format the template
+      // incorrectly.
+      auto next_idx = chunk_idx + 1;
+      if (!sub->consume_after.empty() && next_idx < line.chunks.size() &&
+          !line.chunks[next_idx].is_var) {
+        chunk_idx = next_idx;
+
+        absl::string_view text = line.chunks[chunk_idx].text;
+        for (char c : sub->consume_after) {
+          if (absl::ConsumePrefix(&text, absl::string_view(&c, 1))) {
+            break;
           }
-          PrintRaw(text);
         }
+
+        PrintRaw(text);
       }
 
       if (same_name_record.has_value() &&
diff --git a/src/google/protobuf/io/printer.h b/src/google/protobuf/io/printer.h
index 4896154..45698c2 100644
--- a/src/google/protobuf/io/printer.h
+++ b/src/google/protobuf/io/printer.h
@@ -172,6 +172,26 @@
 // will crash. Callers must statically know that every variable reference is
 // valid, and MUST NOT pass user-provided strings directly into Emit().
 //
+// Substitutions can be configured to "chomp" a single character after them, to
+// help make indentation work out. This can be configured by passing a
+// two-argument io::Printer::Value into Emit's substitution map:
+//
+//   p.Emit({{"var", io::Printer::Value{var_decl, ";"}}}, R"cc(
+//     class $class$ {
+//      public:
+//       $var$;
+//     };
+//   )cc");
+//
+// This will delete the ; after $var$, regardless of whether it was an empty
+// declaration or not. It will also intelligently attempt to clean up
+// empty lines that follow, if it was on an empty line; this promotes cleaner
+// formatting of the output.
+//
+// Any number of different characters can be potentially skipped, but only one
+// will actually be skipped. For example, callback substitutions (see below) use
+// ";," by default as their "chomping set".
+//
 // # Callback Substitution
 //
 // Instead of passing a string into Emit(), it is possible to pass in a callback
@@ -436,31 +456,71 @@
     }
   };
 
-  // Sink type for constructing values to pass to WithVars() and Emit().
-  template <typename K, bool allow_callbacks>
-  struct VarDefinition {
-    using StringOrCallback = absl::variant<std::string, std::function<void()>>;
+  // Helper type for wrapping a variable substitution expansion result.
+  template <bool owned>
+  struct ValueImpl {
+   private:
+    template <typename T>
+    struct IsSubImpl : std::false_type {};
+    template <bool a>
+    struct IsSubImpl<ValueImpl<a>> : std::true_type {};
 
-    template <typename Key, typename Value>
-    VarDefinition(Key&& key, Value&& value)
-        : key(std::forward<Key>(key)),
-          value(ToStringOrCallback(std::forward<Value>(value), Rank2{})),
-          annotation(absl::nullopt) {}
+   public:
+    using StringType =
+        std::conditional_t<owned, std::string, absl::string_view>;
+    // These callbacks return false if this is a recursive call.
+    using Callback = std::function<bool()>;
+    using StringOrCallback = absl::variant<StringType, Callback>;
 
-    // NOTE: This is an overload rather than taking optional<AnnotationRecord>
-    // with a default argument of nullopt, because we want to pick up
-    // AnnotationRecord's user-defined conversions. Because going from
-    // e.g. Descriptor* -> optional<AnnotationRecord> requires two user-defined
-    // conversions, this does not work.
-    template <typename Key, typename Value>
-    VarDefinition(Key&& key, Value&& value, AnnotationRecord annotation)
-        : key(std::forward<Key>(key)),
-          value(ToStringOrCallback(std::forward<Value>(value), Rank2{})),
-          annotation(std::move(annotation)) {}
+    ValueImpl() = default;
 
-    K key;
+    // This is a template to avoid colliding with the copy constructor below.
+    template <typename Value,
+              std::enable_if_t<!IsSubImpl<absl::remove_cvref_t<Value>>::value,
+                               int> = 0>
+    ValueImpl(Value&& value)  // NOLINT
+        : value(ToStringOrCallback(std::forward<Value>(value), Rank2{})) {
+      if (absl::holds_alternative<Callback>(this->value)) {
+        consume_after = ";,";
+      }
+    }
+
+    // Copy ctor/assign allow interconversion of the two template parameters.
+    template <bool that_owned>
+    ValueImpl(const ValueImpl<that_owned>& that) {  // NOLINT
+      *this = that;
+    }
+
+    template <bool that_owned>
+    ValueImpl& operator=(const ValueImpl<that_owned>& that) {
+      // Cast to void* is required, since this and that may potentially be of
+      // different types (due to the `that_owned` parameter).
+      if (static_cast<const void*>(this) == static_cast<const void*>(&that)) {
+        return *this;
+      }
+
+      using ThatStringType = typename ValueImpl<that_owned>::StringType;
+
+      if (auto* str = absl::get_if<ThatStringType>(&that.value)) {
+        value = StringType(*str);
+      } else {
+        value = absl::get<Callback>(that.value);
+      }
+
+      consume_after = that.consume_after;
+      return *this;
+    }
+
+    const StringType* AsString() const {
+      return absl::get_if<StringType>(&value);
+    }
+
+    const Callback* AsCallback() const {
+      return absl::get_if<Callback>(&value);
+    }
+
     StringOrCallback value;
-    absl::optional<AnnotationRecord> annotation;
+    std::string consume_after;
 
    private:
     // go/ranked-overloads
@@ -474,22 +534,33 @@
     //
     // This is done to produce a better error message than the "candidate does
     // not match" SFINAE errors.
-    template <bool allowed = allow_callbacks>
-    StringOrCallback ToStringOrCallback(std::function<void()> cb, Rank2) {
-      static_assert(
-          allowed, "callback-typed variables are not allowed in this location");
-      return cb;
+    template <typename Cb, typename = decltype(std::declval<Cb&&>()())>
+    StringOrCallback ToStringOrCallback(Cb&& cb, Rank2) {
+      return Callback(
+          [cb = std::forward<Cb>(cb), is_called = false]() mutable -> bool {
+            if (is_called) {
+              // Catch whether or not this function is being called recursively.
+              return false;
+            }
+            is_called = true;
+            cb();
+            is_called = false;
+            return true;
+          });
     }
 
     // Separate from the AlphaNum overload to avoid copies when taking strings
-    // by value.
-    StringOrCallback ToStringOrCallback(std::string s, Rank1) { return s; }
+    // by value when in `owned` mode.
+    StringOrCallback ToStringOrCallback(StringType s, Rank1) { return s; }
 
     StringOrCallback ToStringOrCallback(const absl::AlphaNum& s, Rank0) {
-      return std::string(s.Piece());
+      return StringType(s.Piece());
     }
   };
 
+  using ValueView = ValueImpl</*owned=*/false>;
+  using Value = ValueImpl</*owned=*/true>;
+
   // Provide a helper to use heterogeneous lookup when it's available.
   template <class...>
   using void_t = void;
@@ -516,6 +587,44 @@
   static constexpr absl::string_view kProtocCodegenTrace =
       "PROTOC_CODEGEN_TRACE";
 
+  // Sink type for constructing substitutions to pass to WithVars() and Emit().
+  class Sub {
+   public:
+    template <typename Value>
+    Sub(std::string key, Value&& value)
+        : key_(std::move(key)),
+          value_(std::forward<Value>(value)),
+          annotation_(absl::nullopt) {}
+
+    // NOTE: This is an overload rather than taking optional<AnnotationRecord>
+    // with a default argument of nullopt, because we want to pick up
+    // AnnotationRecord's user-defined conversions. Because going from
+    // e.g. Descriptor* -> optional<AnnotationRecord> requires two user-defined
+    // conversions, this does not work.
+    template <typename Key, typename Value>
+    Sub(Key&& key, Value&& value, AnnotationRecord annotation)
+        : key_(std::forward<Key>(key)),
+          value_(std::forward<Value>(value)),
+          annotation_(std::move(annotation)) {}
+
+    Sub Annotate(AnnotationRecord annotation) && {
+      annotation_ = std::move(annotation);
+      return std::move(*this);
+    }
+
+    Sub WithSuffix(std::string sub_suffix) && {
+      value_.consume_after = std::move(sub_suffix);
+      return std::move(*this);
+    }
+
+   private:
+    friend class Printer;
+
+    std::string key_;
+    Value value_;
+    absl::optional<AnnotationRecord> annotation_;
+  };
+
   // Options for controlling how the output of a Printer is formatted.
   struct Options {
     Options() = default;
@@ -572,13 +681,14 @@
   // Returns an RAII object that pops the lookup frame.
   template <typename Map>
   auto WithVars(const Map* vars) {
-    var_lookups_.emplace_back([vars](absl::string_view var) -> LookupResult {
-      auto it = vars->find(ToStringKey<Map>(var));
-      if (it == vars->end()) {
-        return absl::nullopt;
-      }
-      return absl::string_view(it->second);
-    });
+    var_lookups_.emplace_back(
+        [vars](absl::string_view var) -> absl::optional<ValueView> {
+          auto it = vars->find(ToStringKey<Map>(var));
+          if (it == vars->end()) {
+            return absl::nullopt;
+          }
+          return ValueView(it->second);
+        });
     return absl::MakeCleanup([this] { var_lookups_.pop_back(); });
   }
 
@@ -591,20 +701,19 @@
   template <typename Map = absl::flat_hash_map<std::string, std::string>,
             std::enable_if_t<!std::is_pointer<Map>::value, int> = 0>
   auto WithVars(Map&& vars) {
-    var_lookups_.emplace_back([vars = std::forward<Map>(vars)](
-                                  absl::string_view var) -> LookupResult {
-      auto it = vars.find(ToStringKey<Map>(var));
-      if (it == vars.end()) {
-        return absl::nullopt;
-      }
-      return absl::string_view(it->second);
-    });
+    var_lookups_.emplace_back(
+        [vars = std::forward<Map>(vars)](
+            absl::string_view var) -> absl::optional<ValueView> {
+          auto it = vars.find(ToStringKey<Map>(var));
+          if (it == vars.end()) {
+            return absl::nullopt;
+          }
+          return ValueView(it->second);
+        });
     return absl::MakeCleanup([this] { var_lookups_.pop_back(); });
   }
 
-  auto WithVars(std::initializer_list<
-                VarDefinition<std::string, /*allow_callbacks=*/false>>
-                    vars);
+  auto WithVars(std::initializer_list<Sub> vars);
 
   // Looks up a variable set with WithVars().
   //
@@ -673,10 +782,7 @@
   // documentation for more details.
   //
   // `format` MUST be a string constant.
-  void Emit(std::initializer_list<
-                VarDefinition<absl::string_view, /*allow_callbacks=*/true>>
-                vars,
-            absl::string_view format,
+  void Emit(std::initializer_list<Sub> vars, absl::string_view format,
             SourceLocation loc = SourceLocation::current());
 
   // Write a string directly to the underlying output, performing no formatting
@@ -881,10 +987,8 @@
   // Prints a codegen trace, for the given location in the compiler's source.
   void PrintCodegenTrace(absl::optional<SourceLocation> loc);
 
-  // The core implementation for "fully-elaborated" variable definitions. This
-  // is a private function to avoid users being able to set `allow_callbacks`.
-  template <typename K, bool allow_callbacks>
-  auto WithDefs(std::initializer_list<VarDefinition<K, allow_callbacks>> vars);
+  // The core implementation for "fully-elaborated" variable definitions.
+  auto WithDefs(std::initializer_list<Sub> vars, bool allow_callbacks);
 
   // Returns the start and end of the value that was substituted in place of
   // the variable `varname` in the last call to PrintImpl() (with
@@ -899,10 +1003,8 @@
   bool at_start_of_line_ = true;
   bool failed_ = false;
 
-  using LookupResult =
-      absl::optional<absl::variant<absl::string_view, std::function<void()>>>;
-
-  std::vector<std::function<LookupResult(absl::string_view)>> var_lookups_;
+  std::vector<std::function<absl::optional<ValueView>(absl::string_view)>>
+      var_lookups_;
 
   std::vector<
       std::function<absl::optional<AnnotationRecord>(absl::string_view)>>
@@ -918,39 +1020,33 @@
   std::vector<std::string> line_start_variables_;
 };
 
-template <typename K, bool allow_callbacks>
-auto Printer::WithDefs(
-    std::initializer_list<VarDefinition<K, allow_callbacks>> vars) {
-  absl::flat_hash_map<K, absl::variant<std::string, std::function<void()>>>
-      var_map;
+inline auto Printer::WithDefs(std::initializer_list<Sub> vars,
+                              bool allow_callbacks) {
+  absl::flat_hash_map<std::string, Value> var_map;
   var_map.reserve(vars.size());
 
-  absl::flat_hash_map<K, AnnotationRecord> annotation_map;
+  absl::flat_hash_map<std::string, AnnotationRecord> annotation_map;
 
-  for (auto& var : vars) {
-    auto result = var_map.insert({var.key, var.value});
+  for (const auto& var : vars) {
+    GOOGLE_ABSL_CHECK(allow_callbacks || var.value_.AsCallback() == nullptr)
+        << "callback arguments are not permitted in this position";
+    auto result = var_map.insert({var.key_, var.value_});
     GOOGLE_ABSL_CHECK(result.second)
-        << "repeated variable in Emit() or WithVars() call: \"" << var.key
+        << "repeated variable in Emit() or WithVars() call: \"" << var.key_
         << "\"";
-    if (var.annotation.has_value()) {
-      annotation_map.insert({var.key, *var.annotation});
+    if (var.annotation_.has_value()) {
+      annotation_map.insert({var.key_, *var.annotation_});
     }
   }
 
-  var_lookups_.emplace_back(
-      [map = std::move(var_map)](absl::string_view var) -> LookupResult {
-        auto it = map.find(var);
-        if (it == map.end()) {
-          return absl::nullopt;
-        }
-        if (auto* str = absl::get_if<std::string>(&it->second)) {
-          return absl::string_view(*str);
-        }
-
-        auto* f = absl::get_if<std::function<void()>>(&it->second);
-        GOOGLE_ABSL_CHECK(f != nullptr);
-        return *f;
-      });
+  var_lookups_.emplace_back([map = std::move(var_map)](absl::string_view var)
+                                -> absl::optional<ValueView> {
+    auto it = map.find(var);
+    if (it == map.end()) {
+      return absl::nullopt;
+    }
+    return ValueView(it->second);
+  });
 
   bool has_annotations = !annotation_map.empty();
   if (has_annotations) {
@@ -973,10 +1069,8 @@
   });
 }
 
-inline auto Printer::WithVars(
-    std::initializer_list<VarDefinition<std::string, /*allow_callbacks=*/false>>
-        vars) {
-  return WithDefs(vars);
+inline auto Printer::WithVars(std::initializer_list<Sub> vars) {
+  return WithDefs(vars, /*allow_callbacks=*/false);
 }
 }  // namespace io
 }  // namespace protobuf
diff --git a/src/google/protobuf/io/printer_unittest.cc b/src/google/protobuf/io/printer_unittest.cc
index ebf31eb..754426d 100644
--- a/src/google/protobuf/io/printer_unittest.cc
+++ b/src/google/protobuf/io/printer_unittest.cc
@@ -598,6 +598,27 @@
             "};\n");
 }
 
+TEST_F(PrinterTest, EmitConsumeAfter) {
+  {
+    Printer printer(output());
+    printer.Emit(
+        {
+            {"class", "Foo"},
+            Printer::Sub{"var", "int x;"}.WithSuffix(";"),
+        },
+        R"cc(
+          class $class$ {
+            $var$;
+          };
+        )cc");
+  }
+
+  EXPECT_EQ(written(),
+            "class Foo {\n"
+            "  int x;\n"
+            "};\n");
+}
+
 TEST_F(PrinterTest, EmitWithSpacedVars) {
   {
     Printer printer(output());