expand: Derive statements

Expand procedural macros on statements properly.

gcc/rust/ChangeLog:

	* expand/rust-expand-visitor.cc : Change type from pointer to
	reference and expand statements from statements.
	* ast/rust-macro.h: Return macro kind.
	* ast/rust-ast.h: Add Statement kind.
	* ast/rust-item.h: Change module.
	* ast/rust-stmt.h: Return kind.

Co-authored-by: Pierre-Emmanuel Patry <pierre-emmanuel.patry@embecosm.com>
Signed-off-by: Pierre-Emmanuel Patry <pierre-emmanuel.patry@embecosm.com>

5 files changed, 129 insertions, 37 deletions

diff --git a/gcc/rust/ast/rust-ast.h b/gcc/rust/ast/rust-ast.h
index 1e9b5f3..7619690 100644
--- a/gcc/rust/ast/rust-ast.h
+++ b/gcc/rust/ast/rust-ast.h
@@ -884,6 +884,15 @@ class MetaListNameValueStr;
 class Stmt : public Node
 {
 public:
+  enum class Kind
+  {
+    Empty,
+    Item,
+    Let,
+    Expr,
+    MacroInvocation,
+  };
+
   // Unique pointer custom clone function
   std::unique_ptr<Stmt> clone_stmt () const
   {
@@ -900,8 +909,10 @@ public:
   virtual bool is_marked_for_strip () const = 0;
   NodeId get_node_id () const { return node_id; }
 
-  virtual bool is_item () const = 0;
+  virtual Kind get_stmt_kind () = 0;
 
+  // TODO: Can we remove these two?
+  virtual bool is_item () const = 0;
   virtual bool is_expr () const { return false; }
 
   virtual void add_semicolon () {}
@@ -931,6 +942,8 @@ public:
   add_crate_name (std::vector<std::string> &names ATTRIBUTE_UNUSED) const
   {}
 
+  Stmt::Kind get_stmt_kind () final { return Stmt::Kind::Item; }
+
   // FIXME: ARTHUR: Is it okay to have removed that final? Is it *required*
   // behavior that we have items that can also be expressions?
   bool is_item () const override { return true; }
diff --git a/gcc/rust/ast/rust-item.h b/gcc/rust/ast/rust-item.h
index e60c3f0..ff574be 100644
--- a/gcc/rust/ast/rust-item.h
+++ b/gcc/rust/ast/rust-item.h
@@ -1003,7 +1003,7 @@ public:
 
   Identifier get_name () const { return module_name; }
 
-  Kind get_ast_kind () const override { return Kind::MODULE; }
+  AST::Kind get_ast_kind () const override { return AST::Kind::MODULE; }
 
 private:
   Identifier module_name;
diff --git a/gcc/rust/ast/rust-macro.h b/gcc/rust/ast/rust-macro.h
index 075ae63..eeed709 100644
--- a/gcc/rust/ast/rust-macro.h
+++ b/gcc/rust/ast/rust-macro.h
@@ -571,7 +571,10 @@ public:
     is_builtin_rule = true;
   }
 
-  Kind get_ast_kind () const override { return Kind::MACRO_RULES_DEFINITION; }
+  AST::Kind get_ast_kind () const override
+  {
+    return AST::Kind::MACRO_RULES_DEFINITION;
+  }
 
 protected:
   /* Use covariance to implement clone function as returning this object rather
@@ -670,7 +673,10 @@ public:
     return ExprWithoutBlock::get_node_id ();
   }
 
-  Kind get_ast_kind () const override { return Kind::MACRO_INVOCATION; }
+  AST::Kind get_ast_kind () const override
+  {
+    return AST::Kind::MACRO_INVOCATION;
+  }
 
   NodeId get_macro_node_id () const { return node_id; }
 
diff --git a/gcc/rust/ast/rust-stmt.h b/gcc/rust/ast/rust-stmt.h
index d2cace6..db38a8a 100644
--- a/gcc/rust/ast/rust-stmt.h
+++ b/gcc/rust/ast/rust-stmt.h
@@ -48,6 +48,8 @@ public:
 
   bool is_item () const override final { return false; }
 
+  Stmt::Kind get_stmt_kind () final { return Stmt::Kind::Empty; }
+
 protected:
   /* Use covariance to implement clone function as returning this object rather
    * than base */
@@ -172,6 +174,7 @@ public:
   }
 
   bool is_item () const override final { return false; }
