U
    iâËd¦u  ã                   @   sì  d dl Z d dlZd dlZd dlZd dlZd dlZd dlmZ d dl	m
Z
 d dlmZmZ d dlmZmZmZmZ d dlmZmZ d dlmZ d dlmZ d d	lmZ d dlmZ d dlm   m!  m"Z# d d
l$m%Z% d dl&m'Z'm(Z(m)Z)m*Z*m+Z+m,Z, d dl-m.Z. d dl/Z/eƒ Z0de0_e/ 1ej 2d¡ d¡Z3e/ 1e 4¡ dkd¡Z5G dd„ de6ƒZ7G dd„ de/j(ƒZ8G dd„ de/j(ƒZ9G dd„ de/j(ƒZ:G dd„ de(ƒZ;G dd„ de/j(ƒZ<e/ =e#j>d¡G dd „ d e'e(ƒƒZ?G d!d"„ d"e(ƒZ@eAd#krèe/ B¡  dS )$é    N©Únjit)Útypes)Úcompile_isolatedÚFlags)ÚrtsysÚnrtoptÚ_nrt_pythonÚnrt)Ú	intrinsicÚinclude_path)Ú	signature)Úimpl_ret_untracked)Úir©ÚNRT_get_api)ÚEnableNRTStatsMixinÚTestCaseÚtemp_directoryÚimport_dynamicÚskip_if_32bitÚrun_in_subprocess)Ú
cpu_targetTÚlinuxzlinux only test)Úi386Úx86_64zx86 only testc                   @   s    e Zd ZdZdd„ Zdd„ ZdS )ÚDummyr   c                 C   s   t | ƒ jd7  _d S ©Né   ©ÚtypeÚalive©Úself© r$   úM/home/sam/Atlas/atlas_env/lib/python3.8/site-packages/numba/tests/test_nrt.pyÚ__init__,   s    zDummy.__init__c                 C   s   t | ƒ jd8  _d S r   r   r"   r$   r$   r%   Ú__del__/   s    zDummy.__del__N)Ú__name__Ú
__module__Ú__qualname__r!   r&   r'   r$   r$   r$   r%   r   )   s   r   c                   @   s   e Zd ZdZdZdd„ ZdS )ÚTestNrtMemInfoNotInitializedzn
    Unit test for checking the use of the NRT fails if the
    initialization sequence has not been run.
    Fc              
   C   st   ddddœ}|  ¡ D ]Z\}}zH|  t¡}dt_tt|ƒ}||Ž  W 5 Q R X d}|  |t|jƒ¡ W 5 dt_X qd S )Nr$   )r$   r$   )r$   )ZlibraryÚmeminfo_newÚmeminfo_allocTFz'Runtime must be initialized before use.)	Úitemsr   Ú_initÚassertRaisesÚRuntimeErrorÚgetattrÚassertInÚstrÚ	exception)r#   ÚmethodsÚmethÚargsÚraisesÚfnÚmsgr$   r$   r%   Útest_init_fail:   s    þ
z+TestNrtMemInfoNotInitialized.test_init_failN)r(   r)   r*   Ú__doc__Z_numba_parallel_test_r<   r$   r$   r$   r%   r+   3   s   r+   c                       s\   e Zd ZdZ‡ fdd„Zdd„ Zdd„ Zdd	„ Zd
d„ Zdd„ Z	e
dd„ ƒZdd„ Z‡  ZS )ÚTestNrtMemInfoz2
    Unit test for core MemInfo functionality
    c                    s$   dt _t tj¡ tt| ƒ ¡  d S )Nr   )	r   r!   r   Ú
initializer   Útarget_contextÚsuperr>   ÚsetUpr"   ©Ú	__class__r$   r%   rB   R   s    zTestNrtMemInfo.setUpc                 C   s   t ƒ }|  t jd¡ d}t ||¡}|  |jd¡ ~|  t jd¡ | ¡  |  |jd¡ |  t jd¡ | ¡  |  |jd¡ ~|  t jd¡ d S )Nr   ì   þJ[= é   r   )r   ÚassertEqualr!   r   r,   ÚrefcountÚacquireÚrelease)r#   ÚdÚaddrÚmir$   r$   r%   Útest_meminfo_refct_1Y   s    z#TestNrtMemInfo.test_meminfo_refct_1c                 C   s¬   t ƒ }|  t jd¡ d}t ||¡}|  |jd¡ ~|  t jd¡ tdƒD ]}| ¡  qJ|  |jd¡ |  t jd¡ tdƒD ]}| ¡  q||  |jd¡ ~|  t jd¡ d S )Nr   rE   éd   ée   r   )	r   rG   r!   r   r,   rH   ÚrangerI   rJ   )r#   rK   rL   rM   ÚctÚ_r$   r$   r%   Útest_meminfo_refct_2j   s     

