o
    9n0jY                     @   sp  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	 G dd de
Zded	efd
dZdg dg fddZG dd de
ZG dd deZG dd deZe e_G dd deZd3dedefddZdededed efd!d"Zdededed efd#d$Zdededefd%d&Zdededed efd'd(Zdefd)d*Zd+ed,ed-efd.d/Zd4d+ed0ed-efd1d2ZdS )5    )StringIOIllegalStateException)RuleContext)ATN)ATNStatec                   @   sX   e Zd ZdZdZdZeZdefddZdd Z	d	d
 Z
dd ZdefddZdd ZdS )PredictionContextNi   cachedHashCodec                 C   s
   || _ d S Nr
   )selfr
    r   _/home/wildlama/miniconda3/envs/lam_a2e/lib/python3.10/site-packages/antlr4/PredictionContext.py__init__4      
zPredictionContext.__init__c                 C      dS Nr   r   r   r   r   r   __len__7      zPredictionContext.__len__c                 C   s
   | | j u S r   )EMPTYr   r   r   r   isEmpty;   r   zPredictionContext.isEmptyc                 C   s   |  t| d | jkS Nr	   )getReturnStatelenEMPTY_RETURN_STATEr   r   r   r   hasEmptyPath>   s   zPredictionContext.hasEmptyPathindexc                 C   s   t d)Nzillegal!r   r   r   r   r   r   r   A      z PredictionContext.getReturnStatec                 C      | j S r   r   r   r   r   r   __hash__D      zPredictionContext.__hash__)__name__
__module____qualname__r   r   ZglobalNodeCountidintr   r   r   r   r   r"   r   r   r   r   r      s    r   parentreturnStatec                 C   s    | d u rt dS t t | |fS )N )hashr)   r*   r   r   r   calculateHashCodeG   s    r.   parentsreturnStatesc                 C   s.   d}t | |D ]\}}t|t||f}q|S r   )zipr,   r.   )r/   r0   hr)   r*   r   r   r   calculateListsHashCodeJ   s   r3   c                   @   s8   e Zd Zdd ZdefddZdefddZdd	 Zd
S )PredictionContextCachec                 C   s   t  | _d S r   )dictcacher   r   r   r   r   V   s   zPredictionContextCache.__init__ctxc                 C   s8   |t jkrt jS | j|d }|d ur|S || j|< |S r   )r   r   r6   get)r   r7   existingr   r   r   add]   s   

zPredictionContextCache.addc                 C   s   | j |d S r   )r6   r8   )r   r7   r   r   r   r8   f   s   zPredictionContextCache.getc                 C   
   t | jS r   )r   r6   r   r   r   r   r   i   r   zPredictionContextCache.__len__N)r$   r%   r&   r   r   r:   r8   r   r   r   r   r   r4   T   s
    	r4   c                       sx   e Zd ZededefddZdedef fddZdd Zd	efd
dZ	d	efddZ
dd Zdd Zdd Z  ZS )SingletonPredictionContextr)   r*   c                 C   s"   |t jkr| d u rtjS t| |S r   )r   r   r<   r   r-   r   r   r   createo   s   
z!SingletonPredictionContext.createc                    s&   t ||}t | || _|| _d S r   )r.   superr   	parentCtxr*   )r   r)   r*   ZhashCode	__class__r   r   r   w   s   

