# Tests for our minimal LLVM wrappers
import __future__

from test import test_dynamic
from test import test_support

try:
    import _llvm
except ImportError:
    raise test_support.TestSkipped("not built against LLVM")

import __builtin__
import contextlib
import functools
import gc
import sys
import types
import unittest
import weakref


# Calculate default LLVM optimization level.
def _foo():
    pass
DEFAULT_OPT_LEVEL = _foo.__code__.co_optimization
del _foo


# Various constants that feed into the hotness model.
HOTNESS_CALL = 10  # Points for a function entry.
HOTNESS_LOOP = 1   # Points for a taken loop backedge.

JIT_SPIN_COUNT = _llvm.get_hotness_threshold() / HOTNESS_CALL + 1000
JIT_OPT_LEVEL = sys.flags.optimize if sys.flags.optimize > 2 else 2


@contextlib.contextmanager
def set_jit_control(new_level):
    orig_level = _llvm.get_jit_control()
    _llvm.set_jit_control(new_level)
    try:
        yield
    finally:
        _llvm.set_jit_control(orig_level)


def at_each_optimization_level(func):
    """Decorator for test functions, to run them at each optimization level."""
    levels = [None, -1, 0, 1, 2]
    if DEFAULT_OPT_LEVEL != -1:
        levels = [level for level in levels if level >= DEFAULT_OPT_LEVEL]
    @functools.wraps(func)
    def result(self):
        for level in levels:
            func(self, level)
    return result


def compile_for_llvm(function_name, def_string, optimization_level=-1,
                     globals_dict=None):
    """Compiles function_name, defined in def_string to be run through LLVM.

    Compiles and runs def_string in a temporary namespace, pulls the
    function named 'function_name' out of that namespace, optimizes it
    at level 'optimization_level', -1 for the default optimization,
    and marks it to be JITted and run through LLVM.

    """
    namespace = {}
    if globals_dict is None:
        globals_dict = globals()
    exec def_string in globals_dict, namespace
    func = namespace[function_name]
    if optimization_level is not None:
        if optimization_level >= DEFAULT_OPT_LEVEL:
            func.__code__.co_optimization = optimization_level
        func.__code__.co_use_jit = True
    return func


def spin_until_hot(func, *training_args):
    """Compile a function by calling it until it is hot.

    Calls a function JIT_SPIN_COUNT times with each of training_args as
    arguments and returns a list containing the return values from each
    iteration.  When testing feedback directed optimizations, this function is
    used to gather feedback in the code object before performaing JIT
    compilation.

    To inject different training arguments into the function while it is being
    spun up, you can pass multiple tuples of arguments to this function, like
    so:

    spin_until_hot(mul, [1, 3], [1.0, 3.0])

    This will cause the function to be trained on both integer and floating
    point inputs.
    """
    if not training_args:  # An empty training_args isn't what you meant.
        training_args = [[]]

    results = []
    with set_jit_control("whenhot"):
        for _ in xrange(JIT_SPIN_COUNT):
            for args in training_args:
                results.append(func(*args))
    return results


class ExtraAssertsTestCase(unittest.TestCase):
    def assertRaisesWithArgs(self, expected_exception_type,
                             expected_args, f, *args, **kwargs):
        try:
            f(*args, **kwargs)
        except expected_exception_type, real_exception:
            pass
        else:
            self.fail("%r not raised" % expected_exception_type)
        self.assertEquals(real_exception.args, expected_args)

    def assertContains(self, obj, container):
        if obj not in container:
            self.fail("%r not found in %r" % (obj, container))

    def assertNotContains(self, obj, container):
        if obj in container:
            self.fail("%r found in %r" % (obj, container))


class LlvmTestCase(unittest.TestCase):

    """Common base class for LLVM-focused tests.

    Features provided:
        - Assert that the code doesn't bail to the interpreter.
    """

    def setUp(self):
        sys.setbailerror(True)
        self._old_jit_control = _llvm.get_jit_control()
        _llvm.set_jit_control("whenhot")

    def tearDown(self):
        sys.setbailerror(False)
        _llvm.set_jit_control(self._old_jit_control)


