Implement `gn edit "shard"` subcommand

Adds the `shard` subcommand to `gn edit` to split a target into
fine-grained shard targets based on source file stems.

Each shard inherits the original target's compiler flags, defines,
visibility, and conditionals, while stripping existing dependencies.
The original target is converted to a group (or specified group type)
depending on the newly created shards.

Supports optional arguments to customize the sharded target and group types:
  gn edit "shard [sharded_target_type] [group_type]" <labels...>

Change-Id: I5aac2501746d608d5bf1789ec2caf61a6a6a6964
Reviewed-on: https://gn-review.googlesource.com/c/gn/+/25620
Reviewed-by: Takuto Ikuta <tikuta@google.com>
Commit-Queue: Matt Stark <msta@google.com>
diff --git a/docs/reference.md b/docs/reference.md
index 994203e..719bc5d 100644
--- a/docs/reference.md
+++ b/docs/reference.md
@@ -786,6 +786,20 @@
         gn edit "set testonly true" //src/tools:*
         gn edit "set srcs:list foo.cc" //:foo
         gn edit "set deps :bar :baz" //:foo
+
+  shard [sharded_target_type] [group_type]
+      Splits the target's sources into fine-grained shard targets,
+      inheriting compilation flags, and updates the parent
+      target to depend on the newly created shards.
+      Note: Sharding strips existing deps and clones conditional blocks
+      into each shard. For conditional sources (e.g. `if (is_win)`),
+      manual cleanup is often needed, such as:
+        * Moving the condition to wrap the shard instead of the sources
+        * Stripping empty conditions
+
+      Examples:
+        gn edit "shard" //:large_target
+        gn edit "shard source_set static_library" //:large_target
 ```
 ### <a name="cmd_format"></a>**gn format [\--dump-tree] [\--format-width=WIDTH] (\--stdin | &lt;list of build_files...&gt;)**&nbsp;[Back to Top](#gn-reference)
 
diff --git a/src/gn/build_file_editor.cc b/src/gn/build_file_editor.cc
index 8fefdae..6535c43 100644
--- a/src/gn/build_file_editor.cc
+++ b/src/gn/build_file_editor.cc
@@ -42,14 +42,6 @@
   return v;
 }
 
-std::optional<std::string> AsStringLiteral(const ParseNode* node) {
-  auto val = AsLiteralValue(node);
-  if (val && val->type() == Value::STRING) {
-    return std::move(val->string_value());
-  }
-  return std::nullopt;
-}
-
 // Returns true if a node in the tree is a literal node matching the user's
 // request.
 bool Matches(const EditTarget& target,
@@ -168,6 +160,28 @@
 
 }  // namespace
 
+std::optional<std::string> AsStringLiteral(const ParseNode* node) {
+  auto val = AsLiteralValue(node);
+  if (val && val->type() == Value::STRING) {
+    return std::move(val->string_value());
+  }
+  return std::nullopt;
+}
+
+std::vector<TreeNode> FindAllListElements(const TreeNode& assignment) {
+  auto* op = assignment.AsAssignment();
+  if (!op)
+    return {};
+  return FindExpression<TreeNode>(
+      assignment.Descend(op->right()),
+      [](TreeNode& node_ref) -> std::optional<TreeNode> {
+        if (node_ref.parent() && node_ref.parent()->AsList()) {
+          return node_ref;
+        }
+        return std::nullopt;
+      });
+}
+
 std::vector<TreeNode> FindListElementInAssignment(const EditTarget& target,
                                                   const TreeNode& root,
                                                   const Value& value) {
@@ -240,25 +254,24 @@
   return false;
 }
 
-void TreeNode::add_todo(EditState& state, const EditTarget& target) const {
-  const std::vector<std::string> lines = {
-      "# TODO(gn edit: " + state.context + "):",
-      "# This would normally be deleted but is conditional.",
-      "# Manual intervention is required to decide whether it should "
-      "actually be deleted.",
-  };
-  for (const auto& line : lines) {
-    StringAtom atom(line);
-    Token comment_token(node()->GetRange().begin(), Token::LINE_COMMENT,
-                        atom.str());
-    node()->comments_mutable()->append_before(std::move(comment_token));
-  }
+void TreeNode::add_todo(EditState& state,
+                        const EditTarget& target,
+                        std::string_view message) const {
+  std::string line =
+      "# TODO(gn edit: " + state.context + "): " + std::string(message);
+  StringAtom atom(line);
+  Token comment_token(node()->GetRange().begin(), Token::LINE_COMMENT,
+                      atom.str());
+  node()->comments_mutable()->append_before(std::move(comment_token));
   state.needs_manual_review.insert(target.label);
 }
 
 void TreeNode::RemoveSelf(EditState& state, const EditTarget& target) const {
   if (is_conditional()) {
-    add_todo(state, target);
+    add_todo(state, target,
+             "This would normally be deleted but is conditional. Manual "
+             "intervention is required to decide whether it should actually be "
+             "deleted.");
   } else {
     RemoveSelfUnconditionally();
   }
@@ -348,16 +361,19 @@
   return Ok();
 }
 
-std::vector<TreeNode> EditTarget::assignments(std::string_view attr) const {
+std::vector<TreeNode> TreeNode::assignments(
+    std::initializer_list<std::string_view> attrs) const {
   return FindStatement<TreeNode>(
-      block, [attr](TreeNode& node_ref) -> std::optional<TreeNode> {
+      node(), [attrs](TreeNode& node_ref) -> std::optional<TreeNode> {
         if (const auto* op = node_ref->AsBinaryOp()) {
           if (op->op().type() == Token::EQUAL ||
               op->op().type() == Token::PLUS_EQUALS ||
               op->op().type() == Token::MINUS_EQUALS) {
             if (const auto* left = op->left()->AsIdentifier()) {
-              if (left->value().value() == attr) {
-                return node_ref;
+              for (auto attr : attrs) {
+                if (left->value().value() == attr) {
+                  return node_ref;
+                }
               }
             }
           }
@@ -366,6 +382,19 @@
       });
 }
 
+std::vector<TreeNode> TreeNode::assignments(std::string_view attr) const {
+  return assignments({attr});
+}
+
+std::vector<TreeNode> EditTarget::assignments(
+    std::initializer_list<std::string_view> attrs) const {
+  return node.Descend(block).assignments(attrs);
+}
+
+std::vector<TreeNode> EditTarget::assignments(std::string_view attr) const {
+  return assignments({attr});
+}
+
 void EditTarget::add_warning(EditState& state, std::string_view message) const {
   std::string full_message = "Target \"" + label.GetUserVisibleName(false) +
                              "\" " + std::string(message);
diff --git a/src/gn/build_file_editor.h b/src/gn/build_file_editor.h
index ebc0bde..70b27f1 100644
--- a/src/gn/build_file_editor.h
+++ b/src/gn/build_file_editor.h
@@ -6,8 +6,10 @@
 #define TOOLS_GN_BUILD_FILE_EDITOR_H_
 
 #include <functional>
+#include <initializer_list>
 #include <memory>
 #include <optional>
+#include <string_view>
 #include <unordered_map>
 #include <vector>
 
@@ -50,7 +52,9 @@
 
   // Adds a todo comment to the build file to show the user where manual
   // intervention is required.
-  void add_todo(EditState& state, const EditTarget& target) const;
+  void add_todo(EditState& state,
+                const EditTarget& target,
+                std::string_view message) const;
 
   // Removes self from the tree, or adds a TODO suggesting that it should
   // probably be removed.
@@ -76,6 +80,12 @@
 
   TreeNode Descend(ParseNode* child) const;
 
+  // Calculates all =, +=, and -= assignments of the given attributes under this
+  // node.
+  std::vector<TreeNode> assignments(
+      std::initializer_list<std::string_view> attrs) const;
+  std::vector<TreeNode> assignments(std::string_view attr) const;
+
  private:
   std::vector<ParseNode*> stack_;
 };
@@ -124,6 +134,12 @@
   return results;
 }
 
+// Returns the string literal value if the node is a string literal.
+std::optional<std::string> AsStringLiteral(const ParseNode* node);
+
+// Finds all list elements directly in an assignment expression.
+std::vector<TreeNode> FindAllListElements(const TreeNode& assignment);
+
 // Finds an element in an assignment expression ("=" or "+=") whose right-hand
 // side likely evaluates to a list.
 std::vector<TreeNode> FindListElementInAssignment(const EditTarget& target,
@@ -167,7 +183,9 @@
 
 // Represents a build target to be edited.
 struct EditTarget {
-  // Calculates all =, +=, and -= of a given attribute.
+  // Calculates all =, +=, and -= of the given attributes.
+  std::vector<TreeNode> assignments(
+      std::initializer_list<std::string_view> attrs) const;
   std::vector<TreeNode> assignments(std::string_view attr) const;
 
   // Emits a warning to the user.
diff --git a/src/gn/edit_command.cc b/src/gn/edit_command.cc
index fc0c20e..192bd49 100644
--- a/src/gn/edit_command.cc
+++ b/src/gn/edit_command.cc
@@ -91,7 +91,21 @@
     "      Examples:\n"
     "        gn edit \"set testonly true\" //src/tools:*\n"
     "        gn edit \"set srcs:list foo.cc\" //:foo\n"
-    "        gn edit \"set deps :bar :baz\" //:foo\n";
+    "        gn edit \"set deps :bar :baz\" //:foo\n"
+    "\n"
+    "  shard [sharded_target_type] [group_type]\n"
+    "      Splits the target's sources into fine-grained shard targets,\n"
+    "      inheriting compilation flags, and updates the parent\n"
+    "      target to depend on the newly created shards.\n"
+    "      Note: Sharding strips existing deps and clones conditional blocks\n"
+    "      into each shard. For conditional sources (e.g. `if (is_win)`),\n"
+    "      manual cleanup is often needed, such as:\n"
+    "        * Moving the condition to wrap the shard instead of the sources\n"
+    "        * Stripping empty conditions\n"
+    "\n"
+    "      Examples:\n"
+    "        gn edit \"shard\" //:large_target\n"
+    "        gn edit \"shard source_set static_library\" //:large_target\n";
 
 Result<std::pair<std::vector<SourceFile>, EditState>> RunEditImpl(
     const std::vector<std::string>& args,
@@ -186,6 +200,19 @@
         "build file of the form:\n'# TODO(gn edit: <command>): ...'\n",
         DECORATION_DIM);
   }
+  const auto& needs_fix_deps = result->second.needs_fix_deps;
+  if (!needs_fix_deps.empty()) {
+    OutputString(
+        "\nTo automatically resolve and populate dependencies for sharded "
+        "targets, run:\n",
+        DECORATION_YELLOW);
+    std::string check_cmd = "gn check <out_dir>";
+    for (const Label& label : needs_fix_deps) {
+      check_cmd += " " + label.GetUserVisibleName(false);
+    }
+    check_cmd += " --fix\n";
+    OutputString(check_cmd, DECORATION_GREEN);
+  }
   return 0;
 }
 
diff --git a/src/gn/edit_command_unittest.cc b/src/gn/edit_command_unittest.cc
index 7b27d81..9467e92 100644
--- a/src/gn/edit_command_unittest.cc
+++ b/src/gn/edit_command_unittest.cc
@@ -43,6 +43,10 @@
   if (!edited.edit_state_.warnings.empty()) {
     res += "\nWarnings: " + testing::Pretty(edited.edit_state_.warnings);
   }
+  if (!edited.edit_state_.needs_fix_deps.empty()) {
+    res += "\nNeeds check --fix: " +
+           testing::Pretty(edited.edit_state_.needs_fix_deps);
+  }
   return res;
 }
 
@@ -243,16 +247,16 @@
 )"),
                  Edited(R"(
 if (is_win) {
-  # TODO(gn edit: delete):
-  # This would normally be deleted but is conditional.
-  # Manual intervention is required to decide whether it should actually be deleted.
+  # TODO(gn edit: delete): This would normally be deleted but is conditional.
+  # Manual intervention is required to decide whether it should actually be
+  # deleted.
   executable("bar") {
     sources = [ "bar.cc" ]
   }
 } else {
-  # TODO(gn edit: delete):
-  # This would normally be deleted but is conditional.
-  # Manual intervention is required to decide whether it should actually be deleted.
+  # TODO(gn edit: delete): This would normally be deleted but is conditional.
+  # Manual intervention is required to decide whether it should actually be
+  # deleted.
   executable("bar") {
     sources = [ "bar.cc" ]
   }
@@ -580,9 +584,9 @@
   ]
 
   if (is_linux) {
-    # TODO(gn edit: set deps //foo //bar):
-    # This would normally be deleted but is conditional.
-    # Manual intervention is required to decide whether it should actually be deleted.
+    # TODO(gn edit: set deps //foo //bar): This would normally be deleted but is
+    # conditional. Manual intervention is required to decide whether it should
+    # actually be deleted.
     deps += [ "//linux" ]
   }
 }
@@ -603,9 +607,9 @@
                      R"(
 executable("foo") {
   if (is_linux) {
-    # TODO(gn edit: set deps:list //foo):
-    # This would normally be deleted but is conditional.
-    # Manual intervention is required to decide whether it should actually be deleted.
+    # TODO(gn edit: set deps:list //foo): This would normally be deleted but is
+    # conditional. Manual intervention is required to decide whether it should
+    # actually be deleted.
     deps = [ "//linux" ]
     public_deps = [ "//linux" ]
   }
@@ -615,4 +619,150 @@
                      EditState({Label(SourceDir("//"), "foo")})));
 }
 
+TEST_F(EditCommandTest, ShardSubcommand) {
+  // Shards target across sources and subdirectory files, preserving
+  // conditionals, attributes, and visibility.
+  EXPECT_SUCCESS(DoEdit("shard", {"//:foo"},
+                        R"(
+static_library("foo") {
+  sources = [
+    "a.cc",
+    "a.h",
+    "foo.h",
+    "util/foo-bar.cc",
+    "util/foo-bar.h",
+  ]
+  if (is_win) {
+    sources += [ "c_win.cc" ]
+  }
+  defines = [ "ENABLE_FEATURE" ]
+  deps = [ "//base" ]
+  testonly = true
+  visibility = [ "//..." ]
+}
+)"),
+                 Edited(R"(
+group("foo") {
+  public_deps = [
+    ":a",
+    ":c_win",
+    ":foo_foo",
+    ":util_foo_bar",
+  ]
+
+  testonly = true
+  visibility = [ "//..." ]
+}
+static_library("a") {
+  sources = [
+    "a.cc",
+    "a.h",
+  ]
+  if (is_win) {
+    sources += []
+  }
+  defines = [ "ENABLE_FEATURE" ]
+
+  testonly = true
+  visibility = [ "//..." ]
+}
+static_library("c_win") {
+  sources = []
+  if (is_win) {
+    sources += [ "c_win.cc" ]
+  }
+  defines = [ "ENABLE_FEATURE" ]
+
+  testonly = true
+  visibility = [ "//..." ]
+}
+static_library("foo_foo") {
+  sources = [ "foo.h" ]
+  if (is_win) {
+    sources += []
+  }
+  defines = [ "ENABLE_FEATURE" ]
+
+  testonly = true
+  visibility = [ "//..." ]
+}
+static_library("util_foo_bar") {
+  sources = [
+    "util/foo-bar.cc",
+    "util/foo-bar.h",
+  ]
+  if (is_win) {
+    sources += []
+  }
+  defines = [ "ENABLE_FEATURE" ]
+
+  testonly = true
+  visibility = [ "//..." ]
+}
+)",
+                        EditState({}, {},
+                                  {Label(SourceDir("//"), "a"),
+                                   Label(SourceDir("//"), "c_win"),
+                                   Label(SourceDir("//"), "foo_foo"),
+                                   Label(SourceDir("//"), "util_foo_bar")})));
+
+  // Sharding with explicit shard target type and custom group type.
+  EXPECT_SUCCESS(DoEdit("shard source_set static_library", {"//:foo"},
+                        R"(
+static_library("foo") {
+  sources = [
+    "a.cc",
+    "b.cc",
+  ]
+}
+)"),
+                 Edited(R"(
+static_library("foo") {
+  public_deps = [
+    ":a",
+    ":b",
+  ]
+}
+source_set("a") {
+  sources = [ "a.cc" ]
+}
+source_set("b") {
+  sources = [ "b.cc" ]
+}
+)",
+                        EditState({}, {},
+                                  {Label(SourceDir("//"), "a"),
+                                   Label(SourceDir("//"), "b")})));
+
+  // Single source target emits a warning and is not sharded.
+  EXPECT_SUCCESS(
+      DoEdit("shard", {"//:foo"},
+             R"(
+static_library("foo") {
+  sources = [ "foo.cc" ]
+  deps = [ "//base" ]
+}
+)"),
+      Edited(R"(
+static_library("foo") {
+  sources = [ "foo.cc" ]
+  deps = [ "//base" ]
+}
+)",
+             EditState(
+                 {}, {Err(Location(),
+                          "Target \"//:foo\" does not need to be sharded.")})));
+
+  // Absolute source paths produce an error.
+  EXPECT_FAILURE(DoEdit("shard", {"//:foo"},
+                        R"(
+static_library("foo") {
+  sources = [
+    "//base/a.cc",
+    "b.cc",
+  ]
+}
+)"));
+}
+
 }  // namespace commands
diff --git a/src/gn/edit_subcommands.cc b/src/gn/edit_subcommands.cc
index ceddf05..595feda 100644
--- a/src/gn/edit_subcommands.cc
+++ b/src/gn/edit_subcommands.cc
@@ -4,6 +4,12 @@
 
 #include "gn/edit_subcommands.h"
 
+#include <optional>
+#include <set>
+#include <string>
+#include <utility>
+#include <vector>
+
 #include "base/containers/span.h"
 #include "base/strings/string_number_conversions.h"
 #include "gn/build_file_editor.h"
@@ -311,6 +317,130 @@
   });
 }
 
+Result<std::string> GetShardName(const LocationRange& location,
+                                 std::string_view path) {
+  if (path.empty() || path.starts_with("/")) {
+    return Err(
+        location,
+        "Cannot shard target with non-relative source path: " +
+            std::string(path),
+        "Determining a shard name is not supported for non-relative paths.");
+  }
+  if (path.starts_with("./")) {
+    path.remove_prefix(2);
+  }
+
+  size_t last_slash = path.rfind('/');
+  size_t last_dot = path.rfind('.');
+  if (last_dot != std::string_view::npos &&
+      (last_slash == std::string_view::npos || last_dot > last_slash)) {
+    path = path.substr(0, last_dot);
+  }
+
+  std::string shard_name;
+  shard_name.reserve(path.size());
+  for (char c : path) {
+    if (c == '/' || c == '\\' || c == '.' || c == '-') {
+      shard_name.push_back('_');
+    } else {
+      shard_name.push_back(c);
+    }
+  }
+  return shard_name;
+}
+
+EditCommand ShardCommand(
+    std::optional<std::string> shard_target_type = std::nullopt,
+    std::string group_type = "group") {
+  return EditTargetCommand([shard_target_type, group_type](
+                               BuildFile& build_file, const EditTarget& target,
+                               EditState& state) -> Err {
+    std::set<std::string> shards;
+    for (auto assign : target.assignments({"sources", "public"})) {
+      for (const auto& item : FindAllListElements(assign)) {
+        if (auto lit = AsStringLiteral(item.node()); lit) {
+          ASSIGN_OR_RETURN(std::string shard_name,
+                           GetShardName(item->GetRange(), *lit));
+          shards.insert(std::move(shard_name));
+        }
+      }
+    }
+
+    if (shards.size() <= 1) {
+      target.add_warning(state, "does not need to be sharded.");
+      return Ok();
+    }
+
+    for (auto assign : target.assignments({"public_deps", "deps"})) {
+      assign.RemoveSelfUnconditionally();
+    }
+
+    std::vector<Value> target_names;
+    auto [container, it] = target.node.node_location();
+    // Ensure we insert the shard after the group.
+    ++it;
+
+    for (const auto& shard : shards) {
+      // Note: If a target "foo" contains both "foo.cc" and "foo_foo.cc", both
+      // would map to "foo_foo". This is a rare edge case that is not worth
+      // complex disambiguation logic.
+      std::string target_name = (shard == target.label.name())
+                                    ? (target.label.name() + "_" + shard)
+                                    : shard;
+      target_names.push_back(Value(nullptr, ":" + target_name));
+      state.needs_fix_deps.insert(Label(target.label.dir(), target_name));
+
+      auto node = target.node.node()->Clone();
+      auto* func = node->AsFunctionCallMut();
+      if (shard_target_type) {
+        func->set_function(Token(func->function().location(), Token::IDENTIFIER,
+                                 *shard_target_type));
+      }
+      it = container.insert(it, std::move(node));
+      ++it;
+
+      func->args()->contents()[0] =
+          build_file.to_node(Value(nullptr, target_name));
+
+      for (auto assign :
+           TreeNode({func->block()}).assignments({"sources", "public"})) {
+        for (const auto& item : FindAllListElements(assign)) {
+          if (auto lit = AsStringLiteral(item.node()); lit) {
+            auto res = GetShardName(item->GetRange(), *lit);
+            DCHECK(!res.has_error()) << "should have failed above";
+            if (*res != shard) {
+              item.RemoveSelfUnconditionally();
+            }
+          }
+        }
+      }
+    }
+
+    auto* orig_func = target.node.node()->AsFunctionCallMut();
+    orig_func->set_function(
+        Token(orig_func->function().location(), Token::IDENTIFIER, group_type));
+
+    std::vector<std::unique_ptr<ParseNode>> group_stmts;
+    group_stmts.push_back(build_file.create_assignment(
+        "public_deps",
+        build_file.to_node(Value(nullptr, std::move(target_names)))));
+    for (const auto& assign : target.assignments({"testonly", "visibility"})) {
+      auto node = assign.node()->Clone();
+      if (assign.is_conditional()) {
+        TreeNode({node.get()})
+            .add_todo(
+                state, target,
+                "This was conditional in the original target. Manual review is "
+                "required to decide if it applies to this group target.");
+      }
+      group_stmts.push_back(std::move(node));
+    }
+
+    target.block->statements() = std::move(group_stmts);
+    return Ok();
+  });
+}
+
 }  // namespace
 
 Result<EditCommand> ParseCommand(std::vector<std::string> args) {
@@ -411,6 +541,17 @@
     }
 
     return SetCommand(std::string(attribute), std::move(val));
+  } else if (args[0] == "shard") {
+    if (args.size() == 1) {
+      return ShardCommand();
+    } else if (args.size() == 2) {
+      return ShardCommand(args[1]);
+    } else if (args.size() == 3) {
+      return ShardCommand(args[1], args[2]);
+    } else {
+      return Err(Location(), "Invalid shard command.",
+                 "Usage: shard [sharded target type] [group type]");
+    }
   } else {
     return Err(Location(),
                "Unknown edit command: " + std::string(args[0]) +
diff --git a/src/gn/edit_subcommands.h b/src/gn/edit_subcommands.h
index 720be9d..76f0b81 100644
--- a/src/gn/edit_subcommands.h
+++ b/src/gn/edit_subcommands.h
@@ -19,8 +19,12 @@
   explicit EditState(std::string context) : context(std::move(context)) {}
 
   EditState() = default;
-  EditState(std::set<Label> review, std::vector<Err> warn = {})
-      : needs_manual_review(std::move(review)), warnings(std::move(warn)) {}
+  EditState(std::set<Label> review,
+            std::vector<Err> warn = {},
+            std::set<Label> needs_fix_deps = {})
+      : needs_manual_review(std::move(review)),
+        warnings(std::move(warn)),
+        needs_fix_deps(std::move(needs_fix_deps)) {}
 
   // When something needs manual review, gn will output
   // "# TODO(gn edit: <context>)"
@@ -34,6 +38,9 @@
   // runs: `gn edit "remove deps //bar" //foo`, but //bar was not a dependency
   // of //foo.
   std::vector<Err> warnings;
+  // Targets that have had dependencies stripped and need dependency resolution
+  // via `gn check <out_dir> --fix`.
+  std::set<Label> needs_fix_deps;
 };
 
 // EditCommand is a function that modifies a build file.