# Path: soap/datatype.py
# class TypeBase(object):
# class AutoType(TypeBase):
# class BoolType(TypeBase):
# class IntType(TypeBase):
# class FloatingPointType(TypeBase):
# class FloatType(FloatingPointType):
# class DoubleType(FloatingPointType):
# class FunctionType(TypeBase):
# class ArrayType(TypeBase):
# class IntegerArrayType(ArrayType):
# class FloatArrayType(ArrayType):
#     def __repr__(self):
#     def __eq__(self, other):
#     def __hash__(self):
#     def __str__(self):
#     def __str__(self):
#     def __str__(self):
#     def __str__(self):
#     def __str__(self):
#     def __str__(self):
#     def __str__(self):
#     def __init__(self, shape):
#     def __str__(self):
#     def __eq__(self, other):
#     def __hash__(self):
# def type_of(value):
# def type_cast(dtype, value=None, top=False, bottom=False):
# 
# Path: soap/expression/variable.py
# class Variable(ArithmeticMixin, BooleanMixin, BinaryExpression):
#     """The variable class."""
#     __slots__ = ()
#     _str_brackets = False
# 
#     def __init__(self, name, dtype=auto_type):
#         super().__init__(VARIABLE_OP, name, dtype)
# 
#     @property
#     def name(self):
#         return self.args[0]
# 
#     @property
#     def dtype(self):
#         return self.args[1]
# 
#     def format(self):
#         return '{}'.format(self.name)
# 
#     def __repr__(self):
#         return '{cls}({name!r}, {dtype!r})'.format(
#             cls=self.__class__.__name__, name=self.name, dtype=self.dtype)
# 
# Path: soap/parser/program.py
# def parse(program, decl=None):
#     decl = decl or {}
#     visitor = _ProgramVisitor(decl)
#     program = _preprocess(program)
#     flow = visitor.parse(program)
#     return ProgramFlow(flow)
# 
# Path: soap/semantics/schedule/extract.py
# class ForLoopNestExtractor(ForLoopExtractor):
#     def __init__(self, fix_expr):
#         super().__init__(fix_expr)
#         try:
#             self.iter_vars, self.iter_slices, self._kernel = \
#                 self._extract_for_loop_nest(fix_expr)
#             self.is_for_loop_nest = True
#         except ForLoopNestExtractionFailureException as exception:
#             self.iter_vars = self.iter_slices = None
#             self._kernel = fix_expr.loop_state
#             self.is_for_loop_nest = False
#             self.exception = exception
# 
#     @property
#     def kernel(self):
#         return self._kernel
# 
#     def _check_init_sandwich(self, fix_expr):
#         for var, expr in fix_expr.init_state.items():
#             if is_expression(expr):
#                 raise ForLoopNestExtractionFailureException(
#                     'Loop has logic sandwich because its init_state '
#                     'has logic.')
# 
#     def _find_inner_fix_expr(self, fix_expr, iter_var):
#         found_inner_loop = False
#         for var, expr in fix_expr.loop_state.items():
#             if var == iter_var:
#                 pass  # iter_var is already checked by ForLoopExtractor
#             elif var == fix_expr.loop_var:
#                 if isinstance(expr, FixExpr):
#                     # has SIMPLE inner loop
#                     inner_fix_expr = expr
#                     found_inner_loop = True
#                 else:
#                     # has sandwich
#                     raise ForLoopNestExtractionFailureException(
#                         'Loop has logic sandwich somewhere.')
#             elif var == expr:
#                 pass  # no logic
#             else:
#                 raise ForLoopNestExtractionFailureException(
#                     'Loop has logic sandwich on variables other than the '
#                     'iteration variable and the loop variable.')
#         if not found_inner_loop:
#             raise ForLoopNestExtractionFailureException(
#                 'Did not find loop var in loop.')
#         return inner_fix_expr
# 
#     def _extract_for_loop_nest(self, fix_expr):
#         extractor = ForLoopExtractor(fix_expr)
#         if not extractor.is_for_loop:
#             raise ForLoopNestExtractionFailureException(
#                 'Loop is not for loop, reason: ' + str(extractor.exception))
# 
#         iter_var = extractor.iter_var
#         iter_vars = [iter_var]
#         iter_slices = [extractor.iter_slice]
# 
#         if not extractor.has_inner_loops:
#             return iter_vars, iter_slices, fix_expr.loop_state
# 
#         # nested loops
#         inner_fix_expr = self._find_inner_fix_expr(fix_expr, iter_var)
# 
#         # and check inner loop init_state is simple
#         self._check_init_sandwich(inner_fix_expr)
# 
#         inner_iter_vars, inner_iter_slices, inner_kernel = \
#             self._extract_for_loop_nest(inner_fix_expr)
#         iter_vars += inner_iter_vars
#         iter_slices += inner_iter_slices
#         return iter_vars, iter_slices, inner_kernel
# 
# Path: soap/semantics/state/meta.py
# def flow_to_meta_state(flow):
#     id_state = MetaState({v: v for v in flow.vars(output=False)})
#     return id_state.transition(flow)