z#SingletonPredictionContext.__init__c                 C   r   r   r   r   r   r   r   r   }   r   z"SingletonPredictionContext.__len__r   c                 C   r!   r   )r?   r   r   r   r   	getParent   r#   z$SingletonPredictionContext.getParentc                 C   r!   r   )r*   r   r   r   r   r      r#   z)SingletonPredictionContext.getReturnStatec                 C   s>   | |u rdS |d u rdS t |tsdS | j|jko| j|jkS NTF)
isinstancer<   r*   r?   r   otherr   r   r   __eq__   s   
z!SingletonPredictionContext.__eq__c                 C   r!   r   r   r   r   r   r   r"      r#   z#SingletonPredictionContext.__hash__c                 C   sP   | j d u rdnt| j }t|dkr| j| jkrdS t| jS t| jd | S )Nr+   r   $ )r?   strr   r*   r   )r   upr   r   r   __str__   s   
z"SingletonPredictionContext.__str__)r$   r%   r&   staticmethodr   r(   r=   r   r   rB   r   rG   r"   rL   __classcell__r   r   r@   r   r<   m   s    
r<   c                       s<   e Zd Z fddZdd Zdd Zdd Zd	d
 Z  ZS )EmptyPredictionContextc                    s   t  d tj d S r   )r>   r   r   r   r   r@   r   r   r      s   zEmptyPredictionContext.__init__c                 C   r   )NTr   r   r   r   r   r      r   zEmptyPredictionContext.isEmptyc                 C   s   | |u S r   r   rE   r   r   r   rG      r    zEmptyPredictionContext.__eq__c                 C   r!   r   r   r   r   r   r   r"      r#   zEmptyPredictionContext.__hash__c                 C   r   )NrH   r   r   r   r   r   rL      r   zEmptyPredictionContext.__str__)	r$   r%   r&   r   r   rG   r"   rL   rN   r   r   r@   r   rO      s    rO   c                       sj   e Zd Zdedef fddZdd Zdd Zd	efd
dZd	efddZ	dd Z
dd Zdd Z  ZS )ArrayPredictionContextr/   r0   c                    s"   t  t|| || _|| _d S r   )r>   r   r3   r/   r0   )r   r/   r0   r@   r   r   r      s   
zArrayPredictionContext.__init__c                 C   s   | j d tjkS r   )r0   r   r   r   r   r   r   r      s   zArrayPredictionContext.isEmptyc                 C   r;   r   )r   r0   r   r   r   r   r      r   zArrayPredictionContext.__len__r   c                 C   
   | j | S r   )r/   r   r   r   r   rB      r   z ArrayPredictionContext.getParentc                 C   rQ   r   )r0   r   r   r   r   r      r   z%ArrayPredictionContext.getReturnStatec                 C   sF   | |u rdS t |tsdS t| t|krdS | j|jko"| j|jkS rC   )rD   rP   r,   r0   r/   rE   r   r   r   rG      s   
zArrayPredictionContext.__eq__c                 C   s   |   rdS t ^}|d tdt| jD ]?}|dkr"|d | j| tjkr0|d q|t| j|  | j	| d urQ|d |t| j	|  q|d q|d |
 W  d    S 1 sjw   Y  d S )	Nz[][r   z, rH   rI   null])r   r   writeranger   r0   r   r   rJ   r/   getvalue)r   bufir   r   r   rL      s$   




$zArrayPredictionContext.__str__c                 C   r!   r   r   r   r   r   r   r"      r#   zArrayPredictionContext.__hash__)r$   r%   r&   listr   r   r   r(   rB   r   rG   rL   r"   rN   r   r   r@   r   rP      s    
rP   NatnouterContextc                 C   sZ   |d u rt j}|jd u s|t ju rtjS t| |j}| j|j }|jd }t	||j
jS r   )r   r   r?   r    PredictionContextFromRuleContextZstatesZinvokingStateZtransitionsr<   r=   ZfollowStateZstateNumber)r[   r\   r)   stateZ
transitionr   r   r   r]      s   
r]   abrootIsWildcard
mergeCachec                 C   s   | |kr| S t | trt |trt| |||S |r't | tr | S t |tr'|S t | tr5t| jg| jg} t |trCt|jg|jg}t| |||S r   )rD   r<   mergeSingletonsrO   rP   r?   r*   mergeArrays)r_   r`   ra   rb   r   r   r   merge   s   



re   c           
      C   s  |d ur | | |fd }|d ur|S | || fd }|d ur |S t| ||}|d ur6|d ur4||| |f< |S | j|jkrft| j|j||}|| jkrL| S ||jkrS|S t|| j}|d urd||| |f< |S d }| |u sw| jd urz| j|jkrz| j}|d ur| j|jg}| j|jkr|j| jg}||g}	t|	|}|d ur||| |f< |S | j|jg}| j|jg}	| j|jkr|j| jg}|j| jg}	t|	|}|d ur||| |f< |S r   )r8   	mergeRootr*   re   r?   r<   r=   rP   )
r_   r`   ra   rb   previousmergedr)   ZsingleParentpayloadsr/   r   r   r   rc   4  sV   



