# Copyright 2015 Google Inc. All Rights Reserved.
#
# Licensed 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 collections, operator

from ..jvm import arraytypes as arrays
from ..jvm import scalartypes as scalars
from ..jvm import mathops, jvmops
from ..treelist import TreeList
from .. import flags, dalvik


# The two main things we need type inference for are determining the types of
# primative values and arrays. Luckily, we don't care about actual classes in
# these cases, we just need to know whether it is int,float,reference, etc. to
# generate the correct bytecode instructions, which are typed in Java.
#
# One additional problem is that ART's implicit casts narrow the type instead of
# replacing it like regular checkcasts do. This means that there is no way to
# replicate the behavior in Java using normal casts unless you know which class
# is a subclass of another and which classes are interfaces. However, we want to
# be able to translate code without knowing about every other class that could be
# referenced by the application, so we make do with a hack.
#
# Variables subjected to implicit casting are marked as tainted. Whenever a
# tained value is used, it is explcitly checkcasted to the expected type. This
# isn't ideal since it will incorrectly throw in the cast of bad interface casts,
# but it's the best we can do without requiring knowledge of the whole inheritance
# hierarchy.

class TypeInfo:
    def __init__(self, prims, arrs, tainted):
        # copy on write
        self.prims = prims
        self.arrs = arrs
        self.tainted = tainted

    def _copy(self): return TypeInfo(self.prims.copy(), self.arrs.copy(), self.tainted.copy())
    def _get(self, reg): return self.prims[reg], self.arrs[reg], self.tainted[reg]

    def _set(self, reg, st, at, taint=False):
        self.prims[reg] = st
        self.arrs[reg] = at
        self.tainted[reg] = taint
        return self

    def move(self, src, dest, wide):
        new = self._copy()._set(dest, *self._get(src))
        if wide:
            new._set(dest+1, *self._get(src+1))
        return new

    def assign(self, reg, st, at=arrays.INVALID, taint=False):
        assert st is not None
        return self._copy()._set(reg, st, at, taint)

    def assign2(self, reg, st):
        assert st is not None
        at = arrays.INVALID
        return self._copy()._set(reg, st, at)._set(reg+1, scalars.INVALID, at)

    def assignFromDesc(self, reg, desc):
        st = scalars.fromDesc(desc)
        at = arrays.fromDesc(desc)
        if scalars.iswide(st):
            return self.assign2(reg, st)
        else:
            return self.assign(reg, st, at)

    def isSame(self, other):
        return (self.prims.data is other.prims.data and
            self.arrs.data is other.arrs.data and
            self.tainted.data is other.tainted.data)

def merge(old, new):
    temp = old._copy()
    temp.prims.merge(new.prims)
    temp.arrs.merge(new.arrs)
    temp.tainted.merge(new.tainted)
    return old if old.isSame(temp) else temp

def fromParams(method, num_regs):
    isstatic = method.access & flags.ACC_STATIC
    full_ptypes = method.id.getSpacedParamTypes(isstatic)
    offset = num_regs - len(full_ptypes)

    prims = TreeList(scalars.INVALID, operator.__and__)
    arrs = TreeList(arrays.INVALID, arrays.merge)
    tainted = TreeList(False, operator.__or__)

    for i, desc in enumerate(full_ptypes):
        if desc is not None:
            prims[offset + i] = scalars.fromDesc(desc)
            arrs[offset + i] = arrays.fromDesc(desc)
    return TypeInfo(prims, arrs, tainted)

_MATH_THROW_OPS = [jvmops.IDIV, jvmops.IREM, jvmops.LDIV, jvmops.LREM]
def pruneHandlers(all_handlers):
    result = collections.defaultdict(list)
    for instr, handlers in all_handlers.items():
        if not instr.type in dalvik.PRUNED_THROW_TYPES:
            continue
        # if math op, make sure it is int div/rem
        if instr.type == dalvik.BinaryOp:
            if mathops.BINARY[instr.opcode][0] not in _MATH_THROW_OPS:
                continue
        elif instr.type == dalvik.BinaryOpConst:
            if mathops.BINARY_LIT[instr.opcode] not in _MATH_THROW_OPS:
                continue

        types = set()
        for ctype, handler in handlers:
            # if multiple handlers with same catch type, only include the first
            if ctype not in types:
                result[instr].append((ctype, handler))
                types.add(ctype)
            # stop as soon as we reach a catch all handler
            if ctype == b'java/lang/Throwable':
                break
    return dict(result)

################################################################################
# Lots of instructions just return an object or int for type inference purposes
# so we have a single function for these cases
def visitRetObj(dex, instr, cur):
    return cur.assign(instr.args[0], scalars.OBJ)
def visitRetInt(dex, instr, cur):
    return cur.assign(instr.args[0], scalars.INT)

# Instruction specific callbacks
def visitMove(dex, instr, cur):
    return cur.move(instr.args[1], instr.args[0], wide=False)
def visitMoveWide(dex, instr, cur):
    return cur.move(instr.args[1], instr.args[0], wide=True)
def visitMoveResult(dex, instr, cur):
    return cur.assignFromDesc(instr.args[0], instr.prev_result)
def visitConst32(dex, instr, cur):
    val = instr.args[1] % (1<<32)
    if val == 0:
        return cur.assign(instr.args[0], scalars.ZERO, arrays.NULL)
    else:
        return cur.assign(instr.args[0], scalars.C32)
