o
    ۰i                     @   s  d Z ddlZddlZddlZddlmZ ddlmZmZm	Z	m
Z
mZmZ ddlZddlZddlZddlmZ ddlZddlmZmZmZmZmZmZmZmZmZ ddlmZ G d	d
 d
eZ ej!G dd dZ"ej!G dd dZ#de
eef de$fddZ%dS )zFA simple program to transform bytecode into more readable source code.    N)CodeType)ListTupleDictUnionCallableOptional)defaultdict   )	nop_unreachable_bytecodenop_instructionadd_indentationremove_indentationremove_some_temppropagate_line_numsconvert_instructionsimplify_finally_statementInstruction)get_function_signaturec                       s*   e Zd ZdZd fdd	Zdd Z  ZS )DecompilationErrorz)Custom exception class for decompilation. c                    s   || _ t | j  d S N)messagesuper__init__)selfr   	__class__ D/home/ubuntu/.local/lib/python3.10/site-packages/depyf/decompiler.pyr   !   s   zDecompilationError.__init__c                 C   s   d| j  S )NzDecompilationError: )r   r   r   r   r   __str__%   s   zDecompilationError.__str__)r   )__name__
__module____qualname____doc__r   r!   __classcell__r   r   r   r   r      s    r   c                   @   sF   e Zd ZU dZeed< eed< dZeed< dZ	e
ed< dZe
ed< d	S )
DecompilerStatez_State of decompiler, keep track of the evaluation stack, as well as the decompiled source code.source_codestackFinside_looploop_start_indexloop_end_indexN)r"   r#   r$   r%   str__annotations__listr*   boolr,   intr-   r   r   r   r   r'   )   s   
 r'   c                   @   sR  e Zd ZU dZeed< dZeed< dZe	ed< e
jdd d	Zeed
< dZeed< ejdddZdefddZeZe Z Z Z Z Z Z Z Z Z Z Z Z Z!e Z" Z# Z$ Z% Z& Z' Z(Z)eZ*e Z+Z,eZ-e Z. Z/ Z0 Z1 Z2 Z3Z4eZ5defddZ6e6Z7e6 Z8 Z9 Z: Z; Z< Z=Z>e6Z?e6Z@e6ZAdefddZBdefddZCeC ZD ZE ZF ZG ZH ZI ZJZKdefddZLdefddZMeMZNdefddZOdefd d!ZPdefd"d#ZQdefd$d%ZRdefd&d'ZSdefd(d)ZTdefd*d+ZUdefd,d-ZVdefd.d/ZWeW ZX ZY ZZZ[defd0d1Z\defd2d3Z]defd4d5Z^defd6d7Z_defd8d9Z`e` Za ZbZce6Zddefd:d;Zedefd<d=Zfdefd>d?Zgdefd@dAZheh Zi Zj ZkZldefdBdCZmdefdDdEZnen Zo Zp Zq Zr Zs Zt Zu Zv Zw Zx Zy ZzZ{defdFdGZ|defdHdIZ}defdJdKZ~e~ Z Z Z Z Z Z Z Z Z Z Z ZZdefdLdMZdefdNdOZdefdPdQZdefdRdSZdefdTdUZdefdVdWZe Z Z ZZdefdXdYZdefdZd[Zdefd\d]Zdefd^d_Zdefd`daZdefdbdcZe ZZe ZZe ZZe ZZe ZZe ZZeZeZdefdddeZdefdfdgZeZdefdhdiZdefdjdkZe Z Z ZZdefdldmZdefdndoZdefdpdqZdefdrdsZdefdtduZdefdvdwZdefdxdyZdefdzd{Ze ZZdefd|d}Zdefd~dZdefddZdefddZdefddZdefddZdefddZe Z ZZdefddZe ZZdefddZe ZZdefddZe ZZdefddZdefddZdefddZdefddZdefddZdefddZdefddZe Z ZZdefddZdefddZdefddZdefddZdedefddZdefddZedee fddZdeeef fddZdd ZddefddZedd Zejdd			dde	dee	 de	fddZe			ddedede	dee	 def
ddĄZddƄ ZddȄ ZdS )
DecompilerzA decompiler for a code object.coder   
temp_count__temp_temp_prefixc                   C   s   t dg dS )Nr   r(   r)   )r'   r   r   r   r   <lambda>:   s    zDecompiler.<lambda>)default_factorystate   indentationFr+   c                 c   sT    t d||||d}| j}|jr|js|j|_|j|_|j|_|| _dV  || _dS )z"Create a new state for decompiler.r   )r(   r)   r*   r,   r-   N)r'   r;   r*   r,   r-   )r   r)   r*   r,   r-   r;   	old_stater   r   r   	new_state>   s   
zDecompiler.new_stateinstc                 C   s   t d|j )NzUnsupported instruction: )NotImplementedErroropnamer   r@   r   r   r   unimplemented_instructionN      z$Decompiler.unimplemented_instructionc                 C      d S r   r   rC   r   r   r   generic_nopi   s   zDecompiler.generic_nopc                 C   sD  d}zt t|j|jk}W n	 ty   Y nw |r&| jjt|j dS t|jtrZ|jj	}|jj
}| j jd7  _|  }| j j| d| d| d7  _| jj| dS |jdr| j jd7  _|  }| j j| d|j d7  _| jj| dS t|jtr| jj|j dS | jjd	|j d
 dS )zPush a constant onto the stack.
        `inst.argval` is the constant value, we have to use `repr` to get the source code
        Fzimport importlib
z = importlib.import_module("z").
ztorch.zimport torch
 = z'__co_consts[z]'N)evalreprargvalBaseExceptionr;   r)   append