z#TestNrtMemInfo.test_meminfo_refct_2c                 C   s¾   t ƒ }|  t jd¡ d}t ||¡}|  |jd¡ t|ƒ}|  |jd¡ |  ||j¡ |  |j	¡ |  
||j¡ |  |j¡ |  |jd¡ |  |jd¡ ~~|  t jd¡ ~|  t jd¡ d S )Nr   rE   r   )r   rG   r!   r   r,   rH   Ú
memoryviewÚdataÚassertFalseÚreadonlyÚassertIsÚobjÚ
assertTrueÚc_contiguousÚitemsizeÚndim)r#   rK   rL   rM   Úmviewr$   r$   r%   Útest_fake_memoryview}   s$    z#TestNrtMemInfo.test_fake_memoryviewc                 C   sh  ddl m}m}m}m} t tj¡}|jd }t	j
|dd}|j}|||jƒ||d ƒƒ}	tdƒD ]}
|  |	j|
 d¡ qbtdƒD ]}
|
d |	j|
< q‚t|ƒ}|  |j|¡ |  |j¡ |  ||j¡ |  |j¡ |  |jd¡ |  |jd¡ ~tj||j|j |d}~tj t |j¡d |¡ |d7 }tdƒD ]}
|  |	j|
 |
d	 ¡ q4|  |j j|¡ ~d S )
Nr   ©Úc_uint32Úc_void_pÚPOINTERÚcasté
   T©Úsafeì   ËK— r   ©ÚdtypeÚshapeÚbufferrF   )Úctypesrb   rc   rd   re   Únprk   Úuint32r]   r   r-   rV   rQ   rG   ÚcontentsrU   ÚnbytesrW   rX   rY   rZ   r[   r\   r^   ÚndarrayÚtestingÚassert_equalÚarangeÚsize)r#   rb   rc   rd   re   rk   ÚbytesizerM   rL   Úc_arrÚir_   Úarrr$   r$   r%   Útest_memoryview“   s8    
ÿzTestNrtMemInfo.test_memoryviewc                 C   s$  ddl m}m}m}m} t tj¡}|jd }t	j
|dd}|  |jd¡ |j}|||ƒ||d ƒƒ}	tdƒD ]}
|  |	j|
 d¡ qntdƒD ]}
|
d |	j|
< qŽtj|||j |d}|  |jd¡ ~tj t |j¡d |¡ |d7 }tdƒD ]}
|  |	j|
 |
d	 ¡ qò|  |j j|¡ ~d S )
Nr   ra   rf   Trg   r   ri   rj   rF   )rn   rb   rc   rd   re   ro   rk   rp   r]   r   r-   rG   rH   rV   rQ   rq   rs   rt   ru   rv   rw   )r#   rb   rc   rd   re   rk   rx   rM   rL   ry   rz   r{   r$   r$   r%   Útest_buffer¼   s,    
ÿzTestNrtMemInfo.test_bufferc              
   C   s\   t jjd d }dD ]B}|  t¡}tj||d W 5 Q R X |  d|› dt|j	ƒ¡ qd S )Né   rF   ©TFrg   zRequested allocation of z bytes failed.)
r   Úsize_tÚmaxvalr0   ÚMemoryErrorr   r-   r3   r4   r5   )r#   rw   Úpredr9   r$   r$   r%   Útest_allocate_invalid_sizeà   s    ÿz)TestNrtMemInfo.test_allocate_invalid_sizec              
   C   sT   d}dD ]F}|   t¡}tj||d W 5 Q R X d|› d}|  |t|jƒ¡ qd S )Niöÿÿÿr   rg   z,Cannot allocate a negative number of bytes: Ú.)r0   Ú
ValueErrorr   r-   r3   r4   r5   )r#   rw   rƒ   r9   r;   r$   r$   r%   Útest_allocate_negative_sizeê   s    z*TestNrtMemInfo.test_allocate_negative_size)r(   r)   r*   r=   rB   rN   rT   r`   r|   r}   r   r„   r‡   Ú__classcell__r$   r$   rC   r%   r>   M   s   )$
	r>   c                   @   s    e Zd ZdZdd„ Zdd„ ZdS )ÚTestTracemalloczB
    Test NRT-allocated memory can be tracked by tracemalloc.
    c                 C   sl   zdd l }W n tk
r*   |  d¡ Y nX | ¡  z(| ¡ }|ƒ }| ¡ }~| |d¡W ¢S | ¡  X d S )Nr   ztracemalloc not availableÚlineno)ÚtracemallocÚImportErrorZskipTestÚstartÚstopZtake_snapshotZ
compare_to)r#   Úfuncr‹   ÚbeforeÚresÚafterr$   r$   r%   Úmeasure_memory_diffú   s    z#TestTracemalloc.measure_memory_diffc                    sÒ   d‰ t j‰t‡ ‡fdd„ƒ‰‡fdd„}‡fdd„}|jjd }ˆƒ  |  |¡}|d	 }|  |jˆ ¡ | j|jˆ d
 dd |j	d	 }|  
