metal.bzl

""" Rules for organizing and compiling Metal. """

load("@bazel_skylib//lib:dicts.bzl", "dicts")
load("@bazel_skylib//lib:paths.bzl", "paths")
load("@build_bazel_apple_support//lib:apple_support.bzl", "apple_support")
load("@build_bazel_rules_swift//swift:swift.bzl", "swift_library")
load("@rules_xcodeproj//xcodeproj/internal:providers.bzl", "XcodeProjExtraFilesHintInfo")

MetalFilesInfo = provider(
    "Collects Metal AIR files and headers",
    fields = ["transitive_airs", "transitive_headers", "transitive_header_paths"],
)

def get_transitive_airs(airs, deps):
    """Obtain the source files for a target and its transitive dependencies.

    Args:
      srcs: a list of source files
      deps: a list of targets that are direct dependencies
    Returns:
      a collection of the transitive sources
    """
    return depset(
        airs,
        transitive = [dep[MetalFilesInfo].transitive_airs for dep in deps],
    )

def get_transitive_hdrs(hdrs, deps):
    """Obtain the header files for a target and its transitive dependencies.

    Args:
      hdrs: a list of header files
      deps: a list of targets that are direct dependencies
    Returns:
      a collection of the transitive headers
    """
    return depset(
        hdrs,
        transitive = [dep[MetalFilesInfo].transitive_headers for dep in deps],
    )

def get_transitive_hdr_paths(paths, deps):
    """Obtain the header include paths for a target and its transitive dependencies.

    Args:
      paths: a list of header file include paths
      deps: a list of targets that are direct dependencies
    Returns:
      a collection of the transitive header include paths
    """
    return depset(
        paths,
        transitive = [dep[MetalFilesInfo].transitive_header_paths for dep in deps],
    )

def _process_hdrs(ctx, hdrs, include_prefix, strip_include_prefix):
    if strip_include_prefix == "" and include_prefix == "":
        return paths.dirname(ctx.build_file_path), hdrs
    if paths.is_absolute(strip_include_prefix):
        fail("should be relative", attr = "strip_include_prefix")
    if paths.is_absolute(include_prefix):
        fail("should be relative", attr = "include_prefix")
    virtual_hdr_prefix = paths.join("_virtual_hdrs", ctx.label.name)
    virtual_hdr_path = paths.join(
        ctx.genfiles_dir.path,
        paths.dirname(ctx.build_file_path),
        virtual_hdr_prefix,
    )
    virtual_hdr_prefix = paths.join(virtual_hdr_prefix, include_prefix)

    full_strip_prefix = paths.normalize(paths.join(
        ctx.label.package,
        ctx.attr.strip_include_prefix,
    )) + "/"

    virtual_hdrs = []
    for hdr in hdrs:
        if full_strip_prefix != "" and not hdr.short_path.startswith(full_strip_prefix):
            fail("hdr '%s' is not under the specified strip prefix '%s'" %
                 (hdr, strip_include_prefix))
        virtual_hdr_name = paths.join(
            virtual_hdr_prefix,
            hdr.short_path.removeprefix(full_strip_prefix),
        )
        virtual_hdr = ctx.actions.declare_file(virtual_hdr_name)
        ctx.actions.symlink(
            output = virtual_hdr,
            target_file = hdr,
            progress_message = "Symlinking virtual hdr sources for %{label}",
        )
        virtual_hdrs.append(virtual_hdr)

    return virtual_hdr_path, virtual_hdrs

def _is_debug(ctx):
    return ctx.var["COMPILATION_MODE"] == "dbg"

def _get_os_version_flag(ctx):
    if ctx.fragments.apple.apple_platform_type == "ios":
        if ctx.fragments.apple.ios_minimum_os_flag:
            return "-mios-version-min={}".format(ctx.fragments.apple.ios_minimum_os_flag)
        print(ctx.fragments.apple.ios_minimum_os_flag)
    elif ctx.fragments.apple.apple_platform_type == "macos":
        if ctx.fragments.apple.macos_minimum_os_flag:
            return "-mmacos-version-min={}".format(ctx.fragments.apple.macos_minimum_os_flag)
    return None