rc   c                 C   s   |r| t jkr
t jS |t jkrt jS d S | t jkr!|t jkr!t jS | t jkr6|jt jg}|jd g}t||S |t jkrK| jt jg}| jd g}t||S d S r   )r   r   r*   r   r?   rP   )r_   r`   ra   ri   r/   r   r   r   rf     s"   







rf   c                 C   s  |d ur | | |fd }|d ur|S | || fd }|d ur |S d}d}d}d gt| jt|j  }d gt| }	|t| jk r|t|jk r| j| }
|j| }| j| |j| kr| j| }|tjkom|
d u om|d u }|
d uoy|d uoy|
|k}|s~|r|
|	|< |||< nt|
|||}||	|< |||< |d7 }|d7 }n)| j| |j| k r|
|	|< | j| ||< |d7 }n||	|< |j| ||< |d7 }|d7 }|t| jk r|t|jk sH|t| jk rt|t| jD ]}| j| |	|< | j| ||< |d7 }qnt|t|jD ]}|j| |	|< |j| ||< |d7 }q|t|	k rL|dkr@t	|	d |d }|d ur>||| |f< |S |	d| }	|d| }t
|	|}|| krc|d ura| || |f< | S ||kru|d urs||| |f< |S t|	 |d ur||| |f< |S )Nr   r	   )r8   r   r0   r/   r   r   re   rV   r<   r=   rP   combineCommonParents)r_   r`   ra   rb   rg   rY   jkZmergedReturnStatesZmergedParentsZa_parentZb_parentpayloadZbothDollarsZax_axZmergedParentprh   r   r   r   rd     s   















rd   c                 C   sb   t  }tdt| D ]}| | }||d d u r|||< q
tdt| D ]
}|| |  | |< q$d S r   )r5   rV   r   r8   )r/   ZuniqueParentsrn   r)   r   r   r   rj     s   rj   contextcontextCachevisitedc           	         s4     r S | }|d ur|S | }|d ur || < |S d}d gt  }tdt|D ])}t |||}|sD| |urY|sU fddtt D }d}|||< q0|sg|   | <  S d }t|dkrstj}nt|dkrt	
|d  d}nt| j}|| |||< || < |S )NFr   c                    s   g | ]}  |qS r   )rB   ).0rk   ro   r   r   
<listcomp>2  s    z.getCachedPredictionContext.<locals>.<listcomp>Tr	   )r   r8   r   rV   getCachedPredictionContextrB   r:   r   r   r<   r=   r   rP   r0   )	ro   rp   rq   r9   changedr/   rY   r)   updatedr   rs   r   ru   "  sB   



ru   nodesc                 C   s   |d u rt  }t| ||S |d u rt }t| ||S | d u s&|| d d ur(|S || |  ||  tdt| D ]}t| ||| q:|S r   )	rZ   getAllContextNodesr5   r8   putr:   rV   r   rB   )ro   rx   rq   rY   r   r   r   ry   `  s   
ry   r   )NN)ior   Zantlr4.error.Errorsr   Zantlr4.RuleContextr   Zantlr4.atn.ATNr   Zantlr4.atn.ATNStater   objectr   r(   r.   r3   r4   r<   rO   r   rP   r]   boolr5   re   rc   rf   rd   rZ   rj   ru   ry   r   r   r   r   <module>   s*   8
1>4a'[>