Combine kt_forbidden_deps and kt_deps_jdeps into a single aspect, kt_traverse_exports

Both of the original aspects have very similar behaviour, and the same basic logic will be used for plugin and friends propagation in followup CLs.

PiperOrigin-RevId: 463184696
Change-Id: Ic13b20d95e69d3f0b5c634e8694e561c8b3ead12

8 files changed, 181 insertions, 149 deletions

diff --git a/kotlin/direct_jdeps.bzl b/kotlin/direct_jdeps.bzl
new file mode 100644
index 0000000..8111ff3
--- /dev/null
+++ b/kotlin/direct_jdeps.bzl
@@ -0,0 +1,26 @@
+# Copyright 2022 Google LLC. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the License);
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""kt_traverse_exports visitor for exposing jdeps files from direct deps."""
+
+def _get_jdeps(target, _ctx):
+    return [out.compile_jdeps for out in target[JavaInfo].java_outputs if out.compile_jdeps]
+
+kt_direct_jdeps_visitor = struct(
+    name = "direct_jdeps",
+    visit_target = _get_jdeps,
+    filter_export = None,
+    process_unvisited_target = None,
+    finish_expansion = None,
+)
diff --git a/kotlin/forbidden_deps.bzl b/kotlin/forbidden_deps.bzl
index 1239a4a..74eb0ea 100644
--- a/kotlin/forbidden_deps.bzl
+++ b/kotlin/forbidden_deps.bzl
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Identifies and reports forbidden deps of Kotlin rules.
+"""kt_traverse_exports visitor for identifying forbidden deps of Kotlin rules.
 
 Currently this system recognizes:
   - nano protos
@@ -23,69 +23,43 @@ Currently this system recognizes:
 load("@bazel_skylib//lib:sets.bzl", "sets")
 load("@//bazel:stubs.bzl", "EXEMPT_DEPS", "FORBIDDEN_DEP_PACKAGES")
 
-_KtForbiddenDepInfo = provider(
-    doc = "Info about forbiddenness as a Kotlin dep.",
-    fields = dict(
-        cause = "Optional[string|dict]: Why is this target forbidden as a dep. None means not forbidden.",
-    ),
-)
+def _error(target, msg):
+    return (str(target.label), msg)
+
+def _is_exempt(target):
+    return sets.contains(EXEMPT_DEPS, str(target.label))
 
-def _aspect_impl(_unused_target, ctx):
-    if _is_exempt(str(ctx.label)):
-        return _KtForbiddenDepInfo(cause = None)
+def _check_forbidden(target, ctx):
+    if _is_exempt(target):
+        return []
 
     if sets.contains(FORBIDDEN_DEP_PACKAGES, ctx.label.package):
-        return _KtForbiddenDepInfo(cause = "Forbidden package")
+        return [_error(target, "Forbidden package")]
 
     # Identify nano protos using tag (b/122083175)
     for tag in ctx.rule.attr.tags:
         if "nano_proto_library" == tag:
-            return _KtForbiddenDepInfo(cause = "nano_proto_library")
-
-    # Check exports if the visited rule isn't itself a problem
-    export_errors = _merge_errors(getattr(ctx.rule.attr, "exports", []))
-    if len(export_errors) > 0:
-        return _KtForbiddenDepInfo(cause = export_errors)
-
-    return _KtForbiddenDepInfo(cause = None)
-
-_aspect = aspect(
-    implementation = _aspect_impl,
-    # Transitively check exports, since they are effectively directly depended on
-    attr_aspects = ["exports"],
-)
+            return [_error(target, "nano_proto_library")]
 
-def _validate_deps(deps):
-    errors = _merge_errors(deps)
-    if len(errors) > 0:
-        fail("Forbidden deps, see go/kotlin/build-rules#restrictions:\n" + "\n".join(
-            [
-                "  " + k + " : " + v
-                for k, v in errors.items()
-            ],
-        ))
+    return []
 
