
    xKg[t                        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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(jR                  ej                  jU                  d       d      Z+ e(jR                   ejX                         dvd      Z- G d de.      Z/ G d de(j@                        Z0 G d de(j@                        Z1 G d de(j@                        Z2 G d de       Z3 G d de(j@                        Z4e$ G d dee              Z5 G d de       Z6e7d k(  r e(jp                          yy)!    Nnjit)types)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_targetlinuxzlinux only test)i386x86_64zx86 only testc                       e Zd ZdZd Zd Zy)Dummyr   c                 @    t        |       xj                  dz  c_        y N   typealiveselfs    X/home/alanp/www/video.onchill/myenv/lib/python3.12/site-packages/numba/tests/test_nrt.py__init__zDummy.__init__(       T
A    c                 @    t        |       xj                  dz  c_        y r   r!   r$   s    r&   __del__zDummy.__del__+   r(   r)   N)__name__
__module____qualname__r#   r'   r+    r)   r&   r   r   %   s    Er)   r   c                       e Zd ZdZdZd Zy)TestNrtMemInfoNotInitializedzn
    Unit test for checking the use of the NRT fails if the
    initialization sequence has not been run.
    Fc                 Z   dddd}|j                         D ]v  \  }}	 | j                  t              5 }dt        _        t        t        |      } ||  d d d        d}| j                  |t        j                               dt        _        x y # 1 sw Y   >xY w# dt        _        w xY w)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msgs          r&   test_init_failz+TestNrtMemInfoNotInitialized.test_init_fail6   s     "*$)
 "--/JD$	#&&|4"'EK -BI 5
 @c3v'7'7#89" *44 #s"   B!B/BB	BB*N)r,   r-   r.   __doc___numba_parallel_test_rD   r/   r)   r&   r1   r1   /   s     "#r)   r1   c                   V     e Zd ZdZ fdZd Zd Zd Zd Zd Z	e
d        Zd	 Z xZS )
TestNrtMemInfoz2
    Unit test for core MemInfo functionality
    c                     dt         _        t        j                  t        j
                         t        t        | #          y )Nr   )	r   r#   r   
initializer   target_contextsuperrH   setUpr%   	__class__s    r&   rM   zTestNrtMemInfo.setUpN   s+    223nd)+r)   c                 4   t               }| j                  t         j                  d       d}t        j                  ||      }| j                  |j
                  d       ~| j                  t         j                  d       |j                          | j                  |j
                  d       | j                  t         j                  d       |j                          | j                  |j
                  d       ~| j                  t         j                  d       y )Nr       J[=    r   )r   assertEqualr#   r   r4   refcountacquirerelease)r%   daddrmis       r&   test_meminfo_refct_1z#TestNrtMemInfo.test_meminfo_refct_1U   s    Ga(tQ'a(a(


a(a(


a(a(r)   c                 t   t               }| j                  t         j                  d       d}t        j                  ||      }| j                  |j
                  d       ~| j                  t         j                  d       t        d      D ]  }|j                           | j                  |j
                  d       | j                  t         j                  d       t        d      D ]  }|j                           | j                  |j
                  d       ~| j                  t         j                  d       y )Nr    rQ   d   e   r   )	r   rS   r#   r   r4   rT   rangerU   rV   )r%   rW   rX   rY   ct_s         r&   test_meminfo_refct_2z#TestNrtMemInfo.test_meminfo_refct_2f   s    Ga(tQ'a(a(*BJJL g.a(sAJJL a(a(r)   c                    t               }| j                  t         j                  d       d}t        j                  ||      }| j                  |j
                  d       t        |      }| j                  |j
                  d       | j                  ||j                         | j                  |j                         | j                  ||j                         | j                  |j                         | j                  |j                  d       | j                  |j                  d       ~~| j                  t         j                  d       ~| j                  t         j                  d       y )Nr    rQ   r   )r   rS   r#   r   r4   rT   