isinstancetyper#   r"   r(   get_temp_nameargrepr
startswithr   arg)r   r@   can_reprmodulename	temp_namer   r   r   
LOAD_CONST   s.   "zDecompiler.LOAD_CONSTc                 C   sT   d|j v r| jjd |j dr | jj|jdd dS | jj|j dS )z-`inst.argval` is the variable name, in stringzNULL + N.	comp_arg_)rR   r;   r)   rN   rS   rL   replacerC   r   r   r   generic_load   s
   
zDecompiler.generic_loadc                 C   s   | j jd |   d S )Nzlocals())r;   r)   rN   replace_mutable_tos_with_temprC   r   r   r   LOAD_LOCALS   s   zDecompiler.LOAD_LOCALSc                 C   sH   | j j }| j j| d|j d|j d| d|j 	 |   d S )N[z] if 'z' in z else r;   r)   poprN   rL   r^   r   r@   tosr   r   r   LOAD_FROM_DICT_OR_GLOBALS   s
   $z$Decompiler.LOAD_FROM_DICT_OR_GLOBALSc           
      C   sp  t jdk r7| jj }zt|}W n	 ty   Y nw |dd }d|v r6| j jd| d7  _| 	 }n| 	 }| jj }|j
d@ rL| jj  |j
d@ rW| jj  |j
d	@ rb| jj ni }|j
d
@ ro| jj nd}t|syt|r}td | |j}d}| j|d
  jdkr| j|d
  j
}d}t|j|d}	| j j|	7  _|s| jj| d S |d	 S )N      rZ   r+   <z"original function name z is illegal, use a temp name."
   r<      r
   r   zNFunction with default arguments is not supported, ignore the default argumentsF
STORE_FASTT)overwite_fn_name)sysversion_infor;   r)   rb   rJ   	Exceptionsplitr(   rQ   rL   lenprintindex_ofoffsetinstructionsrB   r3   	decompilerN   )
r   r@   	qual_name	func_namer4   kw_defaultsdefaults
this_indeximmediately_used
inner_funcr   r   r   MAKE_FUNCTION   sD   


zDecompiler.MAKE_FUNCTIONc                 C   rF   r   r   rC   r   r   r   COPY_FREE_VARS   s   zDecompiler.COPY_FREE_VARSc                 C   s^   t | jj }|j}| r| jj| d|  d S | jjd| dt| d d S )NrZ   zgetattr(, ))r.   r;   r)   rb   rL   isidentifierrN   rK   )r   r@   lhsrhsr   r   r   	LOAD_ATTR   s
   $zDecompiler.LOAD_ATTRc              	   C   sV   | j j }| j j }| j j }| j j| d| d| d|j  |   d S )N(r   z).ra   )r   r@   self_objcls_obj	super_objr   r   r   LOAD_SUPER_ATTR   s   zDecompiler.LOAD_SUPER_ATTRc                 C   s&   | j j| j j  d|j  d S )NrZ   )r;   r)   rN   rb   rL   rC   r   r   r   LOAD_METHOD   s   &zDecompiler.LOAD_METHODc                 C   s   | j jd d S )NAssertionErrorr;   r)   rN   rC   r   r   r   LOAD_ASSERTION_ERROR   s   zDecompiler.LOAD_ASSERTION_ERRORc                 C      | j jd  d S r   r   rC   r   r   r   	PUSH_NULL     zDecompiler.PUSH_NULLc                 C   &   | j j }| j jd| d d S )Niter(r   r;   r)   rb   rN   rc   r   r   r   GET_ITER     zDecompiler.GET_ITERc                 C   s>   |j }| jj }||kr| j j| d| d7  _d S d S NrI   rH   rL   r;   r)   rb   r(   )r   r@   leftrightr   r   r   generic_store  s
    zDecompiler.generic_storec                 C   sJ   | j j }| j j }| j j }| j  j| d| d| d7  _d S )Nr`   ] = rH   r;   r)   rb   r(   )r   r@   indexxvaluer   r   r   STORE_SUBSCR  s   &zDecompiler.STORE_SUBSCRc              
   C   s\   | j j }| j j }| j j }| j j }| j  j| d| d| d| d7  _d S )Nr`   :r   rH   r   )r   r@   endstart	containerr   r   r   r   STORE_SLICE  s
   ,zDecompiler.STORE_SLICEc                 C   s@   | j j }| j j }| j  j| d|j d| d7  _d S )NrZ   rI   rH   r;   r)   rb   r(   rL   )r   r@   r   r   r   r   r   
