o
    iXu                     @   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 Z  d dl!m"Z"m#Z#m$Z$m%Z%m&Z&m'Z'm(Z( d dl)m*Z* d dl+Z+e Z,de,_e+-ej.d dZ/e+-e0 dvdZ1G dd de2Z3G dd de+j#Z4G dd de+j#Z5G dd de+j#Z6G dd de#Z7G dd de+j#Z8e'G dd de"e#Z9G d d! d!e#Z:e;d"kre+<  dS 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skip_unless_cffirun_in_subprocess)
cpu_targetTlinuxzlinux only test)i386x86_64zx86 only testc                   @       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&   c/var/www/html/eduruby.in/lip-sync/lip-sync-env/lib/python3.10/site-packages/numba/tests/test_nrt.py__init__+      zDummy.__init__c                 C   s   t |  jd8  _d S r   r!   r$   r&   r&   r'   __del__.   r)   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   s   dddd}|  D ]9\}}z1| t}dt_tt|}||  W d    n1 s+w   Y  d}| |t|j W dt_q
dt_w d S )Nr&   )r&   r&   )r&   )librarymeminfo_newmeminfo_allocFz'Runtime must be initialized before use.T)	itemsassertRaisesRuntimeErrorr   _initgetattrassertInstr	exception)r%   methodsmethargsraisesfnmsgr&   r&   r'   test_init_fail9   s   

z+TestNrtMemInfoNotInitialized.test_init_failN)r+   r,   r-   __doc__Z_numba_parallel_test_r@   r&   r&   r&   r'   r.   2   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superrB   setUpr$   	__class__r&   r'   rF   Q   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   r0   refcountacquirerelease)r%   daddrmir&   r&   r'   test_meminfo_refct_1X   s   z#TestNrtMemInfo.test_meminfo_refct_1c                 C   s   t  }| t jd d}t||}| |jd ~| t jd tdD ]}|  q%| |jd | t jd tdD ]}|  q>| |jd ~| t jd d S )Nr    rI   d   e   r   )	r   rK   r#   r   r0   rL   rangerM   rN   )r%   rO   rP   rQ   ct_r&   r&   r'   test_meminfo_refct_2i   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    rI   r   )r   rK   r#   r   r0   rL   
memoryviewdataassertFalsereadonlyassertIsobj
assertTruec_contiguousitemsizendim)r%   rO   rP   rQ   mviewr&   r&   r'   test_fake_memoryview|   s$   z#TestNrtMemInfo.test_fake_memoryviewc                 C   sf  ddl m}m}m}m} ttj}|jd }t	j
|dd}|j}|||j||d }	tdD ]}
| |	j|
 d q1tdD ]	}
