From e25ec536ed2bb98b2048dbe8e8e646f1b91b0082 Mon Sep 17 00:00:00 2001
From: Essien Ita Essien <34972+essiene@users.noreply.github.com>
Date: Fri, 16 Aug 2024 01:10:44 +0100
Subject: [PATCH] next/prev: Fetch descendants with more correctness.

See context in [this discussion](https://github.com/martinvonz/jj/pull/3935#discussion_r1649520967)

Fixes #3947
---
 cli/src/movement_util.rs             | 32 ++++++------
 cli/tests/test_next_prev_commands.rs | 75 ++++++++++++++++++++++++++++
 2 files changed, 91 insertions(+), 16 deletions(-)

diff --git a/cli/src/movement_util.rs b/cli/src/movement_util.rs
index da66fe82bb..ce197f7e6c 100644
--- a/cli/src/movement_util.rs
+++ b/cli/src/movement_util.rs
@@ -121,27 +121,27 @@ impl Direction {
         start_revset: &Rc<RevsetExpression>,
         args: &MovementArgsInternal,
     ) -> Result<Rc<RevsetExpression>, CommandError> {
-        let target_revset = match (self, args.conflict) {
-            (Direction::Next, true) => start_revset
+        let nth = match (self, args.should_edit) {
+            (Direction::Next, true) => start_revset.descendants_at(args.offset),
+            (Direction::Next, false) => start_revset
                 .children()
+                .minus(working_revset)
+                .descendants_at(args.offset - 1),
+            (Direction::Prev, _) => start_revset.ancestors_at(args.offset),
+        };
+
+        let target_revset = match (self, args.conflict) {
+            (_, false) => nth,
+            (Direction::Next, true) => nth
                 .descendants()
                 .filtered(RevsetFilterPredicate::HasConflict)
-                .roots()
-                .minus(working_revset),
-            (Direction::Next, false) => start_revset
-                .descendants_at(args.offset)
-                .minus(working_revset),
-            (Direction::Prev, true) =>
+                .roots(),
             // If people desire to move to the root conflict, replace the `heads()` below
             // with `roots(). But let's wait for feedback.
-            {
-                start_revset
-                    .parents()
-                    .ancestors()
-                    .filtered(RevsetFilterPredicate::HasConflict)
-                    .heads()
-            }
-            (Direction::Prev, false) => start_revset.ancestors_at(args.offset),
+            (Direction::Prev, true) => nth
+                .ancestors()
+                .filtered(RevsetFilterPredicate::HasConflict)
+                .heads(),
         };
 
         Ok(target_revset)
diff --git a/cli/tests/test_next_prev_commands.rs b/cli/tests/test_next_prev_commands.rs
index 84d51b7bc5..6d7f8b5b15 100644
--- a/cli/tests/test_next_prev_commands.rs
+++ b/cli/tests/test_next_prev_commands.rs
@@ -1112,6 +1112,81 @@ fn test_movement_edit_mode_false() {
     "###);
 }
 
+#[test]
+fn test_next_offset_when_wc_has_descendants() {
+    let test_env = TestEnvironment::default();
+    test_env.add_config(r#"ui.movement.edit = false"#);
+
+    test_env.jj_cmd_ok(test_env.env_root(), &["init", "repo", "--git"]);
+    let repo_path = test_env.env_root().join("repo");
+
+    test_env.jj_cmd_ok(&repo_path, &["commit", "-m", "base"]);
+    test_env.jj_cmd_ok(&repo_path, &["commit", "-m", "right-wc"]);
+    test_env.jj_cmd_ok(&repo_path, &["commit", "-m", "right-1"]);
+    test_env.jj_cmd_ok(&repo_path, &["commit", "-m", "right-2"]);
+
+    test_env.jj_cmd_ok(&repo_path, &["new", "description(base)"]);
+    test_env.jj_cmd_ok(&repo_path, &["commit", "-m", "left-wc"]);
+    test_env.jj_cmd_ok(&repo_path, &["commit", "-m", "left-1"]);
+    test_env.jj_cmd_ok(&repo_path, &["commit", "-m", "left-2"]);
+
+    test_env.jj_cmd_ok(&repo_path, &["edit", "description(right-wc)"]);
+    insta::assert_snapshot!(get_log_output(&test_env, &repo_path), @r###"
+    ○  vruxwmqvtpmx left-2
+    ○  yqosqzytrlsw left-1
+    ○  royxmykxtrkr left-wc
+    │ ○  zsuskulnrvyr right-2
+    │ ○  kkmpptxzrspx right-1
+    │ @  rlvkpnrzqnoo right-wc
+    ├─╯
+    ○  qpvuntsmwlqt base
+    ◆  zzzzzzzzzzzz
+    "###);
+
+    test_env.jj_cmd_ok(&repo_path, &["next", "2"]);
+    insta::assert_snapshot!(get_log_output(&test_env, &repo_path), @r###"
+    @  kmkuslswpqwq
+    │ ○  vruxwmqvtpmx left-2
+    ├─╯
+    ○  yqosqzytrlsw left-1
+    ○  royxmykxtrkr left-wc
+    │ ○  zsuskulnrvyr right-2
+    │ ○  kkmpptxzrspx right-1
+    │ ○  rlvkpnrzqnoo right-wc
+    ├─╯
+    ○  qpvuntsmwlqt base
+    ◆  zzzzzzzzzzzz
+    "###);
+
+    test_env.jj_cmd_ok(&repo_path, &["edit", "description(left-wc)"]);
+    insta::assert_snapshot!(get_log_output(&test_env, &repo_path), @r###"
+    ○  vruxwmqvtpmx left-2
+    ○  yqosqzytrlsw left-1
+    @  royxmykxtrkr left-wc
+    │ ○  zsuskulnrvyr right-2
+    │ ○  kkmpptxzrspx right-1
+    │ ○  rlvkpnrzqnoo right-wc
+    ├─╯
+    ○  qpvuntsmwlqt base
+    ◆  zzzzzzzzzzzz
+    "###);
+
+    test_env.jj_cmd_ok(&repo_path, &["next"]);
+    insta::assert_snapshot!(get_log_output(&test_env, &repo_path), @r###"
+    @  nkmrtpmomlro
+    │ ○  zsuskulnrvyr right-2
+    │ ○  kkmpptxzrspx right-1
+    ├─╯
+    ○  rlvkpnrzqnoo right-wc
+    │ ○  vruxwmqvtpmx left-2
+    │ ○  yqosqzytrlsw left-1
+    │ ○  royxmykxtrkr left-wc
+    ├─╯
+    ○  qpvuntsmwlqt base
+    ◆  zzzzzzzzzzzz
+    "###);
+}
+
 fn get_log_output(test_env: &TestEnvironment, cwd: &Path) -> String {
     let template = r#"separate(" ", change_id.short(), local_bookmarks, if(conflict, "conflict"), description)"#;
     test_env.jj_cmd_success(cwd, &["log", "-T", template])