1from collections import defaultdict
2import copy
3import itertools
4import os
5import linecache
6import pprint
7import re
8import sys
9import operator
10from types import FunctionType, BuiltinFunctionType
11from functools import total_ordering
12from io import StringIO
13
14from numba.core import errors, config
15from numba.core.utils import (BINOPS_TO_OPERATORS, INPLACE_BINOPS_TO_OPERATORS,
16                              UNARY_BUITINS_TO_OPERATORS, OPERATORS_TO_BUILTINS)
17from numba.core.errors import (NotDefinedError, RedefinedError,
18                               VerificationError, ConstantInferenceError)
19from numba.core import consts
20
21# terminal color markup
22_termcolor = errors.termcolor()
23
24
25class Loc(object):
26    """Source location
27
28    """
29    _defmatcher = re.compile('def\s+(\w+)\(.*')
30
31    def __init__(self, filename, line, col=None, maybe_decorator=False):
32        """ Arguments:
33        filename - name of the file
34        line - line in file
35        col - column
36        maybe_decorator - Set to True if location is likely a jit decorator
37        """
38        self.filename = filename
39        self.line = line
40        self.col = col
41        self.lines = None # the source lines from the linecache
42        self.maybe_decorator = maybe_decorator
43
44    def __eq__(self, other):
45        # equivalence is solely based on filename, line and col
46        if type(self) is not type(other): return False
47        if self.filename != other.filename: return False
48        if self.line != other.line: return False
49        if self.col != other.col: return False
50        return True
51
52    def __ne__(self, other):
53        return not self.__eq__(other)
54
55    @classmethod
56    def from_function_id(cls, func_id):
57        return cls(func_id.filename, func_id.firstlineno, maybe_decorator=True)
58
59    def __repr__(self):
60        return "Loc(filename=%s, line=%s, col=%s)" % (self.filename,
61                                                      self.line, self.col)
62
63    def __str__(self):
64        if self.col is not None:
65            return "%s (%s:%s)" % (self.filename, self.line, self.col)
66        else:
67            return "%s (%s)" % (self.filename, self.line)
68
69    def _find_definition(self):
70        # try and find a def, go backwards from error line
71        fn_name = None
72        lines = self.get_lines()
73        for x in reversed(lines[:self.line - 1]):
74            # the strip and startswith is to handle user code with commented out
75            # 'def' or use of 'def' in a docstring.
76            if x.strip().startswith('def '):
77                fn_name = x
78                break
79
80        return fn_name
81
82    def _raw_function_name(self):
83        defn = self._find_definition()
84        if defn:
85            return self._defmatcher.match(defn.strip()).groups()[0]
86        else:
87            # Probably exec(<string>) or REPL.
88            return None
89
90    def get_lines(self):
91        if self.lines is None:
92
93            self.lines = linecache.getlines(self._get_path())
94
95        return self.lines
96
97    def _get_path(self):
98        path = None
99        try:
100            # Try to get a relative path
101            # ipython/jupyter input just returns as self.filename
102            path = os.path.relpath(self.filename)
103        except ValueError:
104            # Fallback to absolute path if error occurred in getting the
105            # relative path.
106            # This may happen on windows if the drive is different
107            path = os.path.abspath(self.filename)
108        return path
109
110
111    def strformat(self, nlines_up=2):
112
113        lines = self.get_lines()
114
115        use_line = self.line
116
117        if self.maybe_decorator:
118            # try and sort out a better `loc`, if it's suspected that this loc
119            # points at a jit decorator by virtue of
120            # `__code__.co_firstlineno`
121
122            # get lines, add a dummy entry at the start as lines count from
123            # 1 but list index counts from 0
124            tmplines = [''] + lines
125
126            if lines and use_line and 'def ' not in tmplines[use_line]:
127                # look forward 10 lines, unlikely anyone managed to stretch
128                # a jit call declaration over >10 lines?!
129                min_line = max(0, use_line)
130                max_line = use_line + 10
131                selected = tmplines[min_line : max_line]
132                index = 0
133                for idx, x in enumerate(selected):
134                    if 'def ' in x:
135                        index = idx
136                        break
137                use_line = use_line + index
138
139
140        ret = [] # accumulates output
141        if lines and use_line:
142
143            def count_spaces(string):
144                spaces = 0
145                for x in itertools.takewhile(str.isspace, str(string)):
146                    spaces += 1
147                return spaces
148
149            # A few places in the code still use no `loc` or default to line 1
150            # this is often in places where exceptions are used for the purposes
151            # of flow control. As a result max is in use to prevent slice from
152            # `[negative: positive]`
153            selected = lines[max(0, use_line - nlines_up):use_line]
154
155            # see if selected contains a definition
156            def_found = False
157            for x in selected:
158                if 'def ' in x:
159                    def_found = True
160
161            # no definition found, try and find one
162            if not def_found:
163                # try and find a def, go backwards from error line
164                fn_name = None
165                for x in reversed(lines[:use_line - 1]):
166                    if 'def ' in x:
167                        fn_name = x
168                        break
169                if fn_name:
170                    ret.append(fn_name)
171                    spaces = count_spaces(x)
172                    ret.append(' '*(4 + spaces) + '<source elided>\n')
173
174            if selected:
175                ret.extend(selected[:-1])
176                ret.append(_termcolor.highlight(selected[-1]))
177
178                # point at the problem with a caret
179                spaces = count_spaces(selected[-1])
180                ret.append(' '*(spaces) + _termcolor.indicate("^"))
181
182        # if in the REPL source may not be available
183        if not ret:
184            ret = "<source missing, REPL/exec in use?>"
185
186        err = _termcolor.filename('\nFile "%s", line %d:')+'\n%s'
187        tmp = err % (self._get_path(), use_line, _termcolor.code(''.join(ret)))
188        return tmp
189
190    def with_lineno(self, line, col=None):
191        """
192        Return a new Loc with this line number.
193        """
194        return type(self)(self.filename, line, col)
195
196    def short(self):
197        """
198        Returns a short string
199        """
200        shortfilename = os.path.basename(self.filename)
201        return "%s:%s" % (shortfilename, self.line)
202
203
204# Used for annotating errors when source location is unknown.
205unknown_loc = Loc("unknown location", 0, 0)
206
207
208@total_ordering
209class SlotEqualityCheckMixin(object):
210    # some ir nodes are __dict__ free using __slots__ instead, this mixin
211    # should not trigger the unintended creation of __dict__.
212    __slots__ = tuple()
213
214    def __eq__(self, other):
215        if type(self) is type(other):
216            for name in self.__slots__:
217                if getattr(self, name) != getattr(other, name):
218                    return False
219            else:
220                return True
221        return False
222
223    def __le__(self, other):
224        return str(self) <= str(other)
225
226    def __hash__(self):
227        return id(self)
228
229
230@total_ordering
231class EqualityCheckMixin(object):
232    """ Mixin for basic equality checking """
233
234    def __eq__(self, other):
235        if type(self) is type(other):
236            def fixup(adict):
237                bad = ('loc', 'scope')
238                d = dict(adict)
239                for x in bad:
240                    d.pop(x, None)
241                return d
242            d1 = fixup(self.__dict__)
243            d2 = fixup(other.__dict__)
244            if d1 == d2:
245                return True
246        return False
247
248    def __le__(self, other):
249        return str(self) < str(other)
250
251    def __hash__(self):
252        return id(self)
253
254
255class VarMap(object):
256    def __init__(self):
257        self._con = {}
258
259    def define(self, name, var):
260        if name in self._con:
261            raise RedefinedError(name)
262        else:
263            self._con[name] = var
264
265    def get(self, name):
266        try:
267            return self._con[name]
268        except KeyError:
269            raise NotDefinedError(name)
270
271    def __contains__(self, name):
272        return name in self._con
273
274    def __len__(self):
275        return len(self._con)
276
277    def __repr__(self):
278        return pprint.pformat(self._con)
279
280    def __hash__(self):
281        return hash(self.name)
282
283    def __iter__(self):
284        return self._con.iterkeys()
285
286    def __eq__(self, other):
287        if type(self) is type(other):
288            # check keys only, else __eq__ ref cycles, scope -> varmap -> var
289            return self._con.keys() == other._con.keys()
290        return False
291
292    def __ne__(self, other):
293        return not self.__eq__(other)
294
295
296class AbstractRHS(object):
297    """Abstract base class for anything that can be the RHS of an assignment.
298    This class **does not** define any methods.
299    """
300
301
302class Inst(EqualityCheckMixin, AbstractRHS):
303    """
304    Base class for all IR instructions.
305    """
306
307    def list_vars(self):
308        """
309        List the variables used (read or written) by the instruction.
310        """
311        raise NotImplementedError
312
313    def _rec_list_vars(self, val):
314        """
315        A recursive helper used to implement list_vars() in subclasses.
316        """
317        if isinstance(val, Var):
318            return [val]
319        elif isinstance(val, Inst):
320            return val.list_vars()
321        elif isinstance(val, (list, tuple)):
322            lst = []
323            for v in val:
324                lst.extend(self._rec_list_vars(v))
325            return lst
326        elif isinstance(val, dict):
327            lst = []
328            for v in val.values():
329                lst.extend(self._rec_list_vars(v))
330            return lst
331        else:
332            return []
333
334
335class Stmt(Inst):
336    """
337    Base class for IR statements (instructions which can appear on their
338    own in a Block).
339    """
340    # Whether this statement ends its basic block (i.e. it will either jump
341    # to another block or exit the function).
342    is_terminator = False
343    # Whether this statement exits the function.
344    is_exit = False
345
346    def list_vars(self):
347        return self._rec_list_vars(self.__dict__)
348
349
350class Terminator(Stmt):
351    """
352    IR statements that are terminators: the last statement in a block.
353    A terminator must either:
354    - exit the function
355    - jump to a block
356
357    All subclass of Terminator must override `.get_targets()` to return a list
358    of jump targets.
359    """
360    is_terminator = True
361
362    def get_targets(self):
363        raise NotImplementedError(type(self))
364
365
366class Expr(Inst):
367    """
368    An IR expression (an instruction which can only be part of a larger
369    statement).
370    """
371
372    def __init__(self, op, loc, **kws):
373        assert isinstance(op, str)
374        assert isinstance(loc, Loc)
375        self.op = op
376        self.loc = loc
377        self._kws = kws
378
379    def __getattr__(self, name):
380        if name.startswith('_'):
381            return Inst.__getattr__(self, name)
382        return self._kws[name]
383
384    def __setattr__(self, name, value):
385        if name in ('op', 'loc', '_kws'):
386            self.__dict__[name] = value
387        else:
388            self._kws[name] = value
389
390    @classmethod
391    def binop(cls, fn, lhs, rhs, loc):
392        assert isinstance(fn, BuiltinFunctionType)
393        assert isinstance(lhs, Var)
394        assert isinstance(rhs, Var)
395        assert isinstance(loc, Loc)
396        op = 'binop'
397        return cls(op=op, loc=loc, fn=fn, lhs=lhs, rhs=rhs,
398                   static_lhs=UNDEFINED, static_rhs=UNDEFINED)
399
400    @classmethod
401    def inplace_binop(cls, fn, immutable_fn, lhs, rhs, loc):
402        assert isinstance(fn, BuiltinFunctionType)
403        assert isinstance(immutable_fn, BuiltinFunctionType)
404        assert isinstance(lhs, Var)
405        assert isinstance(rhs, Var)
406        assert isinstance(loc, Loc)
407        op = 'inplace_binop'
408        return cls(op=op, loc=loc, fn=fn, immutable_fn=immutable_fn,
409                   lhs=lhs, rhs=rhs,
410                   static_lhs=UNDEFINED, static_rhs=UNDEFINED)
411
412    @classmethod
413    def unary(cls, fn, value, loc):
414        assert isinstance(value, (str, Var, FunctionType))
415        assert isinstance(loc, Loc)
416        op = 'unary'
417        fn = UNARY_BUITINS_TO_OPERATORS.get(fn, fn)
418        return cls(op=op, loc=loc, fn=fn, value=value)
419
420    @classmethod
421    def call(cls, func, args, kws, loc, vararg=None):
422        assert isinstance(func, (Var, Intrinsic))
423        assert isinstance(loc, Loc)
424        op = 'call'
425        return cls(op=op, loc=loc, func=func, args=args, kws=kws,
426                   vararg=vararg)
427
428    @classmethod
429    def build_tuple(cls, items, loc):
430        assert isinstance(loc, Loc)
431        op = 'build_tuple'
432        return cls(op=op, loc=loc, items=items)
433
434    @classmethod
435    def build_list(cls, items, loc):
436        assert isinstance(loc, Loc)
437        op = 'build_list'
438        return cls(op=op, loc=loc, items=items)
439
440    @classmethod
441    def build_set(cls, items, loc):
442        assert isinstance(loc, Loc)
443        op = 'build_set'
444        return cls(op=op, loc=loc, items=items)
445
446    @classmethod
447    def build_map(cls, items, size, literal_value, value_indexes, loc):
448        assert isinstance(loc, Loc)
449        op = 'build_map'
450        return cls(op=op, loc=loc, items=items, size=size,
451                   literal_value=literal_value, value_indexes=value_indexes)
452
453    @classmethod
454    def pair_first(cls, value, loc):
455        assert isinstance(value, Var)
456        op = 'pair_first'
457        return cls(op=op, loc=loc, value=value)
458
459    @classmethod
460    def pair_second(cls, value, loc):
461        assert isinstance(value, Var)
462        assert isinstance(loc, Loc)
463        op = 'pair_second'
464        return cls(op=op, loc=loc, value=value)
465
466    @classmethod
467    def getiter(cls, value, loc):
468        assert isinstance(value, Var)
469        assert isinstance(loc, Loc)
470        op = 'getiter'
471        return cls(op=op, loc=loc, value=value)
472
473    @classmethod
474    def iternext(cls, value, loc):
475        assert isinstance(value, Var)
476        assert isinstance(loc, Loc)
477        op = 'iternext'
478        return cls(op=op, loc=loc, value=value)
479
480    @classmethod
481    def exhaust_iter(cls, value, count, loc):
482        assert isinstance(value, Var)
483        assert isinstance(count, int)
484        assert isinstance(loc, Loc)
485        op = 'exhaust_iter'
486        return cls(op=op, loc=loc, value=value, count=count)
487
488    @classmethod
489    def getattr(cls, value, attr, loc):
490        assert isinstance(value, Var)
491        assert isinstance(attr, str)
492        assert isinstance(loc, Loc)
493        op = 'getattr'
494        return cls(op=op, loc=loc, value=value, attr=attr)
495
496    @classmethod
497    def getitem(cls, value, index, loc):
498        assert isinstance(value, Var)
499        assert isinstance(index, Var)
500        assert isinstance(loc, Loc)
501        op = 'getitem'
502        return cls(op=op, loc=loc, value=value, index=index)
503
504    @classmethod
505    def typed_getitem(cls, value, dtype, index, loc):
506        assert isinstance(value, Var)
507        assert isinstance(loc, Loc)
508        op = 'typed_getitem'
509        return cls(op=op, loc=loc, value=value, dtype=dtype,
510                   index=index)
511
512    @classmethod
513    def static_getitem(cls, value, index, index_var, loc):
514        assert isinstance(value, Var)
515        assert index_var is None or isinstance(index_var, Var)
516        assert isinstance(loc, Loc)
517        op = 'static_getitem'
518        return cls(op=op, loc=loc, value=value, index=index,
519                   index_var=index_var)
520
521    @classmethod
522    def cast(cls, value, loc):
523        """
524        A node for implicit casting at the return statement
525        """
526        assert isinstance(value, Var)
527        assert isinstance(loc, Loc)
528        op = 'cast'
529        return cls(op=op, value=value, loc=loc)
530
531    @classmethod
532    def phi(cls, loc):
533        """Phi node
534        """
535        assert isinstance(loc, Loc)
536        return cls(op='phi', incoming_values=[], incoming_blocks=[], loc=loc)
537
538    @classmethod
539    def make_function(cls, name, code, closure, defaults, loc):
540        """
541        A node for making a function object.
542        """
543        assert isinstance(loc, Loc)
544        op = 'make_function'
545        return cls(op=op, name=name, code=code, closure=closure, defaults=defaults, loc=loc)
546
547    @classmethod
548    def null(cls, loc):
549        """
550        A node for null value.
551
552        This node is not handled by type inference. It is only added by
553        post-typing passes.
554        """
555        op = 'null'
556        return cls(op=op, loc=loc)
557
558    def __repr__(self):
559        if self.op == 'call':
560            args = ', '.join(str(a) for a in self.args)
561            pres_order = self._kws.items() if config.DIFF_IR == 0 else sorted(self._kws.items())
562            kws = ', '.join('%s=%s' % (k, v) for k, v in pres_order)
563            vararg = '*%s' % (self.vararg,) if self.vararg is not None else ''
564            arglist = ', '.join(filter(None, [args, vararg, kws]))
565            return 'call %s(%s)' % (self.func, arglist)
566        elif self.op == 'binop':
567            lhs, rhs = self.lhs, self.rhs
568            if self.fn == operator.contains:
569                lhs, rhs = rhs, lhs
570            fn = OPERATORS_TO_BUILTINS.get(self.fn, self.fn)
571            return '%s %s %s' % (lhs, fn, rhs)
572        else:
573            pres_order = self._kws.items() if config.DIFF_IR == 0 else sorted(self._kws.items())
574            args = ('%s=%s' % (k, v) for k, v in pres_order)
575            return '%s(%s)' % (self.op, ', '.join(args))
576
577    def list_vars(self):
578        return self._rec_list_vars(self._kws)
579
580    def infer_constant(self):
581        raise ConstantInferenceError('%s' % self, loc=self.loc)
582
583
584class SetItem(Stmt):
585    """
586    target[index] = value
587    """
588
589    def __init__(self, target, index, value, loc):
590        assert isinstance(target, Var)
591        assert isinstance(index, Var)
592        assert isinstance(value, Var)
593        assert isinstance(loc, Loc)
594        self.target = target
595        self.index = index
596        self.value = value
597        self.loc = loc
598
599    def __repr__(self):
600        return '%s[%s] = %s' % (self.target, self.index, self.value)
601
602
603class StaticSetItem(Stmt):
604    """
605    target[constant index] = value
606    """
607
608    def __init__(self, target, index, index_var, value, loc):
609        assert isinstance(target, Var)
610        assert not isinstance(index, Var)
611        assert isinstance(index_var, Var)
612        assert isinstance(value, Var)
613        assert isinstance(loc, Loc)
614        self.target = target
615        self.index = index
616        self.index_var = index_var
617        self.value = value
618        self.loc = loc
619
620    def __repr__(self):
621        return '%s[%r] = %s' % (self.target, self.index, self.value)
622
623
624class DelItem(Stmt):
625    """
626    del target[index]
627    """
628
629    def __init__(self, target, index, loc):
630        assert isinstance(target, Var)
631        assert isinstance(index, Var)
632        assert isinstance(loc, Loc)
633        self.target = target
634        self.index = index
635        self.loc = loc
636
637    def __repr__(self):
638        return 'del %s[%s]' % (self.target, self.index)
639
640
641class SetAttr(Stmt):
642    def __init__(self, target, attr, value, loc):
643        assert isinstance(target, Var)
644        assert isinstance(attr, str)
645        assert isinstance(value, Var)
646        assert isinstance(loc, Loc)
647        self.target = target
648        self.attr = attr
649        self.value = value
650        self.loc = loc
651
652    def __repr__(self):
653        return '(%s).%s = %s' % (self.target, self.attr, self.value)
654
655
656class DelAttr(Stmt):
657    def __init__(self, target, attr, loc):
658        assert isinstance(target, Var)
659        assert isinstance(attr, str)
660        assert isinstance(loc, Loc)
661        self.target = target
662        self.attr = attr
663        self.loc = loc
664
665    def __repr__(self):
666        return 'del (%s).%s' % (self.target, self.attr)
667
668
669class StoreMap(Stmt):
670    def __init__(self, dct, key, value, loc):
671        assert isinstance(dct, Var)
672        assert isinstance(key, Var)
673        assert isinstance(value, Var)
674        assert isinstance(loc, Loc)
675        self.dct = dct
676        self.key = key
677        self.value = value
678        self.loc = loc
679
680    def __repr__(self):
681        return '%s[%s] = %s' % (self.dct, self.key, self.value)
682
683
684class Del(Stmt):
685    def __init__(self, value, loc):
686        assert isinstance(value, str)
687        assert isinstance(loc, Loc)
688        self.value = value
689        self.loc = loc
690
691    def __str__(self):
692        return "del %s" % self.value
693
694
695class Raise(Terminator):
696    is_exit = True
697
698    def __init__(self, exception, loc):
699        assert exception is None or isinstance(exception, Var)
700        assert isinstance(loc, Loc)
701        self.exception = exception
702        self.loc = loc
703
704    def __str__(self):
705        return "raise %s" % self.exception
706
707    def get_targets(self):
708        return []
709
710
711class StaticRaise(Terminator):
712    """
713    Raise an exception class and arguments known at compile-time.
714    Note that if *exc_class* is None, a bare "raise" statement is implied
715    (i.e. re-raise the current exception).
716    """
717    is_exit = True
718
719    def __init__(self, exc_class, exc_args, loc):
720        assert exc_class is None or isinstance(exc_class, type)
721        assert isinstance(loc, Loc)
722        assert exc_args is None or isinstance(exc_args, tuple)
723        self.exc_class = exc_class
724        self.exc_args = exc_args
725        self.loc = loc
726
727    def __str__(self):
728        if self.exc_class is None:
729            return "<static> raise"
730        elif self.exc_args is None:
731            return "<static> raise %s" % (self.exc_class,)
732        else:
733            return "<static> raise %s(%s)" % (self.exc_class,
734                                     ", ".join(map(repr, self.exc_args)))
735
736    def get_targets(self):
737        return []
738
739
740class TryRaise(Stmt):
741    """A raise statement inside a try-block
742    Similar to ``Raise`` but does not terminate.
743    """
744    def __init__(self, exception, loc):
745        assert exception is None or isinstance(exception, Var)
746        assert isinstance(loc, Loc)
747        self.exception = exception
748        self.loc = loc
749
750    def __str__(self):
751        return "try_raise %s" % self.exception
752
753
754class StaticTryRaise(Stmt):
755    """A raise statement inside a try-block.
756    Similar to ``StaticRaise`` but does not terminate.
757    """
758
759    def __init__(self, exc_class, exc_args, loc):
760        assert exc_class is None or isinstance(exc_class, type)
761        assert isinstance(loc, Loc)
762        assert exc_args is None or isinstance(exc_args, tuple)
763        self.exc_class = exc_class
764        self.exc_args = exc_args
765        self.loc = loc
766
767    def __str__(self):
768        if self.exc_class is None:
769            return "static_try_raise"
770        elif self.exc_args is None:
771            return "static_try_raise %s" % (self.exc_class,)
772        else:
773            return "static_try_raise %s(%s)" % (self.exc_class,
774                                     ", ".join(map(repr, self.exc_args)))
775
776
777class Return(Terminator):
778    """
779    Return to caller.
780    """
781    is_exit = True
782
783    def __init__(self, value, loc):
784        assert isinstance(value, Var), type(value)
785        assert isinstance(loc, Loc)
786        self.value = value
787        self.loc = loc
788
789    def __str__(self):
790        return 'return %s' % self.value
791
792    def get_targets(self):
793        return []
794
795
796class Jump(Terminator):
797    """
798    Unconditional branch.
799    """
800
801    def __init__(self, target, loc):
802        assert isinstance(loc, Loc)
803        self.target = target
804        self.loc = loc
805
806    def __str__(self):
807        return 'jump %s' % self.target
808
809    def get_targets(self):
810        return [self.target]
811
812
813class Branch(Terminator):
814    """
815    Conditional branch.
816    """
817
818    def __init__(self, cond, truebr, falsebr, loc):
819        assert isinstance(cond, Var)
820        assert isinstance(loc, Loc)
821        self.cond = cond
822        self.truebr = truebr
823        self.falsebr = falsebr
824        self.loc = loc
825
826    def __str__(self):
827        return 'branch %s, %s, %s' % (self.cond, self.truebr, self.falsebr)
828
829    def get_targets(self):
830        return [self.truebr, self.falsebr]
831
832
833class Assign(Stmt):
834    """
835    Assign to a variable.
836    """
837    def __init__(self, value, target, loc):
838        assert isinstance(value, AbstractRHS)
839        assert isinstance(target, Var)
840        assert isinstance(loc, Loc)
841        self.value = value
842        self.target = target
843        self.loc = loc
844
845    def __str__(self):
846        return '%s = %s' % (self.target, self.value)
847
848
849class Print(Stmt):
850    """
851    Print some values.
852    """
853    def __init__(self, args, vararg, loc):
854        assert all(isinstance(x, Var) for x in args)
855        assert vararg is None or isinstance(vararg, Var)
856        assert isinstance(loc, Loc)
857        self.args = tuple(args)
858        self.vararg = vararg
859        # Constant-inferred arguments
860        self.consts = {}
861        self.loc = loc
862
863    def __str__(self):
864        return 'print(%s)' % ', '.join(str(v) for v in self.args)
865
866
867class Yield(Inst):
868    def __init__(self, value, loc, index):
869        assert isinstance(value, Var)
870        assert isinstance(loc, Loc)
871        self.value = value
872        self.loc = loc
873        self.index = index
874
875    def __str__(self):
876        return 'yield %s' % (self.value,)
877
878    def list_vars(self):
879        return [self.value]
880
881
882class EnterWith(Stmt):
883    """Enter a "with" context
884    """
885    def __init__(self, contextmanager, begin, end, loc):
886        """
887        Parameters
888        ----------
889        contextmanager : IR value
890        begin, end : int
891            The beginning and the ending offset of the with-body.
892        loc : ir.Loc instance
893            Source location
894        """
895        assert isinstance(contextmanager, Var)
896        assert isinstance(loc, Loc)
897        self.contextmanager = contextmanager
898        self.begin = begin
899        self.end = end
900        self.loc = loc
901
902    def __str__(self):
903        return 'enter_with {}'.format(self.contextmanager)
904
905    def list_vars(self):
906        return [self.contextmanager]
907
908
909class Arg(EqualityCheckMixin, AbstractRHS):
910    def __init__(self, name, index, loc):
911        assert isinstance(name, str)
912        assert isinstance(index, int)
913        assert isinstance(loc, Loc)
914        self.name = name
915        self.index = index
916        self.loc = loc
917
918    def __repr__(self):
919        return 'arg(%d, name=%s)' % (self.index, self.name)
920
921    def infer_constant(self):
922        raise ConstantInferenceError('%s' % self, loc=self.loc)
923
924
925class Const(EqualityCheckMixin, AbstractRHS):
926    def __init__(self, value, loc, use_literal_type=True):
927        assert isinstance(loc, Loc)
928        self.value = value
929        self.loc = loc
930        # Note: need better way to tell if this is a literal or not.
931        self.use_literal_type = use_literal_type
932
933    def __repr__(self):
934        return 'const(%s, %s)' % (type(self.value).__name__, self.value)
935
936    def infer_constant(self):
937        return self.value
938
939    def __deepcopy__(self, memo):
940        # Override to not copy constant values in code
941        return Const(
942            value=self.value, loc=self.loc,
943            use_literal_type=self.use_literal_type,
944        )
945
946
947class Global(EqualityCheckMixin, AbstractRHS):
948    def __init__(self, name, value, loc):
949        assert isinstance(loc, Loc)
950        self.name = name
951        self.value = value
952        self.loc = loc
953
954    def __str__(self):
955        return 'global(%s: %s)' % (self.name, self.value)
956
957    def infer_constant(self):
958        return self.value
959
960    def __deepcopy__(self, memo):
961        # don't copy value since it can fail (e.g. modules)
962        # value is readonly and doesn't need copying
963        return Global(self.name, self.value, copy.deepcopy(self.loc))
964
965
966class FreeVar(EqualityCheckMixin, AbstractRHS):
967    """
968    A freevar, as loaded by LOAD_DECREF.
969    (i.e. a variable defined in an enclosing non-global scope)
970    """
971
972    def __init__(self, index, name, value, loc):
973        assert isinstance(index, int)
974        assert isinstance(name, str)
975        assert isinstance(loc, Loc)
976        # index inside __code__.co_freevars
977        self.index = index
978        # variable name
979        self.name = name
980        # frozen value
981        self.value = value
982        self.loc = loc
983
984    def __str__(self):
985        return 'freevar(%s: %s)' % (self.name, self.value)
986
987    def infer_constant(self):
988        return self.value
989
990    def __deepcopy__(self, memo):
991        # Override to not copy constant values in code
992        return FreeVar(index=self.index, name=self.name, value=self.value,
993                       loc=self.loc)
994
995
996
997class Var(EqualityCheckMixin, AbstractRHS):
998    """
999    Attributes
1000    -----------
1001    - scope: Scope
1002
1003    - name: str
1004
1005    - loc: Loc
1006        Definition location
1007    """
1008
1009    def __init__(self, scope, name, loc):
1010        # NOTE: Use of scope=None should be removed.
1011        assert scope is None or isinstance(scope, Scope)
1012        assert isinstance(name, str)
1013        assert isinstance(loc, Loc)
1014        self.scope = scope
1015        self.name = name
1016        self.loc = loc
1017
1018    def __repr__(self):
1019        return 'Var(%s, %s)' % (self.name, self.loc.short())
1020
1021    def __str__(self):
1022        return self.name
1023
1024    @property
1025    def is_temp(self):
1026        return self.name.startswith("$")
1027
1028    @property
1029    def unversioned_name(self):
1030        """The unversioned name of this variable, i.e. SSA renaming removed
1031        """
1032        for k, redef_set in self.scope.var_redefinitions.items():
1033            if self.name in redef_set:
1034                return k
1035        return self.name
1036
1037    @property
1038    def versioned_names(self):
1039        """Known versioned names for this variable, i.e. known variable names in
1040        the scope that have been formed from applying SSA to this variable
1041        """
1042        return self.scope.get_versions_of(self.unversioned_name)
1043
1044    @property
1045    def all_names(self):
1046        """All known versioned and unversioned names for this variable
1047        """
1048        return self.versioned_names | {self.unversioned_name,}
1049
1050class Intrinsic(EqualityCheckMixin):
1051    """
1052    A low-level "intrinsic" function.  Suitable as the callable of a "call"
1053    expression.
1054
1055    The given *name* is backend-defined and will be inserted as-is
1056    in the generated low-level IR.
1057    The *type* is the equivalent Numba signature of calling the intrinsic.
1058    """
1059
1060    def __init__(self, name, type, args, loc=None):
1061        self.name = name
1062        self.type = type
1063        self.loc = loc
1064        self.args = args
1065
1066    def __repr__(self):
1067        return 'Intrinsic(%s, %s, %s)' % (self.name, self.type, self.loc)
1068
1069    def __str__(self):
1070        return self.name
1071
1072
1073class Scope(EqualityCheckMixin):
1074    """
1075    Attributes
1076    -----------
1077    - parent: Scope
1078        Parent scope
1079
1080    - localvars: VarMap
1081        Scope-local variable map
1082
1083    - loc: Loc
1084        Start of scope location
1085
1086    """
1087
1088    def __init__(self, parent, loc):
1089        assert parent is None or isinstance(parent, Scope)
1090        assert isinstance(loc, Loc)
1091        self.parent = parent
1092        self.localvars = VarMap()
1093        self.loc = loc
1094        self.redefined = defaultdict(int)
1095        self.var_redefinitions = defaultdict(set)
1096
1097    def define(self, name, loc):
1098        """
1099        Define a variable
1100        """
1101        v = Var(scope=self, name=name, loc=loc)
1102        self.localvars.define(v.name, v)
1103        return v
1104
1105    def get(self, name):
1106        """
1107        Refer to a variable.  Returns the latest version.
1108        """
1109        if name in self.redefined:
1110            name = "%s.%d" % (name, self.redefined[name])
1111        return self.get_exact(name)
1112
1113    def get_exact(self, name):
1114        """
1115        Refer to a variable.  The returned variable has the exact
1116        name (exact variable version).
1117        """
1118        try:
1119            return self.localvars.get(name)
1120        except NotDefinedError:
1121            if self.has_parent:
1122                return self.parent.get(name)
1123            else:
1124                raise
1125
1126    def get_or_define(self, name, loc):
1127        if name in self.redefined:
1128            name = "%s.%d" % (name, self.redefined[name])
1129
1130        if name not in self.localvars:
1131            return self.define(name, loc)
1132        else:
1133            return self.localvars.get(name)
1134
1135    def redefine(self, name, loc, rename=True):
1136        """
1137        Redefine if the name is already defined
1138        """
1139        if name not in self.localvars:
1140            return self.define(name, loc)
1141        elif not rename:
1142            # Must use the same name if the variable is a cellvar, which
1143            # means it could be captured in a closure.
1144            return self.localvars.get(name)
1145        else:
1146            while True:
1147                ct = self.redefined[name]
1148                self.redefined[name] = ct + 1
1149                newname = "%s.%d" % (name, ct + 1)
1150                try:
1151                    res = self.define(newname, loc)
1152                except RedefinedError:
1153                    continue
1154                else:
1155                    self.var_redefinitions[name].add(newname)
1156                return res
1157
1158    def get_versions_of(self, name):
1159        """
1160        Gets all known versions of a given name
1161        """
1162        vers = set()
1163        def walk(thename):
1164            redefs = self.var_redefinitions.get(thename, None)
1165            if redefs:
1166                for v in redefs:
1167                    vers.add(v)
1168                    walk(v)
1169        walk(name)
1170        return vers
1171
1172    def make_temp(self, loc):
1173        n = len(self.localvars)
1174        v = Var(scope=self, name='$%d' % n, loc=loc)
1175        self.localvars.define(v.name, v)
1176        return v
1177
1178    @property
1179    def has_parent(self):
1180        return self.parent is not None
1181
1182    def __repr__(self):
1183        return "Scope(has_parent=%r, num_vars=%d, %s)" % (self.has_parent,
1184                                                          len(self.localvars),
1185                                                          self.loc)
1186
1187
1188class Block(EqualityCheckMixin):
1189    """A code block
1190
1191    """
1192
1193    def __init__(self, scope, loc):
1194        assert isinstance(scope, Scope)
1195        assert isinstance(loc, Loc)
1196        self.scope = scope
1197        self.body = []
1198        self.loc = loc
1199
1200    def copy(self):
1201        block = Block(self.scope, self.loc)
1202        block.body = self.body[:]
1203        return block
1204
1205    def find_exprs(self, op=None):
1206        """
1207        Iterate over exprs of the given *op* in this block.
1208        """
1209        for inst in self.body:
1210            if isinstance(inst, Assign):
1211                expr = inst.value
1212                if isinstance(expr, Expr):
1213                    if op is None or expr.op == op:
1214                        yield expr
1215
1216    def find_insts(self, cls=None):
1217        """
1218        Iterate over insts of the given class in this block.
1219        """
1220        for inst in self.body:
1221            if isinstance(inst, cls):
1222                yield inst
1223
1224    def find_variable_assignment(self, name):
1225        """
1226        Returns the assignment inst associated with variable "name", None if
1227        it cannot be found.
1228        """
1229        for x in self.find_insts(cls=Assign):
1230            if x.target.name == name:
1231                return x
1232        return None
1233
1234    def prepend(self, inst):
1235        assert isinstance(inst, Stmt)
1236        self.body.insert(0, inst)
1237
1238    def append(self, inst):
1239        assert isinstance(inst, Stmt)
1240        self.body.append(inst)
1241
1242    def remove(self, inst):
1243        assert isinstance(inst, Stmt)
1244        del self.body[self.body.index(inst)]
1245
1246    def clear(self):
1247        del self.body[:]
1248
1249    def dump(self, file=None):
1250        # Avoid early bind of sys.stdout as default value
1251        file = file or sys.stdout
1252        for inst in self.body:
1253            if hasattr(inst, 'dump'):
1254                inst.dump(file)
1255            else:
1256                inst_vars = sorted(str(v) for v in inst.list_vars())
1257                print('    %-40s %s' % (inst, inst_vars), file=file)
1258
1259    @property
1260    def terminator(self):
1261        return self.body[-1]
1262
1263    @property
1264    def is_terminated(self):
1265        return self.body and self.body[-1].is_terminator
1266
1267    def verify(self):
1268        if not self.is_terminated:
1269            raise VerificationError("Missing block terminator")
1270            # Only the last instruction can be a terminator
1271        for inst in self.body[:-1]:
1272            if inst.is_terminator:
1273                raise VerificationError("Terminator before the last "
1274                                        "instruction")
1275
1276    def insert_after(self, stmt, other):
1277        """
1278        Insert *stmt* after *other*.
1279        """
1280        index = self.body.index(other)
1281        self.body.insert(index + 1, stmt)
1282
1283    def insert_before_terminator(self, stmt):
1284        assert isinstance(stmt, Stmt)
1285        assert self.is_terminated
1286        self.body.insert(-1, stmt)
1287
1288    def __repr__(self):
1289        return "<ir.Block at %s>" % (self.loc,)
1290
1291
1292class Loop(SlotEqualityCheckMixin):
1293    """Describes a loop-block
1294    """
1295    __slots__ = "entry", "exit"
1296
1297    def __init__(self, entry, exit):
1298        self.entry = entry
1299        self.exit = exit
1300
1301    def __repr__(self):
1302        args = self.entry, self.exit
1303        return "Loop(entry=%s, exit=%s)" % args
1304
1305
1306class With(SlotEqualityCheckMixin):
1307    """Describes a with-block
1308    """
1309    __slots__ = "entry", "exit"
1310
1311    def __init__(self, entry, exit):
1312        self.entry = entry
1313        self.exit = exit
1314
1315    def __repr__(self):
1316        args = self.entry, self.exit
1317        return "With(entry=%s, exit=%s)" % args
1318
1319
1320class FunctionIR(object):
1321
1322    def __init__(self, blocks, is_generator, func_id, loc,
1323                 definitions, arg_count, arg_names):
1324        self.blocks = blocks
1325        self.is_generator = is_generator
1326        self.func_id = func_id
1327        self.loc = loc
1328        self.arg_count = arg_count
1329        self.arg_names = arg_names
1330
1331        self._definitions = definitions
1332
1333        self._reset_analysis_variables()
1334
1335    def equal_ir(self, other):
1336        """ Checks that the IR contained within is equal to the IR in other.
1337        Equality is defined by being equal in fundamental structure (blocks,
1338        labels, IR node type and the order in which they are defined) and the
1339        IR nodes being equal. IR node equality essentially comes down to
1340        ensuring a node's `.__dict__` or `.__slots__` is equal, with the
1341        exception of ignoring 'loc' and 'scope' entries. The upshot is that the
1342        comparison is essentially location and scope invariant, but otherwise
1343        behaves as unsurprisingly as possible.
1344        """
1345        if type(self) is type(other):
1346            return self.blocks == other.blocks
1347        return False
1348
1349    def diff_str(self, other):
1350        """
1351        Compute a human readable difference in the IR, returns a formatted
1352        string ready for printing.
1353        """
1354        msg = []
1355        for label, block in self.blocks.items():
1356            other_blk = other.blocks.get(label, None)
1357            if other_blk is not None:
1358                if block != other_blk:
1359                    msg.append(("Block %s differs" % label).center(80, '-'))
1360                    # see if the instructions are just a permutation
1361                    block_del = [x for x in block.body if isinstance(x, Del)]
1362                    oth_del = [x for x in other_blk.body if isinstance(x, Del)]
1363                    if block_del != oth_del:
1364                        # this is a common issue, dels are all present, but
1365                        # order shuffled.
1366                        if sorted(block_del) == sorted(oth_del):
1367                            msg.append(("Block %s contains the same dels but "
1368                                        "their order is different") % label)
1369                    if len(block.body) > len(other_blk.body):
1370                        msg.append("This block contains more statements")
1371                    elif len(block.body) < len(other_blk.body):
1372                        msg.append("Other block contains more statements")
1373
1374                    # find the indexes where they don't match
1375                    tmp = []
1376                    for idx, stmts in enumerate(zip(block.body,
1377                                                    other_blk.body)):
1378                        b_s, o_s = stmts
1379                        if b_s != o_s:
1380                            tmp.append(idx)
1381
1382                    def get_pad(ablock, l):
1383                        pointer = '-> '
1384                        sp = len(pointer) * ' '
1385                        pad = []
1386                        nstmt = len(ablock)
1387                        for i in range(nstmt):
1388                            if i in tmp:
1389                                item = pointer
1390                            elif i >= l:
1391                                item = pointer
1392                            else:
1393                                item = sp
1394                            pad.append(item)
1395                        return pad
1396
1397                    min_stmt_len = min(len(block.body), len(other_blk.body))
1398
1399                    with StringIO() as buf:
1400                        it = [("self", block), ("other", other_blk)]
1401                        for name, _block in it:
1402                            buf.truncate(0)
1403                            _block.dump(file=buf)
1404                            stmts = buf.getvalue().splitlines()
1405                            pad = get_pad(_block.body, min_stmt_len)
1406                            title = ("%s: block %s" % (name, label))
1407                            msg.append(title.center(80, '-'))
1408                            msg.extend(["{0}{1}".format(a, b) for a, b in
1409                                        zip(pad, stmts)])
1410        if msg == []:
1411            msg.append("IR is considered equivalent.")
1412        return '\n'.join(msg)
1413
1414    def _reset_analysis_variables(self):
1415
1416        self._consts = consts.ConstantInference(self)
1417
1418        # Will be computed by PostProcessor
1419        self.generator_info = None
1420        self.variable_lifetime = None
1421        # { ir.Block: { variable names (potentially) alive at start of block } }
1422        self.block_entry_vars = {}
1423
1424    def derive(self, blocks, arg_count=None, arg_names=None,
1425               force_non_generator=False):
1426        """
1427        Derive a new function IR from this one, using the given blocks,
1428        and possibly modifying the argument count and generator flag.
1429
1430        Post-processing will have to be run again on the new IR.
1431        """
1432        firstblock = blocks[min(blocks)]
1433
1434        new_ir = copy.copy(self)
1435        new_ir.blocks = blocks
1436        new_ir.loc = firstblock.loc
1437        if force_non_generator:
1438            new_ir.is_generator = False
1439        if arg_count is not None:
1440            new_ir.arg_count = arg_count
1441        if arg_names is not None:
1442            new_ir.arg_names = arg_names
1443        new_ir._reset_analysis_variables()
1444        # Make fresh func_id
1445        new_ir.func_id = new_ir.func_id.derive()
1446        return new_ir
1447
1448    def copy(self):
1449        new_ir = copy.copy(self)
1450        blocks = {}
1451        block_entry_vars = {}
1452        for label, block in self.blocks.items():
1453            new_block = block.copy()
1454            blocks[label] = new_block
1455            if block in self.block_entry_vars:
1456                block_entry_vars[new_block] = self.block_entry_vars[block]
1457        new_ir.blocks = blocks
1458        new_ir.block_entry_vars = block_entry_vars
1459        return new_ir
1460
1461    def get_block_entry_vars(self, block):
1462        """
1463        Return a set of variable names possibly alive at the beginning of
1464        the block.
1465        """
1466        return self.block_entry_vars[block]
1467
1468    def infer_constant(self, name):
1469        """
1470        Try to infer the constant value of a given variable.
1471        """
1472        if isinstance(name, Var):
1473            name = name.name
1474        return self._consts.infer_constant(name)
1475
1476    def get_definition(self, value, lhs_only=False):
1477        """
1478        Get the definition site for the given variable name or instance.
1479        A Expr instance is returned by default, but if lhs_only is set
1480        to True, the left-hand-side variable is returned instead.
1481        """
1482        lhs = value
1483        while True:
1484            if isinstance(value, Var):
1485                lhs = value
1486                name = value.name
1487            elif isinstance(value, str):
1488                lhs = value
1489                name = value
1490            else:
1491                return lhs if lhs_only else value
1492            defs = self._definitions[name]
1493            if len(defs) == 0:
1494                raise KeyError("no definition for %r"
1495                               % (name,))
1496            if len(defs) > 1:
1497                raise KeyError("more than one definition for %r"
1498                               % (name,))
1499            value = defs[0]
1500
1501    def get_assignee(self, rhs_value, in_blocks=None):
1502        """
1503        Finds the assignee for a given RHS value. If in_blocks is given the
1504        search will be limited to the specified blocks.
1505        """
1506        if in_blocks is None:
1507            blocks = self.blocks.values()
1508        elif isinstance(in_blocks, int):
1509            blocks = [self.blocks[in_blocks]]
1510        else:
1511            blocks = [self.blocks[blk] for blk in list(in_blocks)]
1512
1513        assert isinstance(rhs_value, AbstractRHS)
1514
1515        for blk in blocks:
1516            for assign in blk.find_insts(Assign):
1517                if assign.value == rhs_value:
1518                    return assign.target
1519
1520        raise ValueError("Could not find an assignee for %s" % rhs_value)
1521
1522
1523    def dump(self, file=None):
1524        nofile = file is None
1525        # Avoid early bind of sys.stdout as default value
1526        file = file or StringIO()
1527        for offset, block in sorted(self.blocks.items()):
1528            print('label %s:' % (offset,), file=file)
1529            block.dump(file=file)
1530        if nofile:
1531            text = file.getvalue()
1532            if config.HIGHLIGHT_DUMPS:
1533                try:
1534                    import pygments
1535                except ImportError:
1536                    msg = "Please install pygments to see highlighted dumps"
1537                    raise ValueError(msg)
1538                else:
1539                    from pygments import highlight
1540                    from numba.misc.dump_style import NumbaIRLexer as lexer
1541                    from numba.misc.dump_style import by_colorscheme
1542                    from pygments.formatters import Terminal256Formatter
1543                    print(highlight(text, lexer(), Terminal256Formatter(
1544                        style=by_colorscheme())))
1545            else:
1546                print(text)
1547
1548
1549    def dump_to_string(self):
1550        with StringIO() as sb:
1551            self.dump(file=sb)
1552            return sb.getvalue()
1553
1554    def dump_generator_info(self, file=None):
1555        file = file or sys.stdout
1556        gi = self.generator_info
1557        print("generator state variables:", sorted(gi.state_vars), file=file)
1558        for index, yp in sorted(gi.yield_points.items()):
1559            print("yield point #%d: live variables = %s, weak live variables = %s"
1560                  % (index, sorted(yp.live_vars), sorted(yp.weak_live_vars)),
1561                  file=file)
1562
1563    def render_dot(self, filename_prefix="numba_ir", include_ir=True):
1564        """Render the CFG of the IR with GraphViz DOT via the
1565        ``graphviz`` python binding.
1566
1567        Returns
1568        -------
1569        g : graphviz.Digraph
1570            Use `g.view()` to open the graph in the default PDF application.
1571        """
1572
1573        try:
1574            import graphviz as gv
1575        except ImportError:
1576            raise ImportError(
1577                "The feature requires `graphviz` but it is not available. "
1578                "Please install with `pip install graphviz`"
1579            )
1580        g = gv.Digraph(
1581            filename="{}{}.dot".format(
1582                filename_prefix,
1583                self.func_id.unique_name,
1584            )
1585        )
1586        # Populate the nodes
1587        for k, blk in self.blocks.items():
1588            with StringIO() as sb:
1589                blk.dump(sb)
1590                label = sb.getvalue()
1591            if include_ir:
1592                label = ''.join(
1593                    ['  {}\l'.format(x) for x in label.splitlines()],
1594                )
1595                label = "block {}\l".format(k) + label
1596                g.node(str(k), label=label, shape='rect')
1597            else:
1598                label = "{}\l".format(k)
1599                g.node(str(k), label=label, shape='circle')
1600        # Populate the edges
1601        for src, blk in self.blocks.items():
1602            for dst in blk.terminator.get_targets():
1603                g.edge(str(src), str(dst))
1604        return g
1605
1606
1607# A stub for undefined global reference
1608class UndefinedType(EqualityCheckMixin):
1609
1610    _singleton = None
1611
1612    def __new__(cls):
1613        obj = cls._singleton
1614        if obj is not None:
1615            return obj
1616        else:
1617            obj = object.__new__(cls)
1618            cls._singleton = obj
1619        return obj
1620
1621    def __repr__(self):
1622        return "Undefined"
1623
1624
1625UNDEFINED = UndefinedType()
1626