import dis
import re
import sys
from io import StringIO
import unittest
import ast
import types

def disassemble(func):
    f = StringIO()
    tmp = sys.stdout
    sys.stdout = f
    dis.dis(func)
    sys.stdout = tmp
    result = f.getvalue()
    f.close()
    return result

def dis_single(line):
    return disassemble(compile(line, '', 'single'))

def compile_fn(src, fnname='f'):
    """
    Compile src to a Module, return an (ast.FunctionDef, ast.Module) pair
    where the former is the named function
    """
    def _get_fn_single(mod, node, name):
        for f in node.body:
            if isinstance(f, (ast.ClassDef, ast.FunctionDef)):
                if f.name == name:
                    return f
        raise ValueError('Could not find name: %r' % name)
    mod = ast.parse(src)
    path = fnname.split('.')
    node = mod
    for name in path:
        node = _get_fn_single(mod, node, name)
    return node, mod

def get_code_for_fn(co, fnname):
    def _get_code_single(co, fnname):
        for const in co.co_consts:
            if isinstance(const, types.CodeType):
                if const.co_name == fnname:
                    return const
        raise ValueError('code for %r not found' % fnname)

    path = fnname.split('.')
    for name in path:
        co = _get_code_single(co, name)
    return co

class TestOptimization(unittest.TestCase):
    def assertHasLineWith(self, asm, items):
        def _has_items(line, items):
            for item in items:
                if item not in line:
                    return False
            return True

        for line in asm.splitlines():
            if _has_items(line, items):
                return
        raise ValueError('Did not find a line containing all of %r within %s'
                         % (items, asm))

    def test_custom_assert(self):
        with self.assertRaises(ValueError):
            self.assertHasLineWith('Hello world', ('foo', 'bar'))

class TestFramework(TestOptimization):
    # Ensure that invoking the optimizer is working:
    def test_eval(self):
        self.assertEqual(eval('42'), 42)