tj |j¡d¡ |  
|j|¡ |  |¡}|d	 }|  |jˆ d ¡ d S )Ni@B c                      s   t  ˆ ˆ¡S )z<
            Allocate and return a large array.
            )ro   Úemptyr$   )ÚNrk   r$   r%   Úalloc_nrt_memory  s    z7TestTracemalloc.test_snapshot.<locals>.alloc_nrt_memoryc                      s   ˆ ƒ S ©Nr$   r$   ©r–   r$   r%   Úkeep_memory  s    z2TestTracemalloc.test_snapshot.<locals>.keep_memoryc                      s
   ˆ ƒ  d S r—   r$   r$   r˜   r$   r%   Úrelease_memory  s    z5TestTracemalloc.test_snapshot.<locals>.release_memoryr   r   g=
×£p=ð?zpUnexpected allocation overhead encountered. May be due to difference in CPython builds or running under coverage)r;   ztest_nrt.pyg{®Gáz„?)ro   Zint8r   Ú__code__Úco_firstlinenor“   ZassertGreaterEqualrw   Ú
assertLessÚ	tracebackrG   ÚosÚpathÚbasenameÚfilenamerŠ   )r#   r™   rš   Zalloc_linenoÚdiffÚstatÚframer$   )r•   r–   rk   r%   Útest_snapshot
  s(    
ÿ

zTestTracemalloc.test_snapshotN)r(   r)   r*   r=   r“   r¦   r$   r$   r$   r%   r‰   õ   s   r‰   c                   @   s:   e Zd Zdd„ Zdd„ Zdd„ Zdd„ Zejd	d
„ ƒZ	dS )ÚTestNRTIssuec                    sh   t dd„ ƒ‰ t ‡ fdd„ƒ‰t ‡fdd„ƒ}d}tj |df¡}|||ƒ}| ||¡}tj ||¡ d	S )
zO
        GitHub Issue #1244 https://github.com/numba/numba/issues/1244
        c                 S   s   | \}}t  |d |d  ¡S )NrF   )ÚmathÚsqrt)ÚvectorÚxÚyr$   r$   r%   Úcalculate_2D_vector_mag:  s    zNTestNRTIssue.test_issue_with_refct_op_pruning.<locals>.calculate_2D_vector_magc                    s<   t jdt jd}ˆ | ƒ}| \}}|| |d< || |d< |S ©NrF   ©rk   r   r   )ro   r”   Úfloat64)rª   Únormalized_vectorÚmagr«   r¬   )r­   r$   r%   Únormalize_2D_vector@  s    zJTestNRTIssue.test_issue_with_refct_op_pruning.<locals>.normalize_2D_vectorc                    sV   t j| dft jd}t| ƒD ]4}|| }ˆ |ƒ}|d ||df< |d ||df< q|S r®   )ro   r”   r°   rQ   )Únum_vectorsZvectorsZnormalized_vectorsrz   rª   r±   )r³   r$   r%   Únormalize_vectorsL  s    zHTestNRTIssue.test_issue_with_refct_op_pruning.<locals>.normalize_vectorsrf   rF   N)r   ro   ÚrandomÚpy_funcrt   Zassert_almost_equal)r#   rµ   r´   Ztest_vectorsÚgotÚexpectedr$   )r­   r³   r%   Ú test_issue_with_refct_op_pruning6  s    