memoryviewdataassertFalsereadonlyassertIsobj
assertTruec_contiguousitemsizendim)r%   rW   rX   rY   mviews        r&   test_fake_memoryviewz#TestNrtMemInfo.test_fake_memoryviewy   s   Ga(tQ'a(2a(rww'(b%))$**++Q'a(a(r)   c                    ddl m}m}m}m} t        j                  t
        j                        }|j                  dz  }t        j                  |d      }|j                  } | ||j                         ||dz              }	t        d      D ]!  }
| j                  |	j                  |
   d       # t        d      D ]  }
|
dz   |	j                  |
<    t        |      }| j                  |j                   |       | j#                  |j$                         | j'                  ||j(                         | j+                  |j,                         | j                  |j                  d       | j                  |j.                  d       ~t        j0                  ||j                   |j                  z  |      }~t
        j2                  j5                  t        j6                  |j8                        dz   |       |dz  }t        d      D ]$  }
| j                  |	j                  |
   |
d	z          & | j                  |j                   j                  |       ~y )
Nr   c_uint32c_void_pPOINTERcast
   Tsafe   K r    dtypeshapebufferrR   )ctypesrq   rr   rs   rt   nprz   uint32rk   r   r5   rd   r^   rS   contentsrc   nbytesre   rf   rg   rh   ri   rj   rl   ndarraytestingassert_equalarangesize)r%   rq   rr   rs   rt   rz   bytesizerY   rX   c_arrirm   arrs                r&   test_memoryviewzTestNrtMemInfo.test_memoryview   s   <<#>>B&  5wwXbgg&2(>?rAU^^A.
;  rA !AENN1 2x0(b%))$**++Q'jjuELLENN,J %'


		#(( 3a 7=q rAU^^A.A6  	$/r)   c                    ddl m}m}m}m} t        j                  t
        j                        }|j                  dz  }t        j                  |d      }| j                  |j                  d       |j                  } | ||       ||dz              }	t        d      D ]!  }
| j                  |	j                  |
   d       # t        d      D ]  }
|
dz   |	j                  |
<    t        j                   |||j                  z  |      }| j                  |j                  d       ~t
        j"                  j%                  t        j&                  |j(                        dz   |       |dz  }t        d      D ]$  }
| j                  |	j                  |
   |
d	z          & | j                  |j                   j                  |       ~y )
Nr   rp   ru   Trv   r    rx   ry   rR   )r}   rq   rr   rs   rt   r~   rz   r   rk   r   r5   rS   rT   rd   r^   r   r   r   r   r   r   )r%   rq   rr   rs   rt   rz   r   rY   rX   r   r   r   s               r&   test_bufferzTestNrtMemInfo.test_buffer   si   <<#>>B&  5a(wwXd^WX]%;<rAU^^A.
;  rA !AENN1  jjuH,F "$a(


		#(( 3a 7=q rAU^^A.A6  	$/r)   c                 &   t         j                  j                  dz  dz  }dD ]`  }| j                  t              5 }t        j                  ||       d d d        | j                  d| dt        j                               b y # 1 sw Y   5xY w)N   rR   TFrv   zRequested allocation of z bytes failed.)
r   size_tmaxvalr7   MemoryErrorr   r5   r;   r<   r=   )r%   r   predrA   s       r&   test_allocate_invalid_sizez)TestNrtMemInfo.test_allocate_invalid_size   sy     ||""a'1,!D"";/6##Dt4 0MM4TF.If../1 "//s   BB	c                     d}dD ]b  }| j                  t              5 }t        j                  ||       d d d        d| d}| j	                  |t        j                               d y # 1 sw Y   7xY w)Nir   rv   z,Cannot allocate a negative number of bytes: .)r7   
ValueErrorr   r5   r;   r<   r=   )r%   r   r   rA   rC   s        r&   test_allocate_negative_sizez*TestNrtMemInfo.test_allocate_negative_size   sf     !D"":.&##Dt4 /@aHCMM#s6#3#345	 "..s   A++A4	)r,   r-   r.   rE   rM   rZ   ra   rn   r   r   r   r   r   __classcell__rO   s   @r&   rH   rH   I   sA    ,)")&),$RH 1 16r)   rH   c                       e Zd ZdZd Zd Zy)TestTracemalloczB
    Test NRT-allocated memory can be tracked by tracemalloc.
    c                 ,   	 dd l }j                          	 |j	                         } |       }|j	                         }~|j                  |d      |j                          S # t        $ r | j                  d       Y ww xY w# |j                          w xY w)Nr   ztracemalloc not availablelineno)tracemallocImportErrorskipTeststarttake_snapshot
