Handle `third_party/java/` nested source roots in `_java_segments`

When a path starts with `third_party/java/`, try skipping past `third_party/java/` and looking for a nested source root (`java`, `javatests`, `src`, `testsrc`) before falling back to `third_party/java/` as the root.

PiperOrigin-RevId: 999987683
Change-Id: I4f251857af862cafd90a4641a12c015266ec335d
diff --git a/java/common/rules/java_helper.bzl b/java/common/rules/java_helper.bzl
index ee274ad..412292c 100644
--- a/java/common/rules/java_helper.bzl
+++ b/java/common/rules/java_helper.bzl
@@ -22,6 +22,16 @@
 def _java_segments(path):
     if path.startswith("/"):
         fail("path must not be absolute: '%s'" % path)
+    if path.startswith("third_party/java/"):
+        # For third_party/java, try skipping past the first two segments and looking for a nested
+        # source root, to relativize e.g. third_party/java/foo/java/resource into 'resource',
+        # not 'foo/java/resource'.
+        segments = _java_segments_impl(path[len("third_party/java/"):])
+        if segments != None:
+            return segments
+    return _java_segments_impl(path)
+
+def _java_segments_impl(path):
     segments = path.split("/")
     root_idx = -1
     for idx, segment in enumerate(segments):
diff --git a/test/java/common/rules/BUILD b/test/java/common/rules/BUILD
index 44bb6c3..b115dbb 100644
--- a/test/java/common/rules/BUILD
+++ b/test/java/common/rules/BUILD
@@ -5,6 +5,7 @@
 load(":common_launcher_java_library_tests.bzl", "java_library_launcher_tests")
 load(":deploy_archive_builder_tests.bzl", "deploy_archive_builder_test_suite")
 load(":java_binary_tests.bzl", "java_binary_tests")
+load(":java_helper_tests.bzl", "java_helper_test_suite")
 load(":java_import_tests.bzl", "java_import_tests")
 load(":java_launcher_tests.bzl", "java_launcher_tests")
 load(":java_library_tests.bzl", "java_library_tests")
@@ -17,6 +18,8 @@
 
 deploy_archive_builder_test_suite(name = "deploy_archive_builder_tests")
 
+java_helper_test_suite(name = "java_helper_tests")
+
 merge_attrs_test_suite(name = "merge_attrs_tests")
 
 java_binary_tests(name = "java_binary_tests")
diff --git a/test/java/common/rules/java_helper_tests.bzl b/test/java/common/rules/java_helper_tests.bzl
new file mode 100644
index 0000000..5fdbbb6
--- /dev/null
+++ b/test/java/common/rules/java_helper_tests.bzl
@@ -0,0 +1,76 @@
+"""Tests for java_helper functions"""
+
+load("@bazel_skylib//lib:unittest.bzl", "asserts", "unittest")
+load(
+    "//java/common/rules:java_helper.bzl",
+    "helper",
+)
+
+def _java_segments_test_impl(ctx):
+    env = unittest.begin(ctx)
+
+    asserts.equals(
+        env,
+        ["com", "google", "foo", "Foo"],
+        helper.java_segments("java/com/google/foo/Foo"),
+    )
+    asserts.equals(
+        env,
+        ["com", "google", "foo", "FooTest"],
+        helper.java_segments("javatests/com/google/foo/FooTest"),
+    )
+    asserts.equals(
+        env,
+        ["com", "google", "foo", "Foo"],
+        helper.java_segments("src/main/java/com/google/foo/Foo"),
+    )
+    asserts.equals(
+        env,
+        ["com", "google", "foo", "FooTest"],
+        helper.java_segments("src/test/java/com/google/foo/FooTest"),
+    )
+    asserts.equals(
+        env,
+        ["com", "google", "foo", "Foo"],
+        helper.java_segments("third_party/java/foo/java/com/google/foo/Foo"),
+    )
+    asserts.equals(
+        env,
+        ["com", "google", "foo", "FooTest"],
+        helper.java_segments("third_party/java/foo/javatests/com/google/foo/FooTest"),
+    )
+    asserts.equals(
+        env,
+        ["com", "google", "foo", "Foo"],
+        helper.java_segments("third_party/java/foo/src/main/java/com/google/foo/Foo"),
+    )
+    asserts.equals(
+        env,
+        ["com", "google", "foo", "FooTest"],
+        helper.java_segments("third_party/java/foo/src/test/java/com/google/foo/FooTest"),
+    )
+    asserts.equals(
+        env,
+        ["foo", "org", "foo", "Foo"],
+        helper.java_segments("third_party/java/foo/org/foo/Foo"),
+    )
+    asserts.equals(
+        env,
+        None,
+        helper.java_segments("foo/bar/Baz"),
+    )
+
+    return unittest.end(env)
+
+java_segments_test = unittest.make(_java_segments_test_impl)
+
+def java_helper_test_suite(name):
+    """Sets up java_helper test suite
+
+    Args:
+        name: the name of the test suite target
+    """
+    unittest.suite(
+        name,
+        java_segments_test,
+    )