z-TestNRTIssue.test_issue_with_refct_op_pruningc                 C   s`   dd„ }t |dt tjt tjdd¡f¡ƒ}| ¡ \}}|  |d¡ |  |tj	dtjd¡ d S )Nc                   S   s   dt jdt jdfS )Ng        r   r¯   )ro   ÚzerosÚint32r$   r$   r$   r%   Úfd  s    z.TestNRTIssue.test_incref_after_cast.<locals>.fr$   r   ÚCy                r¯   )
r   r   ÚTupleZ
complex128ZArrayr¼   Úentry_pointZassertPreciseEqualro   r»   )r#   r½   ZcresÚzr{   r$   r$   r%   Útest_incref_after_casta  s    ÿÿz#TestNRTIssue.test_incref_after_castc                 C   sB   t dd„ ƒ}|ƒ \}}tj ||¡ tj |tjdtjd¡ d S )Nc                  S   s0   t jdt jd} t jdt jd}| |d d … fS )Nrf   r¯   )ro   Úonesr°   )ÚaÚbr$   r$   r%   r½   s  s    z5TestNRTIssue.test_refct_pruning_issue_1511.<locals>.frf   r¯   )r   ro   rt   ru   rÃ   r°   )r#   r½   rÄ   rÅ   r$   r$   r%   Útest_refct_pruning_issue_1511r  s
    

z*TestNRTIssue.test_refct_pruning_issue_1511c                 C   sV   t dd„ ƒ}t ddddddddddg
g¡}| |dd¡}||ddƒ}|  ||¡ d S )Nc                 S   s`   t  | dk¡}t|d ƒdkr(d}d}n|d d }|d d }|d | }|d | }||fS )Nr   r   )ro   ÚwhereÚlen)Úimager«   r¬   Znext_locZy_offsetZx_offsetZ
next_loc_xZ
next_loc_yr$   r$   r%   Úudt~  s    z7TestNRTIssue.test_refct_pruning_issue_1526.<locals>.udtr   r   é   )r   ro   Úarrayr·   rG   )r#   rÊ   rÄ   Úexpectr¸   r$   r$   r%   Útest_refct_pruning_issue_1526}  s    
 z*TestNRTIssue.test_refct_pruning_issue_1526c                 C   sR   ddl m} |  tj¡ |dd„ ƒ}|  tj¡ |  |ƒ | ¡ ¡ |  tj¡ d S )Nr   r   c                   S   s   dS )Né{   r$   r$   r$   r$   r%   Úfoož  s    z8TestNRTIssue.test_no_nrt_on_njit_decoration.<locals>.foo)Únumbar   rW   r   r/   rG   r·   r[   )r#   r   rÐ   r$   r$   r%   Útest_no_nrt_on_njit_decoration”  s    
z+TestNRTIssue.test_no_nrt_on_njit_decorationN)
r(   r)   r*   rº   rÂ   rÆ   rÎ   r   Zrun_test_in_subprocessrÒ   r$   r$   r$   r%   r§   5  s   +r§   c                   @   s:   e Zd ZdZdd„ Ze d¡dd„ ƒZee	dd„ ƒƒZ
d	S )
ÚTestRefCtPruninga¡  
define i32 @"MyFunction"(i8** noalias nocapture %retptr, { i8*, i32 }** noalias nocapture %excinfo, i8* noalias nocapture readnone %env, double %arg.vt.0, double %arg.vt.1, double %arg.vt.2, double %arg.vt.3, double %arg.bounds.0, double %arg.bounds.1, double %arg.bounds.2, double %arg.bounds.3, i8* %arg.xs.0, i8* nocapture readnone %arg.xs.1, i64 %arg.xs.2, i64 %arg.xs.3, double* nocapture readonly %arg.xs.4, i64 %arg.xs.5.0, i64 %arg.xs.6.0, i8* %arg.ys.0, i8* nocapture readnone %arg.ys.1, i64 %arg.ys.2, i64 %arg.ys.3, double* nocapture readonly %arg.ys.4, i64 %arg.ys.5.0, i64 %arg.ys.6.0, i8* %arg.aggs_and_cols.0.0, i8* nocapture readnone %arg.aggs_and_cols.0.1, i64 %arg.aggs_and_cols.0.2, i64 %arg.aggs_and_cols.0.3, i32* nocapture %arg.aggs_and_cols.0.4, i64 %arg.aggs_and_cols.0.5.0, i64 %arg.aggs_and_cols.0.5.1, i64 %arg.aggs_and_cols.0.6.0, i64 %arg.aggs_and_cols.0.6.1) local_unnamed_addr {
entry:
tail call void @NRT_incref(i8* %arg.xs.0)
tail call void @NRT_incref(i8* %arg.ys.0)
tail call void @NRT_incref(i8* %arg.aggs_and_cols.0.0)
%.251 = icmp sgt i64 %arg.xs.5.0, 0
br i1 %.251, label %B42.preheader, label %B160

B42.preheader:                                    ; preds = %entry
%0 = add i64 %arg.xs.5.0, 1
br label %B42

B42:                                              ; preds = %B40.backedge, %B42.preheader
%lsr.iv3 = phi i64 [ %lsr.iv.next, %B40.backedge ], [ %0, %B42.preheader ]
%lsr.iv1 = phi double* [ %scevgep2, %B40.backedge ], [ %arg.xs.4, %B42.preheader ]
%lsr.iv = phi double* [ %scevgep, %B40.backedge ], [ %arg.ys.4, %B42.preheader ]
%.381 = load double, double* %lsr.iv1, align 8
%.420 = load double, double* %lsr.iv, align 8
%.458 = fcmp ole double %.381, %arg.bounds.1
%not..432 = fcmp oge double %.381, %arg.bounds.0
%"$phi82.1.1" = and i1 %.458, %not..432
br i1 %"$phi82.1.1", label %B84, label %B40.backedge

B84:                                              ; preds = %B42
%.513 = fcmp ole double %.420, %arg.bounds.3
%not..487 = fcmp oge double %.420, %arg.bounds.2
%"$phi106.1.1" = and i1 %.513, %not..487
br i1 %"$phi106.1.1", label %B108.endif.endif.endif, label %B40.backedge

B160:                                             ; preds = %B40.backedge, %entry
tail call void @NRT_decref(i8* %arg.ys.0)
tail call void @NRT_decref(i8* %arg.xs.0)
tail call void @NRT_decref(i8* %arg.aggs_and_cols.0.0)
store i8* null, i8** %retptr, align 8
ret i32 0

B108.endif.endif.endif:                           ; preds = %B84
%.575 = fmul double %.381, %arg.vt.0
%.583 = fadd double %.575, %arg.vt.1
%.590 = fptosi double %.583 to i64
%.630 = fmul double %.420, %arg.vt.2
%.638 = fadd double %.630, %arg.vt.3
%.645 = fptosi double %.638 to i64
tail call void @NRT_incref(i8* %arg.aggs_and_cols.0.0)              ; GONE 1
tail call void @NRT_decref(i8* null)                                ; GONE 2
tail call void @NRT_incref(i8* %arg.aggs_and_cols.0.0), !noalias !0 ; GONE 3
%.62.i.i = icmp slt i64 %.645, 0
%.63.i.i = select i1 %.62.i.i, i64 %arg.aggs_and_cols.0.5.0, i64 0
%.64.i.i = add i64 %.63.i.i, %.645
%.65.i.i = icmp slt i64 %.590, 0
%.66.i.i = select i1 %.65.i.i, i64 %arg.aggs_and_cols.0.5.1, i64 0
%.67.i.i = add i64 %.66.i.i, %.590
%.84.i.i = mul i64 %.64.i.i, %arg.aggs_and_cols.0.5.1
%.87.i.i = add i64 %.67.i.i, %.84.i.i
%.88.i.i = getelementptr i32, i32* %arg.aggs_and_cols.0.4, i64 %.87.i.i
%.89.i.i = load i32, i32* %.88.i.i, align 4, !noalias !3
%.99.i.i = add i32 %.89.i.i, 1
store i32 %.99.i.i, i32* %.88.i.i, align 4, !noalias !3
tail call void @NRT_decref(i8* %arg.aggs_and_cols.0.0), !noalias !0 ; GONE 4
tail call void @NRT_decref(i8* %arg.aggs_and_cols.0.0)              ; GONE 5
br label %B40.backedge

B40.backedge:                                     ; preds = %B108.endif.endif.endif, %B84, %B42
%scevgep = getelementptr double, double* %lsr.iv, i64 1
%scevgep2 = getelementptr double, double* %lsr.iv1, i64 1
%lsr.iv.next = add i64 %lsr.iv3, -1
%.294 = icmp sgt i64 %lsr.iv.next, 1
br i1 %.294, label %B42, label %B160
}
    c                 C   sú   | j }t| ¡ ƒ}dd„ |D ƒ}dd„ |D ƒ}t |¡}t| ¡ ƒ}dd„ |D ƒ}dd„ |D ƒ}|  ||¡ |  ||¡ t|ƒt|ƒ }	t|ƒt|ƒ }
