o
    ۾i                     @   sp   d Z ddlZddlmZmZ ddlmZ ddlm	Z	 e
dZe
dZe
dg d	Zd
d Zdd ZdS )z
NRT specific optimizations
    N)defaultdictdeque)binding)cgutilsz,\s*(?:tail)?\s*call void @NRT_incref\((.*)\)z,\s*(?:tail)?\s*call void @NRT_decref\((.*)\)|)z[0-9]+:z,[\'"]?[-a-zA-Z$._0-9][-a-zA-Z$._0-9]*[\'"]?:z^definez^;\s*<label>c                    s|   dd }fdd}dd fddd	d
   fdddd g }|| D ]\}}|r4||}||7 }q*d |S )Nc                 s   s    g }t |  D ]5}|dr|rJ || q	|dr0|s#J || d|fV  g }q	|r8|| q	d|gfV  q	d S )Ndefine}TF)str
splitlines
startswithappend)modulecurline r   M/home/ubuntu/.local/lib/python3.10/site-packages/numba/core/runtime/nrtopt.py_extract_functions   s   



z7_remove_redundant_nrt_refct.<locals>._extract_functionsc                    s2   g } | D ]\}}|r|r|}||7 }q|S Nr   )
func_linesoutis_bbbb_lines)_extract_basic_blocks_process_basic_blockr   r   _process_function/   s   
z6_remove_redundant_nrt_refct.<locals>._process_functionc                 s   s    | d  ds
J | d  dsJ d| d gfV  g }| dd D ] }t|}|d ur<d|fV  g }d|gfV  q#|rC|| q#d|fV  d| d gfV  d S )Nr   r   r   F   T)r   	_regex_bbmatchr   )r   r   lnmr   r   r   r   7   s    



z:_remove_redundant_nrt_refct.<locals>._extract_basic_blocksc                    s    | } | } | S r   r   )r   )(_move_and_group_decref_after_all_increfs_prune_redundant_refct_opsr   r   r   J   s   z9_remove_redundant_nrt_refct.<locals>._process_basic_blockc                 s   sp    t | D ]0\}}t|}|d ur||dd fV  qt|}|d ur/|d |dfV  q|d d fV  qd S )Nr   )	enumerate_regex_increfr   group_regex_decref)r   numr   r    r   r   r   _examine_refct_opO   s   

z6_remove_redundant_nrt_refct.<locals>._examine_refct_opc                    s   t t}t t}t  | D ]2\}}}|r|rJ |r.|dkr& | q|| | q|rA|dkr: | q|| | q| D ]&\}}|| }tt|t|}	t|	D ]}
 |	   |
  q[qF fddt| D S )Nzi8* nullc                    s   g | ]
\}}| vr|qS r   r   ).0r'   r   	to_remover   r   
<listcomp>u   s    zS_remove_redundant_nrt_refct.<locals>._prune_redundant_refct_ops.<locals>.<listcomp>)r   r   setaddr   itemsminlenrangepoppopleftr#   )r   
incref_map
decref_mapr'   
incref_var
decref_varvardecopsincopsct_)r(   r*   r   r"   ]   s,   z?_remove_redundant_nrt_refct.<locals>._prune_redundant_refct_opsc                 S   s   d}t | D ]\}}t|d ur|d }qd}t | D ]\}}t|d ur+|d }qt||}g }g }| d | D ]}t|d urJ|| q;|| q;|| | |d   S )Nr   r   )r#   r$   r   r&   maxr   )r   last_incref_posposr   last_decref_poslast_posdecrefsheadr   r   r   r!   x   s$   
zM_remove_redundant_nrt_refct.<locals>._move_and_group_decref_after_all_increfs
)join)llvmirr   r   	processedis_funclinesr   )r(   r   r!   r   r"   r   _remove_redundant_nrt_refct   s   

rK   c                 C   sR   z|  d W n ty   |  Y S w | j}tt| }t|}t||_|S )a  
    Remove redundant reference count operations from the
    `llvmlite.binding.ModuleRef`. This parses the ll_module as a string and
    line by line to remove the unnecessary nrt refct pairs within each block.
    Decref calls are moved after the last incref call in the block to avoid
    temporarily decref'ing to zero (which can happen due to hidden decref from
    alias).

    Note: non-threadsafe due to usage of global LLVMcontext
    
NRT_incref)	get_function	NameErrornamerK   r	   llparse_assemblyr   normalize_ir_text)	ll_modulerO   newllnew_modr   r   r   remove_redundant_nrt_refct   s   
rV   )__doc__recollectionsr   r   llvmliter   rP   
numba.corer   compiler$   r&   rF   r   rK   rV   r   r   r   r   <module>   s    

 