compare_tostop)r%   funcr   beforeresafters         r&   measure_memory_diffz#TestTracemalloc.measure_memory_diff   s    	7 		 ..0F&C--/E##FH5  	7MM56	7 s   A! 9B !A>=A>Bc                 x  	 dt         j                  	t        	fd       fd}fd}|j                  j                  dz   }         | j                  |      }|d   }| j                  |j                         | j                  |j                  dz  d	       |j                  d   }| j                  t        j                  j                  |j                        d
       | j                  |j                  |       | j                  |      }|d   }| j                  |j                  dz         y )Ni@B c                  0    t        j                         S )z<
            Allocate and return a large array.
            )r~   empty)Nrz   s   r&   alloc_nrt_memoryz7TestTracemalloc.test_snapshot.<locals>.alloc_nrt_memory
  s    
 88Au%%r)   c                               S Nr/   r   s   r&   keep_memoryz2TestTracemalloc.test_snapshot.<locals>.keep_memory  s    #%%r)   c                                y r   r/   r   s   r&   release_memoryz5TestTracemalloc.test_snapshot.<locals>.release_memory  s	    r)   r    r   g=
ףp=?zpUnexpected allocation overhead encountered. May be due to difference in CPython builds or running under coverage)rC   ztest_nrt.pyg{Gz?)r~   int8r   __code__co_firstlinenor   assertGreaterEqualr   
assertLess	tracebackrS   ospathbasenamefilenamer   )
r%   r   r   alloc_linenodiffstatframer   r   rz   s
          @@@r&   test_snapshotzTestTracemalloc.test_snapshot  s   		& 
	&	&	 #++::Q> 	 ''4Aw		1-		1u9@ 	 	B q!))%..9=I|4 ''7Aw		1t8,r)   N)r,   r-   r.   rE   r   r   r/   r)   r&   r   r      s     (-r)   r   c                   H    e Zd Zd Zd Zd Zd Zej                  d        Z	y)TestNRTIssuec                    t         d        t         fd       t         fd       }d}t        j                  j                  |df      } |||      }|j                  ||      }t        j                  j                  ||       y)zO
        GitHub Issue #1244 https://github.com/numba/numba/issues/1244
        c                 H    | \  }}t        j                  |dz  |dz  z         S )NrR   )mathsqrt)vectorxys      r&   calculate_2D_vector_magzNTestNRTIssue.test_issue_with_refct_op_pruning.<locals>.calculate_2D_vector_mag6  s&    DAq99Q!Va1f_--r)   c                     t        j                  dt         j                        } |       }| \  }}||z  |d<   ||z  |d<   |S NrR   rz   r   r    )r~   r   float64)r   normalized_vectormagr   r   r   s        r&   normalize_2D_vectorzJTestNRTIssue.test_issue_with_refct_op_pruning.<locals>.normalize_2D_vector<  sO     ""** =)&1CDAq#$s7a #$s7a $$r)   c                     t        j                  | dft         j                        }t        |       D ]#  }||   } |      }|d   ||df<   |d   ||df<   % |S r   )r~   r   r   r^   )num_vectorsvectorsnormalized_vectorsr   r   r   r   s         r&   normalize_vectorszHTestNRTIssue.test_issue_with_refct_op_pruning.<locals>.normalize_vectorsH  sp    !#;*:"**!M;' $7$?!+<Q+?"1a4(+<Q+?"1a4( ( &%r)   ru   rR   N)r   r~   randompy_funcr   assert_almost_equal)r%   r   r   test_vectorsgotexpectedr   r   s         @@r&    test_issue_with_refct_op_pruningz-TestNRTIssue.test_issue_with_refct_op_pruning2  s     
	. 
	.
 
		% 
		% 
	& 
	& yy''a(89\:$,,[,G


&&x5r)   c                 d   d } t         t        j                  t        j                  t        j                  t        j
                  dd      f                   |      } |       \  }}| j                  |d       | j                  |t        j                  dt        j
                               y )Nc                  P    dt        j                  dt         j                        fS )Ng        r    r   )r~   zerosint32r/   r)   r&   fz.TestNRTIssue.test_incref_after_cast.<locals>.f`  s    "((333r)   r    Cy                r   )	r   r   Tuple
complex128Arrayr   assertPreciseEqualr~   r   )r%   r   cfunczr   s        r&   test_incref_after_castz#TestNRTIssue.test_incref_after_cast]  s    	4
K Gekk5#3#3#(;;u{{As#C#F G J KKLN32&RXXarxx%@Ar)   c                     t         d        } |       \  }}t        j                  j                  ||       t        j                  j                  |t        j                  dt        j
                               y )Nc                      t        j                  dt         j                        } t        j                  dt         j                        }| |d d  fS )Nru   r   )r~   onesr   )abs     r&   r   z5TestNRTIssue.test_refct_pruning_issue_1511.<locals>.fl  s7    "**-A"**-Aad7Nr)   ru   r   )r   r~   r   r   r   r   )r%   r   r   r   s       r&   test_refct_pruning_issue_1511z*TestNRTIssue.test_refct_pruning_issue_1511k  sX    		 
	
 s1


1%


2772RZZ#@Ar)   c                     t         d        }t        j                  g dg      }|j                  |dd      } ||dd      }| j	                  ||       y )Nc                     t        j                  | dk(        }t        |d         dk(  rd}d}n|d   d   }|d   d   }|dz
  |z   }|dz
  |z   }||fS )Nr    r   )r~   wherelen)imager   r   next_locy_offsetx_offset
next_loc_x
next_loc_ys           r&   udtz7TestNRTIssue.test_refct_pruning_issue_1526.<locals>.udtw  sp    xx
+H8A;1$#A;q>#A;q>a%8+Ja%8+Jz))r)   )
r    r   r    r   r    r   r   r    r   r   r       )r   r~   arrayr   rS   )r%   r   r   expectr   s        r&   test_refct_pruning_issue_1526z*TestNRTIssue.test_refct_pruning_issue_1526v  sW    		* 
	* HH456Q1%!Ql%r)   c                 $   ddl m} | j                  t        j                         |d        }| j                  t        j                         | j                   |       |j                                | j                  t        j                         y )Nr   r   c                       y)N{   r/   r/   r)   r&   fooz8TestNRTIssue.test_no_nrt_on_njit_decoration.<locals>.foo  s    r)   )numbar   re   r   r9   rS   r   ri   )r%   r   r  s      r&   test_no_nrt_on_njit_decorationz+TestNRTIssue.test_no_nrt_on_njit_decoration  so     	 	% 
	 
	 	% 	. 	$r)   N)
r,   r-   r.   r   r   r   r  r   run_test_in_subprocessr
  r/   r)   r&   r   r   1  s3    )6VB	B&. $$% %%r)   r   c                   `    e Zd ZdZd Z ej                  d      d        Zee	d               Z