|	|
B }|  ||	|
A ¡ d |¡}dD ]}d 	|¡}|  
||¡ q¾|  tt| ¡ ƒƒt|ƒ¡ d S )	Nc                 S   s   g | ]}d |kr|‘qS ©Z
NRT_increfr$   ©Ú.0Úlnr$   r$   r%   Ú
<listcomp>ú  s      zDTestRefCtPruning.test_refct_pruning_op_recognize.<locals>.<listcomp>c                 S   s   g | ]}d |kr|‘qS ©Z
NRT_decrefr$   rÕ   r$   r$   r%   rØ   û  s      c                 S   s   g | ]}d |kr|‘qS rÔ   r$   rÕ   r$   r$   r%   rØ      s      c                 S   s   g | ]}d |kr|‘qS rÙ   r$   rÕ   r$   r$   r%   rØ     s      Ú
)r   rF   é   é   é   z	; GONE {})Úsample_llvm_irÚlistÚ
splitlinesr   Z_remove_redundant_nrt_refctZassertNotEqualÚsetrG   ÚjoinÚformatr3   rÈ   )r#   Zinput_irZinput_linesZbefore_increfsZbefore_decrefsZ	output_irÚoutput_linesZafter_increfsZafter_decrefsZpruned_increfsZpruned_decrefsÚcombinedZpruned_linesrz   Úgoner$   r$   r%   Útest_refct_pruning_op_recognize÷  s&    


