# Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved.
""" Tests constant folding with globals. """
import dace
import numpy as np

from dace.frontend.python import astutils
from dace.frontend.python.newast import (GlobalResolver,
                                         ConditionalCodeResolver,
                                         DeadCodeEliminator)
from dace.frontend.python.parser import DaceProgram


class MyConfiguration:
    def __init__(self, parameter):
        self.p = parameter * 2

    @property
    def q(self):
        return self.p * 2

    def get_parameter(self):
        return self.p // 2

    @staticmethod
    def get_random_number():
        return 4

    @property
    def cloned(self):
        return MyConfiguration(self.get_parameter())


N = 2
cfg = MyConfiguration(N)
val = 5

# Confuse AST parser with global of the same name as array
A = 5


@dace.program
def instantiated_global(A):
    A[cfg.q] = (A[cfg.get_parameter()] * MyConfiguration.get_random_number() +
                cfg.p) + val


def test_instantiated_global():
    """
    Tests constant/symbolic values with predetermined global values.
    """
    A = np.random.rand(10)
    reg_A = np.copy(A)
    reg_A[cfg.q] = (reg_A[cfg.get_parameter()] *
                    MyConfiguration.get_random_number() + cfg.p) + val

    instantiated_global(A)

    assert np.allclose(A, reg_A)


def test_nested_globals():
    """
    Tests constant/symbolic values with multiple nesting levels.
    """
    @dace.program
    def instantiated_global2(A):
        A[cfg.q] = cfg.cloned.p

    A = np.random.rand(10)
    reg_A = np.copy(A)
    reg_A[cfg.q] = cfg.cloned.p

    instantiated_global2(A)

    assert np.allclose(A, reg_A)


def _analyze_and_unparse_code(func: DaceProgram) -> str:
    src_ast, _, _, _ = astutils.function_to_ast(func.f)
    resolved = {
        k: v
        for k, v in func.global_vars.items() if k not in func.argnames
    }
    src_ast = GlobalResolver(resolved).visit(src_ast)
    src_ast = ConditionalCodeResolver(resolved).visit(src_ast)
    src_ast = DeadCodeEliminator().visit(src_ast)

    return astutils.unparse(src_ast)


def test_dead_code_elimination_if():
    """
    Tests dead code elimination with compile-time if conditions.
    """
    sym = dace.symbol('sym', positive=True)
    cfg_symbolic = MyConfiguration(sym)

    @dace.program
    def test(A):
        if cfg_symbolic.q > sym:
            return 2 * A
        else:
            return 4 * A

    parsed_code = _analyze_and_unparse_code(test)
    assert '4' not in parsed_code
    assert '2' in parsed_code


def test_dead_code_elimination_ifexp():
    """
    Tests dead code elimination with compile-time ternary expressions.
    """
    sym = dace.symbol('sym', positive=True)
    cfg_symbolic = MyConfiguration(sym)

    @dace.program
    def test(A):
        return 2 * A if cfg_symbolic.q > sym else 4 * A

    parsed_code = _analyze_and_unparse_code(test)
    assert '4' not in parsed_code
    assert '2' in parsed_code


def test_dead_code_elimination_noelse():
    """
    Tests dead code elimination with compile-time if conditions (without else).
    """
    scale = None

    @dace.program
    def test(A):
        if scale is None:
            return 2 * A
        return scale * A

    parsed_code = _analyze_and_unparse_code(test)
    assert 'scale' not in parsed_code
    assert '2' in parsed_code


def test_dead_code_elimination_unreachable():
    """
    Tests dead code elimination with unreachable code.
    """
    @dace.program
    def test(A):
        if A[5] > 1:
            return 3 * A
            return 6 * A
        return 2 * A
        return 4 * A

    parsed_code = _analyze_and_unparse_code(test)
    assert '6' not in parsed_code and '4' not in parsed_code  # Dead code
    assert '5' in parsed_code and '1' in parsed_code  # Condition
    assert '3' in parsed_code and '2' in parsed_code  # Reachable code


