# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

import pytest

import tvm
import tvm.testing
from tvm.support import utils
from tvm.testing import env


@pytest.mark.skipif(not env.has_llvm(), reason="need llvm")
def test_import_static_library():
    from tvm import te

    # Generate two LLVM modules.
    A = te.placeholder((1024,), name="A")
    B = te.compute(A.shape, lambda *i: A(*i) + 1.0, name="B")
    irmod0 = tvm.IRModule.from_expr(
        te.create_prim_func([A, B]).with_attr("global_symbol", "myadd0")
    )
    irmod1 = tvm.IRModule.from_expr(
        te.create_prim_func([A, B]).with_attr("global_symbol", "myadd1")
    )

    mod0 = tvm.tirx.build(irmod0, target="llvm")
    mod1 = tvm.tirx.build(irmod1, target="llvm")

    assert mod0.implements_function("myadd0")
    assert mod1.implements_function("myadd1")
    assert mod1.is_compilation_exportable()

    # mod1 is currently an 'llvm' module.
    # Save and reload it as a vanilla 'static_library'.
    temp = utils.tempdir()
    mod1_o_path = temp.relpath("mod1.o")
    mod1.write_to_file(mod1_o_path)
    mod1_o = tvm.runtime.load_static_library(mod1_o_path, ["myadd1"])
    assert mod1_o.implements_function("myadd1")
    assert mod1_o.is_compilation_exportable()

    # Import mod1 as a static library into mod0 and compile to its own DSO.
    mod0.import_module(mod1_o)
    mod0_dso_path = temp.relpath("mod0.so")
    tvm.runtime.Executable(mod0).export_library(mod0_dso_path)

    # The imported mod1 is statically linked into mod0.
    loaded_lib = tvm.runtime.load_module(mod0_dso_path)
    assert loaded_lib.kind == "library"
    assert len(loaded_lib.imports) == 0
    assert loaded_lib.implements_function("myadd0")
    assert loaded_lib.get_function("myadd0")
    assert loaded_lib.implements_function("myadd1")
    assert loaded_lib.get_function("myadd1")
    assert not loaded_lib.is_compilation_exportable()


if __name__ == "__main__":
    tvm.testing.main()