z0TestRefCtPruning.test_refct_pruning_op_recognizez3Pass removed as it was buggy. Re-enable when fixed.c                    s‚   t dd„ ƒ‰t dd„ ƒ‰ t ‡ ‡fdd„ƒ‰t dd‡fd	d
„ƒ}| d¡ t| |jd ¡ƒ}tt d|¡ƒ}|  t	|ƒd¡ dS )ztestcase from #2350c                 S   s"   t  |¡s||| f  d7  < d S r   ©ro   Úisnan©r«   r¬   ZaggÚfieldr$   r$   r%   Ú_append_non_na  s    
zITestRefCtPruning.test_refct_pruning_with_branches.<locals>._append_non_nac                 S   sB   t  |¡s>t  ||| f ¡r*|||| f< n||| f  |7  < d S r—   rè   rê   r$   r$   r%   Ú_append  s    
zBTestRefCtPruning.test_refct_pruning_with_branches.<locals>._appendc                    s    ˆ| |||ƒ ˆ | |||ƒ d S r—   r$   rê   )rí   rì   r$   r%   Úappend&  s    zATestRefCtPruning.test_refct_pruning_with_branches.<locals>.appendT)Zno_cpython_wrapperc                    s:   t | jd ƒD ]&}t | jd ƒD ]}ˆ ||| |ƒ q qd S )Nr   r   )rQ   rl   )r{   rë   rz   Új)rî   r$   r%   Úextend-  s    zATestRefCtPruning.test_refct_pruning_with_branches.<locals>.extendz(f4[:,::1], f4)r   z!(NRT_incref|NRT_decref)\([^\)]+\)N)
r   Úcompiler4   Zinspect_llvmZ
signaturesrß   ÚreÚfinditerrG   rÈ   )r#   rð   ZllvmirZrefopsr$   )rí   rì   rî   r%   Ú test_refct_pruning_with_branches  s    


z1TestRefCtPruning.test_refct_pruning_with_branchesc                    sV   t  ¡  t  ¡  t  ¡  t  ¡  tdd„ ƒ‰ tdgƒ‡ fdd„ƒ}|  |dƒd¡ dS )	ztThe InlineAsm class from llvmlite.ir has no 'name' attr the refcount
        pruning pass should be tolerant to thisc                 S   s   dd„ }t tjtjtjƒ|fS )Nc                 S   sT   |\}}t  t  d¡t  d¡t  d¡g¡}|j|dd||fddd}t| ||j|ƒS )Né    zmov $2, $0; imul $1, $0z=&r,r,rZasm_mulF)ÚnameZside_effect)r   ÚFunctionTypeZIntTypeÚasmr   Úreturn_type)ZcgctxZbuilderÚsigr8   Zarg_0Zarg_1ZftyÚmulr$   r$   r%   ÚcodegenG  s    ÿ
 þz>TestRefCtPruning.test_inline_asm.<locals>.bar.<locals>.codegen)r   r   r¼   )Ztyctxr«   r¬   rü   r$   r$   r%   ÚbarE  s    z-TestRefCtPruning.test_inline_asm.<locals>.barzint32(int32)c                    s   | d7 } ˆ | dƒ}|S )Nr   rF   r$   )r«   rÁ   ©rý   r$   r%   rÐ   Q  s    
z-TestRefCtPruning.test_inline_asm.<locals>.foorf   é   N)Úllvmr?   Zinitialize_native_targetZinitialize_native_asmprinterZinitialize_native_asmparserr   r   rG   )r#   rÐ   r$   rþ   r%   Útest_inline_asm;  s    
z TestRefCtPruning.test_inline_asmN)r(   r)   r*   rÞ   rç   ÚunittestÚskiprô   Ú
linux_onlyÚx86_onlyr  r$   r$   r$   r%   rÓ   ­  s   H
$rÓ   zcffi requiredc                       sH   e Zd ZdZ‡ fdd„Zdd„ Zdd„ Zdd	„ Zd
d„ Zdd„ Z	‡  Z
S )ÚTestNrtExternalCFFIz?Testing the use of externally compiled C code that use NRT
    c                    s   t j tt| ƒ ¡  d S r—   )r   r@   rA   r  rB   r"   rC   r$   r%   rB   ^  s    zTestNrtExternalCFFI.setUpc              	   C   sz   ddl m} |ƒ }|j||tƒ gd | |¡ td |¡ƒ}|j|d tj	 
|¡ zt|ƒ}W 5 tj	 |¡ X ||fS )Nr   ©ÚFFI)Úinclude_dirszcffi_test_{})Útmpdir)Úcffir  Z
set_sourcer   Úcdefr   rã   rñ   Úsysr    rî   Úremover   )r#   rö   Úsourcer  r  Úffir
  Úmodr$   r$   r%   Úcompile_cffi_modulec  s    