class TestInlining(TestOptimization):
    def assertIsInlinable(self, src, fnname='f'):
        from __optimizer__ import fn_is_inlinable
        fn, mod = compile_fn(src, fnname)
        self.assert_(fn_is_inlinable(fn, mod),
                     msg='Unexpectedly not inlinable: %s\n%r' % (src, ast.dump(fn)))

    def assertIsNotInlinable(self, src, fnname='f'):
        from __optimizer__ import fn_is_inlinable
        fn, mod = compile_fn(src, fnname)
        self.assert_(not fn_is_inlinable(fn, mod),
                     msg='Unexpectedly inlinable: %s\n%r' % (src, ast.dump(fn)))

    def compile_to_code(self, src, fnname):
        # Compile the given source code, returning the code object for the given function name
        co = compile(src, 'optimizable.py', 'exec')
        return get_code_for_fn(co, fnname)

    def test_simple_inlining(self):
        src = (
'''
def function_to_be_inlined(x, y, z):
    # this currently needs to exist to enable the inliner
    pass
sep = '-'
def f(x, y, z):
    return (2 * x * y) + (sep * z)
def callsite():
    f('foo', 3, 16)
''')
        self.assertIsInlinable(src)
        callsite = self.compile_to_code(src, 'callsite')
        asm = disassemble(callsite)
        self.assertIn('JUMP_IF_SPECIALIZABLE', asm)
        self.assertIn("('foofoofoofoofoofoo')", asm)
        self.assertIn("(sep)", asm)



    def test_copy_propagation(self):
        src = (
'''
def function_to_be_inlined(a, b, c, d):
    return(a + b + c + d)
def callsite():
    print(function_to_be_inlined('fee', 'fi', 'fo', 'fum'))
''')
        callsite = self.compile_to_code(src, 'callsite')
        asm = disassemble(callsite)
        self.assertHasLineWith(asm, ('LOAD_GLOBAL', '(__internal__.saved.function_to_be_inlined)'))
        self.assertIn('JUMP_IF_SPECIALIZABLE', asm)
        self.assertHasLineWith(asm, ('LOAD_CONST', "('feefifofum')"))

    def test_keyword_args(self):
        src = (
'''
def f(x, y, z=42):
    return 42
''')
        self.assertIsNotInlinable(src)

    def test_star_args(self):
        src = ('''
def f(x, y, *args):
    return [x, y] + args
''')
        self.assertIsNotInlinable(src)


    def test_kwargs(self):
        src = (
'''
def f(x, y, **kwargs):
    return kwargs
''')
        self.assertIsNotInlinable(src)

    def test_decorators(self):
        src = (
'''
@login_required
def f(x, y, z):
    return 42
''')
        self.assertIsNotInlinable(src)

    def test_imports(self):
        src = (
'''
def f(x, y, z):
    import os
    return 42
''')
        self.assertIsNotInlinable(src)

    def test_generators(self):
        src = (
'''
def f(x, y, z):
    yield 42
''')
        self.assertIsNotInlinable(src)

    def test_complex_return_logic(self):
        # Only simple use of "return" for now
        src = (
'''
def f(x, y, z):
    if x:
        return 42
''')
        self.assertIsNotInlinable(src)


    def test_free_and_cell_variables(self):
        # No free/cell variables
        src = (
'''
def f(x, y, z):
    def g(p, q):
        return x*p + y*q + z
    return g
''')
        self.assertIsNotInlinable(src)

    def test_reassignment_of_function_name(self):
        src = (
'''
def f(x):
    return 42
f = len
''')
        self.assertIsNotInlinable(src)



    def not_test_simple(self):
        src = '''
my_global = 34

def function_to_be_inlined(x, y, z):
    return (2 * x * y) + z + my_global
print(function_to_be_inlined(3, 4, 5))
'''
        asm = disassemble(src)
        print(asm)
        # FIXME: verify that the inlined callsite actually works!

    def test_double(self):
        src = '''
def function_to_be_inlined():
    pass
def f(x, y, z):
    return (2 * x * y) + z
def g(x, y, z):
    return (f(x, y, 3 * z)
            - f(x + 1, y * 5, z -2))
def h(x, y):
    return g(x, y, 0)
'''
        co = compile(src, 'optimizable.py', 'exec')
        g = get_code_for_fn(co, 'g')
        h = get_code_for_fn(co, 'h')
        co_asm = disassemble(co)
        g_asm = disassemble(g)
        h_asm = disassemble(h)
        # Verify co_asm:
        self.assertHasLineWith(co_asm,
                               ('STORE_GLOBAL', '(__internal__.saved.f)'))
        self.assertHasLineWith(co_asm,
                               ('STORE_GLOBAL', '(__internal__.saved.g)'))
        self.assertHasLineWith(co_asm,
                               ('STORE_GLOBAL', '(__internal__.saved.h)'))

        # Verify g_asm:
        self.assertHasLineWith(g_asm,
                               ('LOAD_GLOBAL', '(__internal__.saved.f)'))
        self.assertIn('JUMP_IF_SPECIALIZABLE', g_asm)

        # Verify h_asm:
        self.assertHasLineWith(h_asm,
                               ('LOAD_GLOBAL', '(__internal__.saved.g)'))
        self.assertHasLineWith(h_asm,
                               ('LOAD_GLOBAL', '(__internal__.saved.f)'))
        self.assertIn('JUMP_IF_SPECIALIZABLE', h_asm)

    def test_ignore_implicit_return(self):
        src = '''
my_global = 34
def function_to_be_inlined(x, y, z):
    print(2 * x * y) + z + my_global
def callsite():
    function_to_be_inlined(3, 4, 5)
'''
        callsite = self.compile_to_code(src, 'callsite')
        asm = disassemble(callsite)
        #print(asm)
        # FIXME: actually test it

    def test_call_call(self):
        # seen in copy.py:
        src = '''
def _copy_with_constructor(x):
    return type(x)(x)
'''
        fn = self.compile_to_code(src, '_copy_with_constructor')
        asm = disassemble(fn)

    def test_global(self):
        # Taken from tempfile.py:
        src = '''
tempdir = None

def gettempdir():
    """Accessor for tempfile.tempdir."""
    global tempdir
    if tempdir is None:
        _once_lock.acquire()
        try:
            if tempdir is None:
                tempdir = _get_default_tempdir()
        finally:
            _once_lock.release()
    return tempdir

def callsite():
    return gettempdir()
'''
        fn = self.compile_to_code(src, 'callsite')
        asm = disassemble(fn)
        #print(asm)
        self.assertIn('JUMP_IF_SPECIALIZABLE', asm)
        self.assertIn('(tempdir)', asm)

    def test_del_local(self):
        # Adapted from tempfile.py:
        src = (
'''
def f():
    fd = _os.open(filename, _bin_openflags, 0o600)
    fp = _io.open(fd, 'wb')
    fp.write(b'blat')
    fp.close()
    _os.unlink(filename)
    del fp, fd
    return dir
''')
        fn = self.compile_to_code(src, 'f')
        asm = disassemble(fn)

    def test_inplace(self):
        src = (
'''
def inplace(self):
    a = foo()
    a *= 42
''')
        fn = self.compile_to_code(src, 'inplace')
        asm = disassemble(fn)
        # FIXME: Ensure that the initial write to "a" isn't optimized away:
        #print(asm)

    def test_use_implicit_return(self):
        pass # TODO

    def test_predicates(self):
        src = '''
def function_to_be_inlined():
    pass

def is_crunchy(path):
    return True

def find_file(filename, std_dirs, paths):
    for dir in std_dirs:
        if is_crunchy(dir):
            f = os.path.join(sysroot, dir[1:], filename)
    return None
'''
        fn = self.compile_to_code(src, 'find_file')
        asm = disassemble(fn)
        self.assertIn('JUMP_IF_SPECIALIZABLE', asm)

    def test_useless_args(self):
        src = '''
def function_to_be_inlined():
    pass

def foo(a, b, c):
    bar(a, b)

def bar(a, b):
    baz(a)

def baz(a):
    pass

foo(1, 2, 3)
'''
        asm = disassemble(src)
        #print(asm)
        # FIXME

    def test_simple_recursion(self):
        src = '''
def function_to_be_inlined():
    return function_to_be_inlined()
'''
        fn = self.compile_to_code(src, 'function_to_be_inlined')
        asm = disassemble(fn)
        #print(asm)

    def test_mutual_recursion(self):
        src = '''
def function_to_be_inlined():
    return function_to_be_inlined()

def f():
    return g()

def g():
    return f()
'''
        fn = self.compile_to_code(src, 'f')
        asm = disassemble(fn)
        #print(asm)

    def test_method(self):
        src = '''
class Foo:
    def simple_method(self, a):
         return 'I am Foo.simple_method' + self.bar + a + self.baz

    def user_of_method(self):
         return self.simple_method(10)

    def overridable_method(self, a):
         return 'I am Foo.overridable_method'

class Bar(Foo):
    def user_of_base_class_method(self):
         return self.simple_method(10)

    def overridable_method(self, a):
         return 'I am Bar.overridable_method'

    def user_of_overridden_method(self, a):
         return self.overridable_method(a)
'''

        # Various methods should be inlinable:
        self.assertIsInlinable(src, fnname='Foo.simple_method')
        self.assertIsInlinable(src, fnname='Foo.overridable_method')
        self.assertIsInlinable(src, fnname='Bar.overridable_method')

        saved_foo_simple_method = '(__internal__.saved.Foo.simple_method)'
        saved_foo_overridable_method = '(__internal__.saved.Foo.overridable_method)'
        saved_bar_overridable_method = '(__internal__.saved.Bar.overridable_method)'


        # Ensure that we're saving the inlinable functions as global for use
        # by JUMP_IF_SPECIALIZABLE at callsites:
        fn = self.compile_to_code(src, 'Foo')
        asm = disassemble(fn)
        self.assertHasLineWith(asm, ('STORE_GLOBAL', saved_foo_simple_method))
        self.assertHasLineWith(asm, ('STORE_GLOBAL', saved_foo_overridable_method))
        if 0:
            print('\nFoo')
            print(asm)

        fn = self.compile_to_code(src, 'Bar')
        asm = disassemble(fn)
        self.assertHasLineWith(asm, ('STORE_GLOBAL', saved_bar_overridable_method))
        if 0:
            print('\nBar')
            print(asm)

        # Ensure that the various callsites use these globals:
        fn = self.compile_to_code(src, 'Foo.user_of_method')
        asm = disassemble(fn)
        if 0:
            print('\nFoo.user_of_method')
            print(asm)
        self.assertHasLineWith(asm, ('LOAD_GLOBAL', saved_foo_simple_method))
        self.assertIn('JUMP_IF_SPECIALIZABLE', asm)
        self.assertHasLineWith(asm, ('LOAD_CONST', "('I am Foo.simple_method')"))

        # Verify a subclass method that inlines a call to a base class method:
        fn = self.compile_to_code(src, 'Bar.user_of_base_class_method')
        asm = disassemble(fn)
        if 0:
            print('\nBar.user_of_base_class_method')
            print(asm)
        if 0:
            # FIXME: doesn't inline these yet:
            self.assertHasLineWith(asm, ('LOAD_GLOBAL', saved_foo_simple_method))
            self.assertIn('JUMP_IF_SPECIALIZABLE', asm)

        # Verify a subclass method that inlines a call to an overridden base class method:
        fn = self.compile_to_code(src, 'Bar.user_of_overridden_method')
        asm = disassemble(fn)
        if 0:
            print('\nBar.user_of_overridden_method')
            print(asm)
        self.assertHasLineWith(asm, ('LOAD_GLOBAL', saved_bar_overridable_method))
        self.assertIn('JUMP_IF_SPECIALIZABLE', asm)
        self.assertHasLineWith(asm, ('LOAD_CONST', "('I am Bar.overridable_method')"))

        # FIXME: execute it, verify both paths



    def test_namespaces(self):
        src = '''
class Foo:
    def simple_method(self, a):
         return self.bar + a + self.baz
'''
        fn = self.compile_to_code(src, 'Foo.simple_method')
        asm = disassemble(fn)
        #from __optimizer__ import get_dotted_name