y)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                 6   | j                   }t        |j                               }|D cg c]	  }d|v s| }}|D cg c]	  }d|v s| }}t        j                  |      }t        |j                               }|D cg c]	  }d|v s| }}|D cg c]	  }d|v s| }	}| j                  ||       | j                  ||	       t        |      t        |      z
  }
t        |      t        |	      z
  }|
|z  }| j                  ||
|z         dj                  |      }dD ]%  }dj                  |      }| j                  ||       ' | j                  t        t        |j                                     t        |             y c c}w c c}w c c}w c c}w )N
NRT_incref
NRT_decref
)r    rR            z	; GONE {})sample_llvm_irlist
splitlinesr   _remove_redundant_nrt_refctassertNotEqualsetrS   joinformatr;   r   )r%   input_irinput_lineslnbefore_increfsbefore_decrefs	output_iroutput_linesafter_increfsafter_decrefspruned_increfspruned_decrefscombinedpruned_linesr   gones                   r&   test_refct_pruning_op_recognizez0TestRefCtPruning.test_refct_pruning_op_recognize  s   &&8..01'2I{lb6H"{I'2I{lb6H"{I 66x@	I0023&2Illb6HlI&2Illb6HlI 	NM:NM:^,s=/AA^,s=/AA "N2>N#BCyy* !A%%a(DMM$- ! 	T,"9"9";<=s8}M5 JI
 JIs-   	FF	FF 	F
F	FFz3Pass removed as it was buggy. Re-enable when fixed.c                 d   t         d        t         d        t         fd       t        d      fd       }|j                  d       t        |j                  |j                  d               }t        t        j                  d	|            }| j                  t        |      d       y
)ztestcase from #2350c                 N    t        j                  |      s||| fxx   dz  cc<   y y r   r~   isnanr   r   aggfields       r&   _append_non_nazITestRefCtPruning.test_refct_pruning_with_branches.<locals>._append_non_na  s#    88E?AqD	Q	 #r)   c                     t        j                  |      s2t        j                  ||| f         r|||| f<   y ||| fxx   |z  cc<   y y r   r.  r0  s       r&   _appendzBTestRefCtPruning.test_refct_pruning_with_branches.<locals>._append  sE    88E?88C1I& %C1I1I&I	 #r)   c                 2     | |||        | |||       y r   r/   )r   r   r1  r2  r5  r3  s       r&   appendzATestRefCtPruning.test_refct_pruning_with_branches.<locals>.append  s    1ae,Aq#u%r)   T)no_cpython_wrapperc                     t        | j                  d         D ]*  }t        | j                  d         D ]  } ||| |        , y )Nr   r    )r^   r{   )r   r2  r   jr7  s       r&   extendzATestRefCtPruning.test_refct_pruning_with_branches.<locals>.extend&  s?    399Q<(syy|,A1ae, - )r)   z(f4[:,::1], f4)r   z!(NRT_incref|NRT_decref)\([^\)]+\)N)
r   compiler<   inspect_llvm
signaturesr  refinditerrS   r   )r%   r;  llvmirrefopsr5  r3  r7  s       @@@r&    test_refct_pruning_with_branchesz1TestRefCtPruning.test_refct_pruning_with_branches  s     
	 
	 
	' 
	' 
	& 
	& 
	&	- 