STORE_ATTR%     (zDecompiler.STORE_ATTRc                 C   sV   | j j }| j j }| d| d| j jvr)| j  jd| d| d7  _d S d S )Nr`   ]del z]
r   )r   r@   r   r   r   r   r   DELETE_SUBSCR,  s
   "zDecompiler.DELETE_SUBSCRc                 C      | j  jd|j d7  _d S )Nr   rH   r;   r(   rL   rC   r   r   r   generic_delete5     zDecompiler.generic_deletec                 C   s0   | j j }| j  jd| d|j d7  _d S )Nr   rZ   rH   r   )r   r@   r   r   r   r   DELETE_ATTR=  s   $zDecompiler.DELETE_ATTRc              
   C   sh   |j dd }| jj }| jj }| j j| dt|j  d| d| d7  _| jj| d S )NrZ   r   z = __import__(z, fromlist=z, level=)
)rL   rq   r;   r)   rb   r(   rK   rN   )r   r@   rW   fromlistlevelr   r   r   IMPORT_NAMEB  s
   .zDecompiler.IMPORT_NAMEc                 C   sF   |j }| jjd }| j j| d| d| d7  _| jj| d S )Nr+   rI   rZ   rH   )rL   r;   r)   r(   rN   )r   r@   rW   rV   r   r   r   IMPORT_FROMJ  s   "zDecompiler.IMPORT_FROMc                 C   s<   ddddd|j  }| jjd| d| jj  d d S )	N-+~not)UNARY_NEGATIVEUNARY_POSITIVEUNARY_INVERT	UNARY_NOTr    r   )rB   r;   r)   rN   rb   )r   r@   opr   r   r   generic_unaryR  s   (zDecompiler.generic_unaryc                 C   s"   | j jd| j jd  d d S )Nzlen(r+   r   r   rC   r   r   r   GET_LEN]  s   "zDecompiler.GET_LENc                 C   sd   | j j }| j j }ddddddddd	d
dddd|j }| j jd| d| d| d d S )N*r   r   ///%**&|^<<>>@)BINARY_MULTIPLY
BINARY_ADDBINARY_SUBTRACTBINARY_TRUE_DIVIDEBINARY_FLOOR_DIVIDEBINARY_MODULOBINARY_POWER
BINARY_AND	BINARY_OR
BINARY_XORBINARY_LSHIFTBINARY_RSHIFTBINARY_MATRIX_MULTIPLYr   r   r   )r;   r)   rb   rB   rN   r   r@   r   r   r   r   r   r   generic_binarya  s&   &zDecompiler.generic_binaryc                 C   s6   | j j }| j j }| j j| d| d d S )Nr`   r   r   r   r@   r   r   r   r   r   BINARY_SUBSCRw  s   zDecompiler.BINARY_SUBSCRc                 C   sH   | j j }| j j }| j j }| j j| d| d| d d S )Nr`   r   r   r   )r   r@   r   r   r   r   r   r   BINARY_SLICE|  s   $zDecompiler.BINARY_SLICEc                 C   sr   | j j }| j j }ddddddddd	d
dddd|j }| j  j| d| d| d7  _| j j| d S )Nr   r   r   r   r   r   r   r   r   r   r   r   r   )INPLACE_MULTIPLYINPLACE_ADDINPLACE_SUBTRACTINPLACE_TRUE_DIVIDEINPLACE_FLOOR_DIVIDEINPLACE_MODULOINPLACE_POWERINPLACE_AND
INPLACE_ORINPLACE_XORINPLACE_LSHIFTINPLACE_RSHIFTINPLACE_MATRIX_MULTIPLYr   = rH   )r;   r)   rb   rB   r(   rN   r   r   r   r   generic_inplace_binary  s(   "z!Decompiler.generic_inplace_binaryc              	   C   s   | j j }| j j }d|jv r,| j  j| d|j d| d7  _| j j| d S | j jd| d|j d| d d S )N=r   rH   r   r   )r;   r)   rb   rR   r(   rN   r   r   r   r   	BINARY_OP  s   
$(zDecompiler.BINARY_OPc              	   C   s@   | j j }| j j }| j jd| d|j d| d d S )Nr   r   r   )r;   r)   rb   rN   rL   r   r   r   r   
COMPARE_OP  r   zDecompiler.COMPARE_OPc              	   C   P   | j j }| j j }|jdkrdnd}| j jd| d| d| d d S )Nr   iszis notr   r   r   r;   r)   rb   rL   rN   r   r   r   r   IS_OP     &zDecompiler.IS_OPc              	   C   r   )Nr   inznot inr   r   r   r   r   r   r   r   CONTAINS_OP  r   zDecompiler.CONTAINS_OPc                 C   s   | j  jd7  _d S )Nbreak
)r;   r(   rC   r   r   r   
BREAK_LOOP     zDecompiler.BREAK_LOOPc                 C   sb   |  }| |}| jjr/|| jjkr| j jd7  _d S || jjkr-| j jd7  _d S |S |S )Nr   	continue
)get_jump_targetrt   r;   r*   r-   r(   r,   )r   r@   jump_offset
jump_indexr   r   r   generic_abs_jump  s   
zDecompiler.generic_abs_jumpc                 C   s0   | j  jd| j jd  d7  _| j j  d S )Nreturn r+   rH   )r;   r(   r)   rb   rC   r   r   r   RETURN_VALUE  s    zDecompiler.RETURN_VALUEc                 C   r   )Nr  rH   r   rC   r   r   r   RETURN_CONST  r   zDecompiler.RETURN_CONSTc                 C   s6   t jdkr	td| j jd| jjd  d7  _d S )N)rg      z+YIELD_VALUE is not supported in Python 3.12zyield r+   rH   )rn   ro   rA   r;   r(   r)   rC   r   r   r   YIELD_VALUE  s
   
$zDecompiler.YIELD_VALUEc                 C   r   r   r   rC   r   r   r   RETURN_GENERATOR  r   zDecompiler.RETURN_GENERATORc                 C   s   |j dks	J dd S )Nr   z&Only generator expression is supported)rL   rC   r   r   r   	GEN_START  s   zDecompiler.GEN_STARTc                    s  |   |  }| |j}| jjd }| jj }| jj }d|jv r-d| d}n d|jv r9d| d}nd|jv rEd| d	}nd
|jv rM| }d|jv r[|  |  n	d|jv rd|  t| j	g}| jj
ru|| jj dtf fddfdd| j	|| D }|s| jjd }d|jv rd| d}n d|jv rd| d}nd|jv r| }nd
|jv rd| d	}d| dtd| j }	| j j|	7  _dS t|}
| |
}fdd| j	|| D }| t|}| j	|d  }d|jv sd|jv sd|jv rn;|}|t| j	k rI| j	| j}d|v s(d|v r-|d7 }nd|v r7||ks<d|v r=n|d7 }|t| j	k s|| t|}| |  | |d | | jj}t|| j}| jj }W d   n	1 s{w   Y  d| d| }	| j j|	7  _| | | || | jj}W d   n	1 sw   Y  |rt|| j}d| }| j j|7  _|| j_|S )u  How we support if-else:

        Failed idea: try to paritition the block of instructions into if and else.
        This is not possible, as the if-else block might have overlapping instructions.
        Take this function as an example:

        def f(a):
            b = 1 if a else 2
            print(b)
        
        The bytecode is:
  2           0 LOAD_FAST                0 (a)
              2 POP_JUMP_IF_FALSE        4 (to 8)
              4 LOAD_CONST               1 (1)
              6 JUMP_FORWARD             1 (to 10)
        >>    8 LOAD_CONST               2 (2)
        >>   10 STORE_FAST               1 (b)

  3          12 LOAD_GLOBAL              0 (print)
             14 LOAD_FAST                1 (b)
             16 CALL_FUNCTION            1
             18 POP_TOP
             20 LOAD_CONST               0 (None)
             22 RETURN_VALUE
        
        The instructions for if branch: 2, 4, 6, 10
        The instructions for else branch: 8, 10
        They share the same instruction 10, so we cannot partition the block into if and else.

        Another example:

        def f():
            g(arg1=a if a is not None else b, arg2=2)
            print(1)

        The bytecode is:

  2           0 LOAD_GLOBAL              0 (g)
              2 LOAD_GLOBAL              1 (a)
              4 LOAD_CONST               0 (None)
              6 IS_OP                    1
              8 POP_JUMP_IF_FALSE        7 (to 14)
             10 LOAD_GLOBAL              1 (a)
             12 JUMP_FORWARD             1 (to 16)
        >>   14 LOAD_GLOBAL              2 (b)
        >>   16 LOAD_CONST               1 (2)
             18 LOAD_CONST               2 (('arg1', 'arg2'))
             20 CALL_FUNCTION_KW         2
             22 POP_TOP

  3          24 LOAD_GLOBAL              3 (print)
             26 LOAD_CONST               3 (1)
             28 CALL_FUNCTION            1
             30 POP_TOP
             32 LOAD_CONST               0 (None)
             34 RETURN_VALUE

        The instructions for if branch: 8, 14, 16, 18, 20, 22
        The instructions for else branch: 10, 12, 16, 18, 20, 22
        They share the same instructions 16, 18, 20, 22, so we cannot partition the block into if and else.

        Current idea:

        We take advantage of the following fact:

        This code snippet:

        if cond:
            if-body
        else:
            else-body
        rest-body

        is equivalent to:

        if cond：
            if-body
            rest-body
        else:
            else-body
            rest-body
        
        By duplicating the rest-body, we can decompile the if-else block separately. And they will have some duplicated code.

        Of course, we don't want to duplicate too long code, so we need to find the end of if-else block.
        The current heuristic is to find the first store/return/jump/for-iter instruction after the if-else block (because they are indicators that we will generate meaningful source code).
        r+   IF_NOT_NONEr   z	 is None)IF_NONEz is not None)IF_TRUEz(not r   IF_FALSEPOP_JUMPOR_POPic                    s   |   o	|   kS r   )is_jumpr   )r  )r   r   r   qualified_jumpQ  r   z2Decompiler.generic_jump_if.<locals>.qualified_jumpc                       g | ]
} |r|  qS r   r   .0r  r  r   r   
<listcomp>T      z.Decompiler.generic_jump_if.<locals>.<listcomp>zif :
r   Nc                    r  r   r  r  r  r   r   r  k  r  r
   RAISERETURNSTOREJUMPFOR_ITERzelse:
)r   rt   ru   r;   r)   copyrB   rb   rr   rv   r*   rN   r-   r   r   r=   r(   maxminr?   decompile_range)r   r@   r   r|   condfallthrough_stack
jump_stackend_index_candidatesjump_targetsif_codemax_jumpmax_jump_indexall_jump_targets	last_instold_map_jump_indexrB   	end_indexif_bodyif_end_stack	else_body	else_coder   )r   r  r   generic_jump_if  s   X












$



zDecompiler.generic_jump_ifc                    sN  |  |j|  |   fddt| jD d }d}| | jj | d | | jj	}t
|| j}d| }W d    n1 sFw   Y  d}| | jj? fddt| jD }|rg|d   }| j|d   rv|d8 }| |d | | jj	}t
|| j}d	| }W d    n1 sw   Y  | j j	|| 7  _	 S )
Nc                    s6   g | ]\}}|j d kr|  kr k rn n|qS )	POP_BLOCKrB   r  r  r   r.  start_indexr   r   r    s    "z,Decompiler.SETUP_FINALLY.<locals>.<listcomp>r+   r   r
   ztry:
c                    s&   g | ]\}}|j d kr |kr|qS )END_FINALLYr5  r6  r8  r   r   r    s
    r   z	finally:
)rt   ru   r   	enumeraterv   r?   r;   r)   r"  r(   r   r=   r  )r   r@   pop_block_indextry_codefinally_codeend_finally_indexfinally_end_indexr   r7  r   SETUP_FINALLY  sB   


zDecompiler.SETUP_FINALLYc                    s  |  |jfddt| jD d }| t| j    d } | dkrE| j| }|jds7|jdkr;t| |d8 } | dks' fddt| jD }tjd	kr[|d }n|d
 }t	 |d D ]	}t| j|  qf| j
jd }|  }| j
j| d | j
j| d| d| d}	d}
| | j
jB | d   | j
j}
t|
| j}
|
 }g }|D ]}| d|v sd| v rq|| qddd |D }
W d   n1 sw   Y  | j
 j|	|
 7  _|d S )a  
        with expression as var:
            body

        is equivalent to:

        var = expression
        var.__enter__()
        try:
            body
        finally:
            var.__exit__()

        We find the start of `finally` by `WITH_EXCEPT_START`, and the end of `finally` by `POP_EXCEPT`.
        In early python version, the start is `WITH_CLEANUP_START` and the end is `WITH_CLEANUP_FINISH`.
        c                    &   g | ]\}}|j d v r| kr|qS ))WITH_EXCEPT_STARTWITH_CLEANUP_STARTr5  r6  r:  r   r   r        z)Decompiler.SETUP_WITH.<locals>.<listcomp>r+   r
   rk   r  PUSH_EXC_INFOc                    rB  ))
POP_EXCEPTWITH_CLEANUP_FINISHr5  r6  )r.  r   r   r    rE  rf   r   z	.__exit__zwith z as r  r   zNone(None, None)c                 S   s   g | ]}|d  qS )rH   r   r  r   r   r   r   r    s    N)rt   ru   r;  rv   r   rB   rS   rn   ro   ranger;   r)   rQ   rN   r?   r"  r(   r   r=   
splitlinesstripjoin)r   r@   with_except_indexr  _instpop_except_indicespop_except_indexrd   tempwith_clause	with_bodylinesansliner   r7  r   
SETUP_WITH  sV   


zDecompiler.SETUP_WITHc           	      C   s   |  |j}|  | }|  }d| d| jj  d}| jj| | j| }|	 r9| |jkr9|d7 }| j
| jjd||d" | |d | | jj}|t|| j }| jj }W d    n1 siw   Y  | j j|7  _|| j_|S )Nzfor  in r  r
   T)r*   r,   r-   )rt   ru   r   rQ   r;   r)   rb   rN   rv   r  r?   r"  r(   r   r=   r  )	r   r@   r8  r.  rX   for_coder,  r4   for_end_stackr   r   r   r    s"   
zDecompiler.FOR_ITERc                 C   sb   |j dkr	|j}n	dddd|j  }| jj| d  }|d g|d d  }|| jj| d < d S )NROT_Nrk   rg   r<   )ROT_TWO	ROT_THREEROT_FOURr+   )rB   rL   r;   r)   )r   r@   nvaluesr   r   r   rot_n  s   
zDecompiler.rot_nc                 C   sH   |j }| jjd }| jj|  }||}}|| jjd< || jj| < d S Nr+   )rL   r;   r)   )r   r@   r`  rd   r   r   r   r   SWAP&  s   
zDecompiler.SWAPc                 C   s8   |j }|dkr	d S | jjd|d   }| jj| d S )Nr   r+   r
   )rL   r;   r)   rN   )r   r@   r`  r   r   r   r   COPY.  s
   zDecompiler.COPYc                 C   s   | j j  d S r   r;   r)   rb   rC   r   r   r   POP_TOP8  rE   zDecompiler.POP_TOPc                 C   s   | j j| j jd  d S rc  r   rC   r   r   r   DUP_TOP;  s   zDecompiler.DUP_TOPc                 C   s8   | j jd }| j jd }| j j| | j j| d S )Nr+   r   r   r@   rd   tos1r   r   r   DUP_TOP_TWO?  s   zDecompiler.DUP_TOP_TWOc                 C   s$   | j j|j }| jjt| d S r   )r4   	co_constsrT   r;   r)   rN   rK   )r   r@   namesr   r   r   KW_NAMESG  s   zDecompiler.KW_NAMESc                    sl   fddj D }d}|r+|d jdks)t|dkr+|d jdkr+|d jdkr+d	}t }|r8tjj }fd
dt j	D }|d d d }|d t|t|  }|t|t| d  }g }t
||D ]\}	}
||	 d|
  qjjj }jjrjjd d u rjj  d|v r|g}jj }jj| dd||  d   d S )Nc                    s   g | ]
}|j  j k r|qS r   )ru   rI  )r@   r   r   r  L  r  z#Decompiler.CALL.<locals>.<listcomp>Fr+   ro  r
   ri  PRECALLTc                       g | ]} j j qS r   rf  r  _r    r   r   r  U      r   r   r   r   r   )rv   rB   rr   tuplerJ   r;   r)   rb   rJ  rL   ziprN   rM  r^   )r   r@   r,  has_kw_nameskw_namesargspos_argskwargskwcallsrW   r   funcr   )r@   r   r   CALLK  s8   $zDecompiler.CALLc                    s^    fddt |jD }|d d d } jj } jj| dd| d    d S )Nc                    rq  r   rf  rr  r    r   r   r  i  rt  z+Decompiler.generic_call.<locals>.<listcomp>r+   r   r   r   )rJ  rL   r;   r)   rb   rN   rM  r^   )r   r@   ry  r}  r   r    r   generic_callh  s
    zDecompiler.generic_callc           	         s   t  jj } fddtt|D }|  g }t||D ]\}}|| d|  q  fddt|j	t| D }|d d d } jj } jj| dd
||  d    d S )	Nc                    rq  r   rf  rr  r    r   r   r  s  rt  z/Decompiler.CALL_FUNCTION_KW.<locals>.<listcomp>r   c                    rq  r   rf  rr  r    r   r   r  x  s    r+   r   r   r   )rJ   r;   r)   rb   rJ  rr   reverserv  rN   rL   rM  r^   )	r   r@   kw_argskw_valsr|  rW   valrz  r}  r   r    r   CALL_FUNCTION_KWq  s   
$zDecompiler.CALL_FUNCTION_KWc                 C   s   |j dkr| jj }| jj }| jj| d| d n'|j dkrF| jj }| jj }| jj }| jj| d| d| d |   d S )Nr   z(*r   r
   z, **)rL   r;   r)   rb   rN   r^   )r   r@   ry  r}  r  r   r   r   CALL_FUNCTION_EX  s   

 zDecompiler.CALL_FUNCTION_EXc                 C   s   |j dv rd S |j dv rd S |j dkr,| j jd| jj  d7  _| jjd d S |j dkr@d| jjd	  | jjd	< d S |j d
krJ| |S d S )N)INTRINSIC_1_INVALIDINTRINSIC_IMPORT_STARINTRINSIC_STOPITERATION_ERRORINTRINSIC_ASYNC_GEN_WRAP)INTRINSIC_TYPEVARINTRINSIC_PARAMSPECINTRINSIC_TYPEVARTUPLEINTRINSIC_SUBSCRIPT_GENERICINTRINSIC_TYPEALIASINTRINSIC_PRINTzprint(r   NoneINTRINSIC_UNARY_POSITIVEr   r+   INTRINSIC_LIST_TO_TUPLE)rR   r;   r(   r)   rb   rN   LIST_TO_TUPLErC   r   r   r   CALL_INTRINSIC_1  s   


 


zDecompiler.CALL_INTRINSIC_1c                 C   s   | j j }g }t|jD ]	}||   qddd |D }| j  j|d| d 7  _|d d d D ]	}| j j| q6d S )Nr   c                 S   s   g | ]}| d qS )r   r   rI  r   r   r   r        z.Decompiler.UNPACK_SEQUENCE.<locals>.<listcomp>r   rH   r+   )	r;   r)   rb   rJ  rL   rN   rQ   rM  r(   )r   r@   varname	tmp_namesr  r   rW   r   r   r   UNPACK_SEQUENCE  s   zDecompiler.UNPACK_SEQUENCEc                 C   s   | j j }g }t|jD ]	}||   q|  }| j  jd|d|  d| d 7  _| j j| |d d d D ]	}| j j| q?d S )Nr   z, *rI   rH   r+   )	r;   r)   rb   rJ  rL   rN   rQ   r(   rM  )r   r@   r  r  r  	star_namerW   r   r   r   	UNPACK_EX  s   ,zDecompiler.UNPACK_EXc              	   C   s   | j j }| j j }|jdkr!| j jd| d| d d S |jdkr?| j j }| j jd| d| d| d d S d S )Nrk   zslice(r   r   rg   r   )r   r@   rd   rk  tos2r   r   r   BUILD_SLICE  s   
 
&zDecompiler.BUILD_SLICEc                    s    fddt |jD }|d d d }d|jv rdd |D }|jdkr3 jjd|d  d	 d S  jjdd
| d d S )Nc                    rq  r   rf  rr  r    r   r   r    rt  z*Decompiler.build_tuple.<locals>.<listcomp>r+   UNPACKc                 S      g | ]}d | qS r   r   r  rT   r   r   r   r    r  r
   r   r   z,)r   r   )rJ  rL   rB   r;   r)   rN   rM  r   r@   ry  r   r    r   build_tuple  s   

 zDecompiler.build_tuplec                    sf    fddt |jD }|d d d }d|jv rdd |D } jjdd| d    d S )	Nc                    rq  r   rf  rr  r    r   r   r    rt  z)Decompiler.build_list.<locals>.<listcomp>r+   r  c                 S   r  r  r   r  r   r   r   r    r  r`   r   r   )rJ  rL   rB   r;   r)   rN   rM  r^   r  r   r    r   