def _compile_metals(metals, hdrs, hdr_paths, ctx):
    _get_os_version_flag(ctx)
    air_files = []

    version_flag = _get_os_version_flag(ctx)

    for src_metal in metals:
        air_file = ctx.actions.declare_file(paths.replace_extension(src_metal.basename, ".air"))
        air_files.append(air_file)
        input_files = [src_metal] + [src_hdr for src_hdr in hdrs]

        args = ctx.actions.args()
        args.add("metal")
        args.add("-c")
        if version_flag:
            args.add(version_flag)

        if _is_debug(ctx):
            args.add("-frecord-sources")
            args.add("-gline-tables-only")

        args.add("-o", air_file)
        for path in hdr_paths.to_list():
            args.add("-I", path)
        args.add(src_metal.path)

        args.add_all(ctx.attr.copts)

        apple_support.run(
            actions = ctx.actions,
            xcode_config = ctx.attr._xcode_config[apple_common.XcodeVersionConfig],
            apple_fragment = ctx.fragments.apple,
            inputs = input_files,
            outputs = [air_file],
            executable = "/usr/bin/xcrun",
            arguments = [args],
            mnemonic = "MetalCompile",
        )
    return air_files

def _metal_library_impl(ctx):
    hdr_path, hdrs = _process_hdrs(
        ctx,
        ctx.files.hdrs,
        ctx.attr.include_prefix,
        ctx.attr.strip_include_prefix,
    )
    trans_hdrs = get_transitive_hdrs(hdrs, ctx.attr.deps)
    trans_hdr_paths = get_transitive_hdr_paths([hdr_path], ctx.attr.deps)

    srcs_list = ctx.files.srcs + trans_hdrs.to_list()

    srcs_metal_list = [x for x in srcs_list if x.extension == "metal"]

    srcs_hdrs_list = [x for x in srcs_list if x.extension == "h" or x.extension == "hpp"]

    air_files = _compile_metals(srcs_metal_list, srcs_hdrs_list, trans_hdr_paths, ctx)

    trans_airs = get_transitive_airs(air_files, ctx.attr.deps)

    return [
        MetalFilesInfo(
            transitive_airs = trans_airs,
            transitive_headers = trans_hdrs,
            transitive_header_paths = trans_hdr_paths,
        ),
    ]

_metal_library_rule = rule(
    implementation = _metal_library_impl,
    fragments = ["apple"],
    attrs = dicts.add(
        apple_support.action_required_attrs(),
        {
            "srcs": attr.label_list(allow_files = [".metal", ".h", ".hpp"]),
            "hdrs": attr.label_list(allow_files = [".h", ".hpp"]),
            "deps": attr.label_list(),
            "copts": attr.string_list(),
            "include_prefix": attr.string(),
            "strip_include_prefix": attr.string(),
        },
    ),
)

def metal_library(
        name,
        srcs = [],
        hdrs = [],
        deps = [],
        copts = [],
        include_prefix = "",
        strip_include_prefix = "",
        visibility = None,
        **kwargs):
    """Creates a metal_library target that compiles Metal shading language code into reusable libraries.

    Args:
        name: Name of the library
        srcs: Metal source files (.metal)
        hdrs: Metal header files (.h, .hpp)
        deps: Other metal_library targets this library depends on
        copts: Compiler options passed to metal compiler
        include_prefix: Optional prefix to add to include paths
        strip_include_prefix: Optional prefix to strip from include paths
        visibility: Target visibility
        **kwargs: Additional arguments passed to the underlying rule
    """
    _metal_library_rule(
        name = name,
        srcs = srcs,
        hdrs = hdrs,
        deps = deps,
        copts = copts,
        include_prefix = include_prefix,
        strip_include_prefix = strip_include_prefix,
        visibility = visibility,
        **kwargs
    )