import unittest
from soap.datatype import int_type, FloatArrayType
from soap.expression import Variable
from soap.parser import parse
from soap.semantics.schedule.extract import ForLoopNestExtractor
from soap.semantics.state import flow_to_meta_state



def _extract_from_program(program):
    flow = parse(program)
    meta_state = flow_to_meta_state(flow)
    fix_expr = meta_state[flow.outputs[0]]
    return ForLoopNestExtractor(fix_expr)


class TestExtractor(unittest.TestCase):

    def compare(self, extractor, expect_for_loop):
        for key, expect_val in expect_for_loop.items():
            self.assertEqual(getattr(extractor, key), expect_val)

    def test_simple_loop(self):
        program = """
        #pragma soap input float a[30]
        #pragma soap output a
        for (int i = 1; i < 10; i++)
            a[i] = a[i - 1] + 1;
        """
        expect_for_loop = {

# Path: soap/datatype.py
# class TypeBase(object):
# class AutoType(TypeBase):
# class BoolType(TypeBase):
# class IntType(TypeBase):
# class FloatingPointType(TypeBase):
# class FloatType(FloatingPointType):
# class DoubleType(FloatingPointType):
# class FunctionType(TypeBase):
# class ArrayType(TypeBase):
# class IntegerArrayType(ArrayType):
# class FloatArrayType(ArrayType):
#     def __repr__(self):
#     def __eq__(self, other):
#     def __hash__(self):
#     def __str__(self):
#     def __str__(self):
#     def __str__(self):
#     def __str__(self):
#     def __str__(self):
#     def __str__(self):
#     def __str__(self):
#     def __init__(self, shape):
#     def __str__(self):
#     def __eq__(self, other):
#     def __hash__(self):
# def type_of(value):
# def type_cast(dtype, value=None, top=False, bottom=False):
# 
# Path: soap/expression/variable.py
# class Variable(ArithmeticMixin, BooleanMixin, BinaryExpression):
#     """The variable class."""
#     __slots__ = ()
#     _str_brackets = False
# 
#     def __init__(self, name, dtype=auto_type):
#         super().__init__(VARIABLE_OP, name, dtype)
# 
#     @property
#     def name(self):
#         return self.args[0]
# 
#     @property
#     def dtype(self):
#         return self.args[1]
# 
#     def format(self):
#         return '{}'.format(self.name)
# 
#     def __repr__(self):
#         return '{cls}({name!r}, {dtype!r})'.format(
#             cls=self.__class__.__name__, name=self.name, dtype=self.dtype)
# 
# Path: soap/parser/program.py
# def parse(program, decl=None):
#     decl = decl or {}
#     visitor = _ProgramVisitor(decl)
#     program = _preprocess(program)
#     flow = visitor.parse(program)
#     return ProgramFlow(flow)
# 
# Path: soap/semantics/schedule/extract.py
# class ForLoopNestExtractor(ForLoopExtractor):
#     def __init__(self, fix_expr):
#         super().__init__(fix_expr)
#         try:
#             self.iter_vars, self.iter_slices, self._kernel = \
#                 self._extract_for_loop_nest(fix_expr)
#             self.is_for_loop_nest = True
#         except ForLoopNestExtractionFailureException as exception:
#             self.iter_vars = self.iter_slices = None
#             self._kernel = fix_expr.loop_state
#             self.is_for_loop_nest = False
#             self.exception = exception
# 
#     @property
#     def kernel(self):
#         return self._kernel
# 
#     def _check_init_sandwich(self, fix_expr):
#         for var, expr in fix_expr.init_state.items():
#             if is_expression(expr):
#                 raise ForLoopNestExtractionFailureException(
#                     'Loop has logic sandwich because its init_state '
#                     'has logic.')
# 
#     def _find_inner_fix_expr(self, fix_expr, iter_var):
#         found_inner_loop = False
#         for var, expr in fix_expr.loop_state.items():
#             if var == iter_var:
#                 pass  # iter_var is already checked by ForLoopExtractor
#             elif var == fix_expr.loop_var:
#                 if isinstance(expr, FixExpr):
#                     # has SIMPLE inner loop
#                     inner_fix_expr = expr
#                     found_inner_loop = True
#                 else:
#                     # has sandwich
#                     raise ForLoopNestExtractionFailureException(
#                         'Loop has logic sandwich somewhere.')
#             elif var == expr:
#                 pass  # no logic
#             else:
#                 raise ForLoopNestExtractionFailureException(
#                     'Loop has logic sandwich on variables other than the '
#                     'iteration variable and the loop variable.')
#         if not found_inner_loop:
#             raise ForLoopNestExtractionFailureException(
#                 'Did not find loop var in loop.')
#         return inner_fix_expr
# 
#     def _extract_for_loop_nest(self, fix_expr):
#         extractor = ForLoopExtractor(fix_expr)
#         if not extractor.is_for_loop:
#             raise ForLoopNestExtractionFailureException(
#                 'Loop is not for loop, reason: ' + str(extractor.exception))
# 
#         iter_var = extractor.iter_var
#         iter_vars = [iter_var]
#         iter_slices = [extractor.iter_slice]
# 
#         if not extractor.has_inner_loops:
#             return iter_vars, iter_slices, fix_expr.loop_state
# 
#         # nested loops
#         inner_fix_expr = self._find_inner_fix_expr(fix_expr, iter_var)
# 
#         # and check inner loop init_state is simple
#         self._check_init_sandwich(inner_fix_expr)
# 
#         inner_iter_vars, inner_iter_slices, inner_kernel = \
#             self._extract_for_loop_nest(inner_fix_expr)
#         iter_vars += inner_iter_vars
#         iter_slices += inner_iter_slices
#         return iter_vars, iter_slices, inner_kernel
# 
# Path: soap/semantics/state/meta.py
# def flow_to_meta_state(flow):
#     id_state = MetaState({v: v for v in flow.vars(output=False)})
#     return id_state.transition(flow)


# Path: tests/test_extractor.py
import unittest
from soap.datatype import int_type, FloatArrayType
from soap.expression import Variable
from soap.parser import parse
from soap.semantics.schedule.extract import ForLoopNestExtractor
from soap.semantics.state import flow_to_meta_state




def _extract_from_program(program):
    flow = parse(program)
    meta_state = flow_to_meta_state(flow)
    fix_expr = meta_state[flow.outputs[0]]
    return ForLoopNestExtractor(fix_expr)


class TestExtractor(unittest.TestCase):

    def compare(self, extractor, expect_for_loop):
        for key, expect_val in expect_for_loop.items():
            self.assertEqual(getattr(extractor, key), expect_val)

    def test_simple_loop(self):
        program = """
        #pragma soap input float a[30]
        #pragma soap output a
        for (int i = 1; i < 10; i++)
            a[i] = a[i - 1] + 1;
        """
        expect_for_loop = {

            'iter_var': Variable('i', int_type),