build_list  s   
zDecompiler.build_listc                    s~   d}|j dkr
d}n( fddt|j D }|d d d }d|jv r)dd |D }d	d
| d} jj|    d S )Nr   r   zset()c                    rq  r   rf  rr  r    r   r   r    rt  z(Decompiler.build_set.<locals>.<listcomp>r+   r  c                 S   r  r  r   r  r   r   r   r    r  {r   })rL   rJ  rB   rM  r;   r)   rN   r^   )r   r@   rV  ry  r   r    r   	build_set  s   

zDecompiler.build_setc                    sv   |j dkr jjd n( fddt|j D }|d d d }dd |D } jjdd| d	    d S )
Nr   zdict()c                    rq  r   rf  rr  r    r   r   r    rt  z/Decompiler.build_map_unpack.<locals>.<listcomp>r+   c                 S   r  )r   r   r  r   r   r   r    r  r  r   r  )rL   r;   r)   rN   rJ  rM  r^   r  r   r    r   build_map_unpack  s   
zDecompiler.build_map_unpackc              	      s~    fddt |jd D }|d d d }|d d d }|dd d } jjdddd t||D  d	    d S )
Nc                    rq  r   rf  rr  r    r   r   r    rt  z(Decompiler.BUILD_MAP.<locals>.<listcomp>rk   r+   r
   r  r   c                 S      g | ]\}}| d | qS z: r   r  kvr   r   r   r        r  )rJ  rL   r;   r)   rN   rM  rv  r^   )r   r@   ry  keysra  r   r    r   	BUILD_MAP  s    zDecompiler.BUILD_MAPc              	      sn   t  jj } fddt|jD }|d d d } jjdddd t||D  d  	  d S )Nc                    rq  r   rf  rr  r    r   r   r    rt  z2Decompiler.BUILD_CONST_KEY_MAP.<locals>.<listcomp>r+   r  r   c                 S   r  r  r   r  r   r   r   r    r  r  )