def visitConst64(dex, instr, cur):
    return cur.assign2(instr.args[0], scalars.C64)
def visitCheckCast(dex, instr, cur):
    at = arrays.fromDesc(dex.type(instr.args[1]))
    at = arrays.narrow(cur.arrs[instr.args[0]], at)
    return cur.assign(instr.args[0], scalars.OBJ, at)
def visitNewArray(dex, instr, cur):
    at = arrays.fromDesc(dex.type(instr.args[2]))
    return cur.assign(instr.args[0], scalars.OBJ, at)
def visitArrayGet(dex, instr, cur):
    arr_at = cur.arrs[instr.args[1]]
    if arr_at is arrays.NULL:
        # This is unreachable, so use (ALL, NULL), which can be merged with anything
        return cur.assign(instr.args[0], scalars.ALL, arrays.NULL)
    else:
        st, at = arrays.eletPair(arr_at)
        return cur.assign(instr.args[0], st, at)
def visitInstanceGet(dex, instr, cur):
    field_id = dex.field_id(instr.args[2])
    return cur.assignFromDesc(instr.args[0], field_id.desc)
def visitStaticGet(dex, instr, cur):
    field_id = dex.field_id(instr.args[1])
    return cur.assignFromDesc(instr.args[0], field_id.desc)

def visitUnaryOp(dex, instr, cur):
    _, _, st = mathops.UNARY[instr.opcode]
    if scalars.iswide(st):
        return cur.assign2(instr.args[0], st)
    else:
        return cur.assign(instr.args[0], st)

def visitBinaryOp(dex, instr, cur):
    _, st, _ = mathops.BINARY[instr.opcode]
    if scalars.iswide(st):
        return cur.assign2(instr.args[0], st)
    else:
        return cur.assign(instr.args[0], st)

FUNCS = {
    dalvik.ConstString: visitRetObj,
    dalvik.ConstClass: visitRetObj,
    dalvik.NewInstance: visitRetObj,
    dalvik.InstanceOf: visitRetInt,
    dalvik.ArrayLen: visitRetInt,
    dalvik.Cmp: visitRetInt,
    dalvik.BinaryOpConst: visitRetInt,

    dalvik.Move: visitMove,
    dalvik.MoveWide: visitMoveWide,
    dalvik.MoveResult: visitMoveResult,
    dalvik.Const32: visitConst32,
    dalvik.Const64: visitConst64,
    dalvik.CheckCast: visitCheckCast,
    dalvik.NewArray: visitNewArray,
    dalvik.ArrayGet: visitArrayGet,
    dalvik.InstanceGet: visitInstanceGet,
    dalvik.StaticGet: visitStaticGet,
    dalvik.UnaryOp: visitUnaryOp,
    dalvik.BinaryOp: visitBinaryOp,
}

CONTROL_FLOW_OPS = {dalvik.Goto, dalvik.If, dalvik.IfZ, dalvik.Switch}

def doInference(dex, method, code, bytecode, instr_d):
    # get exception handlers
    all_handlers = collections.defaultdict(list)
    for tryi in code.tries:
        for instr in code.bytecode:
            if tryi.start < instr.pos2 and tryi.end > instr.pos:
                all_handlers[instr] += tryi.catches
    all_handlers = pruneHandlers(all_handlers)

    types = {}
    types[0] = fromParams(method, code.nregs)
    dirty = {0}

    def doMerge(pos, new):
        # prevent infinite loops
        if pos not in instr_d:
            return

        if pos in types:
            old = types[pos]
            new = merge(old, new)
            if new is not old:
                types[pos] = new
                dirty.add(pos)
        else:
            types[pos] = new
            dirty.add(pos)

    while dirty: # iterate until convergence
        for instr in bytecode:
            if instr.pos not in dirty:
                continue

            dirty.remove(instr.pos)
            cur = types[instr.pos]
            itype = instr.type
            if itype in FUNCS:
                after = FUNCS[itype](dex, instr, cur)
            elif itype in CONTROL_FLOW_OPS:
                # control flow - none of these are in FUNCS
                result = after = after2 = cur
                if instr.implicit_casts is not None:
                    desc_ind, regs = instr.implicit_casts
                    for reg in regs:
                        st = cur.prims[reg] # could != OBJ if null
                        at = arrays.narrow(cur.arrs[reg], arrays.fromDesc(dex.type(desc_ind)))
                        result = result.assign(reg, st, at, taint=True)
                    # merge into branch if op = if-nez else merge into fallthrough
                    if instr.opcode == 0x39:
                        after2 = result
                    else:
                        after = result

                if instr.type == dalvik.Goto:
                    doMerge(instr.args[0], after2)
                elif instr.type == dalvik.If:
                    doMerge(instr.args[2], after2)
                elif instr.type == dalvik.IfZ:
                    doMerge(instr.args[1], after2)
                elif instr.type == dalvik.Switch:
                    switchdata = instr_d[instr.args[1]].switchdata
                    for offset in switchdata.values():
                        target = (instr.pos + offset) % (1<<32)
                        doMerge(target, cur)
            else:
                after = cur

            # these instructions don't fallthrough
            if instr.type not in (dalvik.Return, dalvik.Throw, dalvik.Goto):
                doMerge(instr.pos2, after)

            # exception handlers
            if instr in all_handlers:
                for ctype, handler in all_handlers[instr]:
                    doMerge(handler, cur)
    return types, all_handlers
