Allow format_args! with a macro as first argument

Fixes #3896. FormatArgsPieces
isn't great, but I don't see a better alternative for now.

gcc/rust/ChangeLog:

	* ast/rust-ast-visitor.cc (DefaultASTVisitor::visit): Visit
	first argument of FormatArgs.
	* ast/rust-ast.cc (FormatArgs::clone_expr_impl): Remove usage of
	std::cerr for debugging.
	* ast/rust-builtin-ast-nodes.h (enum class FormatArgsNewline):
	New enum.
	(class FormatArgsPieces): New class.
	(FormatArgs::Newline): Replace with using declaration
	referencing FormatArgsNewline.
	(FormatArgs::FormatArgs): Accept a FormatArgsPieces instance.
	(FormatArgs::get_template_arg): New member function.
	(FormatArgs::get_template): Handle changes to class fields.
	(FormatArgs::template_pieces): Rename to...
	(FormatArgs::template_arg): ...here and change type to
	FormatArgsPieces.
	* expand/rust-expand-format-args.cc: Include
	"rust-expand-visitor.h".
	(expand_format_first_arg): New static function.
	(expand_format_args): Change signature.
	* expand/rust-expand-format-args.h: Include "rust-system.h" and
	adjust include order.
	(class ExpandVisitor): Forward declare.
	(struct MacroExpander): Likewise.
	(expand_format_args): Change signature.
	* expand/rust-expand-visitor.cc: Include
	"rust-expand-format-args.h".
	(ExpandVisitor::visit): Attempt to expand FormatArgs.
	* expand/rust-expand-visitor.h: Include "rust-ast-full-decls.h".
	(ExpandVisitor::visit): New overload for FormatArgs.
	* expand/rust-macro-builtins-format-args.cc
	(FormatArgsInput::format_str): Rename to...
	(FormatArgsInput::format_expr): ...here and change type to
	std::unique_ptr<Expr>.
	(format_args_parse_arguments): Handle change to FormatArgsInput.
	(MacroBuiltin::format_args_handler): Likewise and handle changes
	to FormatArgs.

gcc/testsuite/ChangeLog:

	* rust/compile/format_args_concat.rs: New test.

Signed-off-by: Owen Avery <[email protected]>

Author: powerboat9

gcc/rust/ast/rust-ast-visitor.cc

@@ -1500,9 +1500,12 @@ DefaultASTVisitor::visit (AST::BareFunctionType &type)
 }
 
 void
-DefaultASTVisitor::visit (AST::FormatArgs &)
+DefaultASTVisitor::visit (AST::FormatArgs &fmt)
 {
   // FIXME: Do we have anything to do? any subnodes to visit? Probably, right?
+  auto &fmt_arg = fmt.get_template_arg ();
+  if (fmt_arg.has_expr ())
+    visit (fmt_arg.get_expr ());
 }
 
 void

gcc/rust/ast/rust-ast.cc

@@ -5025,8 +5025,6 @@ FormatArgs::set_outer_attrs (std::vector<Attribute>)
 Expr *
 FormatArgs::clone_expr_impl () const
 {
-  std::cerr << "[ARTHUR] cloning FormatArgs! " << std::endl;
-
   return new FormatArgs (*this);
 }
 

gcc/rust/ast/rust-builtin-ast-nodes.h

@@ -168,6 +168,141 @@ class FormatArguments
   std::vector<FormatArgument> args;
 };
 