-def _merge_errors(targets):
-    errors = {}
-    for target in targets:
-        label_str = str(target.label)
-        if _KtForbiddenDepInfo in target:
-            cause = target[_KtForbiddenDepInfo].cause
-            if not cause:
-                continue
-            elif type(cause) == "dict":
-                errors.update(cause)
-            else:
-                errors[label_str] = cause
-        elif not _is_exempt(label_str):
-            errors[label_str] = "It was not checked for forbiddenness"
+def _if_not_checked(target):
+    return [] if _is_exempt(target) else [_error(target, "Not checked")]
 
-    return errors
+def _validate_deps(error_set):
+    if not error_set:
+        return
 
-def _is_exempt(label_str):
-    return sets.contains(EXEMPT_DEPS, label_str)
+    error_lines = [
+        "  " + name + " : " + msg
+        for (name, msg) in error_set.to_list()
+    ]
+    fail("Forbidden deps, see go/kotlin/build-rules#restrictions:\n" + "\n".join(error_lines))
 
-kt_forbidden_deps = struct(
-    aspect = _aspect,
-    validate_deps = _validate_deps,
+kt_forbidden_deps_visitor = struct(
+    name = "forbidden_deps",
+    visit_target = _check_forbidden,
+    filter_export = None,
+    process_unvisited_target = _if_not_checked,
+    finish_expansion = _validate_deps,
 )
diff --git a/kotlin/jvm_compile.bzl b/kotlin/jvm_compile.bzl
index ad276da..ae7003a 100644
--- a/kotlin/jvm_compile.bzl
+++ b/kotlin/jvm_compile.bzl
@@ -15,8 +15,7 @@
 """Compile method that can compile kotlin or java sources"""
 
 load(":common.bzl", "common")
-load(":forbidden_deps.bzl", "kt_forbidden_deps")
-load(":kt_jvm_deps.bzl", "kt_jvm_dep_jdeps")
+load(":traverse_exports.bzl", "kt_traverse_exports")
 
 _RULE_FAMILY = common.RULE_FAMILY
 _PARCELIZE_V2_RUNTIME = "@kotlinc//:parcelize_runtime"
@@ -120,7 +119,7 @@ def kt_jvm_compile(
         if _is_eligible_friend(ctx, dep)
     ])
 
-    kt_forbidden_deps.validate_deps(deps + runtime_deps + exports)
+    kt_traverse_exports.expand_forbidden_deps(deps + runtime_deps + exports)
 
     java_plugin_infos = [plugin[JavaPluginInfo] for plugin in plugins]
 