rJ   r;   r)   rb   rJ  rL   rN   rM  rv  r^   )r   r@   r  ry  ra  r   r    r   BUILD_CONST_KEY_MAP  s    zDecompiler.BUILD_CONST_KEY_MAPc                    sB    fddt |jD }|d d d }d|} jj| d S )Nc                    rq  r   rf  rr  r    r   r   r    rt  z+Decompiler.BUILD_STRING.<locals>.<listcomp>r+   z + )rJ  rL   rM  r;   r)   rN   )r   r@   ry  ra  r   r    r   BUILD_STRING  s   
zDecompiler.BUILD_STRINGc                 C   r   )Nztuple(r   r   )r   r@   itemr   r   r   r    r   zDecompiler.LIST_TO_TUPLEc                 C   s:   | j j }| j|jd}| j  j| d| d7  _d S )N)posz.extend(r   )r;   r)   rb   r^   rL   r(   r   r@   ra  rR  r   r   r   LIST_EXTEND  s    zDecompiler.LIST_EXTENDc                 C   T   |j dkr| j d7  _ | jj|j   }| jj }| j j| d| d7  _d S )Nr
   z.append(r   r   r   r@   r   r   r   r   r   LIST_APPEND  
   
 zDecompiler.LIST_APPENDc                 C   sF   |j dks	J d| jj }|  }| j j| d| d7  _d S )Nr
   zOnly tested for argval==1z.update(r   )rL   r;   r)   rb   r^   r(   r  r   r   r   generic_update%  s    zDecompiler.generic_updatec                 C   r  )Nr
   z.add(r   r   r  r   r   r   SET_ADD-  r  zDecompiler.SET_ADDc                 C   sv   | j j|j d  }tjdkr| j j }| j j }n| j j }| j j }| j  j| d| d| d7  _d S )Nr
   )rg   rj   z.__setitem__(r   r   )r;   r)   rL   rn   ro   rb   r(   )r   r@   r   r   keyr   r   r   MAP_ADD5  s   