'	- 	() V(():):1)=>?bkk"FOPVa(r)   c                    t        j                          t        j                          t        j                          t        j                          t
        d        t        dg      fd       }| j                   |d      d       y)ztThe InlineAsm class from llvmlite.ir has no 'name' attr the refcount
        pruning pass should be tolerant to thisc                 z    d }t        t        j                  t        j                  t        j                        |fS )Nc                    |\  }}t        j                  t        j                  d      t        j                  d      t        j                  d      g      }|j                  |dd||fdd      }t	        | ||j
                  |      S )N    zmov $2, $0; imul $1, $0z=&r,r,rasm_mulF)nameside_effect)r   FunctionTypeIntTypeasmr   return_type)cgctxbuildersigr@   arg_0arg_1ftymuls           r&   codegenz>TestRefCtPruning.test_inline_asm.<locals>.bar.<locals>.codegen@  s|    !%oobjjnrzz"~79zz"~7G Hkk#'@)#(%.y.3 " 5 *%#//3OOr)   )r   r   r   )tyctxr   r   rV  s       r&   barz-TestRefCtPruning.test_inline_asm.<locals>.bar>  s*    P U[[%++u{{CWLLr)   zint32(int32)c                 $    | dz  }  | d      }|S )Nr    rR   r/   )r   r   rX  s     r&   r  z-TestRefCtPruning.test_inline_asm.<locals>.fooJ  s    FAAq	AHr)   ru      N)llvmrJ   initialize_native_targetinitialize_native_asmprinterinitialize_native_asmparserr
   r   rS   )r%   r  rX  s     @r&   test_inline_asmz TestRefCtPruning.test_inline_asm4  sy    
 	%%'))+((*			M 
		M 
~		 
 	
 	R"%r)   N)r,   r-   r.   r  r+  unittestskiprC  