# TODO: dace.constant should signal that argument evaluation is deferred to
#       (nested) call time
# dace.constant = lambda x: None
# def test_constant_parameter():
#     """
#     Tests nested functions with constant parameters passed in as arguments.
#     """
#     @dace.program
#     def nested_func(cfg: dace.constant(MyConfiguration), A: dace.float64[20]):
#         return A[cfg.p]

#     @dace.program
#     def constant_parameter(
#             cfg: dace.constant(MyConfiguration),
#             cfg2: dace.constant(MyConfiguration), A: dace.float64[20]):
#         A[cfg.q] = nested_func(cfg, A)
#         A[MyConfiguration.get_random_number()] = nested_func(cfg2, A)

#     cfg1 = MyConfiguration(3)
#     cfg2 = MyConfiguration(4)
#     A = np.random.rand(20)
#     reg_A = np.copy(A)
#     reg_A[12] = reg_A[6]
#     reg_A[4] = reg_A[8]

#     constant_parameter(cfg1, cfg2, A)
#     assert np.allclose(A, reg_A)

if __name__ == '__main__':
    test_instantiated_global()
    test_nested_globals()
    test_dead_code_elimination_if()
    test_dead_code_elimination_ifexp()
    test_dead_code_elimination_noelse()
    test_dead_code_elimination_unreachable()
    # test_constant_parameter()

# Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved.
""" Tests constant folding with globals. """
import dace
import numpy as np

from dace.frontend.python import astutils
from dace.frontend.python.newast import (GlobalResolver,
                                         ConditionalCodeResolver,
                                         DeadCodeEliminator)
from dace.frontend.python.parser import DaceProgram


class MyConfiguration:
    def __init__(self, parameter):
        self.p = parameter * 2

    @property
    def q(self):
        return self.p * 2

    def get_parameter(self):
        return self.p // 2

    @staticmethod
    def get_random_number():
        return 4

    @property
    def cloned(self):
        return MyConfiguration(self.get_parameter())


N = 2
cfg = MyConfiguration(N)
val = 5

# Confuse AST parser with global of the same name as array
A = 5


@dace.program
def instantiated_global(A):
    A[cfg.q] = (A[cfg.get_parameter()] * MyConfiguration.get_random_number() +
                cfg.p) + val


def test_instantiated_global():
    """
    Tests constant/symbolic values with predetermined global values.
    """
    A = np.random.rand(10)
    reg_A = np.copy(A)
    reg_A[cfg.q] = (reg_A[cfg.get_parameter()] *
                    MyConfiguration.get_random_number() + cfg.p) + val

    instantiated_global(A)

    assert np.allclose(A, reg_A)


def test_nested_globals():
    """
    Tests constant/symbolic values with multiple nesting levels.
    """
    @dace.program
    def instantiated_global2(A):
        A[cfg.q] = cfg.cloned.p

    A = np.random.rand(10)
    reg_A = np.copy(A)
    reg_A[cfg.q] = cfg.cloned.p

    instantiated_global2(A)

    assert np.allclose(A, reg_A)


def _analyze_and_unparse_code(func: DaceProgram) -> str:
    src_ast, _, _, _ = astutils.function_to_ast(func.f)
    resolved = {
        k: v
        for k, v in func.global_vars.items() if k not in func.argnames
    }
    src_ast = GlobalResolver(resolved).visit(src_ast)
    src_ast = ConditionalCodeResolver(resolved).visit(src_ast)
    src_ast = DeadCodeEliminator().visit(src_ast)

    return astutils.unparse(src_ast)


def test_dead_code_elimination_if():
    """
    Tests dead code elimination with compile-time if conditions.
    """
    sym = dace.symbol('sym', positive=True)
    cfg_symbolic = MyConfiguration(sym)

    @dace.program
    def test(A):
        if cfg_symbolic.q > sym:
            return 2 * A
        else:
            return 4 * A

    parsed_code = _analyze_and_unparse_code(test)
    assert '4' not in parsed_code
    assert '2' in parsed_code