&zDecompiler.MAP_ADDc                 C   s   |j dkr| j jd7  _d S |j dkr&| j jd| jj  d7  _d S |j dkrH| jj }| jj }| j jd| d| d7  _d S d S )Nr   zraise
r
   zraise rH   rk   z from )rL   r;   r(   r)   rb   rj  r   r   r   RAISE_VARARGSA  s   

$
"zDecompiler.RAISE_VARARGSc                 C   s   |j \}}|r#| jj }| jj }| jjd| d| d d S | jj }|d u r/tn|}| jj|j d| d d S )Nzformat(r   r   r   )rL   r;   r)   rb   rN   r.   r"   )r   r@   r}  spec	form_specr   r   r   r   FORMAT_VALUEK  s   
  zDecompiler.FORMAT_VALUEr   r   c              
   C   s   z+|}||k r)| j | }tt|jtj}|| |}|r|}n|d7 }||k sW d S W d S  tyD } ztd| d| jj |d }~ww )Nr
   z Failed to decompile instruction rY  )	rv   getattrr3   rB   rD   rp   r   r4   co_name)r   r   r   running_indexr@   methodoutputer   r   r   r"  W  s,   

zDecompiler.decompile_rangeru   c                 C   s4   t | jD ]\}}|j|kr|  S qtd| )Nz$Cannot find instruction with offset )r;  rv   ru   
ValueError)r   ru   idxr@   r   r   r   rt   i  s
   