linux_onlyx86_onlyr_  r/   r)   r&   r  r    sO    FNPN> X]]HI") J")H &  &r)   r  c                   @     e Zd ZdZ fdZd Zd Zd Zd Zd Z	 xZ
S )TestNrtExternalCFFIz?Testing the use of externally compiled C code that use NRT
    c                 J    t         j                   t        t        |           y r   )r   rK   rL   re  rM   rN   s    r&   rM   zTestNrtExternalCFFI.setUpW  s    !!!4.0r)   c                    ddl m}  |       }|j                  ||t               g       |j	                  |       t        dj                  |            }|j                  |       t        j                  j                  |       	 t        |      }t        j                  j                  |       ||fS # t        j                  j                  |       w xY w)Nr   FFI)include_dirszcffi_test_{})tmpdir)cffiri  
set_sourcer   cdefr   r  r<  sysr   r7  r   remove)r%   rI  sourcern  ri  ffirk  mods           r&   compile_cffi_modulez'TestNrtExternalCFFI.compile_cffi_module\  s    etV<>2BC 5 5d ;<6"	$ &CHHOOF#Cx HHOOF#s   B6 6!Cc                 t    ddl m}  |       }|j                  dt        j                  d         } |       }|S )Nr   rh  zvoid* (*)()get_api)rl  ri  rt   r   	c_helpers)r%   ri  rr  nrt_get_apitables        r&   get_nrt_api_tablez%TestNrtExternalCFFI.get_nrt_api_tablel  s3    ehh}k.C.CI.NOr)   c                 V   dj                  | j                  j                        }d}d}| j                  |||      \  }}| j	                  |j
                  j                  d       | j                         }|j
                  j                  |      }| j	                  |j
                  j                  d       t        |j                  d|            }t        j                  |      }	| j	                  |	j                  d       ~	| j	                  |j
                  j                  d       y )	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  rO   r,   rt  rS   libstatusrz  test_nrt_apiintrt   r	   MemInforT   )
r%   rI  rq  rn  rr  rs  ry  outmi_addrrY   s
             r&   test_manage_memoryz&TestNrtExternalCFFI.test_manage_memoryt  s    &--dnn.E.EF(
 ++D&$?S+&&(gg""5)1chhx-.[[!a(0r)   c                    dj                  | j                  j                        }d}d}| j                  |||      \  }}| j	                         }dt        j                  t
        j                        j                  z  }|j                  j                  ||      }t        |j                  d|            }	t        j                  |	      }