def test_dead_code_elimination_ifexp():
    """
    Tests dead code elimination with compile-time ternary expressions.
    """
    sym = dace.symbol('sym', positive=True)
    cfg_symbolic = MyConfiguration(sym)

    @dace.program
    def test(A):
        return 2 * A if cfg_symbolic.q > sym else 4 * A

    parsed_code = _analyze_and_unparse_code(test)
    assert '4' not in parsed_code
    assert '2' in parsed_code


def test_dead_code_elimination_noelse():
    """
    Tests dead code elimination with compile-time if conditions (without else).
    """
    scale = None

    @dace.program
    def test(A):
        if scale is None:
            return 2 * A
        return scale * A

    parsed_code = _analyze_and_unparse_code(test)
    assert 'scale' not in parsed_code
    assert '2' in parsed_code


def test_dead_code_elimination_unreachable():
    """
    Tests dead code elimination with unreachable code.
    """
    @dace.program
    def test(A):
        if A[5] > 1:
            return 3 * A
            return 6 * A
        return 2 * A
        return 4 * A

    parsed_code = _analyze_and_unparse_code(test)
    assert '6' not in parsed_code and '4' not in parsed_code  # Dead code
    assert '5' in parsed_code and '1' in parsed_code  # Condition
    assert '3' in parsed_code and '2' in parsed_code  # Reachable code


# TODO: dace.constant should signal that argument evaluation is deferred to
#       (nested) call time
# dace.constant = lambda x: None
# def test_constant_parameter():
#     """
#     Tests nested functions with constant parameters passed in as arguments.
#     """
#     @dace.program
#     def nested_func(cfg: dace.constant(MyConfiguration), A: dace.float64[20]):
#         return A[cfg.p]

#     @dace.program
#     def constant_parameter(
#             cfg: dace.constant(MyConfiguration),
#             cfg2: dace.constant(MyConfiguration), A: dace.float64[20]):
#         A[cfg.q] = nested_func(cfg, A)
#         A[MyConfiguration.get_random_number()] = nested_func(cfg2, A)

#     cfg1 = MyConfiguration(3)
#     cfg2 = MyConfiguration(4)
#     A = np.random.rand(20)
#     reg_A = np.copy(A)
#     reg_A[12] = reg_A[6]
#     reg_A[4] = reg_A[8]

#     constant_parameter(cfg1, cfg2, A)
#     assert np.allclose(A, reg_A)

if __name__ == '__main__':
    test_instantiated_global()
    test_nested_globals()
    test_dead_code_elimination_if()
    test_dead_code_elimination_ifexp()
    test_dead_code_elimination_noelse()
    test_dead_code_elimination_unreachable()
    # test_constant_parameter()

# Copyright 2019-2021 ETH Zurich and the DaCe authors. All rights reserved. Tests constant folding with globals. # Confuse AST parser with global of the same name as array Tests constant/symbolic values with predetermined global values. Tests constant/symbolic values with multiple nesting levels. Tests dead code elimination with compile-time if conditions. Tests dead code elimination with compile-time ternary expressions. Tests dead code elimination with compile-time if conditions (without else). Tests dead code elimination with unreachable code. # Dead code # Condition # Reachable code # TODO: dace.constant should signal that argument evaluation is deferred to #       (nested) call time # dace.constant = lambda x: None # def test_constant_parameter(): #     """ #     Tests nested functions with constant parameters passed in as arguments. #     """ #     @dace.program #     def nested_func(cfg: dace.constant(MyConfiguration), A: dace.float64[20]): #         return A[cfg.p] #     @dace.program #     def constant_parameter( #             cfg: dace.constant(MyConfiguration), #             cfg2: dace.constant(MyConfiguration), A: dace.float64[20]): #         A[cfg.q] = nested_func(cfg, A) #         A[MyConfiguration.get_random_number()] = nested_func(cfg2, A) #     cfg1 = MyConfiguration(3) #     cfg2 = MyConfiguration(4) #     A = np.random.rand(20) #     reg_A = np.copy(A) #     reg_A[12] = reg_A[6] #     reg_A[4] = reg_A[8] #     constant_parameter(cfg1, cfg2, A) #     assert np.allclose(A, reg_A) # test_constant_parameter()