zDecompiler.index_ofrv   c                 C   s   t | t| t| | d S r   )r   r   r   )r4   rv   r   r   r   cleanup_instructionso  s   zDecompiler.cleanup_instructionsc                 C   s`   t |rddlm} ||j}|| _tdd t|D }t	|| || _
tdg d| _d S )Nr   )get_code_ownerc                 s   s    | ]}t |V  qd S r   )r   rr  r   r   r   	<genexpr>z  s    z&Decompiler.__init__.<locals>.<genexpr>r   r8   )callabledepyf.utilsr  __code__r4   r0   disget_instructionsr3   r  rv   r'   r;   )r   r4   r  rv   r   r   r   r   u  s   
zDecompiler.__init__c                 C   s   t  jd7  _| j t j S )Nr
   )r3   r5   r7   r    r   r   r   rQ     s   zDecompiler.get_temp_namer
   r  c                 C   sD   | j j|  }|  }| j  j| d| d7  _|| j j| < |S r   )r;   r)   rQ   r(   )r   r  rV  rX   r   r   r   r^     s
   z(Decompiler.replace_mutable_tos_with_tempc                  C   s2   g } t jD ]}tt|tjtjur| | q| S r   )r  rB   r  r3   rD   rN   )opnamesr   r   r   r   supported_opnames  s   

zDecompiler.supported_opnamesN)maxsizerm   returnc           
   
   C   s  zd|| _ || _| dt| j | jj}tj	dddkr$t
|| j|}t| j|}tdd t| jD }|rBdd| d nd	}| jjrSd
d| jj d nd	}|| | }|t|| }|W S  tyl     ty }	 z
td| jj |	d }	~	ww )Nr   DEPYF_REMOVE_TEMP1c                 s   s     | ]}|j d kr|jV  qdS )STORE_GLOBALN)rB   rL   )r  r@   r   r   r   r    s    z'Decompiler.decompile.<locals>.<genexpr>zglobal r   rH   r   z	nonlocal zFailed to decompile )r=   r7   r"  rr   rv   r;   r(   osenvirongetr   r   r4   setr  r  rM  co_freevarsr   r   rp   r  )
r   r=   r7   rm   r(   headerglobal_namesglobal_statementsnonlocal_statementr  r   r   r   rw     sP   zDecompiler.decompilecode_to_decompilereference_codefilepath_templatec                    s  t | j||dd}ddlm} |||}|d u r%|j |d }d}n||dd  }|jr:||dd  }d}		 ||	 }tj	
|rot|d }
|
|
dd  }|jrc||dd  }||krj|
}n5|	d	7 }	n/|tj	jd
 dd  |d }t|d}|| W d    n1 sw   Y  nq=|tj	jd
 dd  ddlm} t||dd}||} fdd|D d }g }t|jD ])\}}t|tr|drt|dd
 d d
 }|| j|  q|| q|jt|d}|S )N__place_holder__)r=   r7   rm   r   )fix_irregular_codenonamer   Trr
   r+   rZ   w)collect_all_code_objectsexec)filenamemodec                    s   g | ]	}|j  kr|qS r   )r  rI  ry   r   r   r    s    z9Decompiler.decompile_and_compile_like.<locals>.<listcomp>__co_constsr`   )rm  )r3   rw   depyf.code_transformr  r  r\   findr  r  pathexistsopenreadrq   sepwriter  r  compiler;  rm  rO   r.   rS   r2   rN   ru  )r  r  r=   r7   r  srcr  r  src_bodycountexisting_codeexisting_code_bodyfr  transformed_codetransformed_codes!decompiled_and_compiled_back_codeupdated_constsr  r   r   r   r  r   decompile_and_compile_like  sT   	