@@ -203,7 +202,7 @@ def kt_jvm_compile(
         java_toolchain = java_toolchain,
         javacopts = javacopts,
         kotlincopts = kotlincopts,
-        compile_jdeps = kt_jvm_dep_jdeps.collect_compile_jdeps_depset(deps),
+        compile_jdeps = kt_traverse_exports.expand_direct_jdeps(deps),
         kt_toolchain = kt_toolchain,
         manifest = manifest,
         merged_manifest = merged_manifest,
diff --git a/kotlin/jvm_import.bzl b/kotlin/jvm_import.bzl
index d59b5d3..500e78a 100644
--- a/kotlin/jvm_import.bzl
+++ b/kotlin/jvm_import.bzl
@@ -15,7 +15,7 @@
 """Kotlin kt_jvm_import rule."""
 
 load(":common.bzl", "common")
-load(":forbidden_deps.bzl", "kt_forbidden_deps")
+load(":traverse_exports.bzl", "kt_traverse_exports")
 load("@//toolchains/kotlin_jvm:kt_jvm_toolchains.bzl", "kt_jvm_toolchains")
 load("@//toolchains/kotlin_jvm:java_toolchains.bzl", "java_toolchains")
 load("@bazel_skylib//lib:dicts.bzl", "dicts")
@@ -84,7 +84,7 @@ _KT_JVM_IMPORT_ATTRS = dicts.add(
     deps = attr.label_list(
         # We allow android rule deps to make importing android JARs easier.
         allow_rules = common.ALLOWED_JVM_RULES + common.ALLOWED_ANDROID_RULES,
-        aspects = [kt_forbidden_deps.aspect],
+        aspects = [kt_traverse_exports.aspect],
         providers = [
             # Each provider-set expands on allow_rules
             [JavaInfo],
@@ -121,7 +121,7 @@ _KT_JVM_IMPORT_ATTRS = dicts.add(
             # Each provider-set expands on allow_rules
             [CcInfo],  # for JNI / native dependencies
         ],
-        aspects = [kt_forbidden_deps.aspect],
+        aspects = [kt_traverse_exports.aspect],
         doc = """Runtime-only dependencies.""",
     ),
     srcjar = attr.label(
diff --git a/kotlin/jvm_library.internal.bzl b/kotlin/jvm_library.internal.bzl
index 0aaa1b8..4abb686 100644
--- a/kotlin/jvm_library.internal.bzl
+++ b/kotlin/jvm_library.internal.bzl
@@ -19,9 +19,8 @@ load("@//toolchains/kotlin_jvm:java_toolchains.bzl", "java_toolchains")
 load("@//toolchains/kotlin_jvm:kt_jvm_toolchains.bzl", "kt_jvm_toolchains")
 load("@bazel_skylib//lib:dicts.bzl", "dicts")
 load(":common.bzl", "common")
-load(":forbidden_deps.bzl", "kt_forbidden_deps")
+load(":traverse_exports.bzl", "kt_traverse_exports")
 load(":jvm_compile.bzl", "compile")
-load(":kt_jvm_deps.bzl", "kt_jvm_dep_jdeps")
 
 # TODO: Use this function in all Kotlin rules
 def _make_default_info(ctx, direct_files, propagated_attrs):
@@ -143,8 +142,7 @@ _KT_JVM_LIBRARY_ATTRS = dicts.add(
             # Each provider-set expands on allow_rules
         ],
         aspects = [
-            kt_forbidden_deps.aspect,
-            kt_jvm_dep_jdeps.aspect,
+            kt_traverse_exports.aspect,
         ],
         doc = """The list of libraries this library directly depends on at compile-time. For Java
                      and Kotlin libraries listed, the Jars they build as well as the transitive closure
@@ -177,7 +175,7 @@ _KT_JVM_LIBRARY_ATTRS = dicts.add(
             # Each provider-set expands on allow_rules
         ],
         aspects = [
-            kt_forbidden_deps.aspect,
+            kt_traverse_exports.aspect,
         ],
         doc = """List of libraries treated as if they were part of this library by upstream
                      Java/Kotlin dependencies, see go/be-java#java_library.exports. These libraries
@@ -212,7 +210,7 @@ _KT_JVM_LIBRARY_ATTRS = dicts.add(
             [CcInfo],  # for JNI / native dependencies
         ],
         aspects = [
-            kt_forbidden_deps.aspect,
+            kt_traverse_exports.aspect,
         ],
         doc = """Runtime-only dependencies.""",
     ),
diff --git a/kotlin/kt_jvm_deps.bzl b/kotlin/kt_jvm_deps.bzl
deleted file mode 100644
index 1f95273..0000000
--- a/kotlin/kt_jvm_deps.bzl
+++ /dev/null
@@ -1,79 +0,0 @@
-# Copyright 2022 Google LLC. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the License);
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Aspect for exposing jdeps files from direct deps in a consistent way.
-
-jdeps from exports are considered direct, and expost transitively. This includes
-rules with deps-as-exports behaviours.
-"""
-
-_KtJdepsInfo = provider(
-    doc = "Captures compile_jdeps files for a target including its exported deps.",
-    fields = dict(
-        compile_jdeps = "Depset of compile_jdeps files.",
-    ),
-)
-
-# java_xxx_proto_library don't populate java_outputs but we can get them through
-# required_aspect_providers from their proto_library deps.
-_DEPS_AS_EXPORTS_RULES = [
-    "java_proto_library",
-    "java_lite_proto_library",
-    "java_mutable_proto_library",
-]
-
-_NO_SRCS_DEPS_AS_EXPORTS_RULES = [
-    "android_library",
-    "proto_library",
-]
-
-def _aspect_impl(target, ctx):
-    transitive_compile_jdeps = _collect_compile_jdeps(getattr(ctx.rule.attr, "exports", []))
-
-    # It's ok if these lists are incomplete. Kotlin compilation will retry with all transitive
-    # deps if incomplete jdeps are collected here.
-    if ctx.rule.kind in _DEPS_AS_EXPORTS_RULES:
-        transitive_compile_jdeps.extend(_collect_compile_jdeps(ctx.rule.attr.deps))
-    elif ctx.rule.kind in _NO_SRCS_DEPS_AS_EXPORTS_RULES and not ctx.rule.attr.srcs:
-        transitive_compile_jdeps.extend(_collect_compile_jdeps(ctx.rule.attr.deps))
-
-    if JavaInfo in target:
-        compile_jdeps = depset(
-            direct = [out.compile_jdeps for out in target[JavaInfo].java_outputs if out.compile_jdeps],
-            transitive = transitive_compile_jdeps,
-        )
-        return _KtJdepsInfo(compile_jdeps = compile_jdeps)
-    elif transitive_compile_jdeps:
-        return _KtJdepsInfo(compile_jdeps = depset(transitive = transitive_compile_jdeps))
-    else:
-        return []  # skip if empty in non-Java-like targets for efficiency
-
-def _collect_compile_jdeps(targets):
-    return [target[_KtJdepsInfo].compile_jdeps for target in targets if (_KtJdepsInfo in target)]
-
-def _collect_compile_jdeps_depset(targets):
-    return depset(transitive = _collect_compile_jdeps(targets))
-
-_aspect = aspect(
-    implementation = _aspect_impl,
-    # Transitively check exports, since they are effectively directly depended on.
-    # "deps" needed for rules that treat deps as exports (usually absent srcs).
-    attr_aspects = ["exports", "deps"],
-    required_aspect_providers = [JavaInfo],  # to get at JavaXxxProtoAspects' JavaInfos
-)
-
-kt_jvm_dep_jdeps = struct(
-    aspect = _aspect,
-    collect_compile_jdeps_depset = _collect_compile_jdeps_depset,
-)
diff --git a/kotlin/traverse_exports.bzl b/kotlin/traverse_exports.bzl
new file mode 100644
index 0000000..982883e
--- /dev/null
+++ b/kotlin/traverse_exports.bzl
@@ -0,0 +1,114 @@
+# Copyright 2022 Google LLC. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the License);
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Combined aspect for all rules_kotlin behaviours that need to traverse exports."""
+
+load(":forbidden_deps.bzl", "kt_forbidden_deps_visitor")
+load(":direct_jdeps.bzl", "kt_direct_jdeps_visitor")
+
+# java_xxx_proto_library don't populate java_outputs but we can get them through
+# required_aspect_providers from their proto_library deps.
+_DEPS_AS_EXPORTS_RULES = [
+    "java_proto_library",
+    "java_lite_proto_library",
+    "java_mutable_proto_library",
+]
+
+_NO_SRCS_DEPS_AS_EXPORTS_RULES = [
+    "android_library",
+    "proto_library",
+]
+
+# visitor = struct[T](
+#     name = string,
+#     visit_target = function(Target, ctx): list[T],
+#     filter_export = None|(function(Target): bool),
+#     process_unvisited_target = None|(function(Target): list[T]),
+#     finish_expansion = None|(function(depset[T]): depset[T]),
+# )
+_VISITORS = [
+    kt_forbidden_deps_visitor,
+    kt_direct_jdeps_visitor,
+]
+
+_KtTraverseExportsInfo = provider(
+    doc = "depsets for transitive info about exports",
+    fields = {
+        v.name: ("depset[%s]" % v.name)
+        for v in _VISITORS
+    },
+)
+
+_EMPTY_KT_TRAVERSE_EXPORTS_INFO = _KtTraverseExportsInfo(**{
+    v.name: depset()
+    for v in _VISITORS
+})
+
+def _aspect_impl(target, ctx):
+    if not (JavaInfo in target):
+        # Ignore non-JVM targets. This also chops-up the
+        # traversal domain at these targets.
+        # TODO: Supoprt non-JVM targets for KMP
+        return _EMPTY_KT_TRAVERSE_EXPORTS_INFO
+
+    exports = []
+    exports.extend(getattr(ctx.rule.attr, "exports", []))  # exports list is frozen
+    if ctx.rule.kind in _DEPS_AS_EXPORTS_RULES:
+        exports.extend(ctx.rule.attr.deps)
+    elif ctx.rule.kind in _NO_SRCS_DEPS_AS_EXPORTS_RULES and not ctx.rule.attr.srcs:
+        exports.extend(ctx.rule.attr.deps)
+
+    return _KtTraverseExportsInfo(**{
+        v.name: depset(
+            direct = v.visit_target(target, ctx),
+            transitive = [
+                getattr(t[_KtTraverseExportsInfo], v.name)
+                for t in exports
+                if (not v.filter_export or v.filter_export(t))
+            ],
+        )
+        for v in _VISITORS
+    })
+
+_aspect = aspect(
+    implementation = _aspect_impl,
+    provides = [_KtTraverseExportsInfo],
+    # Transitively check exports, since they are effectively directly depended on.
+    # "deps" needed for rules that treat deps as exports (usually absent srcs).
+    attr_aspects = ["exports", "deps"],
+    required_aspect_providers = [JavaInfo],  # to get at JavaXxxProtoAspects' JavaInfos
+)
+
+def _create_visitor_expand(visitor):
+    def _visitor_expand(targets):
+        direct = []
+        transitive = []
+        for t in targets:
+            if _KtTraverseExportsInfo in t:
+                transitive.append(getattr(t[_KtTraverseExportsInfo], visitor.name))
+            elif visitor.process_unvisited_target:
+                direct.extend(visitor.process_unvisited_target(t))
+
+        expanded_set = depset(direct = direct, transitive = transitive)
+        return visitor.finish_expansion(expanded_set) if visitor.finish_expansion else expanded_set
+
+    return _visitor_expand
+
+kt_traverse_exports = struct(
+    aspect = _aspect,
+    **{
+        "expand_" + v.name: _create_visitor_expand(v)
+        for v in _VISITORS
+    }
+)
diff --git a/tests/analysis/jvm_compile_test.bzl b/tests/analysis/jvm_compile_test.bzl
index 899266d..b04a479 100644
--- a/tests/analysis/jvm_compile_test.bzl
+++ b/tests/analysis/jvm_compile_test.bzl
@@ -14,7 +14,7 @@
 
 """Kotlin kt_jvm_compile API test."""
 
-load("@//kotlin:forbidden_deps.bzl", "kt_forbidden_deps")
+load("@//kotlin:traverse_exports.bzl", "kt_traverse_exports")
 load("@//kotlin:jvm_compile.bzl", "compile")
 load("@//tests/analysis:util.bzl", "ONLY_FOR_ANALYSIS_TEST_TAGS", "create_dir", "create_file")
 load("@//toolchains/kotlin_jvm:java_toolchains.bzl", "java_toolchains")
@@ -57,11 +57,11 @@ _kt_jvm_compile = rule(
             allow_files = True,
         ),
         deps = attr.label_list(
-            aspects = [kt_forbidden_deps.aspect],
+            aspects = [kt_traverse_exports.aspect],
             providers = [JavaInfo],
         ),
         exports = attr.label_list(
-            aspects = [kt_forbidden_deps.aspect],
+            aspects = [kt_traverse_exports.aspect],
             providers = [JavaInfo],
         ),
         r_java = attr.label(