load("@com_google_protobuf//bazel:upb_proto_library.bzl", "upb_c_proto_library")
load("@rules_cc//cc:defs.bzl", "cc_library")
load("@rules_zig//zig:defs.bzl", "zig_library", "zig_test")

upb_c_proto_library(
    name = "xla_data_upb",
    deps = ["@xla//xla:xla_data_proto"],
)

upb_c_proto_library(
    name = "xla_compile_options_upb",
    deps = ["@xla//xla/pjrt/proto:compile_options_proto"],
)

cc_library(
    name = "posix",
    hdrs = ["posix.h"],
)

zig_library(
    name = "zml",
    srcs = glob([
        "*.zig",
        "aio/**/*.zig",
        "nn/**/*.zig",
        # TODO: test_runner.zig should not be here.
        # It's here for now because it seems that test_runner property in zig_test is misbehaving.
        # See https://github.com/zml/rules_zig/issues/2
        "test_runner.zig",
    ]),
    copts = ["-lc"],
    main = "zml.zig",
    visibility = ["//visibility:public"],
    deps = [
        ":posix",
        ":xla_compile_options_upb",
        ":xla_data_upb",
        "//async",
        "//mlir",
        "//mlir/dialects",
        "//pjrt",
        "//runtimes",
        "//stdx",
        "//upb",
        "//zml/tokenizer",
        "//zml/tools",
        "@rules_zig//zig/runfiles",
    ],
)

# All ZML Tests

zig_test(
    name = "test",
    data = [
        "aio/torch/simple.pt",
        "aio/torch/simple_test_4.pickle",
    ],
    test_runner = ":test_runner",
    deps = [":zml"],
)

filegroup(
    name = "test_runner",
    srcs = ["test_runner.zig"],
    visibility = ["//visibility:public"],
)