z%Decompiler.decompile_and_compile_likec                 C   s
   t | jS r   )idr4   r    r   r   r   __hash__  s   
zDecompiler.__hash__c                 C   s   t | t |kS r   )hash)r   otherr   r   r   __eq__  rE   zDecompiler.__eq__)Fr+   r+   )r
   )r<   r6   N)r"   r#   r$   r%   r   r/   r5   r2   r7   r.   dataclassesfieldr;   r'   r=   
contextlibcontextmanagerr?   r   rD   GET_YIELD_FROM_ITERrG  rC  JUMP_IF_NOT_EXC_MATCHCHECK_EG_MATCHrF  PREP_RERAISE_STARrH  CALL_FINALLYPOP_FINALLYrD  SETUP_EXCEPTCHECK_EXC_MATCHCLEANUP_THROWGET_AWAITABLE	GET_AITER	GET_ANEXTEND_ASYNC_FORBEFORE_ASYNC_WITHSETUP_ASYNC_WITHSENDASYNC_GEN_WRAPCACHE
PRINT_EXPRCOPY_DICT_WITHOUT_KEYSIMPORT_STAR
YIELD_FROMSETUP_ANNOTATIONSLOAD_BUILD_CLASSMATCH_MAPPINGMATCH_SEQUENCE
MATCH_KEYSMATCH_CLASSCALL_INTRINSIC_2rG   EXTENDED_ARGNOPRESUME
SETUP_LOOPr4  rp  BEGIN_FINALLYr9  	MAKE_CELLRERAISEEND_FORrY   r]   	LOAD_FASTLOAD_FAST_AND_CLEARLOAD_FAST_CHECKLOAD_GLOBAL
LOAD_DEREF	LOAD_NAMELOAD_CLASSDEREFLOAD_CLOSUREr_   re   LOAD_FROM_DICT_OR_DEREFr   r   r   r   r   r   r   r   r   rl   r  STORE_DEREF
STORE_NAMEr   r   r   r   r   DELETE_NAMEDELETE_GLOBALDELETE_DEREFDELETE_FASTr   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r   r  JUMP_ABSOLUTEJUMP_FORWARDJUMP_BACKWARDJUMP_BACKWARD_NO_INTERRUPTr  r  r  r  r  r3  POP_JUMP_IF_TRUEPOP_JUMP_IF_FALSEPOP_JUMP_FORWARD_IF_TRUEPOP_JUMP_FORWARD_IF_FALSEPOP_JUMP_BACKWARD_IF_TRUEPOP_JUMP_BACKWARD_IF_FALSEPOP_JUMP_FORWARD_IF_NONEPOP_JUMP_FORWARD_IF_NOT_NONEPOP_JUMP_BACKWARD_IF_NONEPOP_JUMP_BACKWARD_IF_NOT_NONEJUMP_IF_TRUE_OR_POPJUMP_IF_FALSE_OR_POPPOP_JUMP_IF_NOT_NONEPOP_JUMP_IF_NONErA  rX  BEFORE_WITHr  rb  r\  r]  r^  r_  rd  re  rg  rh  rl  ro  r~  r  CALL_FUNCTIONCALL_METHODr  r  r  r  r  r  r  BUILD_TUPLEBUILD_TUPLE_UNPACKBUILD_TUPLE_UNPACK_WITH_CALLr  
BUILD_LISTBUILD_LIST_UNPACKr  	BUILD_SETBUILD_SET_UNPACKr  BUILD_MAP_UNPACKBUILD_MAP_UNPACK_WITH_CALLr  r  r  r  r  r  r  
SET_UPDATEDICT_UPDATE
DICT_MERGEr  r  r  r  r"  rt   staticmethodr   r  r   r   r   rQ   r^   r  	functools	lru_cacher   rw   r  r  r  r   r   r   r   r3   3   sB  
 4 " ,				44
 8B
	

	



#Er3   r4   r  c                 C   s   t |  S )aS  Decompile any callable or code object into Python source code.
    It is especially useful for some dynamically generated code, like ``torch.compile``,
    or ``dataclasses``.

    Example usage:

    .. code-block:: python

        from dataclasses import dataclass
        @dataclass
        class Data:
            x: int
            y: float

        import depyf
        print(depyf.decompile(Data.__init__))
        print(depyf.decompile(Data.__eq__))
    
    Output:

    .. code-block:: python

        def __init__(self, x, y):
            self.x = x
            self.y = y
            return None

        def __eq__(self, other):
            if other.__class__ is self.__class__:
                return (self.x, self.y) == (other.x, other.y)
            return NotImplemented
    
    The output source code is semantically equivalent to the function, but not syntactically the same. It verbosely adds many details that are hidden in the Python code. For example, the above output code of ``__init__`` explicitly returns ``None``, which is typically ignored.

    Another detail is that the output code of ``__eq__`` returns ``NotImplemented`` instead of raising ``NotImplemented`` exception when the types are different. At the first glance, it seems to be a bug. However, it is actually the correct behavior. The ``__eq__`` method should return ``NotImplemented`` when the types are different, so that the other object can try to compare with the current object. See `the Python documentation <https://docs.python.org/3/library/numbers.html#implementing-the-arithmetic-operations>`_ for more details.
    )r3   rw   )r4   r   r   r   rw     s   %rw   )&r%   rn   r  r  typesr   typingr   r   r   r   r   r   r	  inspectrc  collectionsr	   r  code_transformr   r   r   r   r   r   r   r   r   utilsr   rp   r   	dataclassr'   r3   r.   rw   r   r   r   r   <module>   s8     ,	         ]