def function_to_be_inlined():
    return 'I am the original implementation'

def new_version():
    return 'I am the new version'

def call_inlinable_function():
    return function_to_be_inlined()

class TestRebinding(TestOptimization):

    def test_rebinding(self):
        # "call_inlinable_function" should have an inlined copy
        # of "function_to_be_inlined":
        asm = disassemble(call_inlinable_function)
        #print(asm)
        # Should have logic for detecting if it can specialize:
        self.assertIn('JUMP_IF_SPECIALIZABLE', asm)
        self.assertHasLineWith(asm,
                               ('LOAD_GLOBAL', '(__internal__.saved.function_to_be_inlined)'))
        # Should have inlined constant value:
        self.assertHasLineWith(asm,
                               ('LOAD_CONST', "('I am the original implementation')"))

        # Try calling it:
        self.assertEquals(call_inlinable_function(),
                          'I am the original implementation')

        # Now rebind the wrapped function.
        # The easy way to do this:
        #   function_to_be_inlined = new_version
        # doesn't work: the optimizer spots it, and turns off inlining
        # for "function_to_be_inlined" throughout this module
        # Similarly, using "globals()" directly is spotted.

        __builtins__.globals()['function_to_be_inlined'] = new_version

        # Verify that we get the correct result:
        self.assertEquals(call_inlinable_function(),
                          'I am the new version')
        # FIXME; should also check bytecode instrumentation to verify that we
        # took the generalized implementation




def test_main(verbose=None):
    import sys
    from test import support
    test_classes = (TestFramework, TestInlining, TestRebinding)
    support.run_unittest(*test_classes)

    # verify reference counting
    if verbose and hasattr(sys, "gettotalrefcount"):
        import gc
        counts = [None] * 5
        for i in range(len(counts)):
            support.run_unittest(*test_classes)
            gc.collect()
            counts[i] = sys.gettotalrefcount()
        print(counts)

if __name__ == "__main__":
    test_main(verbose=False)
    pass