z'TestNrtExternalCFFI.compile_cffi_modulec                 C   s.   ddl m} |ƒ }| dtjd ¡}|ƒ }|S )Nr   r  zvoid* (*)()Zget_api)r  r  re   r	   Z	c_helpers)r#   r  r  Znrt_get_apiÚtabler$   r$   r%   Úget_nrt_api_tables  s
    z%TestNrtExternalCFFI.get_nrt_api_tablec           
      C   sš   d  | jj¡}d}d}|  |||¡\}}|  |jjd¡ |  ¡ }|j |¡}|  |jjd¡ t	| 
d|¡ƒ}t |¡}	|  |	jd¡ ~	|  |jjd¡ d S )	Nz{}_test_manage_memoryaz  
#include <stdio.h>
#include "numba/core/runtime/nrt_external.h"

int status = 0;

void my_dtor(void *ptr) {
    free(ptr);
    status = 0xdead;
}

NRT_MemInfo* test_nrt_api(NRT_api_functions *nrt) {
    void * data = malloc(10);
    NRT_MemInfo *mi = nrt->manage_memory(data, my_dtor);
    nrt->acquire(mi);
    nrt->release(mi);
    status = 0xa110c;
    return mi;
}
        z;
void* test_nrt_api(void *nrt);
extern int status;
        r   i
 r€   r   i­Þ  )rã   rD   r(   r  rG   ÚlibÚstatusr  Útest_nrt_apiÚintre   r
   ÚMemInforH   )
r#   rö   r  r  r  r  r  ÚoutÚmi_addrrM   r$   r$   r%   Útest_manage_memory{  s    
z&TestNrtExternalCFFI.test_manage_memoryc                 C   s¼   d  | jj¡}d}d}|  |||¡\}}|  ¡ }dt tj¡j }|j	 
||¡}t| d|¡ƒ}	t |	¡}
|  |
jd¡ | | d  |¡|
j¡¡}tjdtj|d	}tj |d
ddg¡ d S )Nz{}_test_allocatea7  
#include <stdio.h>
#include "numba/core/runtime/nrt_external.h"

NRT_MemInfo* test_nrt_api(NRT_api_functions *nrt, size_t n) {
    size_t *data = NULL;
    NRT_MemInfo *mi = nrt->allocate(n);
    data = nrt->get_data(mi);
    data[0] = 0xded;
    data[1] = 0xabc;
    data[2] = 0xdef;
    return mi;
}
        z(void* test_nrt_api(void *nrt, size_t n);rÛ   r€   r   z	char [{}])rÛ   )rl   rk   rm   ií  i¼
  iï  )rã   rD   r(   r  r  ro   rk   Zintpr]   r  r  r  re   r
   r  rG   rH   rm   rV   rs   rt   ru   )r#   rö   r  r  r  r  r  Únumbytesr  r  rM   rm   r{   r$   r$   r%   Útest_allocate¤  s    
z!TestNrtExternalCFFI.test_allocatec                 C   sH   ddl m} tdd„ ƒ}|ƒ }t| d|  ¡ ¡ƒ}|ƒ }|  ||¡ d S )Nr   r  c                   S   s   t ƒ S r—   r   r$   r$   r$   r%   r  Ç  s    z6TestNrtExternalCFFI.test_get_api.<locals>.test_nrt_apir€   )r  r  r   r  re   r  rG   )r#   r  r  r  rÍ   r¸   r$   r$   r%   Útest_get_apiÄ  s    
z TestNrtExternalCFFI.test_get_api)r(   r)   r*   r=   rB   r  r  r  r  r  rˆ   r$   r$   rC   r%   r  Z  s   ) r  c                   @   sT   e Zd Zdd„ Zdd„ Zdd„ Zdd„ Zd	d
„ Zdd„ Zdd„ Z	dd„ Z
dd„ ZdS )ÚTestNrtStatisticsc                 C   s   t  ¡ | _d S r—   )r	   Úmemsys_stats_enabledÚ_TestNrtStatistics__stats_stater"   r$   r$   r%   rB   Ó  s    zTestNrtStatistics.setUpc                 C   s   | j rt ¡  nt ¡  d S r—   )r"  r	   Úmemsys_enable_statsÚmemsys_disable_statsr"   r$   r$   r%   ÚtearDown×  s    
zTestNrtStatistics.tearDownc                 C   s&   d}t j ¡ }d|d< t||d d S )Na¡  if 1:
        from numba import njit
        import numpy as np
        from numba.core.runtime import rtsys, _nrt_python
        from numba.core.registry import cpu_target

        @njit
        def foo():
            return np.arange(10)[0]

        # initialize the NRT before use
        rtsys.initialize(cpu_target.target_context)
        assert _nrt_python.memsys_stats_enabled()
        orig_stats = rtsys.get_allocation_stats()
        foo()
        new_stats = rtsys.get_allocation_stats()
        total_alloc = new_stats.alloc - orig_stats.alloc
        total_free = new_stats.free - orig_stats.free
        total_mi_alloc = new_stats.mi_alloc - orig_stats.mi_alloc
        total_mi_free = new_stats.mi_free - orig_stats.mi_free

        expected = 1
        assert total_alloc == expected
        assert total_free == expected
        assert total_mi_alloc == expected
        assert total_mi_free == expected
        Ú1ÚNUMBA_NRT_STATS©Úenv)rŸ   ÚenvironÚcopyr   )r#   Úsrcr)  r$   r$   r%   Útest_stats_env_var_explicit_onÞ  s    