|
d |	j|
< qAt|}| |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	  q| |j j| ~d S )
Nr   c_uint32c_void_pPOINTERcast
   Tsafe   K r    dtypeshapebufferrJ   )ctypesrf   rg   rh   ri   npro   uint32ra   r   r1   rZ   rU   rK   contentsrY   nbytesr[   r\   r]   r^   r_   r`   rb   ndarraytestingassert_equalarangesize)r%   rf   rg   rh   ri   ro   bytesizerQ   rP   c_arrirc   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 q7tdD ]	}
|
d |	j|
< qGtj|||j |d}| |jd ~tjt|jd | |d7 }tdD ]}
| |	j|
 |
d	  qy| |j j| ~d S )
Nr   re   rj   Trk   r    rm   rn   rJ   )rr   rf   rg   rh   ri   rs   ro   rt   ra   r   r1   rK   rL   rZ   rU   ru   rw   rx   ry   rz   r{   )r%   rf   rg   rh   ri   ro   r|   rQ   rP   r}   r~   r   r&   r&   r'   test_buffer   s,   
zTestNrtMemInfo.test_bufferc              	   C   sp   t jjd d }dD ]+}| t}tj||d W d    n1 s#w   Y  | d| dt|j	 q
d S )N   rJ   TFrk   zRequested allocation of z bytes failed.)
r   size_tmaxvalr3   MemoryErrorr   r1   r7   r8   r9   )r%   r{   predr=   r&   r&   r'   test_allocate_invalid_size   s   z)TestNrtMemInfo.test_allocate_invalid_sizec              	   C   sh   d}dD ]-}|  t}tj||d W d    n1 sw   Y  d| d}| |t|j qd S )Nir   rk   z,Cannot allocate a negative number of bytes: .)r3   
ValueErrorr   r1   r7   r8   r9   )r%   r{   r   r=   r?   r&   r&   r'   test_allocate_negative_size   s   z*TestNrtMemInfo.test_allocate_negative_size)r+   r,   r-   rA   rF   rR   rX   rd   r   r   r   r   r   __classcell__r&   r&   rG   r'   rB   L   s    )$
	rB   c                   @   r   )TestTracemalloczB
    Test NRT-allocated memory can be tracked by tracemalloc.
    c                 C   sl   zdd l }W n ty   | d Y nw |  z| }| }| }~||dW |  S |  w )Nr   ztracemalloc not availablelineno)tracemallocImportErrorZskipTeststartZtake_snapshotZ
compare_tostop)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.
            )rs   emptyr&   )Nro   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{Gz?)rs   Zint8r   __code__co_firstlinenor   ZassertGreaterEqualr{   
assertLess	tracebackrK   ospathbasenamefilenamer   )r%   r   r   Zalloc_linenodiffstatframer&   )r   r   ro   r'   test_snapshot	  s(   


zTestTracemalloc.test_snapshotN)r+   r,   r-   rA   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 )NrJ   )mathsqrt)vectorxyr&   r&   r'   calculate_2D_vector_mag9  s   zNTestNRTIssue.test_issue_with_refct_op_pruning.<locals>.calculate_2D_vector_magc                    s<   t jdt jd} | }| \}}|| |d< || |d< |S NrJ   ro   r   r    )rs   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 ]}|| } |}|d ||df< |d ||df< q|S r   )rs   r   r   rU   )num_vectorsZvectorsZnormalized_vectorsr~   r   r   )r   r&   r'   normalize_vectorsK  s   zHTestNRTIssue.test_issue_with_refct_op_pruning.<locals>.normalize_vectorsrj   rJ   N)r   rs   randompy_funcrx   Zassert_almost_equal)r%   r   r   Ztest_vectorsgotexpectedr&   )r   r   r'    test_issue_with_refct_op_pruning5  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   )rs   zerosint32r&   r&   r&   r'   fc  s   z.TestNRTIssue.test_incref_after_cast.<locals>.fr&   r    Cy                r   )
r   r   TupleZ
complex128ZArrayr   entry_pointZassertPreciseEqualrs   r   )r%   r   Zcreszr   r&   r&   r'   test_incref_after_cast`  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 )Nrj   r   )rs   onesr   )abr&   r&   r'   r   r  s   z5TestNRTIssue.test_refct_pruning_issue_1511.<locals>.frj   r   )r   rs   rx   ry   r   r   )r%   r   r   r   r&   r&   r'   test_refct_pruning_issue_1511q  s
   

z*TestNRTIssue.test_refct_pruning_issue_1511c                 C   sF   t dd }tg d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   )rs   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>.udt)
r    r   r    r   r    r   r   r    r   r   r       )r   rs   arrayr   rK   )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   r[   r   r5   rK   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   4  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      g | ]}d |v r|qS Z
NRT_increfr&   .0lnr&   r&   r'   
<listcomp>      zDTestRefCtPruning.test_refct_pruning_op_recognize.<locals>.<listcomp>c                 S   r   Z
NRT_decrefr&   r   r&   r&   r'   r     r   c                 S   r   r   r&   r   r&   r&   r'   r     r   c                 S   r   r   r&   r   r&   r&   r'   r      r   
)r    rJ            z	; GONE {})sample_llvm_irlist
splitlinesr   Z_remove_redundant_nrt_refctZassertNotEqualsetrK   joinformatr7   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_linesr~   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 d S r   rs   isnanr   r   Zaggfieldr&   r&   r'   _append_non_na  s   
zITestRefCtPruning.test_refct_pruning_with_branches.<locals>._append_non_nac                 S   sH   t |s"t ||| f r|||| f< d S ||| f  |7  < d S 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    )rU   rp   )r   r   r~   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   compiler8   Zinspect_llvmZ
signaturesr   refinditerrK   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sigr<   Zarg_0Zarg_1Zftymulr&   r&   r'   codegenF  s   
z>TestRefCtPruning.test_inline_asm.<locals>.bar.<locals>.codegen)r   r   r   )Ztyctxr   r   r  r&   r&   r'   barD  s   z-TestRefCtPruning.test_inline_asm.<locals>.barzint32(int32)c                    s   | d7 }  | d}|S )Nr    rJ   r&   )r   r   r  r&   r'   r   P  s   
z-TestRefCtPruning.test_inline_asm.<locals>.foorj      N)llvmrC   Zinitialize_native_targetZinitialize_native_asmprinterZinitialize_native_asmparserr   r   rK   )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   c                       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   rD   rE   r  rF   r$   rG   r&   r'   rF   ]  s   zTestNrtExternalCFFI.setUpc              	   C   s   ddl m} | }|j||t gd || td|}|j|d tj	
| zt|}W tj	| ||fS tj	| w )Nr   FFI)Zinclude_dirszcffi_test_{})tmpdir)cffir  Z
set_sourcer   cdefr   r   r   sysr   r   r   remove)r%   r   sourcer  r  ffir  modr&   r&   r'   compile_cffi_moduleb  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  ri   r	   Z	c_helpers)r%   r  r  Znrt_get_apitabler&   r&   r'   get_nrt_api_tabler  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   rH   r+   r  rK   libstatusr  test_nrt_apiintri   r
   MemInforL   )
