o 8ήc@sddlmZmZmZddlmZmZedGdddeZedGdddeZedGd d d eZ edGd d d eZ d S))errorsirtypes)register_rewriteRewritezbefore-inferencec@ eZdZdZddZddZdS)RewriteConstGetitemsz Rewrite IR expressions of the kind `getitem(value=arr, index=$constXX)` where `$constXX` is a known constant as `static_getitem(value=arr, index=)`. c Csfi|_}||_|jddD]}|jdkr,z||j}Wn tjy'Yqw|||<qt|dkS)Ngetitemopr) getitemsblock find_exprsr infer_constantindexrConstantInferenceErrorlen)selffunc_irr typemap calltypesr exprconstrI/tmp/pip-target-vg8gfxp4/lib/python/numba/core/rewrites/static_getitem.pymatch s   zRewriteConstGetitems.matchcCs|j}||jjD]1}t|tjr9|j}||jvr9|j|}tj j |j||j |j d}tj||j |j d}||q |S)zC Rewrite all matching getitems as static_getitems. valuer index_varlocrtargetr)r copyclearbody isinstancerAssignrr Exprstatic_getitemrrr!append)r new_blockinstrrnew_exprrrrapplys"        zRewriteConstGetitems.applyN__name__ __module__ __qualname____doc__rr-rrrrr rzafter-inferencec@r)RewriteStringLiteralGetitemsz Rewrite IR expressions of the kind `getitem(value=arr, index=$XX)` where `$XX` is a StringLiteral value as `static_getitem(value=arr, index=)`. cCsfi|_}||_||_|jddD]}|jdkr,||jj}t|tj r,|j|j f||<qt |dkS)zh Detect all getitem expressions and find which ones have string literal indexes r r r) r r rrr rnamer%r StringLiteral literal_valuer)rrr rrr rindex_tyrrrr9s     z"RewriteStringLiteralGetitems.matchcCst|jj|jj}|jjD];}t|tjrD|j}||j vrD|j |\}}tj j |j||j |jd}|j ||j |<tj||j|jd}||q|S)z~ Rewrite all matching getitems as static_getitems where the index is the literal value of the string. rr )rBlockr scoperr$r%r&rr r'r(rrr!r))rr*r+rrlit_valr,rrrr-Is"      z"RewriteStringLiteralGetitems.applyNr.rrrrr41s r4c@r)RewriteStringLiteralSetitemsz Rewrite IR expressions of the kind `setitem(value=arr, index=$XX, value=)` where `$XX` is a StringLiteral value as `static_setitem(value=arr, index=, value=)`. cCs\i|_}||_||_|tjD]}||jj}t|t j r'|j|j f||<qt |dkS)zh Detect all setitem expressions and find which ones have string literal indexes r) setitemsr r find_instsrSetItemrr5r%rr6r7r)rrr rrr=r+r8rrrrgs    z"RewriteStringLiteralSetitems.matchcCst|jj|jj}|jjD]1}t|tjr:||jvr:|j|\}}tj |j ||j |j |jd}|j ||j |<|}||q|S)z~ Rewrite all matching setitems as static_setitems where the index is the literal value of the string. )r!rrrr)rr9r r:rr$r%r?r= StaticSetItemr!rrrr))rr*r+rr;new_instrrrr-vs    z"RewriteStringLiteralSetitems.applyNr.rrrrr<_r3r<c@r)RewriteConstSetitemsz Rewrite IR statements of the kind `setitem(target=arr, index=$constXX, ...)` where `$constXX` is a known constant as `static_setitem(target=arr, index=, ...)`. c Cs\i|_}||_|tjD]}z||j}Wn tjy"Yqw|||<qt |dkS)Nr) r=r r>rr?rrrrr)rrr rrr=r+rrrrrs   zRewriteConstSetitems.matchcCsh|j}||jjD]$}||jvr,|j|}t|j||j|j |j }| |q | |q |S)zC Rewrite all matching setitems as static_setitems. ) r r"r#r$r=rr@r!rrrr))rr*r+rrArrrr-s        zRewriteConstSetitems.applyNr.rrrrrBs rBN) numba.corerrrnumba.core.rewritesrrrr4r<rBrrrrs+-+