def _metal_binary_impl(ctx):
    metallib_file = ctx.outputs.out
    if metallib_file == None:
        metallib_file = ctx.actions.declare_file(ctx.label.name + ".metallib")

    trans_hdrs = get_transitive_hdrs([], ctx.attr.deps)
    trans_hdr_paths = get_transitive_hdr_paths([], ctx.attr.deps)
    srcs_list = ctx.files.srcs + trans_hdrs.to_list()

    srcs_metal_list = [x for x in srcs_list if x.extension == "metal"]

    srcs_hdrs_list = [x for x in srcs_list if x.extension == "h" or x.extension == "hpp"]

    air_files = _compile_metals(srcs_metal_list, srcs_hdrs_list, trans_hdr_paths, ctx)

    trans_airs = get_transitive_airs(air_files, ctx.attr.deps)

    args = ctx.actions.args()
    args.add("metal")
    if _is_debug(ctx):
        args.add("-frecord-sources")
        args.add("-gline-tables-only")

    version_flag = _get_os_version_flag(ctx)
    if version_flag:
        args.add(version_flag)

    args.add("-o", metallib_file)
    args.add_all(trans_airs)

    apple_support.run(
        actions = ctx.actions,
        xcode_config = ctx.attr._xcode_config[apple_common.XcodeVersionConfig],
        apple_fragment = ctx.fragments.apple,
        inputs = trans_airs,
        outputs = [metallib_file],
        executable = "/usr/bin/xcrun",
        arguments = [args],
        mnemonic = "MetallibLink",
    )

    return [DefaultInfo(files = depset([metallib_file]))]

_metal_binary_rule = rule(
    implementation = _metal_binary_impl,
    fragments = ["apple"],
    attrs = dicts.add(
        apple_support.action_required_attrs(),
        {
            "srcs": attr.label_list(allow_files = [".metal", ".h", ".hpp"]),
            "deps": attr.label_list(),
            "copts": attr.string_list(),
            "out": attr.output(),
        },
    ),
)

def metal_binary(
        name,
        srcs = [],
        deps = [],
        copts = [],
        out = None,
        visibility = None,
        **kwargs):
    """Creates a metal_binary target that compiles and links Metal code into a .metallib file.

    Args:
        name: Name of the binary
        srcs: Metal source files (.metal) and headers (.h, .hpp)
        deps: metal_library targets this binary depends on
        copts: Compiler options passed to metal compiler
        out: Output .metallib filename (defaults to name.metallib)
        visibility: Target visibility
        **kwargs: Additional arguments passed to the underlying rule
    """
    _metal_binary_rule(
        name = name,
        srcs = srcs,
        deps = deps,
        copts = copts,
        out = out,
        visibility = visibility,
        **kwargs
    )

MetalDebugInfo = provider(
    doc = "Provider for collecting Metal binary targets",
    fields = {
        "targets": "depset of target label strings",
    },
)

def _metal_debug_aspect_impl(target, ctx):
    """Aspect that collects all metal_binary targets."""
    metal_binaries = []

    # Check if this target itself is a metal_binary (produces a .metallib file)
    if hasattr(target, "files"):
        for f in target.files.to_list():
            if f.extension == "metallib":
                metal_binaries.append(str(target.label))
                break

    # Collect metal binaries from dependencies (aspect handles recursion)
    transitive_metal_binaries = []
    if hasattr(ctx.rule.attr, "deps"):
        for dep in ctx.rule.attr.deps:
            if MetalDebugInfo in dep:
                transitive_metal_binaries.append(dep[MetalDebugInfo].targets)

    if hasattr(ctx.rule.attr, "data"):
        for dep in ctx.rule.attr.data:
            if MetalDebugInfo in dep:
                transitive_metal_binaries.append(dep[MetalDebugInfo].targets)

    if hasattr(ctx.rule.attr, "implementation_deps"):
        for dep in ctx.rule.attr.implementation_deps:
            if MetalDebugInfo in dep:
                transitive_metal_binaries.append(dep[MetalDebugInfo].targets)

    # Return a depset combining our findings with transitive ones
    return [
        MetalDebugInfo(
            targets = depset(
                direct = metal_binaries,
                transitive = transitive_metal_binaries,
            ),
        ),
    ]

