o
    Û¾iC  ã                   @   sH   d dl mZ d dlmZ G dd„ deƒZG dd„ deƒZeƒ ZejZ	dS )é    )Údefaultdict)Úconfigc                   @   s*   e Zd ZdZd	dd„Zdd„ Zdd„ ZdS )
ÚRewritez8Defines the abstract base class for Numba rewrites.
    Nc                 C   ó   dS )z+Constructor for the Rewrite class.
        N© )ÚselfÚstater   r   úP/home/ubuntu/.local/lib/python3.10/site-packages/numba/core/rewrites/registry.pyÚ__init__
   s   zRewrite.__init__c                 C   r   )z]Overload this method to check an IR block for matching terms in the
        rewrite.
        Fr   )r   Úfunc_irÚblockÚtypemapÚ	calltypesr   r   r	   Úmatch   s   zRewrite.matchc                 C   s   t dƒ‚)zgOverload this method to return a rewritten IR basic block when a
        match has been found.
        z Abstract Rewrite.apply() called!)ÚNotImplementedError©r   r   r   r	   Úapply   s   zRewrite.apply)N)Ú__name__Ú
__module__Ú__qualname__Ú__doc__r
   r   r   r   r   r   r	   r      s
    
r   c                   @   s4   e Zd ZdZeddgƒZdd„ Zdd„ Zdd	„ Zd
S )ÚRewriteRegistryz+Defines a registry for Numba rewrites.
    zbefore-inferencezafter-inferencec                 C   s   t tƒ| _dS )ziConstructor for the rewrite registry.  Initializes the rewrites
        member to an empty list.
        N)r   ÚlistÚrewritesr   r   r   r	   r
   !   s   zRewriteRegistry.__init__c                    s*   ˆ ˆj vrtdˆ f ƒ‚‡ ‡fdd„}|S )zf
        Decorator adding a subclass of Rewrite to the registry for
        the given *kind*.
        zinvalid kind %rc                    s,   t | tƒstd | ¡ƒ‚ˆjˆ   | ¡ | S )Nz {0} is not a subclass of Rewrite)Ú
issubclassr   Ú	TypeErrorÚformatr   Úappend)Úrewrite_cls©Úkindr   r   r	   Údo_register.   s   
ÿz-RewriteRegistry.register.<locals>.do_register)Ú_kindsÚKeyError)r   r    r!   r   r   r	   Úregister'   s   
zRewriteRegistry.registerc                 C   s&  || j v sJ ‚|jj}| ¡ }| j| D ]Y}||ƒ}t| ¡ ƒ}|rm| ¡ \}}	| |j|	|j	|j
¡}
|
rktjs;tjrNtdƒ td|j ƒ |	 ¡  tdƒ | ¡ }|||< | ||f¡ tjsctjrk| ¡  tdƒ |s"q| ¡ D ]\}}	|	|| kr€|	 ¡  qrddlm} | |j¡}| ¡  dS )zŽGiven a pipeline and a dictionary of basic blocks, exhaustively
        attempt to apply all registered rewrites to all basic blocks.
        ÚF______________________________________________________________________zREWRITING (%s):Ú<____________________________________________________________r   )ÚpostprocN)r"   r   ÚblocksÚcopyr   r   ÚitemsÚpopr   r   r   r   ÚDEBUGÚDUMP_IRÚprintr   Údumpr   r   ÚverifyÚ
numba.corer'   ÚPostProcessorÚrun)r   r    r   r(   Ú
old_blocksr   ÚrewriteÚ	work_listÚkeyr   ÚmatchesÚ	new_blockr'   Ú	post_procr   r   r	   r   6   s@   ÿñ€€zRewriteRegistry.applyN)	r   r   r   r   Ú	frozensetr"   r
   r$   r   r   r   r   r	   r      s    r   N)
Úcollectionsr   r1   r   Úobjectr   r   Úrewrite_registryr$   Úregister_rewriter   r   r   r	   Ú<module>   s    E