+enum class FormatArgsNewline
+{
+  Yes,
+  No
+};
+
+// can be either an expression or Fmt::Pieces
+// used to handle `format_args!(some_macro!(...), ...)`
+class FormatArgsPieces final
+{
+public:
+  using Newline = FormatArgsNewline;
+
+  enum class Kind
+  {
+    EXPR,
+    PIECES
+  };
+
+  FormatArgsPieces (std::unique_ptr<Expr> expr, Newline nl)
+    : kind (Kind::EXPR), expr_nl (std::move (expr), nl)
+  {}
+
+  FormatArgsPieces (Fmt::Pieces template_str)
+    : kind (Kind::PIECES), template_str (std::move (template_str))
+  {}
+
+  FormatArgsPieces (const FormatArgsPieces &other) noexcept : kind (other.kind)
+  {
+    switch (kind)
+      {
+      case Kind::EXPR:
+	new (&expr_nl) std::pair<std::unique_ptr<Expr>, Newline> (
+	  other.expr_nl.first->clone_expr (), other.expr_nl.second);
+	break;
+      case Kind::PIECES:
+	new (&template_str) Fmt::Pieces (other.template_str);
+	break;
+      default:
+	rust_unreachable ();
+      }
+  }
+
+  FormatArgsPieces (FormatArgsPieces &&other) noexcept : kind (other.kind)
+  {
+    switch (kind)
+      {
+      case Kind::EXPR:
+	new (&expr_nl) std::pair<std::unique_ptr<Expr>, Newline> (
+	  std::move (other.expr_nl.first), other.expr_nl.second);
+	break;
+      case Kind::PIECES:
+	new (&template_str) Fmt::Pieces (std::move (other.template_str));
+	break;
+      default:
+	rust_unreachable ();
+      }
+  }
+
+  FormatArgsPieces &operator= (const FormatArgsPieces &other)
+  {
+    this->~FormatArgsPieces ();
+    new (this) FormatArgsPieces (other);
+    return *this;
+  }
+
+  FormatArgsPieces &operator= (FormatArgsPieces &&other)
+  {
+    this->~FormatArgsPieces ();
+    new (this) FormatArgsPieces (std::move (other));
+    return *this;
+  }
+
+  bool has_expr () const { return kind == Kind::EXPR; }
+
+  Expr &get_expr ()
+  {
+    rust_assert (has_expr ());
+    return *expr_nl.first;
+  }
+
+  const Expr &get_expr () const
+  {
+    rust_assert (has_expr ());
+    return *expr_nl.first;
+  }
+
+  std::unique_ptr<Expr> &get_expr_ptr ()
+  {
+    rust_assert (has_expr ());
+    return expr_nl.first;
+  }
+
+  bool should_add_newline () const
+  {
+    rust_assert (has_expr ());
+    return expr_nl.second == Newline::Yes;
+  }
+
+  bool has_pieces () const { return kind == Kind::PIECES; }
+
+  const Fmt::Pieces &get_pieces () const
+  {
+    rust_assert (has_pieces ());
+    return template_str;
+  }
+
+  void try_simplify ();
+
+  Kind get_kind () const { return kind; }
+
+  ~FormatArgsPieces ()
+  {
+    switch (kind)
+      {
+      case Kind::EXPR:
+	expr_nl.~pair ();
+	break;
+      case Kind::PIECES:
+	template_str.~Pieces ();
+	break;
+      default:
+	rust_unreachable ();
+      }
+  }
+
+private:
+  Kind kind;
+  union
+  {
+    std::pair<std::unique_ptr<Expr>, Newline> expr_nl;
+    Fmt::Pieces template_str;
+  };
+};
+
 // TODO: Format documentation better
 // Having a separate AST node for `format_args!()` expansion allows some
 // important optimizations which help reduce generated code a lot. For example,