metal_debug_aspect = aspect(
    implementation = _metal_debug_aspect_impl,
    attr_aspects = ["deps", "data", "implementation_deps"],
)

def _xcode_metal_debug_helper_impl(ctx):
    """Generates a script to create symlinks for Xcode GPU shader debugging."""

    # Get all metallib targets from dependencies via the aspect
    metallib_targets = []
    for dep in ctx.attr.deps:
        if MetalDebugInfo in dep:
            metallib_targets.extend(dep[MetalDebugInfo].targets.to_list())

    # Generate the script
    script = ctx.actions.declare_file(ctx.label.name + ".sh")

    script_content = """#!/bin/bash
# Generated script to create symlinks for Xcode GPU shader debugging
#
# Usage:
#   {script_name}                      # Interactive mode (prompts for temp dir)
#   {script_name} <xcode_temp_directory>  # Direct mode (uses provided path)
#
# When Xcode captures a GPU frame and extracts Metal sources, run this script
# with the temp directory path to create necessary symlinks for includes to resolve.
#
# Example:
#   {script_name} /var/folders/.../T/780970914.71
#
# Metallibs used by this app:
{metallib_list}

set -e

# Check if temp dir was provided as argument
if [ -n "$1" ]; then
    TEMP_DIR="$1"
else
    # Interactive mode - prompt for temp dir
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "  Xcode Metal GPU Shader Debugging Setup"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""
    echo "Steps:"
    echo "  1. Run your app from Xcode"
    echo "  2. Capture GPU frame: Debug → Capture GPU Frame (⌘⌥G)"
    echo "  3. Click on a shader to view it"
    echo "  4. Copy the temp directory from any compile error"
    echo ""
    echo "The temp directory path looks like:"
    echo "  /var/folders/.../T/780970914.71"
    echo "            (copy up to here ^)"
    echo ""
    echo -n "Enter the temp directory path: "
    read TEMP_DIR
fi

if [ -z "$TEMP_DIR" ]; then
    echo "Error: No path provided"
    exit 1
fi

if [ ! -d "$TEMP_DIR" ]; then
    echo "Error: Directory does not exist: $TEMP_DIR"
    echo ""
    echo "Make sure you:"
    echo "  - Captured a GPU frame first"
    echo "  - Copied the full path up to /T/XXXXXX.XX"
    exit 1
fi

echo "Creating symlinks for Xcode GPU shader debugging in: $TEMP_DIR"

# Function to create symlinks in a directory
# Prefers include directories in the same numbered subdirectory
create_symlinks_in_dir() {{
    local target_dir="$1"

    # Extract the numbered subdirectory from target_dir (e.g., /path/src/HASH/0/... -> 0)
    local target_subdir=$(echo "$target_dir" | sed -n 's#.*/src/[^/]*/\\([0-9]*\\)/.*#\\1#p')

    # Find all include directories and create symlinks
    # Process them so that include dirs in the same numbered subdir come first
    find "$TEMP_DIR" -type d -name "include" | sort | while read -r include_path; do
        lib_base=$(dirname "$include_path")
        lib_name=$(basename "$lib_base")

        # Extract numbered subdirectory from this include path
        local include_subdir=$(echo "$include_path" | sed -n 's#.*/src/[^/]*/\\([0-9]*\\)/.*#\\1#p')

        # Determine virtual prefix (e.g., "applelib" -> "apple", "drawlib" -> "draw")
        case "$lib_name" in
            *lib) virtual_name="${{lib_name%lib}}" ;;
            *) virtual_name="$lib_name" ;;
        esac

        # If symlink doesn't exist, create it
        # If it exists but points to different subdir and we have matching subdir, replace it
        if [ ! -e "$target_dir/$virtual_name" ]; then
            ln -sf "$include_path" "$target_dir/$virtual_name"
            echo "  Created symlink: $target_dir/$virtual_name -> $include_path"
        elif [ "$target_subdir" = "$include_subdir" ]; then
            # This include is in the same numbered subdirectory, prefer it
            current_target=$(readlink "$target_dir/$virtual_name" 2>/dev/null || echo "")
            current_subdir=$(echo "$current_target" | sed -n 's#.*/src/[^/]*/\\([0-9]*\\)/.*#\\1#p')
            if [ "$current_subdir" != "$target_subdir" ]; then
                rm -f "$target_dir/$virtual_name"
                ln -sf "$include_path" "$target_dir/$virtual_name"
                echo "  Updated symlink: $target_dir/$virtual_name -> $include_path (same subdirectory)"
            fi
        fi
    done
}}

# Find all directories containing .metal files and create symlinks
find "$TEMP_DIR" -type f -name "*.metal" | while read -r metalfile; do
    msldir=$(dirname "$metalfile")

    # Skip if we've already processed this directory
    if [ -f "$msldir/.symlinks_created" ]; then
        continue
    fi

    echo "Processing: $msldir"

    # Create symlinks in this directory
    create_symlinks_in_dir "$msldir"

    # Also create symlinks in the shared/ subdirectory if it exists
    if [ -d "$msldir/shared" ]; then
        create_symlinks_in_dir "$msldir/shared"
    fi

    # Mark this directory as processed
    touch "$msldir/.symlinks_created"
done

echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "  Setup complete! Go back to Xcode and view your shader."
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo ""
""".format(
        script_name = script.short_path,
        metallib_list = "\n".join(["#   - " + t for t in metallib_targets]),
    )

    ctx.actions.write(
        output = script,
        content = script_content,
        is_executable = True,
    )

    return [DefaultInfo(
        files = depset([script]),
        executable = script,
    )]

