# TensorFlow Lite for Swift

load("@build_bazel_rules_apple//apple:apple_xcframework.bzl", "apple_static_xcframework")
load("@build_bazel_rules_apple//apple:ios.bzl", "ios_static_framework", "ios_unit_test")
load("@build_bazel_rules_swift//swift:swift.bzl", "swift_library")
load("@rules_cc//cc:objc_library.bzl", "objc_library")
load("//tensorflow/lite:special_rules.bzl", "ios_visibility_allowlist", "tflite_ios_lab_runner")
load("//tensorflow/lite/ios:ios.bzl", "TFL_DEFAULT_TAGS", "TFL_DISABLED_SANITIZER_TAGS", "TFL_MINIMUM_OS_VERSION")

package(
    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
    default_visibility = ["//visibility:private"],
    licenses = ["notice"],
)

config_setting(
    name = "use_coreml_delegate",
    define_values = {"use_coreml_delegate": "1"},
)

config_setting(
    name = "use_metal_delegate",
    define_values = {"use_metal_delegate": "1"},
)

# By default this builds with no delegates.
# To build with the Metal delegate pass --define=use_metal_delegate=1
# To build with the CoreML delegate pass --define=use_coreml_delegate=1
swift_library(
    name = "TensorFlowLite",
    srcs = glob(
        [
            "Sources/*.swift",
        ],
        exclude = [
            "Sources/CoreMLDelegate.swift",
            "Sources/MetalDelegate.swift",
        ],
    ) + select({
        ":use_coreml_delegate": [
            "Sources/CoreMLDelegate.swift",
        ],
        "//conditions:default": [],
    }) + select({
        ":use_metal_delegate": [
            "Sources/MetalDelegate.swift",
        ],
        "//conditions:default": [],
    }),
    library_evolution = True,
    linkopts = select({
        ":use_coreml_delegate": [
            "-Wl,-weak_framework,CoreML",
        ],
        "//conditions:default": [],
    }) + select({
        ":use_metal_delegate": [
            "-Wl,-weak_framework,Metal",
        ],
        "//conditions:default": [],
    }),
    module_name = "TensorFlowLite",
    tags = TFL_DEFAULT_TAGS + ["nobuilder"],
    visibility = ios_visibility_allowlist(),
    # Do not sort: these targets sort differently internally vs open source
    deps = ["//tensorflow/lite/ios:tensorflow_lite_c"] + select({
        ":use_coreml_delegate": [
            "//tensorflow/lite/delegates/coreml:coreml_delegate",
        ],
        "//conditions:default": [],
    }) + select({
        ":use_metal_delegate": [
            "//tensorflow/lite/delegates/gpu:metal_delegate",
        ],
        "//conditions:default": [],
    }),
)

swift_library(
    name = "TensorFlowLiteAllDelegates",
    testonly = 1,
    srcs = glob(["Sources/*.swift"]),
    linkopts = [
        "-Wl,-weak_framework,CoreML",
        "-Wl,-weak_framework,Metal",
    ],
    module_name = "TensorFlowLite",
    tags = TFL_DEFAULT_TAGS + ["builder_default_ios_arm64"],
    deps = [
        "//tensorflow/lite/delegates/coreml:coreml_delegate",
        "//tensorflow/lite/delegates/gpu:metal_delegate",
        "//tensorflow/lite/ios:tensorflow_lite_c",
    ],
)

swift_library(
    name = "TensorFlowLiteMetalDelegate",
    testonly = 1,
    srcs = glob(
        [
            "Sources/*.swift",
        ],
        exclude = [
            "Sources/CoreMLDelegate.swift",
        ],
    ),
    linkopts = [
        "-Wl,-weak_framework,Metal",
    ],
    module_name = "TensorFlowLite",
    tags = TFL_DEFAULT_TAGS + ["builder_default_ios_arm64"],
    deps = [
        "//tensorflow/lite/delegates/gpu:metal_delegate",
        "//tensorflow/lite/ios:tensorflow_lite_c",
        "//third_party/apple_frameworks:Darwin",
        "//third_party/apple_frameworks:Foundation",
    ],
)

# bazel build -c opt --config=ios_fat //tensorflow/lite/swift:TensorFlowLite_framework
ios_static_framework(
    name = "TensorFlowLite_framework",
    avoid_deps = [
        "//tensorflow/lite/ios:tensorflow_lite_c",
    ],
    bundle_name = "TensorFlowLite",
    minimum_os_version = TFL_MINIMUM_OS_VERSION,
    deps = [
        ":TensorFlowLite",
    ],
)

# bazel build -c opt --config=ios //tensorflow/lite/swift:TensorFlowLite_xcframework
apple_static_xcframework(
    name = "TensorFlowLite_xcframework",
    avoid_deps = [
        "//tensorflow/lite/ios:tensorflow_lite_c",
    ],
    bundle_name = "TensorFlowLite",
    ios = {
        "simulator": [
            "x86_64",
            "arm64",
        ],
        "device": ["arm64"],
    },
    minimum_os_versions = {"ios": TFL_MINIMUM_OS_VERSION},
    tags = ["manual"],  # TODO: Remove once tf is on bazel 6.x+
    deps = [
        ":TensorFlowLite",
    ],
)

ios_unit_test(
    name = "Tests",
    size = "small",
    timeout = "moderate",
    minimum_os_version = TFL_MINIMUM_OS_VERSION,
    runner = tflite_ios_lab_runner("IOS_LATEST"),
    tags = TFL_DEFAULT_TAGS + TFL_DISABLED_SANITIZER_TAGS,
    deps = [
        ":TestsLibrary",
    ],
)

swift_library(
    name = "TestsLibrary",
    testonly = 1,
    srcs = glob(["Tests/*.swift"]),
    tags = TFL_DEFAULT_TAGS + ["nobuilder"],
    deps = [
        ":Resources",
        ":TensorFlowLiteAllDelegates",
    ],
)

objc_library(
    name = "Resources",
    data = [
        "//tensorflow/lite:testdata/add.bin",
        "//tensorflow/lite:testdata/add_quantized.bin",
        "//tensorflow/lite:testdata/multi_add.bin",
        "//tensorflow/lite:testdata/multi_signatures.bin",
    ],
    tags = TFL_DEFAULT_TAGS,
)