@@ -180,15 +315,11 @@ class FormatArguments
 class FormatArgs : public Expr
 {
 public:
-  enum class Newline
-  {
-    Yes,
-    No
-  };
+  using Newline = FormatArgsNewline;
 
-  FormatArgs (location_t loc, Fmt::Pieces &&template_str,
+  FormatArgs (location_t loc, FormatArgsPieces template_arg,
 	      FormatArguments &&arguments)
-    : loc (loc), template_pieces (std::move (template_str)),
+    : loc (loc), template_arg (std::move (template_arg)),
       arguments (std::move (arguments))
   {}
 
@@ -198,7 +329,12 @@ class FormatArgs : public Expr
 
   void accept_vis (AST::ASTVisitor &vis) override;
 
-  const Fmt::Pieces &get_template () const { return template_pieces; }
+  FormatArgsPieces &get_template_arg () { return template_arg; }
+  const Fmt::Pieces &get_template () const
+  {
+    rust_assert (template_arg.has_pieces ());
+    return template_arg.get_pieces ();
+  }
   const FormatArguments &get_arguments () const { return arguments; }
   virtual location_t get_locus () const override;
 
@@ -210,7 +346,7 @@ class FormatArgs : public Expr
   // expansion of format_args!(). There is extra handling associated with it.
   // we can maybe do that in rust-fmt.cc? in collect_pieces()? like do the
   // transformation into something we can handle better
-  Fmt::Pieces template_pieces;
+  FormatArgsPieces template_arg;
   FormatArguments arguments;
 
   bool marked_for_strip = false;

gcc/rust/expand/rust-expand-format-args.cc

@@ -27,6 +27,7 @@
 #include "rust-path.h"
 #include "rust-system.h"
 #include "rust-token.h"
+#include "rust-expand-visitor.h"
 
 namespace Rust {
 namespace Fmt {
@@ -65,10 +66,58 @@ get_trait_name (ffi::FormatSpec format_specifier)
   return it->second;
 }
 
-tl::optional<AST::Fragment>
-expand_format_args (AST::FormatArgs &fmt,
-		    std::vector<std::unique_ptr<AST::Token>> &&tokens)
+static bool
+expand_format_first_arg (ExpandVisitor &expand_visitor,
+			 AST::FormatArgsPieces &fmt_arg)
 {
+  if (fmt_arg.has_pieces ())
+    return true;
+
+  expand_visitor.maybe_expand_expr (fmt_arg.get_expr_ptr ());
+
+  AST::Expr &expr = fmt_arg.get_expr ();
+
+  auto emit_err = [&] () {
+    rust_error_at (expr.get_locus (),
+		   "format argument must be a string literal");
+  };
+
+  if (!expr.is_literal ())
+    {
+      if (expr.get_expr_kind () != AST::Expr::Kind::MacroInvocation)
+	emit_err ();
+      return false;
+    }
+
+  auto &lit_expr = static_cast<AST::LiteralExpr &> (expr);
+
+  if (lit_expr.get_lit_type () != AST::Literal::STRING
+      && lit_expr.get_lit_type () != AST::Literal::RAW_STRING)
+    {
+      emit_err ();
+      return false;
+    }
+
+  auto fmt_str = lit_expr.get_literal ().as_string ();
+
+  // TODO: does Fmt::Pieces::collect do this?
+  if (fmt_arg.should_add_newline ())
+    fmt_str += '\n';
+
+  auto pieces = Fmt::Pieces::collect (fmt_str, fmt_arg.should_add_newline (),
+				      Fmt::ffi::ParseMode::Format);
+
+  fmt_arg = pieces;
+  return true;
+}
+
+void
+expand_format_args (ExpandVisitor &expand_visitor, MacroExpander &expander,
+		    AST::FormatArgs &fmt)
+{
+  if (!expand_format_first_arg (expand_visitor, fmt.get_template_arg ()))
+    return;
+
   auto loc = fmt.get_locus ();
   auto builder = AST::Builder (loc);
   auto &arguments = fmt.get_arguments ();
@@ -131,9 +180,12 @@ expand_format_args (AST::FormatArgs &fmt,
   auto final_call
     = builder.call (std::move (final_path), std::move (final_args));
 
-  auto node = AST::SingleASTNode (std::move (final_call));
+  // TODO: is it alright to have an empty token list?
+  std::vector<std::unique_ptr<AST::Token>> fragment_tokens;
 
-  return AST::Fragment ({node}, std::move (tokens));
+  expander.set_expanded_fragment (
+    AST::Fragment ({AST::SingleASTNode (std::move (final_call))},
+		   std::move (fragment_tokens)));
 }
 
 } // namespace Fmt

gcc/rust/expand/rust-expand-format-args.h

@@ -19,15 +19,20 @@
 #ifndef RUST_EXPAND_FORMAT_ARGS_H
 #define RUST_EXPAND_FORMAT_ARGS_H
 
-#include "optional.h"
+#include "rust-system.h"
 #include "rust-ast-fragment.h"
+#include "optional.h"
 
 namespace Rust {
+
+// forward declare
+class ExpandVisitor;
+struct MacroExpander;
+
 namespace Fmt {
 
-tl::optional<AST::Fragment>
-expand_format_args (AST::FormatArgs &fmt,
-		    std::vector<std::unique_ptr<AST::Token>> &&tokens);
+void expand_format_args (ExpandVisitor &expand_visitor, MacroExpander &expander,
+			 AST::FormatArgs &fmt);
 
 } // namespace Fmt
 } // namespace Rust

gcc/rust/expand/rust-expand-visitor.cc

@@ -23,6 +23,7 @@
 #include "rust-ast.h"
 #include "rust-type.h"
 #include "rust-derive.h"
+#include "rust-expand-format-args.h"
 
 namespace Rust {
 
@@ -1061,6 +1062,21 @@ ExpandVisitor::visit (AST::SelfParam &param)
     maybe_expand_type (param.get_type_ptr ());
 }
 
+void
+ExpandVisitor::visit (AST::FormatArgs &fmt)
+{
+  if (macro_invoc_expect_id != fmt.get_node_id ())
+    {
+      rust_internal_error_at (fmt.get_locus (),
+			      "attempting to expand FormatArgs node with id %d "
+			      "into position with node id %d",
+			      (int) fmt.get_node_id (),
+			      (int) macro_invoc_expect_id);
+    }
+
+  Fmt::expand_format_args (*this, expander, fmt);
+}
+
 template <typename T>
 void
 ExpandVisitor::expand_inner_attribute (T &item, AST::SimplePath &path)

gcc/rust/expand/rust-expand-visitor.h

@@ -22,6 +22,7 @@
 #include "rust-ast-visitor.h"
 #include "rust-macro-expand.h"
 #include "rust-proc-macro.h"
+#include "rust-ast-full-decls.h"
 
 namespace Rust {
 
@@ -284,6 +285,9 @@ class ExpandVisitor : public AST::DefaultASTVisitor
   void visit (AST::FunctionParam &type) override;
   void visit (AST::SelfParam &type) override;
 
+  // expand these if possible
+  void visit (AST::FormatArgs &fmt) override;
+
   template <typename T>
   void expand_inner_attribute (T &item, AST::SimplePath &Path);