_xcode_metal_debug_helper = rule(
    implementation = _xcode_metal_debug_helper_impl,
    attrs = {
        "deps": attr.label_list(aspects = [metal_debug_aspect]),
    },
    executable = True,
)

def _add_metal_xcode_integration(name, kwargs):
    """Helper function to add automatic Metal source visibility for Xcode projects.

    Scans data and deps for Metal targets, collects their transitive sources,
    and attaches them as aspect_hints for rules_xcodeproj integration.

    Args:
        name: Name of the target
        kwargs: Target's kwargs dict (will be modified to add aspect_hints)

    Returns:
        Modified kwargs dict
    """
    # Collect both data and regular dependencies to scan for Metal binaries
    # Filter out source files from data (only include rule targets)
    targets_to_scan = []
    for item in kwargs.get("data", []):
        # Only include if it looks like a target (starts with : or //) and doesn't look like a file
        # Skip items that look like files (contain a dot in the last component)
        if type(item) == "string":
            item_str = str(item)
            # Get the last component after / or :
            last_component = item_str.split("/")[-1].split(":")[-1]
            # If it doesn't contain a dot (likely not a file), include it
            if "." not in last_component:
                targets_to_scan.append(item)
    targets_to_scan.extend(kwargs.get("deps", []))

    # If there are any dependencies, create a transitive_metal_extra_files target
    # and attach it as an aspect_hint for Xcode project integration
    if targets_to_scan:
        metal_hint_name = name + "_metal_sources"
        transitive_metal_extra_files(
            name = metal_hint_name,
            targets = targets_to_scan,
        )

        # Get existing aspect_hints and add our metal sources hint
        aspect_hints = kwargs.pop("aspect_hints", [])
        kwargs["aspect_hints"] = aspect_hints + [":" + metal_hint_name]

    return kwargs

def cc_binary_with_metal_debug(name, **kwargs):
    """Wrapper for cc_binary that adds Xcode GPU shader debugging support and automatic Metal source visibility.

    This macro automatically collects all transitive Metal sources from data and regular dependencies
    and makes them visible in the Xcode project navigator.

    Args:
        name: Name of the binary
        **kwargs: All other arguments passed to cc_binary

    Usage:
        cc_binary_with_metal_debug(
            name = "my_app",
            srcs = ["main.cpp"],
            data = [":shaders"],  # Metal sources will automatically appear in Xcode
            deps = [...],  # Will also scan deps for Metal sources
        )

        # Then run: bazel run //path/to:my_app_copy_metal_hdrs_to_trace
    """
    kwargs = _add_metal_xcode_integration(name, kwargs)
    native.cc_binary(name = name, **kwargs)
    metal_debug_script(
        name = name + "_copy_metal_hdrs_to_trace",
        deps = kwargs.get("deps", []),
    )