class GeneralCompilationTests(ExtraAssertsTestCase, LlvmTestCase):

    def test_uncreatable(self):
        # Functions can only be created by their static factories.
        self.assertRaises(TypeError, _llvm._function)

    def test_use_llvm(self):
        # Regression test: setting co_use_jit without setting an optimization
        # level used to segfault when the function was called.
        def foo():
            return 5
        foo.__code__.co_use_jit = True
        foo()

    @at_each_optimization_level
    def test_llvm_compile(self, level):
        # Makes no sense at level None
        if level is None:
            return
        def f(x):
            pass
        f = _llvm.compile(f.func_code, level)
        self.assertTrue(isinstance(f, _llvm._function))
        self.assertRaises(TypeError, _llvm.compile, f, level)

    def test_co_llvm(self):
        def f():
            return 1 + 2
        spin_until_hot(f, [])
        str_co_llvm = str(f.__code__.co_llvm)
        # After code objects are compiled to machine code, the IR form
        # is mostly cleared out. However, we want to be able to print
        # the IR generated for any particular function, so this tests
        # that it gets regenerated when we ask for its string
        # representation.  The cleared-out form is around 5 lines
        # long, while all full functions are much longer.
        self.assertTrue(10 < len(str_co_llvm.split('\n')), msg=str_co_llvm)

    @at_each_optimization_level
    def test_run_simple_function(self, level):
        foo = compile_for_llvm("foo", """
def foo():
    pass
""", level)
        self.assertEquals(None, foo())

    @at_each_optimization_level
    def test_constants(self, level):
        foo = compile_for_llvm("foo", """
def foo(x):
    if x: x[...]
    return [1, 2L, 3.2, 7j, (), "Hello", u"Hello", None]
""", level)
        self.assertEquals([1, 2L, 3.2, 7j, (), "Hello", u"Hello", None],
                          foo(False))

    def test_same_named_functions_coexist(self):
        foo1 = compile_for_llvm("foo", """
def foo(a):
    return a
""")
        foo2 = compile_for_llvm("foo", """
def foo():
    return 7
""")
        self.assertEquals("Hello", foo1("Hello"))
        self.assertEquals(7, foo2())

    def test_stack_pointer_optimized_to_register(self):
        def test_func():
            # We may have to add opcode uses to here as we find things
            # that break the stack pointer optimization.
            return sum(range(*[1, 10, 3]))
        # Run mem2reg.
        test_func.__code__.co_optimization = 2
        self.assertFalse("%stack_pointer_addr = alloca"
                         in str(test_func.__code__.co_llvm))

    # -Xjit=always will cause this test to always fail.
    if _llvm.get_jit_control() != "always":
        def test_fetch_unset_co_llvm(self):
            def test_func():
                pass
            test_func.__code__.co_use_jit = True
            # Just setting co_use_jit doesn't force code generation.
            self.assertEqual(str(test_func.__code__.co_llvm), "None")

    @at_each_optimization_level
    def test_return_arg(self, level):
        foo = compile_for_llvm("foo", """
def foo(a):
    return a
""", level)
        self.assertEquals(3, foo(3))
        self.assertEquals("Hello", foo("Hello"))

    @at_each_optimization_level
    def test_unbound_local(self, level):
        foo = compile_for_llvm("foo", """
def foo():
    a = a
""", level)
        try:
            foo()
        except UnboundLocalError as e:
            self.assertEquals(
                str(e), "local variable 'a' referenced before assignment")
        else:
            self.fail("Expected UnboundLocalError")

    @at_each_optimization_level
    def test_assign(self, level):
        foo = compile_for_llvm("foo", """
def foo(a):
    b = a
    return b
""", level)
        self.assertEquals(3, foo(3))
        self.assertEquals("Hello", foo("Hello"))

    @at_each_optimization_level
    def test_raising_getiter(self, level):
        class RaisingIter(object):
            def __iter__(self):
                raise RuntimeError
        loop = compile_for_llvm("loop", """
def loop(range):
    for i in range:
        pass
""", level)
        self.assertRaises(RuntimeError, loop, RaisingIter())

    @at_each_optimization_level
    def test_raising_next(self, level):
        class RaisingNext(object):
            def __iter__(self):
                return self
            def next(self):
                raise RuntimeError
        loop = compile_for_llvm("loop", """
def loop(range):
    for i in range:
        pass
""", level)
        self.assertRaises(RuntimeError, loop, RaisingNext())

    @at_each_optimization_level
    def test_import_name(self, level):
        importer = compile_for_llvm("importer", """
def importer():
    import os
    return os
""", level)
        import os
        self.assertEqual(importer(), os)

    @at_each_optimization_level
    def test_loop(self, level):
        loop = compile_for_llvm("loop", """
def loop(range):
    for i in range:
        pass
""", level)
        r = iter(range(12))
        self.assertEquals(None, loop(r))
        self.assertRaises(StopIteration, next, r)

    @at_each_optimization_level
    def test_return_from_loop(self, level):
        loop = compile_for_llvm("loop", """
def loop(range):
    for i in range:
        return i
""", level)
        self.assertEquals(1, loop([1,2,3]))

    @at_each_optimization_level
    def test_finally(self, level):
        cleanup = compile_for_llvm("cleanup", """
def cleanup(obj):
    try:
        return 3
    finally:
        obj['x'] = 2
""", level)
        obj = {}
        self.assertEquals(3, cleanup(obj))
        self.assertEquals({'x': 2}, obj)

    @at_each_optimization_level
    def test_nested_finally(self, level):
        cleanup = compile_for_llvm("cleanup", """
def cleanup(obj):
    try:
        try:
            return 3
        finally:
            obj['x'] = 2
    finally:
        obj['y'] = 3
""", level)
        obj = {}
        self.assertEquals(3, cleanup(obj))
        self.assertEquals({'x': 2, 'y': 3}, obj)

    @at_each_optimization_level
    def test_finally_fallthrough(self, level):
        cleanup = compile_for_llvm("cleanup", """
def cleanup(obj):
    try:
        obj['y'] = 3
    finally:
        obj['x'] = 2
    return 3
""", level)
        obj = {}
        self.assertEquals(3, cleanup(obj))
        self.assertEquals({'x': 2, 'y': 3}, obj)

    @at_each_optimization_level
    def test_exception_out_of_finally(self, level):
        cleanup = compile_for_llvm("cleanup", """
def cleanup(obj):
    try:
        obj['x'] = 2
    finally:
        a = a
        obj['y'] = 3
    return 3
""", level)
        obj = {}
        self.assertRaises(UnboundLocalError, cleanup, obj)
        self.assertEquals({'x': 2}, obj)

    @at_each_optimization_level
    def test_exception_through_finally(self, level):
        cleanup = compile_for_llvm("cleanup", """
def cleanup(obj):
    try:
        a = a
    finally:
        obj['x'] = 2
""", level)
        obj = {}
        self.assertRaises(UnboundLocalError, cleanup, obj)
        self.assertEquals({'x': 2}, obj)

    @at_each_optimization_level
    def test_return_in_finally_overrides(self, level):
        cleanup = compile_for_llvm("cleanup", """
def cleanup():
    try:
        return 3
    finally:
        return 2
""", level)
        self.assertEquals(2, cleanup())

    @at_each_optimization_level
    def test_except(self, level):
        catch = compile_for_llvm("catch", """
def catch(obj):
    try:
        raise ZeroDivisionError
    except:
        obj["x"] = 2
""", level)
        obj = {}
        self.assertEquals(None, catch(obj))
        self.assertEquals({"x": 2}, obj)

    @at_each_optimization_level
    def test_filtered_except(self, level):
        catch = compile_for_llvm("catch", """
def catch(exc_type, obj):
    try:
        1 / 0
    except exc_type:
        obj["x"] = 2
""", level)
        obj = {}
        self.assertEquals(None, catch(ZeroDivisionError, obj))
        self.assertEquals({"x": 2}, obj)
        obj = {}
        self.assertRaises(ZeroDivisionError, catch, UnboundLocalError, obj)
        self.assertEquals({}, obj)

    @at_each_optimization_level
    def test_filtered_except_var(self, level):
        catch = compile_for_llvm("catch", """
def catch():
    try:
        1 / 0
    except ZeroDivisionError, exc:
        return exc
""", level)
        exc = catch()
        self.assertEquals(ZeroDivisionError, type(exc))
        self.assertEquals(('integer division or modulo by zero',), exc.args)

    @at_each_optimization_level
    def test_except_skipped_on_fallthrough(self, level):
        catch = compile_for_llvm("catch", """
def catch(obj):
    try:
        obj["x"] = 2
    except:
        obj["y"] = 3
    return 7
""", level)
        obj = {}
        self.assertEquals(7, catch(obj))
        self.assertEquals({"x": 2}, obj)

    @at_each_optimization_level
    def test_else_hit_on_fallthrough(self, level):
        catch = compile_for_llvm("catch", """
def catch(obj):
    try:
        obj["x"] = 2
    except:
        obj["y"] = 3
    else:
        obj["z"] = 4
    return 7
""", level)
        obj = {}
        self.assertEquals(7, catch(obj))
        self.assertEquals({"x": 2, "z": 4}, obj)

    @at_each_optimization_level
    def test_else_skipped_on_catch(self, level):
        catch = compile_for_llvm("catch", """
def catch(obj):
    try:
        a = a
    except:
        obj["y"] = 3
    else:
        obj["z"] = 4
    return 7
""", level)
        obj = {}
        self.assertEquals(7, catch(obj))
        self.assertEquals({"y": 3}, obj)

    @at_each_optimization_level
    def test_raise_from_except(self, level):
        catch = compile_for_llvm("catch", """
def catch(obj):
    try:
        raise ZeroDivisionError
    except:
        a = a
    obj["x"] = 2
""", level)
        obj = {}
        self.assertRaises(UnboundLocalError, catch, obj)
        self.assertEquals({}, obj)

    @at_each_optimization_level
    def test_nested_except(self, level):
        catch = compile_for_llvm("catch", """
def catch(obj):
    try:
        try:
            1 / 0
        except:
            obj["x"] = 2
            a = a
    except:
        obj["y"] = 3
""", level)
        obj = {}
        self.assertEquals(None, catch(obj))
        self.assertEquals({"x": 2, "y": 3}, obj)

    @at_each_optimization_level
    def test_nested_except_skipped_on_fallthrough(self, level):
        catch = compile_for_llvm("catch", """
def catch(obj):
    try:
        try:
            raise ZeroDivisionError
        except:
            obj["x"] = 2
    except:
        obj["y"] = 3
""", level)
        obj = {}
        self.assertEquals(None, catch(obj))
        self.assertEquals({"x": 2}, obj)

    @at_each_optimization_level
    def test_nested_finally_doesnt_block_catch(self, level):
        # Raise from the try.
        catch = compile_for_llvm("catch", """
def catch(obj):
    try:
        try:
            1 / 0
        finally:
            obj["x"] = 2
    except:
        obj["y"] = 3
""", level)
        obj = {}
        self.assertEquals(None, catch(obj))
        self.assertEquals({"x": 2, "y": 3}, obj)

        # And raise from the finally.
        catch = compile_for_llvm("catch", """
def catch(obj):
    try:
        try:
            obj["x"] = 2
        finally:
            raise ZeroDivisionError
    except:
        obj["y"] = 3
""", level)
        obj = {}
        self.assertEquals(None, catch(obj))
        self.assertEquals({"x": 2, "y": 3}, obj)

    @at_each_optimization_level
    def test_nested_except_goes_through_finally(self, level):
        # Raise from the try.
        catch = compile_for_llvm("catch", """
def catch(obj):
    try:
        try:
            1 / 0
        except:
            obj["x"] = 2
    finally:
        obj["y"] = 3
""", level)
        obj = {}
        self.assertEquals(None, catch(obj))
        self.assertEquals({"x": 2, "y": 3}, obj)

        # And raise from the except.
        catch = compile_for_llvm("catch", """
def catch(obj):
    try:
        try:
            raise UnboundLocalError
        except:
            1 / 0
    finally:
        obj["y"] = 3
""", level)
        obj = {}
        self.assertRaises(ZeroDivisionError, catch, obj)
        self.assertEquals({"y": 3}, obj)

    @at_each_optimization_level
    def test_finally_in_finally(self, level):
        catch = compile_for_llvm("catch", """
def catch(obj):
    try:
        raise ZeroDivisionError
    finally:
        try:
            a = a
        finally:
            obj["x"] = 2
""", level)
        obj = {}
        self.assertRaises(UnboundLocalError, catch, obj)
        self.assertEquals({"x": 2}, obj)

    @at_each_optimization_level
    def test_subexception_caught_in_finally(self, level):
        catch = compile_for_llvm("catch", """
def catch(obj):
    try:
        1 / 0
    finally:
        try:
            a = a
        except:
            obj["x"] = 2
""", level)
        obj = {}
        self.assertRaises(ZeroDivisionError, catch, obj)
        self.assertEquals({"x": 2}, obj)

    @at_each_optimization_level
    def test_delete_fast(self, level):
        delit = compile_for_llvm('delit', """
def delit(x):
    y = 2
    z = 3
    del y
    del x
    return z
""", level)
        self.assertEquals(delit(1), 3)

        useit = compile_for_llvm('useit', """
def useit(x):
    del x
    return x
""", level)
        self.assertRaises(UnboundLocalError, useit, 1)

        misuseit = compile_for_llvm('misuseit', 'def misuseit(x): del y',
                                    level)
        self.assertRaises(UnboundLocalError, misuseit, 1)

        reuseit = compile_for_llvm('reuseit', """
def reuseit(x):
    del x
    x = 3
    return x
""", level)
        self.assertEquals(reuseit(1), 3)

    @at_each_optimization_level
    def test_call_function(self, level):
        f1 = compile_for_llvm("f1", "def f1(x): return x()", level)
        self.assertEquals(f1(lambda: 5), 5)
        def raise_exc():
            raise ValueError
        self.assertRaises(ValueError, f1, raise_exc)

        f2 = compile_for_llvm("f2", "def f2(x, y, z): return x(y, 2, z)",
                              level)
        self.assertEquals(f2(lambda *args: args, 1, 3), (1, 2, 3))

        f3 = compile_for_llvm("f3", "def f3(x, y, z): return x(y(z()))",
                              level)
        self.assertEquals(f3(lambda x: x+1, lambda x: x+2, lambda: 0), 3)

    @at_each_optimization_level
    def test_load_global(self, level):
        our_globals = dict(globals())  # Isolate the test's global effects.
        testvalue = 'test global value'
        loadglobal = compile_for_llvm('loadglobal',
                                      'def loadglobal(): return testvalue',
                                      level, our_globals)
        our_globals['testvalue'] = testvalue
        self.assertEquals(loadglobal(), testvalue)

        loadbuiltin = compile_for_llvm('loadbuiltin',
                                       'def loadbuiltin(): return str',
                                       level)
        self.assertEquals(loadbuiltin(), str)

        nosuchglobal = compile_for_llvm('nosuchglobal', '''
def nosuchglobal():
    return there_better_be_no_such_global
''', level)
        self.assertRaises(NameError, nosuchglobal)

    @at_each_optimization_level
    def test_store_global(self, level):
        our_globals = dict(globals())  # Isolate the test's global effects.
        setglobal = compile_for_llvm('setglobal', '''
def setglobal(x):
    global _test_global
    _test_global = x
''', level, our_globals)
        testvalue = "test global value"
        self.assertTrue('_test_global' not in our_globals)
        setglobal(testvalue)
        self.assertContains('_test_global', our_globals)
        self.assertEquals(our_globals['_test_global'], testvalue)

    @at_each_optimization_level
    def test_delete_global(self, level):
        our_globals = dict(globals())  # Isolate the test's global effects.
        delglobal = compile_for_llvm('delglobal', '''
def delglobal():
    global _test_global
    del _test_global
''', level, our_globals)
        our_globals['_test_global'] = 'test global value'
        self.assertContains('_test_global', our_globals)
        delglobal()
        self.assertTrue('_test_global' not in our_globals)

    @at_each_optimization_level
    def test_load_name(self, level):
        our_globals = dict(globals())  # Isolate the test's global effects.
        testvalue = 'test name value'
        loadlocal = compile_for_llvm('loadlocal', '''
def loadlocal():
    exec 'testvalue = "Hello"'
    return testvalue
''', level, our_globals)
        our_globals['testvalue'] = testvalue
        self.assertEquals(loadlocal(), 'Hello')

        our_globals = dict(globals())
        loadglobal = compile_for_llvm('loadglobal', '''
def loadglobal():
    exec ''
    return testvalue
''', level, our_globals)
        our_globals['testvalue'] = testvalue
        self.assertEquals(loadglobal(), testvalue)

        loadbuiltin = compile_for_llvm('loadbuiltin', '''
def loadbuiltin():
    exec ''
    return str
''', level)
        self.assertEquals(loadbuiltin(), str)

        nosuchname = compile_for_llvm('nosuchname', '''
def nosuchname():
    exec ''
    return there_better_be_no_such_name
''', level)
        self.assertRaises(NameError, nosuchname)

    @at_each_optimization_level
    def test_store_name(self, level):
        set_local = compile('a = 3', '<string>', 'exec')
        if level is not None:
            set_local.co_optimization = level
            set_local.co_use_jit = True
        exec set_local
        self.assertEquals(a, 3)

    @at_each_optimization_level
    def test_delete_name(self, level):
        do_del = compile('del a', '<string>', 'exec')
        if level is not None:
            do_del.co_optimization = level
            do_del.co_use_jit = True
        exec 'a = 3'
        self.assertEquals(a, 3)
        exec do_del
        try:
            a
        except NameError:
            pass
        else:
            self.fail('Expected "a" to be deleted')

        try:
            exec compile('del nonexistent', '<string>', 'exec')
        except NameError, e:
            self.assertEquals(e.args, ('name \'nonexistent\' is not defined',))
        else:
            self.fail('Expected not to find "nonexistent"')

    @at_each_optimization_level
    def test_simple_if_stmt(self, level):
        simple_if = compile_for_llvm("simple_if", """
def simple_if(x):
    if x:
        return "true"
""", level)
        self.assertEquals(simple_if(True), "true")
        self.assertEquals(simple_if(False), None)

        simple_if_else = compile_for_llvm("simple_if_else", """
def simple_if_else(x):
    if x:
        return "true"
    else:
        return "false"
""", level)
        self.assertEquals(simple_if_else("not false"), "true")
        self.assertEquals(simple_if_else(""), "false")

    @at_each_optimization_level
    def test_if_stmt_exceptions(self, level):
        if_exception = compile_for_llvm("if_exception", """
def if_exception(x):
    if x:
        return 1
""", level)
        class Unboolable(object):
            def __nonzero__(self):
                raise RuntimeError
        self.assertRaises(RuntimeError, if_exception, Unboolable())

    @at_each_optimization_level
    def test_complex_if(self, level):
        complex_if = compile_for_llvm("complex_if", """
def complex_if(x, y, z):
    if x:
        if y:
            return 1
        else:
            if z:
                return 2
            return 1 / 0
    else:
        return 3
""", level)
        self.assertEquals(complex_if(True, True, False), 1)
        self.assertEquals(complex_if(True, False, True), 2)
        self.assertEquals(complex_if(False, True, True), 3)
        self.assertRaises(ZeroDivisionError, complex_if, True, False, False)

    # Asserts aren't compiled when -O is passed.
    if sys.flags.optimize < 1:
        @at_each_optimization_level
        def test_assert(self, level):
            f = compile_for_llvm("f", "def f(x): assert x", level)
            self.assertEquals(f(1), None)
            self.assertRaises(AssertionError, f, 0)

    @at_each_optimization_level
    def test_and(self, level):
        and1 = compile_for_llvm("and1",
                                "def and1(x, y): return x and y", level)
        self.assertEquals(and1("x", "y"), "y")
        self.assertEquals(and1((), "y"), ())
        self.assertEquals(and1((), ""), ())

        and2 = compile_for_llvm("and2",
                                "def and2(x, y): return x+1 and y+1",
                                level)
        self.assertEquals(and2(-1, "5"), 0)
        self.assertRaises(TypeError, and2, "5", 5)
        self.assertRaises(TypeError, and2,  5, "5")

    @at_each_optimization_level
    def test_or(self, level):
        or1 = compile_for_llvm("or1", "def or1(x, y): return x or y", level)
        self.assertEquals(or1("x", "y"), "x")
        self.assertEquals(or1((), "y"), "y")
        self.assertEquals(or1((), ""), "")

        or2 = compile_for_llvm("or2",
                               "def or2(x, y): return x+1 or y+1", level)
        self.assertEquals(or2(5, "5"), 6)
        self.assertRaises(TypeError, or2, "5", 5)
        self.assertRaises(TypeError, or2, -1, "5")

    @at_each_optimization_level
    def test_complex_or(self, level):
        complex_or = compile_for_llvm('complex_or', '''
def complex_or(a, b, c):
    return a or b or c
''', level)
        self.assertEquals(complex_or(1, 2, 0), 1)
        self.assertEquals(complex_or(1, 2, 3), 1)
        self.assertEquals(complex_or(3, 0, 0), 3)
        self.assertEquals(complex_or(0, 3, 0), 3)
        self.assertEquals(complex_or(0, 0, 1), 1)
        self.assertEquals(complex_or(0, 0, 0), 0)

        complex_or_and = compile_for_llvm('complex_or_and', '''
def complex_or_and(a, b, c):
    return a or b and c
''', level)
        self.assertEquals(complex_or_and(3, 0, 0), 3)
        self.assertEquals(complex_or_and("", 3, 0), 0)
        self.assertEquals(complex_or_and("", 0, 1), 0)
        self.assertEquals(complex_or_and(0, 3, 1), 1)

    @at_each_optimization_level
    def test_complex_and(self, level):
        complex_and = compile_for_llvm('complex_and', '''
def complex_and(a, b, c):
    return a and b and c
''', level)
        self.assertEquals(complex_and(3, 0, ""), 0)
        self.assertEquals(complex_and(3, 2, 0), 0)
        self.assertEquals(complex_and(3, 2, 1), 1)
        self.assertEquals(complex_and(0, 3, 2), 0)
        self.assertEquals(complex_and(3, 0, 2), 0)

        complex_and_or = compile_for_llvm('complex_and_or', '''
def complex_and_or(a, b, c):
    return a and b or c
''', level)
        self.assertEquals(complex_and_or(3, "", 0), 0)
        self.assertEquals(complex_and_or(1, 3, 0), 3)
        self.assertEquals(complex_and_or(1, 3, 2), 3)
        self.assertEquals(complex_and_or(0, 3, 1), 1)

    @at_each_optimization_level
    def test_break(self, level):
        break_one = compile_for_llvm("break_one", """
def break_one(x):
    for y in [1, 2]:
        x["break"] = y
        break
        x["post break"] = y
    else:
        x["else"] = True
    return x
""", level)
        self.assertEqual(break_one({}), {"break": 1})

        nested = compile_for_llvm("nested", """
def nested(x):
    for y in [1, 2]:
        for z in [3, 4]:
            x["break"] = z
            break
            x["post break"] = z
        else:
            x["inner else"] = True
        x["outer"] = y
    else:
        x["else"] = True
    return x
""", level)
        self.assertEqual(nested({}), {"break": 3, "outer": 2, "else": True})

    @at_each_optimization_level
    def test_continue(self, level):
        # CONTINUE_LOOP is only used inside a try/except block. Otherwise,
        # the continue statement is lowered to a JUMP_ABSOLUTE.
        continue_one = compile_for_llvm("continue_one", """
def continue_one(x):
    for y in [1, 2]:
        if y:
            x["continue"] = y
            try:
                continue
            except:
                pass
            finally:
                x["finally"] = y
        1 / 0
    else:
        x["else"] = True
    return x
""", level)
        self.assertEqual(continue_one({}), {"continue": 2, "else": True,
                                            "finally": 2})

        nested = compile_for_llvm("nested", """
def nested(x):
    for y in [1, 2]:
        for z in [3, 4]:
            if z:
                x["continue"] = z
                try:
                    continue
                except:
                    pass
            1 / 0
        else:
            x["inner else"] = True
        x["outer"] = y
    else:
        x["else"] = True
    return x
""", level)
        self.assertEqual(nested({}), {"continue": 4, "outer": 2,
                                      "inner else": True, "else": True})

    @at_each_optimization_level
    def test_load_attr(self, level):
        load_attr = compile_for_llvm('load_attr',
                                     'def load_attr(o): return o.attr',
                                     level)
        load_attr.attr = 1
        self.assertEquals(load_attr(load_attr), 1)
        self.assertRaises(AttributeError, load_attr, object())

    @at_each_optimization_level
    def test_store_attr(self, level):
        store_attr = compile_for_llvm('store_attr',
                                      'def store_attr(o): o.attr = 2',
                                      level)
        store_attr(store_attr)
        self.assertEquals(store_attr.attr, 2)
        self.assertRaises(AttributeError, store_attr, object())

    @at_each_optimization_level
    def test_delete_attr(self, level):
        delete_attr = compile_for_llvm('delete_attr',
                                       'def delete_attr(o): del o.attr',
                                       level)
        delete_attr.attr = 3
        delete_attr(delete_attr)
        self.assertFalse(hasattr(delete_attr, 'attr'))
        self.assertRaises(AttributeError, delete_attr, object())

    @at_each_optimization_level
    def test_call_varargs(self, level):
        f1 = compile_for_llvm("f1", "def f1(x, args): return x(*args)",
                              level)
        def receiver1(a, b):
            return a, b
        self.assertEquals(f1(receiver1, (1, 2)), (1, 2))
        self.assertRaises(TypeError, f1, receiver1, None)
        self.assertRaises(TypeError, f1, None, (1, 2))

        f2 = compile_for_llvm("f2",
                              "def f2(x, args): return x(1, 2, *args)",
                              level)
        def receiver2(a, *args):
            return a, args
        self.assertEquals(f2(receiver2, (3, 4, 5)), (1, (2, 3, 4, 5)))

    @at_each_optimization_level
    def test_call_kwargs(self, level):
        f = compile_for_llvm("f",
                             "def f(x, kwargs): return x(a=1, **kwargs)",
                             level)
        def receiver(**kwargs):
            return kwargs
        self.assertEquals(f(receiver, {'b': 2, 'c': 3}),
                          {'a': 1, 'b': 2, 'c': 3})

    @at_each_optimization_level
    def test_call_args_kwargs(self, level):
        f = compile_for_llvm("f", """
def f(x, args, kwargs):
    return x(1, d=4, *args, **kwargs)
""", level)
        def receiver(*args, **kwargs):
            return args, kwargs
        self.assertEquals(f(receiver, (2, 3), {'e': 5, 'f': 6}),
                          ((1, 2, 3), {'d': 4, 'e': 5, 'f': 6}))

    @at_each_optimization_level
    def test_varargs_func(self, level):
        f = compile_for_llvm("f", """
def f(*args):
    return zip(*args)
""", level)
        self.assertEqual(f([1], [2]), [(1, 2)])

    @at_each_optimization_level
    def test_kwargs_func(self, level):
        f = compile_for_llvm("f", """
def f(**kwargs):
    return dict(**kwargs)
""", level)
        self.assertEqual(f(a=3), {"a": 3})

    @at_each_optimization_level
    def test_build_slice(self, level):
        class Sliceable(object):
            def __getitem__(self, item):
                return item
        # Test BUILD_SLICE_TWO; make sure we didn't swap arguments.
        slice_two = compile_for_llvm('slice_two',
                                      'def slice_two(o): return o[1:2:]',
                                      level)
        self.assertEquals(slice_two(Sliceable()), slice(1, 2, None))
        # Test BUILD_SLICE_THREE.
        slice_three = compile_for_llvm('slice_three',
                                     'def slice_three(o): return o[1:2:3]',
                                     level)
        self.assertEquals(slice_three(Sliceable()), slice(1, 2, 3))
        # No way to make BUILD_SLICE_* raise exceptions.

    @at_each_optimization_level
    def test_with(self, level):
        class SimpleCM(object):
            def __enter__(self):
                self.enter = True

            def __exit__(self, *exc_info):
                self.exit = True

        with_simple = compile_for_llvm('with_simple', '''
def with_simple(self, x):
    with x:
        self.assertEqual(x.__dict__, {"enter": True})
    self.assertEqual(x.__dict__, {"enter": True, "exit": True})
''', level)
        with_simple(self, SimpleCM())

        with_raise = compile_for_llvm('with_raise', '''
def with_raise(self, x):
    with x:
        self.assertEqual(x.__dict__, {"enter": True})
        raise ArithmeticError
''', level)
        x = SimpleCM()
        self.assertRaises(ArithmeticError, with_raise, self, x)
        self.assertEqual(x.__dict__, {'enter': True, 'exit': True})

        with_return = compile_for_llvm('with_return', '''
def with_return(self, x):
    with x:
        self.assertEqual(x.__dict__, {"enter": True})
        return 55
''', level)
        x = SimpleCM()
        self.assertEqual(with_return(self, x), 55)
        self.assertEqual(x.__dict__, {'enter': True, 'exit': True})

        class SwallowingCM(object):
            def __enter__(self):
                self.enter = True

            def __exit__(self, *exc_info):
                self.exit = True
                return True  # Swallow the raised exception

        with_swallow = compile_for_llvm('with_swallow', '''
def with_swallow(self, x):
    with x:
        self.assertEqual(x.__dict__, {"enter": True})
        raise ArithmeticError
    return 55
''', level)
        x = SwallowingCM()
        self.assertEqual(with_swallow(self, x), 55)
        self.assertEqual(x.__dict__, {'enter': True, 'exit': True})

        class BustedExitCM(object):
            def __enter__(self):
                self.enter = True

            def __exit__(self, *exc_info):
                self.exit = True
                class A(object):
                    def __nonzero__(self):
                        raise ArithmeticError
                return A()  # Test error paths in WITH_CLEANUP

        with_ignored_bad_exit = compile_for_llvm('with_ignored_bad_exit', '''
def with_ignored_bad_exit(self, x):
    with x:
        self.assertEqual(x.__dict__, {"enter": True})
    return 55
''', level)
        x = BustedExitCM()
        self.assertEqual(with_ignored_bad_exit(self, x), 55)
        self.assertEqual(x.__dict__, {'enter': True, 'exit': True})

        with_bad_exit = compile_for_llvm('with_bad_exit', '''
def with_bad_exit(self, x):
    with x:
        self.assertEqual(x.__dict__, {"enter": True})
        raise KeyError
    return 55
''', level)
        x = BustedExitCM()
        self.assertRaises(ArithmeticError, with_bad_exit, self, x)
        self.assertEqual(x.__dict__, {'enter': True, 'exit': True})

        class RaisingCM(object):
            def __enter__(self):
                self.enter = True

            def __exit__(self, *exc_info):
                self.exit = True
                raise ArithmeticError

        with_error = compile_for_llvm('with_error', '''
def with_error(self, x):
    with x:
        return 55
''', level)
        x = RaisingCM()
        self.assertRaises(ArithmeticError, with_error, self, x)
        self.assertEqual(x.__dict__, {'enter': True, 'exit': True})

        class NestedCM(object):
            def __init__(self):
                self.enter = 0
                self.exit = 0

            def __enter__(self):
                self.enter += 1

            def __exit__(self, *exc_info):
                self.exit += 1

        with_yield = compile_for_llvm('with_yield', '''
def with_yield(self, x):
    with x:
        self.assertEqual(x.__dict__, {"enter": 1, "exit": 0})
        yield 7
        self.assertEqual(x.__dict__, {"enter": 1, "exit": 0})
        yield 8
    self.assertEqual(x.__dict__, {"enter": 1, "exit": 1})
''', level)
        x = NestedCM()
        self.assertEqual(list(with_yield(self, x)), [7, 8])

        with_nested = compile_for_llvm('with_nested', '''
def with_nested(self, x):
    with x:
        self.assertEqual(x.__dict__, {"enter": 1, "exit": 0})
        with x:
            self.assertEqual(x.__dict__, {"enter": 2, "exit": 0})
        self.assertEqual(x.__dict__, {"enter": 2, "exit": 1})
    self.assertEqual(x.__dict__, {"enter": 2, "exit": 2})
''', level)
        x = NestedCM()
        with_nested(self, x)

    @at_each_optimization_level
    def test_raise(self, level):
        raise_onearg = compile_for_llvm('raise_onearg', '''
def raise_onearg(x):
    raise x
''', level)
        self.assertRaises(OpExc, raise_onearg, OpExc);

        raise_twoargs = compile_for_llvm('raise_twoargs', '''
def raise_twoargs(x, y):
    raise x, y
''', level)
        self.assertRaisesWithArgs(OpExc, ('twoarg',),
                                  raise_twoargs, OpExc, OpExc('twoarg'));

    @at_each_optimization_level
    def test_reraise(self, level):
        raise_noargs = compile_for_llvm('raise_noargs', '''
def raise_noargs():
    raise
''', level)
        exc = OpExc('exc')
        def setup_traceback(e):
            raise e
        try:
            setup_traceback(exc)
        except OpExc:
            # orig_tb and exc re-used for raise_threeargs.
            orig_tb = sys.exc_info()[2]
            try:
                raise_noargs()
            except OpExc, e:
                new_tb = sys.exc_info()[2]
                # Test that we got the right exception and the right
                # traceback. Test both equality and identity for more
                # convenient error displays when things aren't as expected.
                self.assertEquals(e, exc)
                self.assertTrue(e is exc)
                self.assertEquals(new_tb.tb_next, orig_tb)
                self.assertTrue(new_tb.tb_next is orig_tb)
            else:
                self.fail('expected OpExc exception')
        else:
            self.fail('expected OpExc exception')

        raise_threeargs = compile_for_llvm('raise_threeargs', '''
def raise_threeargs(x, y, z):
    raise x, y, z
''', level)
        # Explicit version of the no-args raise.
        try:
            # Re-using exc and orig_tb from raise_noargs.
            raise_threeargs(OpExc, exc, orig_tb)
        except OpExc, e:
            new_tb = sys.exc_info()[2]
            self.assertEquals(e, exc)
            self.assertTrue(e is exc)
            self.assertEquals(new_tb.tb_next, orig_tb)
            self.assertTrue(new_tb.tb_next is orig_tb)
        else:
            self.fail('expected OpExc exception')

    @at_each_optimization_level
    def test_complex_reraise(self, level):
        reraise = compile_for_llvm('reraise', '''
def reraise(raiser, exctype):
    try:
        raiser()
    except:
        try:
            raise
        except exctype:
            return "inner"
        return "middle"
    return "outer"
''', level)
        def raiser():
            raise ZeroDivisionError
        self.assertEquals(reraise(raiser, ZeroDivisionError),
                          "inner")
        self.assertRaises(ZeroDivisionError, reraise, raiser, TypeError)

    @at_each_optimization_level
    def test_simple_yield(self, level):
        generator = compile_for_llvm("generator", """
def generator():
    yield 1
    yield 2
    yield 3
""", level)
        g = generator()
        self.assertEquals(1, g.next())
        self.assertEquals(2, g.next())
        self.assertEquals(3, g.next())
        self.assertRaises(StopIteration, g.next)

    @at_each_optimization_level
    def test_yield_in_loop(self, level):
        generator = compile_for_llvm("generator", """
def generator(x):
    for i in x:
        yield i
""", level)
        g = generator([1, 2, 3, 4])
        self.assertEquals([1, 2, 3, 4], list(g))

        cross_product = compile_for_llvm("cross_product", """
def cross_product(x, y):
    for i in x:
        for j in y:
            yield (i, j)
""", level)
        g = cross_product([1, 2], [3, 4])
        self.assertEquals([(1,3), (1,4), (2,3), (2,4)], list(g))

    @at_each_optimization_level
    def test_yield_saves_block_stack(self, level):
        generator = compile_for_llvm("generator", """
def generator(x):
    yield "starting"
    for i in x:
        try:
            try:
                1 / i
                yield ("survived", i)
            finally:
                yield ("finally", i)
        except ZeroDivisionError:
            yield "caught exception"
    yield "done looping"
""", level)
        self.assertEquals(list(generator([0, 1, 2])),
                          ["starting",
                           ("finally", 0),
                           "caught exception",
                           ("survived", 1),
                           ("finally", 1),
                           ("survived", 2),
                           ("finally", 2),
                           "done looping"])

    @at_each_optimization_level
    def test_generator_send(self, level):
        generator = compile_for_llvm("generator", """
def generator():
    yield (yield 1)
""", level)
        g = generator()
        self.assertEquals(1, g.next())
        self.assertEquals("Hello world", g.send("Hello world"))
        self.assertRaises(StopIteration, g.send, 3)

    @at_each_optimization_level
    def test_generator_throw(self, level):
        generator = compile_for_llvm("generator", """
def generator(obj):
    try:
        yield "starting"
    except ArithmeticError:
        obj["caught"] = 1
    finally:
        obj["finally"] = 1
    yield "done"
""", level)
        obj = {}
        g = generator(obj)
        self.assertEquals("starting", g.next())
        self.assertEquals("done", g.throw(ArithmeticError))
        self.assertEquals(None, g.close())
        self.assertEquals({"caught": 1, "finally": 1}, obj)

        obj = {}
        g = generator(obj)
        self.assertEquals("starting", g.next())
        self.assertRaises(UnboundLocalError, g.throw, UnboundLocalError)
        self.assertRaises(StopIteration, g.next)
        self.assertEquals({"finally": 1}, obj)

    # Getting this to work under -Xjit=always is a pain in the ass, and not
    # worth the effort IMHO.
    if _llvm.get_jit_control() != "always":
        def test_toggle_generator(self):
            # Toggling between native code and the interpreter between yields
            # used to cause crashes because f_lasti doesn't get translated
            # between the scheme used for LLVM and the scheme used for the
            # interpreter.  Currently, due to our generator pseudo
            # on-stack-replacement, these assignments take effect on generator
            # reentry.
            def generator():
                yield 1
                generator.func_code.co_use_jit = True
                yield 2
                generator.func_code.co_use_jit = False
                yield 3
            self.assertEqual(list(generator()), [1, 2, 3])

    @at_each_optimization_level
    def test_closure(self, level):
        make_closure = compile_for_llvm('make_closure', '''
def make_closure(a, level):
    b = 5
    c = 3
    def inner(d, e=5):
        c = d + 1
        return a, b, c, d, e
    if level is not None:
        inner.__code__.co_use_jit = True
        inner.__code__.co_optimization = level
    b = 2
    return inner
''', level)
        inner = make_closure(1, level)
        self.assertEquals(inner(4), (1, 2, 5, 4, 5))
        self.assertRaises(TypeError, inner, "5")

    @at_each_optimization_level
    def test_closure_unbound_freevar(self, level):
        unbound_freevar = compile_for_llvm('unbound_freevar', '''
def unbound_freevar(level):
    if 0:
        b = 2
    def inner():
        return b
    if level is not None:
        inner.__code__.co_use_jit = True
        inner.__code__.co_optimization = level
    return inner
''', level)
        inner = unbound_freevar(level)
        self.assertRaisesWithArgs(NameError,
            ("free variable 'b' referenced before "
             "assignment in enclosing scope",), inner)

    @at_each_optimization_level
    def test_closure_unbound_local(self, level):
        unbound_local = compile_for_llvm('unbound_local', '''
def unbound_local(level):
    def inner():
        if 0:
            b = 3
        return b
    if level is not None:
        inner.__code__.co_use_jit = True
        inner.__code__.co_optimization = level
    return inner
''', level)
        inner = unbound_local(level)
        self.assertRaisesWithArgs(UnboundLocalError,
            ("local variable 'b' referenced before assignment",), inner)

    @at_each_optimization_level
    def test_ends_with_unconditional_jump(self, level):
        foo = compile_for_llvm('foo', '''
from opcode import opmap
from types import CodeType, FunctionType
foo_bytecode = [
    opmap["JUMP_FORWARD"],  4, 0,  #  0  JUMP_FORWARD   4 (to 7)
    opmap["LOAD_CONST"],    0, 0,  #  3  LOAD_CONST     0 (1)
    opmap["RETURN_VALUE"],         #  6  RETURN_VALUE
    opmap["JUMP_ABSOLUTE"], 3, 0,  #  7  JUMP_ABSOLUTE  3
]
foo_code = CodeType(0, 0, 1, 0, "".join(chr(x) for x in foo_bytecode),
                    (1,), (), (), "<string>", "foo", 1, "")
foo = FunctionType(foo_code, globals())
''', level)
        self.assertEquals(1, foo())