z0TestNrtStatistics.test_stats_env_var_explicit_onc                 C   s   d}t ||d d S )Na“  if 1:
        from numba import njit
        import numpy as np
        from numba.core.runtime import rtsys, _nrt_python

        @njit
        def foo():
            return np.arange(10)[0]

        assert _nrt_python.memsys_stats_enabled() == False
        try:
            rtsys.get_allocation_stats()
        except RuntimeError as e:
            assert "NRT stats are disabled." in str(e)
        r(  )r   )r#   r)  r,  r$   r$   r%   Úcheck_env_var_off   s    z#TestNrtStatistics.check_env_var_offc                 C   s    t j ¡ }d|d< |  |¡ d S )NÚ0r'  )rŸ   r*  r+  r.  ©r#   r)  r$   r$   r%   Útest_stats_env_var_explicit_off  s    
z1TestNrtStatistics.test_stats_env_var_explicit_offc                 C   s$   t j ¡ }| dd ¡ |  |¡ d S )Nr'  )rŸ   r*  r+  Úpopr.  r0  r$   r$   r%   Útest_stats_env_var_default_off  s    
z0TestNrtStatistics.test_stats_env_var_default_offc                 C   sœ   t dd„ ƒ}t ¡  |  t ¡ ¡ tdƒD ]l}t ¡ }t ¡  |  	t ¡ ¡ |ƒ  t ¡  |  t ¡ ¡ t ¡ }|ƒ  t ¡ }|  
||¡ |  ||¡ q*d S )Nc                  S   s   t  d¡} t  d| d  ¡S )NrÛ   rÝ   r   )ro   rÃ   rv   )Útmpr$   r$   r%   rÐ   "  s    
z7TestNrtStatistics.test_stats_status_toggle.<locals>.foorF   )r   r	   r#  r[   r!  rQ   r   Úget_allocation_statsr$  rW   rG   r   )r#   rÐ   rz   Zstats_1Zstats_2Zstats_3r$   r$   r%   Útest_stats_status_toggle   s     
z*TestNrtStatistics.test_stats_status_togglec              	   C   sJ   t  ¡  |  t  ¡ ¡ |  t¡}t ¡  W 5 Q R X |  dt	|j
ƒ¡ d S )NúNRT stats are disabled.)r	   r$  rW   r!  r0   r1   r   r5  r3   r4   r5   )r#   r9   r$   r$   r%   Ú5test_rtsys_stats_query_raises_exception_when_disabledD  s
    zGTestNrtStatistics.test_rtsys_stats_query_raises_exception_when_disabledc                 C   s~   d}|D ]p}t td|› ƒ}| j|dJ t ¡  |  t ¡ ¡ |  t¡}|ƒ  W 5 Q R X |  dt	|j
ƒ¡ W 5 Q R X qd S )N)ZallocÚfreeZmi_allocZmi_freeZmemsys_get_stats_)Ú
stats_funcr7  )r2   r	   ZsubTestr$  rW   r!  r0   r1   r3   r4   r5   )r#   Zmethod_variationsr7   r:  r9   r$   r$   r%   Ú<test_nrt_explicit_stats_query_raises_exception_when_disabledP  s    zNTestNrtStatistics.test_nrt_explicit_stats_query_raises_exception_when_disabledN)r(   r)   r*   rB   r%  r-  r.  r1  r3  r6  r8  r;  r$   r$   r$   r%   r   Ñ  s   "$r   Ú__main__)Cr¨   rŸ   Úplatformr  rò   Únumpyro   rÑ   r   Z
numba.corer   Znumba.core.compilerr   r   Znumba.core.runtimer   r   r	   r
   Znumba.core.extendingr   r   Znumba.core.typingr   Znumba.core.imputilsr   Zllvmliter   Zllvmlite.bindingZbindingr   Znumba.core.typing.cffi_utilsÚcoreÚtypingZ
cffi_utilsZcffi_supportZnumba.core.unsafe.nrtr   Znumba.tests.supportr   r   r   r   r   r   Znumba.core.registryr   r  Zenable_nrt_flagsZskipIfÚ
startswithr  Úmachiner  Úobjectr   r+   r>   r‰   r§   rÓ   Z
skipUnlessZ	SUPPORTEDr  r   r(   Úmainr$   r$   r$   r%   Ú<module>   sR    ÿÿ
 )@x .v 