def cc_library_with_metal_debug(name, **kwargs):
    """Wrapper for cc_library that adds Xcode GPU shader debugging support and automatic Metal source visibility.

    This macro automatically collects all transitive Metal sources from data and regular dependencies
    and makes them visible in the Xcode project navigator.

    Args:
        name: Name of the library
        **kwargs: All other arguments passed to cc_library

    Usage:
        cc_library_with_metal_debug(
            name = "my_lib",
            srcs = ["lib.cpp"],
            data = [":shaders"],  # Metal sources will automatically appear in Xcode
            deps = [...],  # Will also scan deps for Metal sources
        )

        # Then run: bazel run //path/to:my_lib_copy_metal_hdrs_to_trace
    """
    kwargs = _add_metal_xcode_integration(name, kwargs)
    native.cc_library(name = name, **kwargs)
    metal_debug_script(
        name = name + "_copy_metal_hdrs_to_trace",
        deps = kwargs.get("deps", []),
    )

def swift_library_with_metal_debug(name, **kwargs):
    """Wrapper for swift_library that adds Xcode GPU shader debugging support and automatic Metal source visibility.

    This macro automatically collects all transitive Metal sources from data and regular dependencies
    and makes them visible in the Xcode project navigator.

    Args:
        name: Name of the library
        **kwargs: All other arguments passed to swift_library

    Usage:
        load("@rules_metal//:metal.bzl", "swift_library_with_metal_debug")

        swift_library_with_metal_debug(
            name = "MyLib",
            srcs = ["MyLib.swift"],
            data = [":shaders"],  # Metal sources will automatically appear in Xcode
            deps = [...],  # Will also scan deps for Metal sources
        )

        # Then run: bazel run //path/to:MyLib_copy_metal_hdrs_to_trace
    """
    kwargs = _add_metal_xcode_integration(name, kwargs)
    swift_library(name = name, **kwargs)
    metal_debug_script(
        name = name + "_copy_metal_hdrs_to_trace",
        deps = kwargs.get("deps", []),
    )

def metal_debug_script(name, deps):
    """Creates Xcode GPU debug helper script(s) for existing targets.

    This is useful for Apple application targets (macos_application, ios_application, etc.)
    that can't use cc_binary_with_metal_debug.

    Args:
        name: Base name for the debug targets (e.g., "MyApp_copy_metal_hdrs_to_trace")
        deps: Dependencies to scan for metal binaries

    Creates:
        - {name}: The debug script (supports both interactive and direct modes)

    Usage:
        macos_application(
            name = "MyApp",
            deps = ["MyAppLib"],
        )

        metal_debug_script(
            name = "MyApp_copy_metal_hdrs_to_trace",
            deps = ["MyAppLib"],
        )

        # Run from command line (auto-rebuilds and prompts for temp dir):
        bazel run //path/to:MyApp_copy_metal_hdrs_to_trace
    """
    _xcode_metal_debug_helper(
        name = name,
        deps = deps,
    )

# Rule for exposing extra files to rules_xcodeproj via aspect_hints
# This is compatible with rules_xcodeproj's xcodeproj_extra_files feature

def _metal_extra_files_impl(ctx):
    """Implementation for metal_extra_files rule."""
    return [XcodeProjExtraFilesHintInfo(files = depset(ctx.files.files))]

metal_extra_files = rule(
    doc = """Surfaces extra Metal source files for inclusion in Xcode project navigator.

    This rule is compatible with rules_xcodeproj's xcodeproj_extra_files feature.
    Attach to metal_library or metal_binary targets via aspect_hints attribute.

    Example:
        metal_extra_files(
            name = "shaders_xcode_hint",
            files = ["shader.metal", "types.h"],
        )

        metal_binary(
            name = "shaders",
            srcs = ["shader.metal", "types.h"],
            aspect_hints = [":shaders_xcode_hint"],
        )
    """,
    implementation = _metal_extra_files_impl,
    attrs = {
        "files": attr.label_list(
            doc = "The list of Metal source files to surface in the Xcode navigator",
            allow_files = [".metal", ".h", ".hpp"],
        ),
    },
    provides = [XcodeProjExtraFilesHintInfo],
)