+  Stmt::Kind get_stmt_kind () final { return Stmt::Kind::Let; }
 
 protected:
   /* Use covariance to implement clone function as returning this object rather
@@ -199,6 +202,8 @@ public:
 
   std::string as_string () const override;
 
+  Stmt::Kind get_stmt_kind () final { return Stmt::Kind::Expr; }
+
   std::vector<LetStmt *> locals;
 
   ExprStmt (std::unique_ptr<Expr> &&expr, Location locus,
diff --git a/gcc/rust/expand/rust-expand-visitor.cc b/gcc/rust/expand/rust-expand-visitor.cc
index 9183a63..e9eefdb 100644
--- a/gcc/rust/expand/rust-expand-visitor.cc
+++ b/gcc/rust/expand/rust-expand-visitor.cc
@@ -97,14 +97,14 @@ get_traits_to_derive (AST::Attribute &attr)
 }
 
 static std::unique_ptr<AST::Item>
-builtin_derive_item (std::unique_ptr<AST::Item> &item,
-		     const AST::Attribute &derive, BuiltinMacro to_derive)
+builtin_derive_item (AST::Item &item, const AST::Attribute &derive,
+		     BuiltinMacro to_derive)
 {
-  return AST::DeriveVisitor::derive (*item, derive, to_derive);
+  return AST::DeriveVisitor::derive (item, derive, to_derive);
 }
 
 static std::vector<std::unique_ptr<AST::Item>>
-derive_item (std::unique_ptr<AST::Item> &item, const std::string &to_derive,
+derive_item (AST::Item &item, const std::string &to_derive,
 	     MacroExpander &expander)
 {
   std::vector<std::unique_ptr<AST::Item>> result;
@@ -127,8 +127,8 @@ derive_item (std::unique_ptr<AST::Item> &item, const std::string &to_derive,
 }
 
 static std::vector<std::unique_ptr<AST::Item>>
-expand_attribute (std::unique_ptr<AST::Item> &item, AST::SimplePath &name,
-		  MacroExpander &expander)
+expand_item_attribute (AST::Item &item, AST::SimplePath &name,
+		       MacroExpander &expander)
 {
   std::vector<std::unique_ptr<AST::Item>> result;
   auto frag = expander.expand_attribute_proc_macro (item, name);
@@ -149,6 +149,29 @@ expand_attribute (std::unique_ptr<AST::Item> &item, AST::SimplePath &name,
   return result;
 }
 
+static std::vector<std::unique_ptr<AST::Stmt>>
+expand_stmt_attribute (AST::Item &item, AST::SimplePath &name,
+		       MacroExpander &expander)
+{
+  std::vector<std::unique_ptr<AST::Stmt>> result;
+  auto frag = expander.expand_attribute_proc_macro (item, name);
+  if (!frag.is_error ())
+    {
+      for (auto &node : frag.get_nodes ())
+	{
+	  switch (node.get_kind ())
+	    {
+	    case AST::SingleASTNode::STMT:
+	      result.push_back (node.take_stmt ());
+	      break;
+	    default:
+	      gcc_unreachable ();
+	    }
+	}
+    }
+  return result;
+}
+
 void
 ExpandVisitor::expand_inner_items (
   std::vector<std::unique_ptr<AST::Item>> &items)
@@ -179,7 +202,7 @@ ExpandVisitor::expand_inner_items (
 		      if (MacroBuiltin::builtins.is_iter_ok (maybe_builtin))
 			{
 			  auto new_item
-			    = builtin_derive_item (item, current,
+			    = builtin_derive_item (*item, current,
 						   maybe_builtin->second);
 			  // this inserts the derive *before* the item - is it a
 			  // problem?
@@ -188,7 +211,7 @@ ExpandVisitor::expand_inner_items (
 		      else
 			{
 			  auto new_items
-			    = derive_item (item, to_derive, expander);
+			    = derive_item (*item, to_derive, expander);
 			  std::move (new_items.begin (), new_items.end (),
 				     std::inserter (items, it));
 			}
@@ -204,8 +227,8 @@ ExpandVisitor::expand_inner_items (
 		    {
 		      attr_it = attrs.erase (attr_it);
 		      auto new_items
-			= expand_attribute (item, current.get_path (),
-					    expander);
+			= expand_item_attribute (*item, current.get_path (),
+						 expander);
 		      it = items.erase (it);
 		      std::move (new_items.begin (), new_items.end (),
 				 std::inserter (items, it));
@@ -236,29 +259,74 @@ ExpandVisitor::expand_inner_stmts (
 
   for (auto it = stmts.begin (); it != stmts.end (); it++)
     {
-      // TODO: Eventually we need to derive here as well
-
-      // auto &stmt = *it;
-
-      // if (stmt->has_outer_attrs ())
-      // {
-      //   auto traits_to_derive
-      //     = get_traits_to_derive (stmt->get_outer_attrs ());
-
-      //   // FIXME: This needs to be reworked absolutely
-      //   static const std::set<std::string> builtin_derives
-      //     = {"Clone", "Copy", "Eq", "PartialEq", "Ord", "PartialOrd"};
-
-      //   for (auto &to_derive : traits_to_derive)
-      //     if (builtin_derives.find (to_derive) != builtin_derives.end ())
-      //       {
-      // 	auto new_item = derive_item (
-      // 	  item, item->get_outer_attrs ()[0] /* FIXME: This is wrong */,
-      // 	  to_derive);
-      // 	// this inserts the derive *before* the item - is it a problem?
-      // 	it = items.insert (it, std::move (new_item));
-      //       }
-      // }
+      auto &stmt = *it;
+
+      // skip all non-item statements
+      if (stmt->get_stmt_kind () != AST::Stmt::Kind::Item)
+	continue;
+
+      auto &item = static_cast<AST::Item &> (*stmt.get ());
+
+      if (item.has_outer_attrs ())
+	{
+	  auto &attrs = item.get_outer_attrs ();
+
+	  for (auto attr_it = attrs.begin (); attr_it != attrs.end ();
+	       /* erase => No increment*/)
+	    {
+	      auto current = *attr_it;
+
+	      if (is_derive (current))
+		{
+		  attr_it = attrs.erase (attr_it);
+		  // Get traits to derive in the current attribute
+		  auto traits_to_derive = get_traits_to_derive (current);
+		  for (auto &to_derive : traits_to_derive)
+		    {
+		      auto maybe_builtin
+			= MacroBuiltin::builtins.lookup (to_derive);
+		      if (MacroBuiltin::builtins.is_iter_ok (maybe_builtin))
+			{
+			  auto new_item
+			    = builtin_derive_item (item, current,
+						   maybe_builtin->second);
+			  // this inserts the derive *before* the item - is it a
+			  // problem?
+			  it = stmts.insert (it, std::move (new_item));
+			}
+		      else
+			{
+			  auto new_items
+			    = derive_item (item, to_derive, expander);
+			  std::move (new_items.begin (), new_items.end (),
+				     std::inserter (stmts, it));
+			}
+		    }
+		}
+	      else /* Attribute */
+		{
+		  if (is_builtin (current))
+		    {
+		      attr_it++;
+		    }
+		  else
+		    {
+		      attr_it = attrs.erase (attr_it);
+		      auto new_items
+			= expand_stmt_attribute (item, current.get_path (),
+						 expander);
+		      it = stmts.erase (it);
+		      std::move (new_items.begin (), new_items.end (),
+				 std::inserter (stmts, it));
+		      // TODO: Improve this ?
+		      // item is invalid since it refers to now deleted,
+		      // cancel the loop increment and break.
+		      it--;
+		      break;
+		    }
+		}
+	    }
+	}
     }
 
   std::function<std::unique_ptr<AST::Stmt> (AST::SingleASTNode)> extractor