class LoopExceptionInteractionTests(LlvmTestCase):
    @at_each_optimization_level
    def test_except_through_loop_caught(self, level):
        nested = compile_for_llvm('nested', '''
def nested(lst, obj):
    try:
        for x in lst:
            raise UnboundLocalError
    except:
        obj["x"] = 2
    # Make sure the block stack is ok.
    try:
        for x in lst:
            return x
    finally:
        obj["y"] = 3
''', level)
        obj = {}
        self.assertEquals(1, nested([1,2,3], obj))
        self.assertEquals({"x": 2, "y": 3}, obj)

    @at_each_optimization_level
    def test_except_in_loop(self, level):
        nested = compile_for_llvm('nested', '''
def nested(lst, obj):
    try:
        for x in lst:
            try:
                a = a
            except ZeroDivisionError:
                obj["x"] = 2
    except UnboundLocalError:
        obj["z"] = 4
    # Make sure the block stack is ok.
    try:
        for x in lst:
            return x
    finally:
        obj["y"] = 3
''', level)
        obj = {}
        self.assertEquals(1, nested([1,2,3], obj))
        self.assertEquals({"z": 4, "y": 3}, obj)

    @at_each_optimization_level
    def test_except_through_loop_finally(self, level):
        nested = compile_for_llvm('nested', '''
def nested(lst, obj):
    try:
        for x in lst:
            a = a
    finally:
        obj["x"] = 2
''', level)
        obj = {}
        self.assertRaises(UnboundLocalError, nested, [1,2,3], obj)
        self.assertEquals({"x": 2}, obj)

    @at_each_optimization_level
    def test_break_in_try(self, level):
        break_one = compile_for_llvm("break_one", """
def break_one(x):
    for y in [1, 2]:
        try:
            x["break"] = y
            break
            x["post break"] = y
        except ZeroDivisionError:
            x["except"] = 77
        finally:
            x["finally"] = y
    else:
        x["else"] = True

    # Make sure the block stack is ok.
    try:
        1 / 0
    except ZeroDivisionError:
        x["except"] = ZeroDivisionError
    return x
""", level)
        self.assertEqual(break_one({}), {"break": 1, "finally": 1,
                                         "except": ZeroDivisionError})


def cause_bail_on_next_line():
    sys.settrace(lambda *args: None)


class BailoutTests(ExtraAssertsTestCase):
    @at_each_optimization_level
    def test_bail_inside_loop(self, level):
        # We had a bug where the block stack in the compiled version
        # of the below function used contiguous small integers to
        # identify block handlers.  When we bailed out to the
        # interpreter, it expected an handler's identifier to be the
        # index of the opcode that started the code for the handler.
        # This mismatch caused a crash.
        loop = compile_for_llvm("loop", """
def loop():
    for i in [1, 2]:
        # Use try/finally to get "break" to emit a BREAK_LOOP opcode
        # instead of just jumping out of the loop.
        try:
            cause_bail_on_next_line()
            break
        finally:
            pass
    return i
""", level)
        orig_trace = sys.gettrace()
        try:
            self.assertEquals(loop(), 1)
        finally:
            sys.settrace(orig_trace)

    def test_use_correct_unwind_reason_when_bailing(self):…