# Provider for collecting transitive Metal sources
_MetalSourcesInfo = provider(
    doc = "Provider for collecting transitive Metal source files",
    fields = {
        "srcs": "Depset of Metal source files",
        "hdrs": "Depset of Metal header files",
    },
)

def _metal_sources_aspect_impl(target, ctx):
    """Aspect that collects transitive Metal sources from metal_library and metal_binary targets."""
    srcs = []
    hdrs = []
    transitive_srcs = []
    transitive_hdrs = []

    # Collect from this target if it's a metal_library or metal_binary
    if hasattr(ctx.rule.attr, "srcs"):
        for src in ctx.rule.attr.srcs:
            if hasattr(src, "files"):
                for f in src.files.to_list():
                    if f.extension in ["metal"]:
                        srcs.append(f)
                    elif f.extension in ["h", "hpp"]:
                        hdrs.append(f)

    if hasattr(ctx.rule.attr, "hdrs"):
        for hdr in ctx.rule.attr.hdrs:
            if hasattr(hdr, "files"):
                for f in hdr.files.to_list():
                    if f.extension in ["h", "hpp"]:
                        hdrs.append(f)

    # Collect from deps
    if hasattr(ctx.rule.attr, "deps"):
        for dep in ctx.rule.attr.deps:
            if _MetalSourcesInfo in dep:
                transitive_srcs.append(dep[_MetalSourcesInfo].srcs)
                transitive_hdrs.append(dep[_MetalSourcesInfo].hdrs)

    # Collect from data dependencies
    if hasattr(ctx.rule.attr, "data"):
        for dep in ctx.rule.attr.data:
            if _MetalSourcesInfo in dep:
                transitive_srcs.append(dep[_MetalSourcesInfo].srcs)
                transitive_hdrs.append(dep[_MetalSourcesInfo].hdrs)

    # Collect from implementation_deps (used by cc_library targets)
    if hasattr(ctx.rule.attr, "implementation_deps"):
        for dep in ctx.rule.attr.implementation_deps:
            if _MetalSourcesInfo in dep:
                transitive_srcs.append(dep[_MetalSourcesInfo].srcs)
                transitive_hdrs.append(dep[_MetalSourcesInfo].hdrs)

    return [_MetalSourcesInfo(
        srcs = depset(direct = srcs, transitive = transitive_srcs),
        hdrs = depset(direct = hdrs, transitive = transitive_hdrs),
    )]

_metal_sources_aspect = aspect(
    implementation = _metal_sources_aspect_impl,
    attr_aspects = ["deps", "data", "implementation_deps"],
)

def _transitive_metal_extra_files_impl(ctx):
    """Collects all transitive Metal sources from a metal_binary and creates XcodeProjExtraFilesHintInfo."""
    all_files = []

    for target in ctx.attr.targets:
        if _MetalSourcesInfo in target:
            all_files.extend(target[_MetalSourcesInfo].srcs.to_list())
            all_files.extend(target[_MetalSourcesInfo].hdrs.to_list())

    return [XcodeProjExtraFilesHintInfo(files = depset(all_files))]

transitive_metal_extra_files = rule(
    doc = """Collects transitive Metal sources and exposes them for Xcode project navigator.

    Use this to surface all Metal sources from a metal_binary and its dependencies.

    Example:
        transitive_metal_extra_files(
            name = "all_metal_sources",
            targets = [":shaders"],
        )

        swift_library(
            name = "MyLib",
            ...
            aspect_hints = [":all_metal_sources"],
        )
    """,
    implementation = _transitive_metal_extra_files_impl,
    attrs = {
        "targets": attr.label_list(
            doc = "Metal targets to collect sources from (typically metal_binary targets)",
            aspects = [_metal_sources_aspect],
        ),
    },
    provides = [XcodeProjExtraFilesHintInfo],
)