r%   r   r  r  r  r  r  outmi_addrrQ   r&   r&   r'   test_manage_memoryz  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|g d
 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   )rp   ro   rq   )i  i
  i  )r   rH   r+   r  r  rs   ro   Zintpra   r  r  r  ri   r
   r  rK   rL   rq   rZ   rw   rx   ry   )r%   r   r  r  r  r  r  numbytesr  r   rQ   rq   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  ri   r  rK   )r%   r  r  r  r   r   r&   r&   r'   test_get_api  s   
z TestNrtExternalCFFI.test_get_api)r+   r,   r-   rA   rF   r  r  r!  r#  r$  r   r&   r&   rG   r'   r  Y  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'   rF     s   zTestNrtStatistics.setUpc                 C   s   | j r	t  d S 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.  r1  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/  r0  r3  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/  r0  popr3  r5  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 ]6}t }t  | 	t  |  t  | t  t }|  t }| 
|| | || qd S )Nc                  S   s   t d} t d| d  S )Nr   r   r   )rs   r   rz   )tmpr&   r&   r'   r   !  s   
z7TestNrtStatistics.test_stats_status_toggle.<locals>.foorJ   )r   r	   r(  r_   r&  rU   r   get_allocation_statsr)  r[   rK   r   )r%   r   r~   Zstats_1Zstats_2Zstats_3r&   r&   r'   test_stats_status_toggle  s"   
z*TestNrtStatistics.test_stats_status_togglec                 C   s^   t   | t   | t}t  W d    n1 sw   Y  | dt	|j
 d S )NNRT stats are disabled.)r	   r)  r[   r&  r3   r4   r   r:  r7   r8   r9   )r%   r=   r&   r&   r'   5test_rtsys_stats_query_raises_exception_when_disabledC  s   
zGTestNrtStatistics.test_rtsys_stats_query_raises_exception_when_disabledc              
   C   s   d}|D ]L}t td| }| j|d4 t  | t  | t}|  W d    n1 s3w   Y  | dt	|j
 W d    n1 sKw   Y  qd S )N)ZallocfreeZmi_allocZmi_freeZmemsys_get_stats_)
stats_funcr<  )r6   r	   ZsubTestr)  r[   r&  r3   r4   r7   r8   r9   )r%   Zmethod_variationsr;   r?  r=   r&   r&   r'   <test_nrt_explicit_stats_query_raises_exception_when_disabledO  s   zNTestNrtStatistics.test_nrt_explicit_stats_query_raises_exception_when_disabledN)r+   r,   r-   rF   r*  r2  r3  r6  r8  r;  r=  r@  r&   r&   r&   r'   r%    s    "$r%  __main__)=r   r   platformr  r   numpyrs   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.unsafe.nrtr   Znumba.tests.supportr   r   r   r   r   r   r   Znumba.core.registryr   r  Zenable_nrt_flagsZskipIf
startswithr
  machiner  objectr   r.   rB   r   r   r   r  r%  r+   mainr&   r&   r&   r'   <module>   sT    $
 )@x .v 