| j                  |
j                   d       |j#                  |j                  dj                  |      |
j$                              }t        j&                  dt
        j                  |	      }t
        j(                  j+                  |g d
       y )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  )r{   rz   r|   )i  i
  i  )r  rO   r,   rt  rz  r~   rz   intprk   r|  r~  r  rt   r	   r  rS   rT   r|   rd   r   r   r   )r%   rI  rq  rn  rr  rs  ry  numbytesr  r  rY   r|   r   s                r&   test_allocatez!TestNrtExternalCFFI.test_allocate  s   !(()@)@A :++D&$?S&&(rxx(111gg""5(3chhx-.[[!a(CHH[%7%7%A277KLjjt2776B


%:;r)   c                     ddl m} t        d        } |       }t        |j	                  d| j                                     } |       }| j                  ||       y )Nr   rh  c                      t               S r   r   r/   r)   r&   r~  z6TestNrtExternalCFFI.test_get_api.<locals>.test_nrt_api  s
    = r)   r   )rl  ri  r   r  rt   rz  rS   )r%   ri  r~  rr  r  r   s         r&   test_get_apiz TestNrtExternalCFFI.test_get_api  sU    		! 
	! eSXXh(>(>(@ABn%r)   )r,   r-   r.   rE   rM   rt  rz  r  r  r  r   r   s   @r&   re  re  S  s(    1
 '1R<@
&r)   re  c                   B    e Zd Zd Zd Zd Zd Zd Zd Zd Z	d Z
d	 Zy
)TestNrtStatisticsc                 6    t        j                         | _        y r   )r   memsys_stats_enabled_TestNrtStatistics__stats_stater$   s    r&   rM   zTestNrtStatistics.setUp  s    (==?r)   c                 n    | j                   rt        j                          y t        j                          y r   )r  r   memsys_enable_statsmemsys_disable_statsr$   s    r&   tearDownzTestNrtStatistics.tearDown  s#    ++-,,.r)   c                 h    d}t         j                  j                         }d|d<   t        ||       y )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  s      r&   test_stats_env_var_explicit_onz0TestNrtStatistics.test_stats_env_var_explicit_on  s0    8 jjoo!$#3'r)   c                 "    d}t        ||       y )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  s      r&   check_env_var_offz#TestNrtStatistics.check_env_var_off  s     	#3'r)   c                 l    t         j                  j                         }d|d<   | j                  |       y )N0r  )r   r  r  r  r%   r  s     r&   test_stats_env_var_explicit_offz1TestNrtStatistics.test_stats_env_var_explicit_off  s+    jjoo!$s#r)   c                     t         j                  j                         }|j                  dd        | j	                  |       y )Nr  )r   r  r  popr  r  s     r&   test_stats_env_var_default_offz0TestNrtStatistics.test_stats_env_var_default_off  s1     jjoo!4(s#r)   c                 b   t         d        }t        j                          | j                  t        j                                t        d      D ]  }t        j                         }t        j                          | j                  t        j                                 |        t        j                          | j                  t        j                                t        j                         } |        t        j                         }| j                  ||       | j                  ||        y )Nc                  b    t        j                  d      } t        j                  d| d   z        S )Nr  r  r   )r~   r   r   )tmps    r&   r  z7TestNrtStatistics.test_stats_status_toggle.<locals>.foo  s&    ''!*C99QQZ((r)   rR   )r   r   r  ri   r  r^   r   get_allocation_statsr  re   rS   r   )r%   r  r   stats_1stats_2stats_3s         r&   test_stats_status_togglez*TestNrtStatistics.test_stats_status_toggle  s    		) 
	)
 	'')88:;qA002G,,.[==?@E++-OOK<<>?002GE002GWg.OOGW-- r)   c                 8   t        j                          | j                  t        j                                | j	                  t
              5 }t        j                          d d d        | j                  dt        j                               y # 1 sw Y   /xY w)NNRT stats are disabled.)r   r  re   r  r7   r8   r   r  r;   r<   r=   )r%   rA   s     r&   5test_rtsys_stats_query_raises_exception_when_disabledzGTestNrtStatistics.test_rtsys_stats_query_raises_exception_when_disabled=  sm     	((*99;<|,&&( - 	/V5E5E1FG -,s   BBc                    d}|D ]  }t        t        d|       }| j                  |      5  t        j                          | j	                  t        j
                                | j                  t              5 } |        d d d        | j                  dt        j                               d d d         y # 1 sw Y   9xY w# 1 sw Y   xY w)N)allocfreemi_allocmi_freememsys_get_stats_)
stats_funcr  )r:   r   subTestr  re   r  r7   r8   r;   r<   r=   )r%   method_variationsr?   r  rA   s        r&   <test_nrt_explicit_stats_query_raises_exception_when_disabledzNTestNrtStatistics.test_nrt_explicit_stats_query_raises_exception_when_disabledI  s     E%D 0A$.HIJ4002  !A!A!CD&&|4L 57V=M=M9NO 54 & 54	 54s$   AC:B:-C:C?CC	N)r,   r-   r.   rM   r  r  r  r  r  r  r  r  r/   r)   r&   r  r    s4    @/ (D(&$$".H
HPr)   r  __main__)9r   r   platformro  r?  numpyr~   r	  r   
numba.corer   numba.core.runtimer   r   r   r	   numba.core.extendingr
   r   numba.core.typingr   numba.core.imputilsr   llvmliter   llvmlite.bindingbindingr[  numba.core.unsafe.nrtr   numba.tests.supportr   r   r   r   r   r   r   numba.core.registryr   r`  skipIf
startswithrb  machinerc  objectr   r1   rH   r   r   r  re  r  r,   mainr/   r)   r&   <module>r     s\    	  
 	     9 ' 2   -F F F +  X__!8!8!AA.0
8??+8++-5GG*,F #8#4#4 #4e6X&& e6P=-h'' =-@r%8 r%jj&x(( j&Z s&-x s& s&lKP KP\ zHMMO r)   