Optimize.py 174 KB
Newer Older
1 2 3 4
from __future__ import absolute_import

from . import TypeSlots
from .ExprNodes import not_a_constant
5 6 7
import cython
cython.declare(UtilityCode=object, EncodedString=object, BytesLiteral=object,
               Nodes=object, ExprNodes=object, PyrexTypes=object, Builtin=object,
8 9 10 11 12 13 14 15 16 17
               UtilNodes=object)

from . import Nodes
from . import ExprNodes
from . import PyrexTypes
from . import Visitor
from . import Builtin
from . import UtilNodes
from . import Options

18
from .Code import UtilityCode, TempitaUtilityCode
19 20 21
from .StringEncoding import EncodedString, BytesLiteral
from .Errors import error
from .ParseTreeTransforms import SkipDeclarations
22

23
import copy
24 25
import codecs

26
try:
27 28
    from __builtin__ import reduce
except ImportError:
29 30
    from functools import reduce

31 32 33 34 35
try:
    from __builtin__ import basestring
except ImportError:
    basestring = str # Python 3

36
def load_c_utility(name):
37
    return UtilityCode.load_cached(name, "Optimize.c")
38

39 40 41 42 43
def unwrap_coerced_node(node, coercion_nodes=(ExprNodes.CoerceToPyTypeNode, ExprNodes.CoerceFromPyTypeNode)):
    if isinstance(node, coercion_nodes):
        return node.arg
    return node

44
def unwrap_node(node):
45 46
    while isinstance(node, UtilNodes.ResultRefNode):
        node = node.expression
47
    return node
48 49

def is_common_value(a, b):
50 51
    a = unwrap_node(a)
    b = unwrap_node(b)
52 53 54
    if isinstance(a, ExprNodes.NameNode) and isinstance(b, ExprNodes.NameNode):
        return a.name == b.name
    if isinstance(a, ExprNodes.AttributeNode) and isinstance(b, ExprNodes.AttributeNode):
55
        return not a.is_py_attr and is_common_value(a.obj, b.obj) and a.attribute == b.attribute
56 57
    return False

58 59 60 61 62
def filter_none_node(node):
    if node is not None and node.constant_result is None:
        return None
    return node

63
class IterationTransform(Visitor.EnvTransform):
64 65 66
    """Transform some common for-in loop patterns into efficient C loops:

    - for-in-dict loop becomes a while loop calling PyDict_Next()
Stefan Behnel's avatar
Stefan Behnel committed
67
    - for-in-enumerate is replaced by an external counter variable
68
    - for-in-range loop becomes a plain C for loop
69
    """
70 71
    def visit_PrimaryCmpNode(self, node):
        if node.is_ptr_contains():
72

73 74 75 76 77 78
            # for t in operand2:
            #     if operand1 == t:
            #         res = True
            #         break
            # else:
            #     res = False
79

80 81 82 83 84 85 86 87 88 89 90
            pos = node.pos
            result_ref = UtilNodes.ResultRefNode(node)
            if isinstance(node.operand2, ExprNodes.IndexNode):
                base_type = node.operand2.base.type.base_type
            else:
                base_type = node.operand2.type.base_type
            target_handle = UtilNodes.TempHandle(base_type)
            target = target_handle.ref(pos)
            cmp_node = ExprNodes.PrimaryCmpNode(
                pos, operator=u'==', operand1=node.operand1, operand2=target)
            if_body = Nodes.StatListNode(
91
                pos,
92 93 94 95 96 97 98 99 100 101 102 103 104 105 106
                stats = [Nodes.SingleAssignmentNode(pos, lhs=result_ref, rhs=ExprNodes.BoolNode(pos, value=1)),
                         Nodes.BreakStatNode(pos)])
            if_node = Nodes.IfStatNode(
                pos,
                if_clauses=[Nodes.IfClauseNode(pos, condition=cmp_node, body=if_body)],
                else_clause=None)
            for_loop = UtilNodes.TempsBlockNode(
                pos,
                temps = [target_handle],
                body = Nodes.ForInStatNode(
                    pos,
                    target=target,
                    iterator=ExprNodes.IteratorNode(node.operand2.pos, sequence=node.operand2),
                    body=if_node,
                    else_clause=Nodes.SingleAssignmentNode(pos, lhs=result_ref, rhs=ExprNodes.BoolNode(pos, value=0))))
107
            for_loop = for_loop.analyse_expressions(self.current_env())
108
            for_loop = self.visit(for_loop)
109
            new_node = UtilNodes.TempResultFromStatNode(result_ref, for_loop)
110

111 112 113 114 115 116 117
            if node.operator == 'not_in':
                new_node = ExprNodes.NotNode(pos, operand=new_node)
            return new_node

        else:
            self.visitchildren(node)
            return node
118

119 120
    def visit_ForInStatNode(self, node):
        self.visitchildren(node)
121
        return self._optimise_for_loop(node, node.iterator.sequence)
122

123
    def _optimise_for_loop(self, node, iterator, reversed=False):
124 125
        if iterator.type is Builtin.dict_type:
            # like iterating over dict.keys()
126
            if reversed:
Stefan Behnel's avatar
Stefan Behnel committed
127
                # CPython raises an error here: not a sequence
128
                return node
Stefan Behnel's avatar
Stefan Behnel committed
129
            return self._transform_dict_iteration(
130
                node, dict_obj=iterator, method=None, keys=True, values=False)
131

132
        # C array (slice) iteration?
133
        if iterator.type.is_ptr or iterator.type.is_array:
134
            return self._transform_carray_iteration(node, iterator, reversed=reversed)
135 136 137 138
        if iterator.type is Builtin.bytes_type:
            return self._transform_bytes_iteration(node, iterator, reversed=reversed)
        if iterator.type is Builtin.unicode_type:
            return self._transform_unicode_iteration(node, iterator, reversed=reversed)
139 140 141

        # the rest is based on function calls
        if not isinstance(iterator, ExprNodes.SimpleCallNode):
Stefan Behnel's avatar
Stefan Behnel committed
142 143
            return node

144 145 146 147 148 149 150
        if iterator.args is None:
            arg_count = iterator.arg_tuple and len(iterator.arg_tuple.args) or 0
        else:
            arg_count = len(iterator.args)
            if arg_count and iterator.self is not None:
                arg_count -= 1

Stefan Behnel's avatar
Stefan Behnel committed
151
        function = iterator.function
152
        # dict iteration?
153
        if function.is_attribute and not reversed and not arg_count:
154
            base_obj = iterator.self or function.obj
155
            method = function.attribute
156
            # in Py3, items() is equivalent to Py2's iteritems()
157
            is_safe_iter = self.global_scope().context.language_level >= 3
158 159 160 161 162 163 164 165 166 167

            if not is_safe_iter and method in ('keys', 'values', 'items'):
                # try to reduce this to the corresponding .iter*() methods
                if isinstance(base_obj, ExprNodes.SimpleCallNode):
                    inner_function = base_obj.function
                    if (inner_function.is_name and inner_function.name == 'dict'
                            and inner_function.entry
                            and inner_function.entry.is_builtin):
                        # e.g. dict(something).items() => safe to use .iter*()
                        is_safe_iter = True
168 169

            keys = values = False
170
            if method == 'iterkeys' or (is_safe_iter and method == 'keys'):
171
                keys = True
172
            elif method == 'itervalues' or (is_safe_iter and method == 'values'):
173
                values = True
174
            elif method == 'iteritems' or (is_safe_iter and method == 'items'):
175
                keys = values = True
176 177 178 179

            if keys or values:
                return self._transform_dict_iteration(
                    node, base_obj, method, keys, values)
180

181
        # enumerate/reversed ?
Stefan Behnel's avatar
Stefan Behnel committed
182
        if iterator.self is None and function.is_name and \
183 184 185
               function.entry and function.entry.is_builtin:
            if function.name == 'enumerate':
                if reversed:
Stefan Behnel's avatar
Stefan Behnel committed
186
                    # CPython raises an error here: not a sequence
187 188 189 190
                    return node
                return self._transform_enumerate_iteration(node, iterator)
            elif function.name == 'reversed':
                if reversed:
Stefan Behnel's avatar
Stefan Behnel committed
191
                    # CPython raises an error here: not a sequence
192 193
                    return node
                return self._transform_reversed_iteration(node, iterator)
194

195 196
        # range() iteration?
        if Options.convert_range and node.target.type.is_int:
Stefan Behnel's avatar
Stefan Behnel committed
197 198 199
            if iterator.self is None and function.is_name and \
                   function.entry and function.entry.is_builtin and \
                   function.name in ('range', 'xrange'):
200
                return self._transform_range_iteration(node, iterator, reversed=reversed)
201

Stefan Behnel's avatar
Stefan Behnel committed
202
        return node
203

204 205 206 207 208 209 210 211 212 213
    def _transform_reversed_iteration(self, node, reversed_function):
        args = reversed_function.arg_tuple.args
        if len(args) == 0:
            error(reversed_function.pos,
                  "reversed() requires an iterable argument")
            return node
        elif len(args) > 1:
            error(reversed_function.pos,
                  "reversed() takes exactly 1 argument")
            return node
214 215 216 217 218 219 220 221 222
        arg = args[0]

        # reversed(list/tuple) ?
        if arg.type in (Builtin.tuple_type, Builtin.list_type):
            node.iterator.sequence = arg.as_none_safe_node("'NoneType' object is not iterable")
            node.iterator.reversed = True
            return node

        return self._optimise_for_loop(node, arg, reversed=True)
223

224 225 226 227 228 229 230 231 232 233
    PyBytes_AS_STRING_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_char_ptr_type, [
            PyrexTypes.CFuncTypeArg("s", Builtin.bytes_type, None)
            ])

    PyBytes_GET_SIZE_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_py_ssize_t_type, [
            PyrexTypes.CFuncTypeArg("s", Builtin.bytes_type, None)
            ])

234 235
    def _transform_bytes_iteration(self, node, slice_node, reversed=False):
        target_type = node.target.type
236
        if not target_type.is_int and target_type is not Builtin.bytes_type:
237 238
            # bytes iteration returns bytes objects in Py2, but
            # integers in Py3
239 240 241
            return node

        unpack_temp_node = UtilNodes.LetRefNode(
242
            slice_node.as_none_safe_node("'NoneType' is not iterable"))
243 244

        slice_base_node = ExprNodes.PythonCapiCallNode(
245 246
            slice_node.pos, "PyBytes_AS_STRING",
            self.PyBytes_AS_STRING_func_type,
247 248 249 250
            args = [unpack_temp_node],
            is_temp = 0,
            )
        len_node = ExprNodes.PythonCapiCallNode(
251 252
            slice_node.pos, "PyBytes_GET_SIZE",
            self.PyBytes_GET_SIZE_func_type,
253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268
            args = [unpack_temp_node],
            is_temp = 0,
            )

        return UtilNodes.LetNode(
            unpack_temp_node,
            self._transform_carray_iteration(
                node,
                ExprNodes.SliceIndexNode(
                    slice_node.pos,
                    base = slice_base_node,
                    start = None,
                    step = None,
                    stop = len_node,
                    type = slice_base_node.type,
                    is_temp = 1,
269 270
                    ),
                reversed = reversed))
271

272 273 274 275 276 277 278
    PyUnicode_READ_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_py_ucs4_type, [
            PyrexTypes.CFuncTypeArg("kind", PyrexTypes.c_int_type, None),
            PyrexTypes.CFuncTypeArg("data", PyrexTypes.c_void_ptr_type, None),
            PyrexTypes.CFuncTypeArg("index", PyrexTypes.c_py_ssize_t_type, None)
        ])

279 280 281 282 283 284 285 286
    init_unicode_iteration_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_int_type, [
            PyrexTypes.CFuncTypeArg("s", PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("length", PyrexTypes.c_py_ssize_t_ptr_type, None),
            PyrexTypes.CFuncTypeArg("data", PyrexTypes.c_void_ptr_ptr_type, None),
            PyrexTypes.CFuncTypeArg("kind", PyrexTypes.c_int_ptr_type, None)
        ],
        exception_value = '-1')
287 288

    def _transform_unicode_iteration(self, node, slice_node, reversed=False):
289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304
        if slice_node.is_literal:
            # try to reduce to byte iteration for plain Latin-1 strings
            try:
                bytes_value = BytesLiteral(slice_node.value.encode('latin1'))
            except UnicodeEncodeError:
                pass
            else:
                bytes_slice = ExprNodes.SliceIndexNode(
                    slice_node.pos,
                    base=ExprNodes.BytesNode(
                        slice_node.pos, value=bytes_value,
                        constant_result=bytes_value,
                        type=PyrexTypes.c_char_ptr_type).coerce_to(
                            PyrexTypes.c_uchar_ptr_type, self.current_env()),
                    start=None,
                    stop=ExprNodes.IntNode(
305
                        slice_node.pos, value=str(len(bytes_value)),
306 307 308 309 310 311
                        constant_result=len(bytes_value),
                        type=PyrexTypes.c_py_ssize_t_type),
                    type=Builtin.unicode_type,  # hint for Python conversion
                )
                return self._transform_carray_iteration(node, bytes_slice, reversed)

312 313 314 315 316
        unpack_temp_node = UtilNodes.LetRefNode(
            slice_node.as_none_safe_node("'NoneType' is not iterable"))

        start_node = ExprNodes.IntNode(
            node.pos, value='0', constant_result=0, type=PyrexTypes.c_py_ssize_t_type)
317 318
        length_temp = UtilNodes.TempHandle(PyrexTypes.c_py_ssize_t_type)
        end_node = length_temp.ref(node.pos)
319 320 321 322 323 324
        if reversed:
            relation1, relation2 = '>', '>='
            start_node, end_node = end_node, start_node
        else:
            relation1, relation2 = '<=', '<'

325 326 327
        kind_temp = UtilNodes.TempHandle(PyrexTypes.c_int_type)
        data_temp = UtilNodes.TempHandle(PyrexTypes.c_void_ptr_type)
        counter_temp = UtilNodes.TempHandle(PyrexTypes.c_py_ssize_t_type)
328

329 330 331
        target_value = ExprNodes.PythonCapiCallNode(
            slice_node.pos, "__Pyx_PyUnicode_READ",
            self.PyUnicode_READ_func_type,
332 333 334
            args = [kind_temp.ref(slice_node.pos),
                    data_temp.ref(slice_node.pos),
                    counter_temp.ref(node.target.pos)],
335 336 337 338
            is_temp = False,
            )
        if target_value.type != node.target.type:
            target_value = target_value.coerce_to(node.target.type,
339
                                                  self.current_env())
340 341 342
        target_assign = Nodes.SingleAssignmentNode(
            pos = node.target.pos,
            lhs = node.target,
343
            rhs = target_value)
344 345 346 347 348 349 350
        body = Nodes.StatListNode(
            node.pos,
            stats = [target_assign, node.body])

        loop_node = Nodes.ForFromStatNode(
            node.pos,
            bound1=start_node, relation1=relation1,
351
            target=counter_temp.ref(node.target.pos),
352 353 354 355 356
            relation2=relation2, bound2=end_node,
            step=None, body=body,
            else_clause=node.else_clause,
            from_range=True)

357 358 359
        setup_node = Nodes.ExprStatNode(
            node.pos,
            expr = ExprNodes.PythonCapiCallNode(
360 361 362 363 364 365 366 367 368 369
                slice_node.pos, "__Pyx_init_unicode_iteration",
                self.init_unicode_iteration_func_type,
                args = [unpack_temp_node,
                        ExprNodes.AmpersandNode(slice_node.pos, operand=length_temp.ref(slice_node.pos),
                                                type=PyrexTypes.c_py_ssize_t_ptr_type),
                        ExprNodes.AmpersandNode(slice_node.pos, operand=data_temp.ref(slice_node.pos),
                                                type=PyrexTypes.c_void_ptr_ptr_type),
                        ExprNodes.AmpersandNode(slice_node.pos, operand=kind_temp.ref(slice_node.pos),
                                                type=PyrexTypes.c_int_ptr_type),
                        ],
370 371
                is_temp = True,
                result_is_used = False,
372
                utility_code=UtilityCode.load_cached("unicode_iter", "Optimize.c"),
373 374 375
                ))
        return UtilNodes.LetNode(
            unpack_temp_node,
376 377 378
            UtilNodes.TempsBlockNode(
                node.pos, temps=[counter_temp, length_temp, data_temp, kind_temp],
                body=Nodes.StatListNode(node.pos, stats=[setup_node, loop_node])))
379

380
    def _transform_carray_iteration(self, node, slice_node, reversed=False):
381
        neg_step = False
382 383
        if isinstance(slice_node, ExprNodes.SliceIndexNode):
            slice_base = slice_node.base
384 385
            start = filter_none_node(slice_node.start)
            stop = filter_none_node(slice_node.stop)
386 387
            step = None
            if not stop:
388 389
                if not slice_base.type.is_pyobject:
                    error(slice_node.pos, "C array iteration requires known end index")
390
                return node
391

392
        elif isinstance(slice_node, ExprNodes.IndexNode):
393
            assert isinstance(slice_node.index, ExprNodes.SliceNode)
394 395
            slice_base = slice_node.base
            index = slice_node.index
396 397 398
            start = filter_none_node(index.start)
            stop = filter_none_node(index.stop)
            step = filter_none_node(index.step)
399
            if step:
400
                if not isinstance(step.constant_result, (int,long)) \
401 402 403
                       or step.constant_result == 0 \
                       or step.constant_result > 0 and not stop \
                       or step.constant_result < 0 and not start:
404 405
                    if not slice_base.type.is_pyobject:
                        error(step.pos, "C array iteration requires known step size and end index")
406 407 408
                    return node
                else:
                    # step sign is handled internally by ForFromStatNode
409 410 411 412
                    step_value = step.constant_result
                    if reversed:
                        step_value = -step_value
                    neg_step = step_value < 0
413
                    step = ExprNodes.IntNode(step.pos, type=PyrexTypes.c_py_ssize_t_type,
414 415
                                             value=str(abs(step_value)),
                                             constant_result=abs(step_value))
416

417 418
        elif slice_node.type.is_array:
            if slice_node.type.size is None:
Stefan Behnel's avatar
Stefan Behnel committed
419
                error(slice_node.pos, "C array iteration requires known end index")
420
                return node
421 422 423
            slice_base = slice_node
            start = None
            stop = ExprNodes.IntNode(
424 425
                slice_node.pos, value=str(slice_node.type.size),
                type=PyrexTypes.c_py_ssize_t_type, constant_result=slice_node.type.size)
426
            step = None
427

428
        else:
429
            if not slice_node.type.is_pyobject:
430
                error(slice_node.pos, "C array iteration requires known end index")
431 432
            return node

433
        if start:
434
            start = start.coerce_to(PyrexTypes.c_py_ssize_t_type, self.current_env())
435
        if stop:
436
            stop = stop.coerce_to(PyrexTypes.c_py_ssize_t_type, self.current_env())
437 438 439 440 441 442 443
        if stop is None:
            if neg_step:
                stop = ExprNodes.IntNode(
                    slice_node.pos, value='-1', type=PyrexTypes.c_py_ssize_t_type, constant_result=-1)
            else:
                error(slice_node.pos, "C array iteration requires known step size and end index")
                return node
444

445 446 447 448 449 450 451
        if reversed:
            if not start:
                start = ExprNodes.IntNode(slice_node.pos, value="0",  constant_result=0,
                                          type=PyrexTypes.c_py_ssize_t_type)
            # if step was provided, it was already negated above
            start, stop = stop, start

452 453 454
        ptr_type = slice_base.type
        if ptr_type.is_array:
            ptr_type = ptr_type.element_ptr_type()
455
        carray_ptr = slice_base.coerce_to_simple(self.current_env())
456

457
        if start and start.constant_result != 0:
458 459 460 461 462
            start_ptr_node = ExprNodes.AddNode(
                start.pos,
                operand1=carray_ptr,
                operator='+',
                operand2=start,
463
                type=ptr_type)
464
        else:
465
            start_ptr_node = carray_ptr
466

467 468 469 470 471 472 473
        if stop and stop.constant_result != 0:
            stop_ptr_node = ExprNodes.AddNode(
                stop.pos,
                operand1=ExprNodes.CloneNode(carray_ptr),
                operator='+',
                operand2=stop,
                type=ptr_type
474
                ).coerce_to_simple(self.current_env())
475 476
        else:
            stop_ptr_node = ExprNodes.CloneNode(carray_ptr)
477

478
        counter = UtilNodes.TempHandle(ptr_type)
479 480
        counter_temp = counter.ref(node.target.pos)

481
        if slice_base.type.is_string and node.target.type.is_pyobject:
482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502
            # special case: char* -> bytes/unicode
            if slice_node.type is Builtin.unicode_type:
                target_value = ExprNodes.CastNode(
                    ExprNodes.DereferenceNode(
                        node.target.pos, operand=counter_temp,
                        type=ptr_type.base_type),
                    PyrexTypes.c_py_ucs4_type).coerce_to(
                        node.target.type, self.current_env())
            else:
                # char* -> bytes coercion requires slicing, not indexing
                target_value = ExprNodes.SliceIndexNode(
                    node.target.pos,
                    start=ExprNodes.IntNode(node.target.pos, value='0',
                                            constant_result=0,
                                            type=PyrexTypes.c_int_type),
                    stop=ExprNodes.IntNode(node.target.pos, value='1',
                                           constant_result=1,
                                           type=PyrexTypes.c_int_type),
                    base=counter_temp,
                    type=Builtin.bytes_type,
                    is_temp=1)
503 504 505
        elif node.target.type.is_ptr and not node.target.type.assignable_from(ptr_type.base_type):
            # Allow iteration with pointer target to avoid copy.
            target_value = counter_temp
506
        else:
507
            # TODO: can this safely be replaced with DereferenceNode() as above?
508 509
            target_value = ExprNodes.IndexNode(
                node.target.pos,
510 511 512 513
                index=ExprNodes.IntNode(node.target.pos, value='0',
                                        constant_result=0,
                                        type=PyrexTypes.c_int_type),
                base=counter_temp,
514
                is_buffer_access=False,
515
                type=ptr_type.base_type)
516 517 518

        if target_value.type != node.target.type:
            target_value = target_value.coerce_to(node.target.type,
519
                                                  self.current_env())
520 521 522 523 524 525 526 527 528 529

        target_assign = Nodes.SingleAssignmentNode(
            pos = node.target.pos,
            lhs = node.target,
            rhs = target_value)

        body = Nodes.StatListNode(
            node.pos,
            stats = [target_assign, node.body])

530 531
        relation1, relation2 = self._find_for_from_node_relations(neg_step, reversed)

532 533
        for_node = Nodes.ForFromStatNode(
            node.pos,
534
            bound1=start_ptr_node, relation1=relation1,
535
            target=counter_temp,
536
            relation2=relation2, bound2=stop_ptr_node,
537 538 539 540 541 542 543 544
            step=step, body=body,
            else_clause=node.else_clause,
            from_range=True)

        return UtilNodes.TempsBlockNode(
            node.pos, temps=[counter],
            body=for_node)

545 546 547 548 549 550
    def _transform_enumerate_iteration(self, node, enumerate_function):
        args = enumerate_function.arg_tuple.args
        if len(args) == 0:
            error(enumerate_function.pos,
                  "enumerate() requires an iterable argument")
            return node
551
        elif len(args) > 2:
552
            error(enumerate_function.pos,
553
                  "enumerate() takes at most 2 arguments")
554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570
            return node

        if not node.target.is_sequence_constructor:
            # leave this untouched for now
            return node
        targets = node.target.args
        if len(targets) != 2:
            # leave this untouched for now
            return node

        enumerate_target, iterable_target = targets
        counter_type = enumerate_target.type

        if not counter_type.is_pyobject and not counter_type.is_int:
            # nothing we can do here, I guess
            return node

571
        if len(args) == 2:
572
            start = unwrap_coerced_node(args[1]).coerce_to(counter_type, self.current_env())
573 574 575 576 577 578 579
        else:
            start = ExprNodes.IntNode(enumerate_function.pos,
                                      value='0',
                                      type=counter_type,
                                      constant_result=0)
        temp = UtilNodes.LetRefNode(start)

580 581
        inc_expression = ExprNodes.AddNode(
            enumerate_function.pos,
582
            operand1 = temp,
583
            operand2 = ExprNodes.IntNode(node.pos, value='1',
584 585
                                         type=counter_type,
                                         constant_result=1),
586 587
            operator = '+',
            type = counter_type,
Stefan Behnel's avatar
Stefan Behnel committed
588
            #inplace = True,   # not worth using in-place operation for Py ints
589 590 591
            is_temp = counter_type.is_pyobject
            )

592 593 594 595
        loop_body = [
            Nodes.SingleAssignmentNode(
                pos = enumerate_target.pos,
                lhs = enumerate_target,
596
                rhs = temp),
597 598
            Nodes.SingleAssignmentNode(
                pos = enumerate_target.pos,
599
                lhs = temp,
600 601
                rhs = inc_expression)
            ]
602

603 604 605 606 607 608 609
        if isinstance(node.body, Nodes.StatListNode):
            node.body.stats = loop_body + node.body.stats
        else:
            loop_body.append(node.body)
            node.body = Nodes.StatListNode(
                node.body.pos,
                stats = loop_body)
610 611

        node.target = iterable_target
612
        node.item = node.item.coerce_to(iterable_target.type, self.current_env())
613
        node.iterator.sequence = args[0]
614 615

        # recurse into loop to check for further optimisations
616
        return UtilNodes.LetNode(temp, self._optimise_for_loop(node, node.iterator.sequence))
617

618 619 620 621 622 623 624 625 626 627 628
    def _find_for_from_node_relations(self, neg_step_value, reversed):
        if reversed:
            if neg_step_value:
                return '<', '<='
            else:
                return '>', '>='
        else:
            if neg_step_value:
                return '>=', '>'
            else:
                return '<=', '<'
629

630
    def _transform_range_iteration(self, node, range_function, reversed=False):
631 632 633 634
        args = range_function.arg_tuple.args
        if len(args) < 3:
            step_pos = range_function.pos
            step_value = 1
635
            step = ExprNodes.IntNode(step_pos, value='1', constant_result=1)
636 637 638
        else:
            step = args[2]
            step_pos = step.pos
639
            if not isinstance(step.constant_result, (int, long)):
640 641
                # cannot determine step direction
                return node
642 643 644
            step_value = step.constant_result
            if step_value == 0:
                # will lead to an error elsewhere
645
                return node
646 647
            step = ExprNodes.IntNode(step_pos, value=str(step_value),
                                     constant_result=step_value)
648 649

        if len(args) == 1:
650 651
            bound1 = ExprNodes.IntNode(range_function.pos, value='0',
                                       constant_result=0)
652
            bound2 = args[0].coerce_to_integer(self.current_env())
653
        else:
654 655
            bound1 = args[0].coerce_to_integer(self.current_env())
            bound2 = args[1].coerce_to_integer(self.current_env())
656

657 658
        relation1, relation2 = self._find_for_from_node_relations(step_value < 0, reversed)

659
        bound2_ref_node = None
660 661
        if reversed:
            bound1, bound2 = bound2, bound1
Stefan Behnel's avatar
Stefan Behnel committed
662 663
            abs_step = abs(step_value)
            if abs_step != 1:
664
                if (isinstance(bound1.constant_result, (int, long)) and
Stefan Behnel's avatar
Stefan Behnel committed
665
                        isinstance(bound2.constant_result, (int, long))):
666
                    # calculate final bounds now
667 668 669 670 671 672 673 674 675 676 677 678
                    if step_value < 0:
                        begin_value = bound2.constant_result
                        end_value = bound1.constant_result
                        bound1_value = begin_value - abs_step * ((begin_value - end_value - 1) // abs_step) - 1
                    else:
                        begin_value = bound1.constant_result
                        end_value = bound2.constant_result
                        bound1_value = end_value + abs_step * ((begin_value - end_value - 1) // abs_step) + 1

                    bound1 = ExprNodes.IntNode(
                        bound1.pos, value=str(bound1_value), constant_result=bound1_value,
                        type=PyrexTypes.spanning_type(bound1.type, bound2.type))
679
                else:
680
                    # evaluate the same expression as above at runtime
681
                    bound2_ref_node = UtilNodes.LetRefNode(bound2)
682
                    spanning_type = PyrexTypes.spanning_type(bound1.type, bound2.type)
683 684 685 686 687
                    if step.type.is_int and abs(step_value) < 0x7FFF:
                        # Avoid loss of integer precision warnings.
                        spanning_step_type = PyrexTypes.spanning_type(spanning_type, PyrexTypes.c_int_type)
                    else:
                        spanning_step_type = PyrexTypes.spanning_type(spanning_type, step.type)
688 689

                    if step_value < 0:
690 691
                        begin_value = bound2_ref_node
                        end_value = bound1
692 693
                        final_op = '-'
                    else:
694 695
                        begin_value = bound1
                        end_value = bound2_ref_node
696 697
                        final_op = '+'

698
                    bound1 = ExprNodes.binop_node(
699
                        bound1.pos,
700
                        operand1=ExprNodes.binop_node(
701
                            bound1.pos,
702
                            operand1=bound2_ref_node,
703
                            operator=final_op,  # +/-
704 705 706 707 708 709 710 711 712
                            operand2=ExprNodes.MulNode(
                                bound1.pos,
                                operand1=ExprNodes.IntNode(
                                    bound1.pos,
                                    value=str(abs_step),
                                    constant_value=abs_step,
                                    type=spanning_step_type),
                                operator='*',
                                operand2=ExprNodes.DivNode(
713
                                    bound1.pos,
714 715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736
                                    operand1=ExprNodes.SubNode(
                                        bound1.pos,
                                        operand1=ExprNodes.SubNode(
                                            bound1.pos,
                                            operand1=begin_value,
                                            operator='-',
                                            operand2=end_value,
                                            type=spanning_type),
                                        operator='-',
                                        operand2=ExprNodes.IntNode(
                                            bound1.pos,
                                            value='1',
                                            constant_result=1),
                                        type=spanning_step_type),
                                    operator='//',
                                    operand2=ExprNodes.IntNode(
                                        bound1.pos,
                                        value=str(abs_step),
                                        constant_value=abs_step,
                                        type=spanning_step_type),
                                    type=spanning_step_type),
                                type=spanning_step_type),
                            type=spanning_step_type),
737
                        operator=final_op,  # +/-
738 739 740 741 742
                        operand2=ExprNodes.IntNode(
                            bound1.pos,
                            value='1',
                            constant_result=1),
                        type=spanning_type)
Stefan Behnel's avatar
Stefan Behnel committed
743 744 745

        if step_value < 0:
            step_value = -step_value
746 747
        step.value = str(step_value)
        step.constant_result = step_value
748
        step = step.coerce_to_integer(self.current_env())
749

750
        if not bound2.is_literal:
751 752
            # stop bound must be immutable => keep it in a temp var
            bound2_is_temp = True
753
            bound2 = bound2_ref_node or UtilNodes.LetRefNode(bound2)
754 755 756
        else:
            bound2_is_temp = False

757 758 759 760 761 762 763
        for_node = Nodes.ForFromStatNode(
            node.pos,
            target=node.target,
            bound1=bound1, relation1=relation1,
            relation2=relation2, bound2=bound2,
            step=step, body=node.body,
            else_clause=node.else_clause,
Magnus Lie Hetland's avatar
Magnus Lie Hetland committed
764
            from_range=True)
765 766 767 768

        if bound2_is_temp:
            for_node = UtilNodes.LetNode(bound2, for_node)

769 770
        return for_node

771
    def _transform_dict_iteration(self, node, dict_obj, method, keys, values):
772
        temps = []
773 774 775
        temp = UtilNodes.TempHandle(PyrexTypes.py_object_type)
        temps.append(temp)
        dict_temp = temp.ref(dict_obj.pos)
776 777
        temp = UtilNodes.TempHandle(PyrexTypes.c_py_ssize_t_type)
        temps.append(temp)
778
        pos_temp = temp.ref(node.pos)
779

780
        key_target = value_target = tuple_target = None
781 782 783 784 785
        if keys and values:
            if node.target.is_sequence_constructor:
                if len(node.target.args) == 2:
                    key_target, value_target = node.target.args
                else:
Stefan Behnel's avatar
Stefan Behnel committed
786
                    # unusual case that may or may not lead to an error
787 788 789
                    return node
            else:
                tuple_target = node.target
790 791
        elif keys:
            key_target = node.target
792
        else:
793
            value_target = node.target
794 795 796 797 798 799 800 801 802 803

        if isinstance(node.body, Nodes.StatListNode):
            body = node.body
        else:
            body = Nodes.StatListNode(pos = node.body.pos,
                                      stats = [node.body])

        # keep original length to guard against dict modification
        dict_len_temp = UtilNodes.TempHandle(PyrexTypes.c_py_ssize_t_type)
        temps.append(dict_len_temp)
804 805 806 807 808 809 810 811 812 813 814 815 816 817
        dict_len_temp_addr = ExprNodes.AmpersandNode(
            node.pos, operand=dict_len_temp.ref(dict_obj.pos),
            type=PyrexTypes.c_ptr_type(dict_len_temp.type))
        temp = UtilNodes.TempHandle(PyrexTypes.c_int_type)
        temps.append(temp)
        is_dict_temp = temp.ref(node.pos)
        is_dict_temp_addr = ExprNodes.AmpersandNode(
            node.pos, operand=is_dict_temp,
            type=PyrexTypes.c_ptr_type(temp.type))

        iter_next_node = Nodes.DictIterationNextNode(
            dict_temp, dict_len_temp.ref(dict_obj.pos), pos_temp,
            key_target, value_target, tuple_target,
            is_dict_temp)
818
        iter_next_node = iter_next_node.analyse_expressions(self.current_env())
819 820 821 822 823 824 825 826 827 828 829 830
        body.stats[0:0] = [iter_next_node]

        if method:
            method_node = ExprNodes.StringNode(
                dict_obj.pos, is_identifier=True, value=method)
            dict_obj = dict_obj.as_none_safe_node(
                "'NoneType' object has no attribute '%s'",
                error = "PyExc_AttributeError",
                format_args = [method])
        else:
            method_node = ExprNodes.NullNode(dict_obj.pos)
            dict_obj = dict_obj.as_none_safe_node("'NoneType' object is not iterable")
831

832 833 834
        def flag_node(value):
            value = value and 1 or 0
            return ExprNodes.IntNode(node.pos, value=str(value), constant_result=value)
835 836

        result_code = [
837
            Nodes.SingleAssignmentNode(
Stefan Behnel's avatar
Stefan Behnel committed
838
                node.pos,
839
                lhs = pos_temp,
840 841
                rhs = ExprNodes.IntNode(node.pos, value='0',
                                        constant_result=0)),
842
            Nodes.SingleAssignmentNode(
Stefan Behnel's avatar
Stefan Behnel committed
843
                dict_obj.pos,
844 845 846 847 848 849
                lhs = dict_temp,
                rhs = ExprNodes.PythonCapiCallNode(
                    dict_obj.pos,
                    "__Pyx_dict_iterator",
                    self.PyDict_Iterator_func_type,
                    utility_code = UtilityCode.load_cached("dict_iter", "Optimize.c"),
850
                    args = [dict_obj, flag_node(dict_obj.type is Builtin.dict_type),
851 852 853
                            method_node, dict_len_temp_addr, is_dict_temp_addr,
                            ],
                    is_temp=True,
854 855
                )),
            Nodes.WhileStatNode(
Stefan Behnel's avatar
Stefan Behnel committed
856
                node.pos,
857
                condition = None,
858 859 860 861 862 863 864 865
                body = body,
                else_clause = node.else_clause
                )
            ]

        return UtilNodes.TempsBlockNode(
            node.pos, temps=temps,
            body=Nodes.StatListNode(
866
                node.pos,
867 868 869
                stats = result_code
                ))

870 871 872 873 874 875 876 877 878
    PyDict_Iterator_func_type = PyrexTypes.CFuncType(
        PyrexTypes.py_object_type, [
            PyrexTypes.CFuncTypeArg("dict",  PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("is_dict",  PyrexTypes.c_int_type, None),
            PyrexTypes.CFuncTypeArg("method_name",  PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("p_orig_length",  PyrexTypes.c_py_ssize_t_ptr_type, None),
            PyrexTypes.CFuncTypeArg("p_is_dict",  PyrexTypes.c_int_ptr_type, None),
            ])

879

880
class SwitchTransform(Visitor.EnvTransform):
881
    """
882
    This transformation tries to turn long if statements into C switch statements.
883
    The requirement is that every clause be an (or of) var == value, where the var
884
    is common among all clauses and both var and value are ints.
885
    """
886 887 888
    NO_MATCH = (None, None, None)

    def extract_conditions(self, cond, allow_not_in):
889
        while True:
890 891
            if isinstance(cond, (ExprNodes.CoerceToTempNode,
                                 ExprNodes.CoerceToBooleanNode)):
892
                cond = cond.arg
893 894
            elif isinstance(cond, ExprNodes.BoolBinopResultNode):
                cond = cond.arg.arg
895 896 897 898 899 900 901
            elif isinstance(cond, UtilNodes.EvalWithTempExprNode):
                # this is what we get from the FlattenInListTransform
                cond = cond.subexpression
            elif isinstance(cond, ExprNodes.TypecastNode):
                cond = cond.operand
            else:
                break
902

903
        if isinstance(cond, ExprNodes.PrimaryCmpNode):
904 905 906 907 908 909 910 911 912 913 914 915 916 917 918
            if cond.cascade is not None:
                return self.NO_MATCH
            elif cond.is_c_string_contains() and \
                   isinstance(cond.operand2, (ExprNodes.UnicodeNode, ExprNodes.BytesNode)):
                not_in = cond.operator == 'not_in'
                if not_in and not allow_not_in:
                    return self.NO_MATCH
                if isinstance(cond.operand2, ExprNodes.UnicodeNode) and \
                       cond.operand2.contains_surrogates():
                    # dealing with surrogates leads to different
                    # behaviour on wide and narrow Unicode
                    # platforms => refuse to optimise this case
                    return self.NO_MATCH
                return not_in, cond.operand1, self.extract_in_string_conditions(cond.operand2)
            elif not cond.is_python_comparison():
919 920 921 922 923 924 925 926 927 928 929 930 931 932 933 934 935 936 937 938
                if cond.operator == '==':
                    not_in = False
                elif allow_not_in and cond.operator == '!=':
                    not_in = True
                else:
                    return self.NO_MATCH
                # this looks somewhat silly, but it does the right
                # checks for NameNode and AttributeNode
                if is_common_value(cond.operand1, cond.operand1):
                    if cond.operand2.is_literal:
                        return not_in, cond.operand1, [cond.operand2]
                    elif getattr(cond.operand2, 'entry', None) \
                             and cond.operand2.entry.is_const:
                        return not_in, cond.operand1, [cond.operand2]
                if is_common_value(cond.operand2, cond.operand2):
                    if cond.operand1.is_literal:
                        return not_in, cond.operand2, [cond.operand1]
                    elif getattr(cond.operand1, 'entry', None) \
                             and cond.operand1.entry.is_const:
                        return not_in, cond.operand2, [cond.operand1]
Stefan Behnel's avatar
Stefan Behnel committed
939
        elif isinstance(cond, ExprNodes.BoolBinopNode):
940 941 942 943 944 945 946 947 948
            if cond.operator == 'or' or (allow_not_in and cond.operator == 'and'):
                allow_not_in = (cond.operator == 'and')
                not_in_1, t1, c1 = self.extract_conditions(cond.operand1, allow_not_in)
                not_in_2, t2, c2 = self.extract_conditions(cond.operand2, allow_not_in)
                if t1 is not None and not_in_1 == not_in_2 and is_common_value(t1, t2):
                    if (not not_in_1) or allow_not_in:
                        return not_in_1, t1, c1+c2
        return self.NO_MATCH

949 950
    def extract_in_string_conditions(self, string_literal):
        if isinstance(string_literal, ExprNodes.UnicodeNode):
951
            charvals = list(map(ord, set(string_literal.value)))
952 953 954 955 956 957 958 959 960
            charvals.sort()
            return [ ExprNodes.IntNode(string_literal.pos, value=str(charval),
                                       constant_result=charval)
                     for charval in charvals ]
        else:
            # this is a bit tricky as Py3's bytes type returns
            # integers on iteration, whereas Py2 returns 1-char byte
            # strings
            characters = string_literal.value
961 962
            characters = list(set([ characters[i:i+1] for i in range(len(characters)) ]))
            characters.sort()
963 964 965 966
            return [ ExprNodes.CharNode(string_literal.pos, value=charval,
                                        constant_result=charval)
                     for charval in characters ]

967 968
    def extract_common_conditions(self, common_var, condition, allow_not_in):
        not_in, var, conditions = self.extract_conditions(condition, allow_not_in)
969
        if var is None:
970
            return self.NO_MATCH
971
        elif common_var is not None and not is_common_value(var, common_var):
972
            return self.NO_MATCH
973
        elif not (var.type.is_int or var.type.is_enum) or sum([not (cond.type.is_int or cond.type.is_enum) for cond in conditions]):
974 975 976 977 978 979 980
            return self.NO_MATCH
        return not_in, var, conditions

    def has_duplicate_values(self, condition_values):
        # duplicated values don't work in a switch statement
        seen = set()
        for value in condition_values:
981
            if value.has_constant_result():
982 983 984 985 986 987
                if value.constant_result in seen:
                    return True
                seen.add(value.constant_result)
            else:
                # this isn't completely safe as we don't know the
                # final C value, but this is about the best we can do
988 989 990 991 992 993
                try:
                    if value.entry.cname in seen:
                        return True
                except AttributeError:
                    return True  # play safe
                seen.add(value.entry.cname)
994
        return False
995

996
    def visit_IfStatNode(self, node):
997
        if not self.current_directives.get('optimize.use_switch'):
998 999 1000
            self.visitchildren(node)
            return node

1001 1002 1003
        common_var = None
        cases = []
        for if_clause in node.if_clauses:
1004 1005
            _, common_var, conditions = self.extract_common_conditions(
                common_var, if_clause.condition, False)
1006
            if common_var is None:
1007
                self.visitchildren(node)
1008
                return node
1009 1010 1011 1012
            cases.append(Nodes.SwitchCaseNode(pos = if_clause.pos,
                                              conditions = conditions,
                                              body = if_clause.body))

1013 1014 1015
        condition_values = [
            cond for case in cases for cond in case.conditions]
        if len(condition_values) < 2:
1016 1017
            self.visitchildren(node)
            return node
1018
        if self.has_duplicate_values(condition_values):
1019
            self.visitchildren(node)
1020
            return node
1021

Robert Bradshaw's avatar
Robert Bradshaw committed
1022
        common_var = unwrap_node(common_var)
1023 1024 1025 1026 1027 1028 1029
        switch_node = Nodes.SwitchStatNode(pos = node.pos,
                                           test = common_var,
                                           cases = cases,
                                           else_clause = node.else_clause)
        return switch_node

    def visit_CondExprNode(self, node):
1030
        if not self.current_directives.get('optimize.use_switch'):
1031 1032 1033
            self.visitchildren(node)
            return node

1034 1035 1036 1037 1038 1039
        not_in, common_var, conditions = self.extract_common_conditions(
            None, node.test, True)
        if common_var is None \
               or len(conditions) < 2 \
               or self.has_duplicate_values(conditions):
            self.visitchildren(node)
1040
            return node
1041 1042 1043
        return self.build_simple_switch_statement(
            node, common_var, conditions, not_in,
            node.true_val, node.false_val)
1044 1045

    def visit_BoolBinopNode(self, node):
1046
        if not self.current_directives.get('optimize.use_switch'):
1047 1048 1049
            self.visitchildren(node)
            return node

1050 1051 1052 1053 1054 1055
        not_in, common_var, conditions = self.extract_common_conditions(
            None, node, True)
        if common_var is None \
               or len(conditions) < 2 \
               or self.has_duplicate_values(conditions):
            self.visitchildren(node)
1056
            node.wrap_operands(self.current_env())  # in case we changed the operands
1057 1058
            return node

1059 1060
        return self.build_simple_switch_statement(
            node, common_var, conditions, not_in,
1061 1062 1063 1064
            ExprNodes.BoolNode(node.pos, value=True, constant_result=True),
            ExprNodes.BoolNode(node.pos, value=False, constant_result=False))

    def visit_PrimaryCmpNode(self, node):
1065
        if not self.current_directives.get('optimize.use_switch'):
1066 1067 1068
            self.visitchildren(node)
            return node

1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079 1080
        not_in, common_var, conditions = self.extract_common_conditions(
            None, node, True)
        if common_var is None \
               or len(conditions) < 2 \
               or self.has_duplicate_values(conditions):
            self.visitchildren(node)
            return node

        return self.build_simple_switch_statement(
            node, common_var, conditions, not_in,
            ExprNodes.BoolNode(node.pos, value=True, constant_result=True),
            ExprNodes.BoolNode(node.pos, value=False, constant_result=False))
1081 1082 1083

    def build_simple_switch_statement(self, node, common_var, conditions,
                                      not_in, true_val, false_val):
1084 1085 1086
        result_ref = UtilNodes.ResultRefNode(node)
        true_body = Nodes.SingleAssignmentNode(
            node.pos,
1087 1088 1089
            lhs=result_ref,
            rhs=true_val.coerce_to(node.type, self.current_env()),
            first=True)
1090 1091
        false_body = Nodes.SingleAssignmentNode(
            node.pos,
1092 1093 1094
            lhs=result_ref,
            rhs=false_val.coerce_to(node.type, self.current_env()),
            first=True)
1095

1096 1097 1098
        if not_in:
            true_body, false_body = false_body, true_body

1099 1100 1101 1102 1103 1104 1105 1106 1107
        cases = [Nodes.SwitchCaseNode(pos = node.pos,
                                      conditions = conditions,
                                      body = true_body)]

        common_var = unwrap_node(common_var)
        switch_node = Nodes.SwitchStatNode(pos = node.pos,
                                           test = common_var,
                                           cases = cases,
                                           else_clause = false_body)
1108 1109 1110 1111
        replacement = UtilNodes.TempResultFromStatNode(result_ref, switch_node)
        return replacement

    def visit_EvalWithTempExprNode(self, node):
1112
        if not self.current_directives.get('optimize.use_switch'):
1113 1114 1115
            self.visitchildren(node)
            return node

1116 1117 1118 1119 1120 1121 1122 1123 1124
        # drop unused expression temp from FlattenInListTransform
        orig_expr = node.subexpression
        temp_ref = node.lazy_temp
        self.visitchildren(node)
        if node.subexpression is not orig_expr:
            # node was restructured => check if temp is still used
            if not Visitor.tree_contains(node.subexpression, temp_ref):
                return node.subexpression
        return node
1125

1126
    visit_Node = Visitor.VisitorTransform.recurse_to_children
1127

1128

1129
class FlattenInListTransform(Visitor.VisitorTransform, SkipDeclarations):
1130 1131
    """
    This transformation flattens "x in [val1, ..., valn]" into a sequential list
1132
    of comparisons.
1133
    """
1134

1135 1136 1137 1138 1139 1140 1141 1142 1143 1144 1145 1146
    def visit_PrimaryCmpNode(self, node):
        self.visitchildren(node)
        if node.cascade is not None:
            return node
        elif node.operator == 'in':
            conjunction = 'or'
            eq_or_neq = '=='
        elif node.operator == 'not_in':
            conjunction = 'and'
            eq_or_neq = '!='
        else:
            return node
1147

1148 1149 1150
        if not isinstance(node.operand2, (ExprNodes.TupleNode,
                                          ExprNodes.ListNode,
                                          ExprNodes.SetNode)):
Stefan Behnel's avatar
Stefan Behnel committed
1151
            return node
1152

Stefan Behnel's avatar
Stefan Behnel committed
1153 1154
        args = node.operand2.args
        if len(args) == 0:
1155 1156
            # note: lhs may have side effects
            return node
1157

1158
        lhs = UtilNodes.ResultRefNode(node.operand1)
Stefan Behnel's avatar
Stefan Behnel committed
1159 1160

        conds = []
1161
        temps = []
Stefan Behnel's avatar
Stefan Behnel committed
1162
        for arg in args:
1163 1164 1165 1166 1167 1168 1169 1170 1171
            try:
                # Trial optimisation to avoid redundant temp
                # assignments.  However, since is_simple() is meant to
                # be called after type analysis, we ignore any errors
                # and just play safe in that case.
                is_simple_arg = arg.is_simple()
            except Exception:
                is_simple_arg = False
            if not is_simple_arg:
1172 1173 1174
                # must evaluate all non-simple RHS before doing the comparisons
                arg = UtilNodes.LetRefNode(arg)
                temps.append(arg)
Stefan Behnel's avatar
Stefan Behnel committed
1175 1176 1177 1178 1179 1180 1181
            cond = ExprNodes.PrimaryCmpNode(
                                pos = node.pos,
                                operand1 = lhs,
                                operator = eq_or_neq,
                                operand2 = arg,
                                cascade = None)
            conds.append(ExprNodes.TypecastNode(
1182
                                pos = node.pos,
Stefan Behnel's avatar
Stefan Behnel committed
1183 1184 1185 1186
                                operand = cond,
                                type = PyrexTypes.c_bint_type))
        def concat(left, right):
            return ExprNodes.BoolBinopNode(
1187
                                pos = node.pos,
Stefan Behnel's avatar
Stefan Behnel committed
1188 1189 1190 1191
                                operator = conjunction,
                                operand1 = left,
                                operand2 = right)

1192
        condition = reduce(concat, conds)
1193 1194 1195 1196
        new_node = UtilNodes.EvalWithTempExprNode(lhs, condition)
        for temp in temps[::-1]:
            new_node = UtilNodes.EvalWithTempExprNode(temp, new_node)
        return new_node
1197

1198
    visit_Node = Visitor.VisitorTransform.recurse_to_children
1199 1200


1201 1202 1203 1204 1205 1206
class DropRefcountingTransform(Visitor.VisitorTransform):
    """Drop ref-counting in safe places.
    """
    visit_Node = Visitor.VisitorTransform.recurse_to_children

    def visit_ParallelAssignmentNode(self, node):
Stefan Behnel's avatar
Stefan Behnel committed
1207 1208 1209
        """
        Parallel swap assignments like 'a,b = b,a' are safe.
        """
1210 1211 1212 1213
        left_names, right_names = [], []
        left_indices, right_indices = [], []
        temps = []

1214 1215
        for stat in node.stats:
            if isinstance(stat, Nodes.SingleAssignmentNode):
1216 1217
                if not self._extract_operand(stat.lhs, left_names,
                                             left_indices, temps):
1218
                    return node
1219 1220
                if not self._extract_operand(stat.rhs, right_names,
                                             right_indices, temps):
1221
                    return node
1222 1223 1224
            elif isinstance(stat, Nodes.CascadedAssignmentNode):
                # FIXME
                return node
1225 1226 1227
            else:
                return node

1228 1229
        if left_names or right_names:
            # lhs/rhs names must be a non-redundant permutation
1230 1231
            lnames = [ path for path, n in left_names ]
            rnames = [ path for path, n in right_names ]
1232 1233 1234
            if set(lnames) != set(rnames):
                return node
            if len(set(lnames)) != len(right_names):
1235 1236
                return node

1237 1238 1239 1240 1241 1242 1243 1244 1245 1246 1247 1248 1249 1250 1251
        if left_indices or right_indices:
            # base name and index of index nodes must be a
            # non-redundant permutation
            lindices = []
            for lhs_node in left_indices:
                index_id = self._extract_index_id(lhs_node)
                if not index_id:
                    return node
                lindices.append(index_id)
            rindices = []
            for rhs_node in right_indices:
                index_id = self._extract_index_id(rhs_node)
                if not index_id:
                    return node
                rindices.append(index_id)
1252

1253 1254 1255 1256 1257 1258 1259
            if set(lindices) != set(rindices):
                return node
            if len(set(lindices)) != len(right_indices):
                return node

            # really supporting IndexNode requires support in
            # __Pyx_GetItemInt(), so let's stop short for now
1260 1261
            return node

1262 1263 1264 1265
        temp_args = [t.arg for t in temps]
        for temp in temps:
            temp.use_managed_ref = False

1266
        for _, name_node in left_names + right_names:
1267 1268 1269 1270 1271
            if name_node not in temp_args:
                name_node.use_managed_ref = False

        for index_node in left_indices + right_indices:
            index_node.use_managed_ref = False
1272 1273 1274

        return node

1275 1276 1277 1278 1279 1280 1281
    def _extract_operand(self, node, names, indices, temps):
        node = unwrap_node(node)
        if not node.type.is_pyobject:
            return False
        if isinstance(node, ExprNodes.CoerceToTempNode):
            temps.append(node)
            node = node.arg
1282 1283 1284 1285
        name_path = []
        obj_node = node
        while isinstance(obj_node, ExprNodes.AttributeNode):
            if obj_node.is_py_attr:
1286
                return False
1287 1288 1289 1290 1291
            name_path.append(obj_node.member)
            obj_node = obj_node.obj
        if isinstance(obj_node, ExprNodes.NameNode):
            name_path.append(obj_node.name)
            names.append( ('.'.join(name_path[::-1]), node) )
1292 1293 1294 1295 1296 1297 1298 1299 1300 1301 1302 1303 1304 1305 1306 1307 1308 1309 1310 1311 1312 1313 1314 1315
        elif isinstance(node, ExprNodes.IndexNode):
            if node.base.type != Builtin.list_type:
                return False
            if not node.index.type.is_int:
                return False
            if not isinstance(node.base, ExprNodes.NameNode):
                return False
            indices.append(node)
        else:
            return False
        return True

    def _extract_index_id(self, index_node):
        base = index_node.base
        index = index_node.index
        if isinstance(index, ExprNodes.NameNode):
            index_val = index.name
        elif isinstance(index, ExprNodes.ConstNode):
            # FIXME:
            return None
        else:
            return None
        return (base.name, index_val)

1316

1317 1318 1319 1320 1321 1322 1323
class EarlyReplaceBuiltinCalls(Visitor.EnvTransform):
    """Optimize some common calls to builtin types *before* the type
    analysis phase and *after* the declarations analysis phase.

    This transform cannot make use of any argument types, but it can
    restructure the tree in a way that the type analysis phase can
    respond to.
Stefan Behnel's avatar
Stefan Behnel committed
1324 1325 1326

    Introducing C function calls here may not be a good idea.  Move
    them to the OptimizeBuiltinCalls transform instead, which runs
Stefan Behnel's avatar
Stefan Behnel committed
1327
    after type analysis.
1328
    """
1329 1330
    # only intercept on call nodes
    visit_Node = Visitor.VisitorTransform.recurse_to_children
1331

1332 1333 1334 1335 1336 1337 1338
    def visit_SimpleCallNode(self, node):
        self.visitchildren(node)
        function = node.function
        if not self._function_is_builtin_name(function):
            return node
        return self._dispatch_to_handler(node, function, node.args)

1339
    def visit_GeneralCallNode(self, node):
1340
        self.visitchildren(node)
1341
        function = node.function
1342
        if not self._function_is_builtin_name(function):
1343 1344 1345 1346
            return node
        arg_tuple = node.positional_args
        if not isinstance(arg_tuple, ExprNodes.TupleNode):
            return node
1347
        args = arg_tuple.args
1348
        return self._dispatch_to_handler(
1349
            node, function, args, node.keyword_args)
1350

1351 1352 1353
    def _function_is_builtin_name(self, function):
        if not function.is_name:
            return False
1354
        env = self.current_env()
1355
        entry = env.lookup(function.name)
1356
        if entry is not env.builtin_scope().lookup_here(function.name):
1357
            return False
1358
        # if entry is None, it's at least an undeclared name, so likely builtin
1359 1360 1361 1362 1363 1364 1365 1366 1367 1368 1369 1370 1371 1372 1373 1374 1375 1376 1377 1378 1379 1380 1381 1382 1383 1384 1385 1386 1387 1388 1389 1390 1391 1392 1393 1394 1395 1396
        return True

    def _dispatch_to_handler(self, node, function, args, kwargs=None):
        if kwargs is None:
            handler_name = '_handle_simple_function_%s' % function.name
        else:
            handler_name = '_handle_general_function_%s' % function.name
        handle_call = getattr(self, handler_name, None)
        if handle_call is not None:
            if kwargs is None:
                return handle_call(node, args)
            else:
                return handle_call(node, args, kwargs)
        return node

    def _inject_capi_function(self, node, cname, func_type, utility_code=None):
        node.function = ExprNodes.PythonCapiFunctionNode(
            node.function.pos, node.function.name, cname, func_type,
            utility_code = utility_code)

    def _error_wrong_arg_count(self, function_name, node, args, expected=None):
        if not expected: # None or 0
            arg_str = ''
        elif isinstance(expected, basestring) or expected > 1:
            arg_str = '...'
        elif expected == 1:
            arg_str = 'x'
        else:
            arg_str = ''
        if expected is not None:
            expected_str = 'expected %s, ' % expected
        else:
            expected_str = ''
        error(node.pos, "%s(%s) called with wrong number of args, %sfound %d" % (
            function_name, arg_str, expected_str, len(args)))

    # specific handlers for simple call nodes

1397
    def _handle_simple_function_float(self, node, pos_args):
Stefan Behnel's avatar
Stefan Behnel committed
1398
        if not pos_args:
1399 1400 1401
            return ExprNodes.FloatNode(node.pos, value='0.0')
        if len(pos_args) > 1:
            self._error_wrong_arg_count('float', node, pos_args, 1)
1402 1403 1404
        arg_type = getattr(pos_args[0], 'type', None)
        if arg_type in (PyrexTypes.c_double_type, Builtin.float_type):
            return pos_args[0]
1405 1406
        return node

Stefan Behnel's avatar
Stefan Behnel committed
1407 1408 1409 1410 1411 1412 1413 1414 1415 1416 1417 1418 1419 1420 1421 1422 1423 1424
    def _handle_simple_function_slice(self, node, pos_args):
        arg_count = len(pos_args)
        start = step = None
        if arg_count == 1:
            stop, = pos_args
        elif arg_count == 2:
            start, stop = pos_args
        elif arg_count == 3:
            start, stop, step = pos_args
        else:
            self._error_wrong_arg_count('slice', node, pos_args)
            return node
        return ExprNodes.SliceNode(
            node.pos,
            start=start or ExprNodes.NoneNode(node.pos),
            stop=stop,
            step=step or ExprNodes.NoneNode(node.pos))

1425 1426 1427 1428 1429 1430 1431 1432 1433 1434 1435 1436 1437 1438 1439 1440 1441 1442 1443 1444 1445 1446 1447 1448 1449
    def _handle_simple_function_ord(self, node, pos_args):
        """Unpack ord('X').
        """
        if len(pos_args) != 1:
            return node
        arg = pos_args[0]
        if isinstance(arg, (ExprNodes.UnicodeNode, ExprNodes.BytesNode)):
            if len(arg.value) == 1:
                return ExprNodes.IntNode(
                    arg.pos, type=PyrexTypes.c_long_type,
                    value=str(ord(arg.value)),
                    constant_result=ord(arg.value)
                )
        elif isinstance(arg, ExprNodes.StringNode):
            if arg.unicode_value and len(arg.unicode_value) == 1 \
                    and ord(arg.unicode_value) <= 255:  # Py2/3 portability
                return ExprNodes.IntNode(
                    arg.pos, type=PyrexTypes.c_int_type,
                    value=str(ord(arg.unicode_value)),
                    constant_result=ord(arg.unicode_value)
                )
        return node

    # sequence processing

1450 1451 1452
    class YieldNodeCollector(Visitor.TreeVisitor):
        def __init__(self):
            Visitor.TreeVisitor.__init__(self)
1453
            self.yield_stat_nodes = {}
1454 1455 1456
            self.yield_nodes = []

        visit_Node = Visitor.TreeVisitor.visitchildren
1457 1458
        # XXX: disable inlining while it's not back supported
        def __visit_YieldExprNode(self, node):
1459 1460 1461
            self.yield_nodes.append(node)
            self.visitchildren(node)

1462
        def __visit_ExprStatNode(self, node):
1463 1464 1465 1466
            self.visitchildren(node)
            if node.expr in self.yield_nodes:
                self.yield_stat_nodes[node.expr] = node

1467 1468 1469 1470 1471 1472
        def __visit_GeneratorExpressionNode(self, node):
            # enable when we support generic generator expressions
            #
            # everything below this node is out of scope
            pass

1473
    def _find_single_yield_expression(self, node):
1474 1475 1476
        collector = self.YieldNodeCollector()
        collector.visitchildren(node)
        if len(collector.yield_nodes) != 1:
1477 1478
            return None, None
        yield_node = collector.yield_nodes[0]
1479 1480 1481 1482
        try:
            return (yield_node.arg, collector.yield_stat_nodes[yield_node])
        except KeyError:
            return None, None
1483

1484 1485 1486 1487 1488 1489 1490 1491 1492 1493 1494 1495 1496 1497 1498 1499 1500 1501 1502 1503 1504 1505 1506 1507 1508 1509 1510 1511 1512 1513 1514 1515 1516 1517 1518 1519 1520 1521 1522 1523 1524 1525 1526 1527 1528
    def _handle_simple_function_all(self, node, pos_args):
        """Transform

        _result = all(x for L in LL for x in L)

        into

        for L in LL:
            for x in L:
                if not x:
                    _result = False
                    break
            else:
                continue
            break
        else:
            _result = True
        """
        return self._transform_any_all(node, pos_args, False)

    def _handle_simple_function_any(self, node, pos_args):
        """Transform

        _result = any(x for L in LL for x in L)

        into

        for L in LL:
            for x in L:
                if x:
                    _result = True
                    break
            else:
                continue
            break
        else:
            _result = False
        """
        return self._transform_any_all(node, pos_args, True)

    def _transform_any_all(self, node, pos_args, is_any):
        if len(pos_args) != 1:
            return node
        if not isinstance(pos_args[0], ExprNodes.GeneratorExpressionNode):
            return node
1529 1530
        gen_expr_node = pos_args[0]
        loop_node = gen_expr_node.loop
1531 1532
        yield_expression, yield_stat_node = self._find_single_yield_expression(loop_node)
        if yield_expression is None:
1533 1534 1535 1536 1537 1538 1539
            return node

        if is_any:
            condition = yield_expression
        else:
            condition = ExprNodes.NotNode(yield_expression.pos, operand = yield_expression)

1540
        result_ref = UtilNodes.ResultRefNode(pos=node.pos, type=PyrexTypes.c_bint_type)
1541
        test_node = Nodes.IfStatNode(
1542
            yield_expression.pos,
1543 1544
            else_clause = None,
            if_clauses = [ Nodes.IfClauseNode(
1545
                yield_expression.pos,
1546 1547 1548 1549 1550 1551 1552
                condition = condition,
                body = Nodes.StatListNode(
                    node.pos,
                    stats = [
                        Nodes.SingleAssignmentNode(
                            node.pos,
                            lhs = result_ref,
1553
                            rhs = ExprNodes.BoolNode(yield_expression.pos, value = is_any,
1554 1555 1556 1557 1558 1559 1560 1561 1562
                                                     constant_result = is_any)),
                        Nodes.BreakStatNode(node.pos)
                        ])) ]
            )
        loop = loop_node
        while isinstance(loop.body, Nodes.LoopNode):
            next_loop = loop.body
            loop.body = Nodes.StatListNode(loop.body.pos, stats = [
                loop.body,
1563
                Nodes.BreakStatNode(yield_expression.pos)
1564
                ])
1565
            next_loop.else_clause = Nodes.ContinueStatNode(yield_expression.pos)
1566 1567 1568 1569
            loop = next_loop
        loop_node.else_clause = Nodes.SingleAssignmentNode(
            node.pos,
            lhs = result_ref,
1570
            rhs = ExprNodes.BoolNode(yield_expression.pos, value = not is_any,
1571 1572
                                     constant_result = not is_any))

1573
        Visitor.recursively_replace_node(loop_node, yield_stat_node, test_node)
1574

1575 1576
        return ExprNodes.InlinedGeneratorExpressionNode(
            gen_expr_node.pos, loop = loop_node, result_node = result_ref,
1577
            expr_scope = gen_expr_node.expr_scope, orig_func = is_any and 'any' or 'all')
1578

1579 1580 1581 1582
    PySequence_List_func_type = PyrexTypes.CFuncType(
        Builtin.list_type,
        [PyrexTypes.CFuncTypeArg("it", PyrexTypes.py_object_type, None)])

1583
    def _handle_simple_function_sorted(self, node, pos_args):
1584 1585 1586 1587 1588
        """Transform sorted(genexpr) and sorted([listcomp]) into
        [listcomp].sort().  CPython just reads the iterable into a
        list and calls .sort() on it.  Expanding the iterable in a
        listcomp is still faster and the result can be sorted in
        place.
1589 1590 1591
        """
        if len(pos_args) != 1:
            return node
1592
        if isinstance(pos_args[0], ExprNodes.ComprehensionNode) \
1593
               and pos_args[0].type is Builtin.list_type:
1594 1595 1596 1597 1598 1599 1600 1601
            listcomp_node = pos_args[0]
            loop_node = listcomp_node.loop
        elif isinstance(pos_args[0], ExprNodes.GeneratorExpressionNode):
            gen_expr_node = pos_args[0]
            loop_node = gen_expr_node.loop
            yield_expression, yield_stat_node = self._find_single_yield_expression(loop_node)
            if yield_expression is None:
                return node
1602

1603
            append_node = ExprNodes.ComprehensionAppendNode(
1604
                yield_expression.pos, expr = yield_expression)
1605

1606
            Visitor.recursively_replace_node(loop_node, yield_stat_node, append_node)
1607

1608
            listcomp_node = ExprNodes.ComprehensionNode(
1609
                gen_expr_node.pos, loop = loop_node,
1610 1611 1612
                append = append_node, type = Builtin.list_type,
                expr_scope = gen_expr_node.expr_scope,
                has_local_scope = True)
1613
            append_node.target = listcomp_node
1614 1615 1616 1617 1618
        elif isinstance(pos_args[0], (ExprNodes.ListNode, ExprNodes.TupleNode)):
            # sorted([a, b, c]) or sorted((a, b, c)). The result of the latter
            # is a list in CPython, so change it into one.
            expr = pos_args[0].as_list()
            listcomp_node = loop_node = expr
1619
        else:
1620 1621 1622 1623 1624
            # Interestingly, PySequence_List works on a lot of non-sequence
            # things as well.
            listcomp_node = loop_node = ExprNodes.PythonCapiCallNode(
                node.pos, "PySequence_List", self.PySequence_List_func_type,
                args=pos_args, is_temp=True)
1625

1626 1627
        result_node = UtilNodes.ResultRefNode(
            pos = loop_node.pos, type = Builtin.list_type, may_hold_none=False)
1628 1629 1630 1631 1632 1633 1634 1635 1636 1637 1638 1639 1640 1641 1642 1643 1644
        listcomp_assign_node = Nodes.SingleAssignmentNode(
            node.pos, lhs = result_node, rhs = listcomp_node, first = True)

        sort_method = ExprNodes.AttributeNode(
            node.pos, obj = result_node, attribute = EncodedString('sort'),
            # entry ? type ?
            needs_none_check = False)
        sort_node = Nodes.ExprStatNode(
            node.pos, expr = ExprNodes.SimpleCallNode(
                node.pos, function = sort_method, args = []))

        sort_node.analyse_declarations(self.current_env())

        return UtilNodes.TempResultFromStatNode(
            result_node,
            Nodes.StatListNode(node.pos, stats = [ listcomp_assign_node, sort_node ]))

1645
    def _handle_simple_function_sum(self, node, pos_args):
Stefan Behnel's avatar
Stefan Behnel committed
1646 1647
        """Transform sum(genexpr) into an equivalent inlined aggregation loop.
        """
1648 1649
        if len(pos_args) not in (1,2):
            return node
1650 1651
        if not isinstance(pos_args[0], (ExprNodes.GeneratorExpressionNode,
                                        ExprNodes.ComprehensionNode)):
1652 1653 1654 1655
            return node
        gen_expr_node = pos_args[0]
        loop_node = gen_expr_node.loop

1656 1657 1658 1659
        if isinstance(gen_expr_node, ExprNodes.GeneratorExpressionNode):
            yield_expression, yield_stat_node = self._find_single_yield_expression(loop_node)
            if yield_expression is None:
                return node
1660
        else:  # ComprehensionNode
1661 1662 1663 1664 1665 1666 1667 1668 1669
            yield_stat_node = gen_expr_node.append
            yield_expression = yield_stat_node.expr
            try:
                if not yield_expression.is_literal or not yield_expression.type.is_int:
                    return node
            except AttributeError:
                return node # in case we don't have a type yet
            # special case: old Py2 backwards compatible "sum([int_const for ...])"
            # can safely be unpacked into a genexpr
1670 1671 1672 1673 1674 1675 1676 1677

        if len(pos_args) == 1:
            start = ExprNodes.IntNode(node.pos, value='0', constant_result=0)
        else:
            start = pos_args[1]

        result_ref = UtilNodes.ResultRefNode(pos=node.pos, type=PyrexTypes.py_object_type)
        add_node = Nodes.SingleAssignmentNode(
1678
            yield_expression.pos,
1679 1680 1681 1682
            lhs = result_ref,
            rhs = ExprNodes.binop_node(node.pos, '+', result_ref, yield_expression)
            )

1683
        Visitor.recursively_replace_node(loop_node, yield_stat_node, add_node)
1684 1685 1686 1687 1688 1689 1690 1691 1692 1693 1694 1695 1696 1697

        exec_code = Nodes.StatListNode(
            node.pos,
            stats = [
                Nodes.SingleAssignmentNode(
                    start.pos,
                    lhs = UtilNodes.ResultRefNode(pos=node.pos, expression=result_ref),
                    rhs = start,
                    first = True),
                loop_node
                ])

        return ExprNodes.InlinedGeneratorExpressionNode(
            gen_expr_node.pos, loop = exec_code, result_node = result_ref,
1698 1699
            expr_scope = gen_expr_node.expr_scope, orig_func = 'sum',
            has_local_scope = gen_expr_node.has_local_scope)
1700

1701 1702 1703 1704 1705 1706 1707 1708 1709 1710
    def _handle_simple_function_min(self, node, pos_args):
        return self._optimise_min_max(node, pos_args, '<')

    def _handle_simple_function_max(self, node, pos_args):
        return self._optimise_min_max(node, pos_args, '>')

    def _optimise_min_max(self, node, args, operator):
        """Replace min(a,b,...) and max(a,b,...) by explicit comparison code.
        """
        if len(args) <= 1:
1711 1712 1713 1714 1715
            if len(args) == 1 and args[0].is_sequence_constructor:
                args = args[0].args
            else:
                # leave this to Python
                return node
1716

1717
        cascaded_nodes = list(map(UtilNodes.ResultRefNode, args[1:]))
1718 1719 1720 1721 1722 1723 1724 1725 1726 1727 1728 1729 1730 1731 1732 1733 1734 1735 1736 1737 1738 1739

        last_result = args[0]
        for arg_node in cascaded_nodes:
            result_ref = UtilNodes.ResultRefNode(last_result)
            last_result = ExprNodes.CondExprNode(
                arg_node.pos,
                true_val = arg_node,
                false_val = result_ref,
                test = ExprNodes.PrimaryCmpNode(
                    arg_node.pos,
                    operand1 = arg_node,
                    operator = operator,
                    operand2 = result_ref,
                    )
                )
            last_result = UtilNodes.EvalWithTempExprNode(result_ref, last_result)

        for ref_node in cascaded_nodes[::-1]:
            last_result = UtilNodes.EvalWithTempExprNode(ref_node, last_result)

        return last_result

1740 1741
    # builtin type creation

1742
    def _DISABLED_handle_simple_function_tuple(self, node, pos_args):
Stefan Behnel's avatar
Stefan Behnel committed
1743
        if not pos_args:
1744 1745 1746 1747 1748 1749 1750 1751
            return ExprNodes.TupleNode(node.pos, args=[], constant_result=())
        # This is a bit special - for iterables (including genexps),
        # Python actually overallocates and resizes a newly created
        # tuple incrementally while reading items, which we can't
        # easily do without explicit node support. Instead, we read
        # the items into a list and then copy them into a tuple of the
        # final size.  This takes up to twice as much memory, but will
        # have to do until we have real support for genexps.
1752
        result = self._transform_list_set_genexpr(node, pos_args, Builtin.list_type)
1753 1754 1755 1756
        if result is not node:
            return ExprNodes.AsTupleNode(node.pos, arg=result)
        return node

1757 1758 1759 1760 1761
    def _handle_simple_function_frozenset(self, node, pos_args):
        """Replace frozenset([...]) by frozenset((...)) as tuples are more efficient.
        """
        if len(pos_args) != 1:
            return node
1762 1763 1764
        if pos_args[0].is_sequence_constructor and not pos_args[0].args:
            del pos_args[0]
        elif isinstance(pos_args[0], ExprNodes.ListNode):
1765 1766 1767
            pos_args[0] = pos_args[0].as_tuple()
        return node

1768
    def _handle_simple_function_list(self, node, pos_args):
Stefan Behnel's avatar
Stefan Behnel committed
1769
        if not pos_args:
1770
            return ExprNodes.ListNode(node.pos, args=[], constant_result=[])
1771
        return self._transform_list_set_genexpr(node, pos_args, Builtin.list_type)
1772 1773

    def _handle_simple_function_set(self, node, pos_args):
Stefan Behnel's avatar
Stefan Behnel committed
1774
        if not pos_args:
1775
            return ExprNodes.SetNode(node.pos, args=[], constant_result=set())
1776
        return self._transform_list_set_genexpr(node, pos_args, Builtin.set_type)
1777

1778
    def _transform_list_set_genexpr(self, node, pos_args, target_type):
1779 1780 1781 1782 1783 1784 1785 1786 1787
        """Replace set(genexpr) and list(genexpr) by a literal comprehension.
        """
        if len(pos_args) > 1:
            return node
        if not isinstance(pos_args[0], ExprNodes.GeneratorExpressionNode):
            return node
        gen_expr_node = pos_args[0]
        loop_node = gen_expr_node.loop

1788 1789
        yield_expression, yield_stat_node = self._find_single_yield_expression(loop_node)
        if yield_expression is None:
1790 1791 1792
            return node

        append_node = ExprNodes.ComprehensionAppendNode(
1793
            yield_expression.pos,
1794
            expr = yield_expression)
1795

1796
        Visitor.recursively_replace_node(loop_node, yield_stat_node, append_node)
1797

1798
        comp = ExprNodes.ComprehensionNode(
1799 1800 1801 1802 1803
            node.pos,
            has_local_scope = True,
            expr_scope = gen_expr_node.expr_scope,
            loop = loop_node,
            append = append_node,
1804 1805 1806
            type = target_type)
        append_node.target = comp
        return comp
1807 1808 1809 1810 1811 1812 1813 1814 1815 1816 1817 1818 1819

    def _handle_simple_function_dict(self, node, pos_args):
        """Replace dict( (a,b) for ... ) by a literal { a:b for ... }.
        """
        if len(pos_args) == 0:
            return ExprNodes.DictNode(node.pos, key_value_pairs=[], constant_result={})
        if len(pos_args) > 1:
            return node
        if not isinstance(pos_args[0], ExprNodes.GeneratorExpressionNode):
            return node
        gen_expr_node = pos_args[0]
        loop_node = gen_expr_node.loop

1820 1821
        yield_expression, yield_stat_node = self._find_single_yield_expression(loop_node)
        if yield_expression is None:
1822 1823 1824 1825 1826 1827 1828 1829
            return node

        if not isinstance(yield_expression, ExprNodes.TupleNode):
            return node
        if len(yield_expression.args) != 2:
            return node

        append_node = ExprNodes.DictComprehensionAppendNode(
1830
            yield_expression.pos,
1831
            key_expr = yield_expression.args[0],
1832
            value_expr = yield_expression.args[1])
1833

1834
        Visitor.recursively_replace_node(loop_node, yield_stat_node, append_node)
1835 1836 1837 1838 1839 1840 1841

        dictcomp = ExprNodes.ComprehensionNode(
            node.pos,
            has_local_scope = True,
            expr_scope = gen_expr_node.expr_scope,
            loop = loop_node,
            append = append_node,
1842
            type = Builtin.dict_type)
1843 1844 1845
        append_node.target = dictcomp
        return dictcomp

1846 1847 1848 1849 1850 1851 1852 1853 1854 1855 1856 1857
    # specific handlers for general call nodes

    def _handle_general_function_dict(self, node, pos_args, kwargs):
        """Replace dict(a=b,c=d,...) by the underlying keyword dict
        construction which is done anyway.
        """
        if len(pos_args) > 0:
            return node
        if not isinstance(kwargs, ExprNodes.DictNode):
            return node
        return kwargs

1858 1859

class InlineDefNodeCalls(Visitor.NodeRefCleanupMixin, Visitor.EnvTransform):
1860 1861
    visit_Node = Visitor.VisitorTransform.recurse_to_children

1862
    def get_constant_value_node(self, name_node):
1863 1864 1865
        if name_node.cf_state is None:
            return None
        if name_node.cf_state.cf_is_null:
1866 1867 1868 1869 1870 1871
            return None
        entry = self.current_env().lookup(name_node.name)
        if not entry or (not entry.cf_assignments
                         or len(entry.cf_assignments) != 1):
            # not just a single assignment in all closures
            return None
1872
        return entry.cf_assignments[0].rhs
1873

1874 1875 1876 1877 1878 1879 1880
    def visit_SimpleCallNode(self, node):
        self.visitchildren(node)
        if not self.current_directives.get('optimize.inline_defnode_calls'):
            return node
        function_name = node.function
        if not function_name.is_name:
            return node
1881
        function = self.get_constant_value_node(function_name)
1882 1883 1884 1885 1886 1887
        if not isinstance(function, ExprNodes.PyCFunctionNode):
            return node
        inlined = ExprNodes.InlinedDefNodeCallNode(
            node.pos, function_name=function_name,
            function=function, args=node.args)
        if inlined.can_be_inlined():
1888
            return self.replace(node, inlined)
1889 1890
        return node

1891

1892 1893
class OptimizeBuiltinCalls(Visitor.NodeRefCleanupMixin,
                           Visitor.MethodDispatcherTransform):
Stefan Behnel's avatar
Stefan Behnel committed
1894
    """Optimize some common methods calls and instantiation patterns
1895 1896 1897 1898 1899
    for builtin types *after* the type analysis phase.

    Running after type analysis, this transform can only perform
    function replacements that do not alter the function return type
    in a way that was not anticipated by the type analysis.
1900
    """
1901 1902
    ### cleanup to avoid redundant coercions to/from Python types

1903 1904 1905
    def _visit_PyTypeTestNode(self, node):
        # disabled - appears to break assignments in some cases, and
        # also drops a None check, which might still be required
1906 1907 1908 1909 1910 1911 1912 1913
        """Flatten redundant type checks after tree changes.
        """
        old_arg = node.arg
        self.visitchildren(node)
        if old_arg is node.arg or node.arg.type != node.type:
            return node
        return node.arg

1914 1915 1916
    def _visit_TypecastNode(self, node):
        # disabled - the user may have had a reason to put a type
        # cast, even if it looks redundant to Cython
1917 1918 1919 1920 1921 1922 1923 1924
        """
        Drop redundant type casts.
        """
        self.visitchildren(node)
        if node.type == node.operand.type:
            return node.operand
        return node

1925 1926 1927 1928 1929 1930 1931 1932 1933
    def visit_ExprStatNode(self, node):
        """
        Drop useless coercions.
        """
        self.visitchildren(node)
        if isinstance(node.expr, ExprNodes.CoerceToPyTypeNode):
            node.expr = node.expr.arg
        return node

1934 1935 1936 1937 1938
    def visit_CoerceToBooleanNode(self, node):
        """Drop redundant conversion nodes after tree changes.
        """
        self.visitchildren(node)
        arg = node.arg
Stefan Behnel's avatar
Stefan Behnel committed
1939 1940
        if isinstance(arg, ExprNodes.PyTypeTestNode):
            arg = arg.arg
1941 1942
        if isinstance(arg, ExprNodes.CoerceToPyTypeNode):
            if arg.type in (PyrexTypes.py_object_type, Builtin.bool_type):
1943
                return arg.arg.coerce_to_boolean(self.current_env())
1944 1945
        return node

1946 1947 1948 1949 1950 1951 1952 1953 1954 1955 1956 1957 1958 1959
    def visit_CoerceFromPyTypeNode(self, node):
        """Drop redundant conversion nodes after tree changes.

        Also, optimise away calls to Python's builtin int() and
        float() if the result is going to be coerced back into a C
        type anyway.
        """
        self.visitchildren(node)
        arg = node.arg
        if not arg.type.is_pyobject:
            # no Python conversion left at all, just do a C coercion instead
            if node.type == arg.type:
                return arg
            else:
1960
                return arg.coerce_to(node.type, self.current_env())
1961 1962
        if isinstance(arg, ExprNodes.PyTypeTestNode):
            arg = arg.arg
1963 1964
        if arg.is_literal:
            if (node.type.is_int and isinstance(arg, ExprNodes.IntNode) or
1965 1966
                    node.type.is_float and isinstance(arg, ExprNodes.FloatNode) or
                    node.type.is_int and isinstance(arg, ExprNodes.BoolNode)):
1967 1968
                return arg.coerce_to(node.type, self.current_env())
        elif isinstance(arg, ExprNodes.CoerceToPyTypeNode):
1969 1970 1971
            if arg.type is PyrexTypes.py_object_type:
                if node.type.assignable_from(arg.arg.type):
                    # completely redundant C->Py->C coercion
1972
                    return arg.arg.coerce_to(node.type, self.current_env())
Stefan Behnel's avatar
Stefan Behnel committed
1973
        elif isinstance(arg, ExprNodes.SimpleCallNode):
Stefan Behnel's avatar
Stefan Behnel committed
1974 1975
            if node.type.is_int or node.type.is_float:
                return self._optimise_numeric_cast_call(node, arg)
1976 1977 1978 1979 1980 1981
        elif isinstance(arg, ExprNodes.IndexNode) and not arg.is_buffer_access:
            index_node = arg.index
            if isinstance(index_node, ExprNodes.CoerceToPyTypeNode):
                index_node = index_node.arg
            if index_node.type.is_int:
                return self._optimise_int_indexing(node, arg, index_node)
Stefan Behnel's avatar
Stefan Behnel committed
1982 1983
        return node

1984 1985 1986 1987 1988 1989 1990 1991 1992 1993 1994 1995
    PyBytes_GetItemInt_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_char_type, [
            PyrexTypes.CFuncTypeArg("bytes", Builtin.bytes_type, None),
            PyrexTypes.CFuncTypeArg("index", PyrexTypes.c_py_ssize_t_type, None),
            PyrexTypes.CFuncTypeArg("check_bounds", PyrexTypes.c_int_type, None),
            ],
        exception_value = "((char)-1)",
        exception_check = True)

    def _optimise_int_indexing(self, coerce_node, arg, index_node):
        env = self.current_env()
        bound_check_bool = env.directives['boundscheck'] and 1 or 0
1996
        if arg.base.type is Builtin.bytes_type:
1997 1998 1999 2000 2001 2002 2003 2004
            if coerce_node.type in (PyrexTypes.c_char_type, PyrexTypes.c_uchar_type):
                # bytes[index] -> char
                bound_check_node = ExprNodes.IntNode(
                    coerce_node.pos, value=str(bound_check_bool),
                    constant_result=bound_check_bool)
                node = ExprNodes.PythonCapiCallNode(
                    coerce_node.pos, "__Pyx_PyBytes_GetItemInt",
                    self.PyBytes_GetItemInt_func_type,
2005
                    args=[
Stefan Behnel's avatar
Stefan Behnel committed
2006
                        arg.base.as_none_safe_node("'NoneType' object is not subscriptable"),
2007 2008 2009
                        index_node.coerce_to(PyrexTypes.c_py_ssize_t_type, env),
                        bound_check_node,
                        ],
2010 2011 2012
                    is_temp=True,
                    utility_code=UtilityCode.load_cached(
                        'bytes_index', 'StringTools.c'))
2013 2014 2015 2016 2017
                if coerce_node.type is not PyrexTypes.c_char_type:
                    node = node.coerce_to(coerce_node.type, env)
                return node
        return coerce_node

Stefan Behnel's avatar
Stefan Behnel committed
2018
    def _optimise_numeric_cast_call(self, node, arg):
2019 2020 2021 2022 2023 2024 2025 2026 2027 2028 2029 2030 2031 2032 2033 2034 2035 2036 2037
        function = arg.function
        if not isinstance(function, ExprNodes.NameNode) \
               or not function.type.is_builtin_type \
               or not isinstance(arg.arg_tuple, ExprNodes.TupleNode):
            return node
        args = arg.arg_tuple.args
        if len(args) != 1:
            return node
        func_arg = args[0]
        if isinstance(func_arg, ExprNodes.CoerceToPyTypeNode):
            func_arg = func_arg.arg
        elif func_arg.type.is_pyobject:
            # play safe: Python conversion might work on all sorts of things
            return node
        if function.name == 'int':
            if func_arg.type.is_int or node.type.is_int:
                if func_arg.type == node.type:
                    return func_arg
                elif node.type.assignable_from(func_arg.type) or func_arg.type.is_float:
2038 2039
                    return ExprNodes.TypecastNode(
                        node.pos, operand=func_arg, type=node.type)
2040 2041 2042 2043 2044
        elif function.name == 'float':
            if func_arg.type.is_float or node.type.is_float:
                if func_arg.type == node.type:
                    return func_arg
                elif node.type.assignable_from(func_arg.type) or func_arg.type.is_float:
2045 2046
                    return ExprNodes.TypecastNode(
                        node.pos, operand=func_arg, type=node.type)
2047 2048
        return node

2049 2050 2051 2052 2053 2054 2055 2056 2057 2058 2059 2060 2061 2062 2063 2064
    def _error_wrong_arg_count(self, function_name, node, args, expected=None):
        if not expected: # None or 0
            arg_str = ''
        elif isinstance(expected, basestring) or expected > 1:
            arg_str = '...'
        elif expected == 1:
            arg_str = 'x'
        else:
            arg_str = ''
        if expected is not None:
            expected_str = 'expected %s, ' % expected
        else:
            expected_str = ''
        error(node.pos, "%s(%s) called with wrong number of args, %sfound %d" % (
            function_name, arg_str, expected_str, len(args)))

2065 2066 2067 2068 2069 2070 2071 2072 2073 2074 2075 2076 2077 2078 2079 2080 2081 2082 2083 2084 2085 2086 2087 2088 2089 2090 2091 2092 2093 2094 2095 2096 2097 2098 2099 2100 2101 2102 2103 2104 2105 2106 2107 2108 2109 2110
    ### generic fallbacks

    def _handle_function(self, node, function_name, function, arg_list, kwargs):
        return node

    def _handle_method(self, node, type_name, attr_name, function,
                       arg_list, is_unbound_method, kwargs):
        """
        Try to inject C-API calls for unbound method calls to builtin types.
        While the method declarations in Builtin.py already handle this, we
        can additionally resolve bound and unbound methods here that were
        assigned to variables ahead of time.
        """
        if kwargs:
            return node
        if not function or not function.is_attribute or not function.obj.is_name:
            # cannot track unbound method calls over more than one indirection as
            # the names might have been reassigned in the meantime
            return node
        type_entry = self.current_env().lookup(type_name)
        if not type_entry:
            return node
        method = ExprNodes.AttributeNode(
            node.function.pos,
            obj=ExprNodes.NameNode(
                function.pos,
                name=type_name,
                entry=type_entry,
                type=type_entry.type),
            attribute=attr_name,
            is_called=True).analyse_as_unbound_cmethod_node(self.current_env())
        if method is None:
            return node
        args = node.args
        if args is None and node.arg_tuple:
            args = node.arg_tuple.args
        call_node = ExprNodes.SimpleCallNode(
            node.pos,
            function=method,
            args=args)
        if not is_unbound_method:
            call_node.self = function.obj
        call_node.analyse_c_function_call(self.current_env())
        call_node.analysed = True
        return call_node.coerce_to(node.type, self.current_env())

2111 2112
    ### builtin types

2113 2114 2115 2116 2117
    PyDict_Copy_func_type = PyrexTypes.CFuncType(
        Builtin.dict_type, [
            PyrexTypes.CFuncTypeArg("dict", Builtin.dict_type, None)
            ])

2118
    def _handle_simple_function_dict(self, node, function, pos_args):
2119
        """Replace dict(some_dict) by PyDict_Copy(some_dict).
2120
        """
2121
        if len(pos_args) != 1:
2122
            return node
2123
        arg = pos_args[0]
2124
        if arg.type is Builtin.dict_type:
2125
            arg = arg.as_none_safe_node("'NoneType' is not iterable")
2126 2127
            return ExprNodes.PythonCapiCallNode(
                node.pos, "PyDict_Copy", self.PyDict_Copy_func_type,
2128
                args = [arg],
2129 2130 2131
                is_temp = node.is_temp
                )
        return node
2132

2133 2134 2135 2136 2137 2138 2139 2140 2141 2142 2143 2144 2145 2146
    PySequence_List_func_type = PyrexTypes.CFuncType(
        Builtin.list_type,
        [PyrexTypes.CFuncTypeArg("it", PyrexTypes.py_object_type, None)])

    def _handle_simple_function_list(self, node, function, pos_args):
        """Turn list(ob) into PySequence_List(ob).
        """
        if len(pos_args) != 1:
            return node
        arg = pos_args[0]
        return ExprNodes.PythonCapiCallNode(
            node.pos, "PySequence_List", self.PySequence_List_func_type,
            args=pos_args, is_temp=node.is_temp)

2147 2148 2149 2150 2151
    PyList_AsTuple_func_type = PyrexTypes.CFuncType(
        Builtin.tuple_type, [
            PyrexTypes.CFuncTypeArg("list", Builtin.list_type, None)
            ])

2152 2153 2154 2155
    PySequence_Tuple_func_type = PyrexTypes.CFuncType(
        Builtin.tuple_type,
        [PyrexTypes.CFuncTypeArg("it", PyrexTypes.py_object_type, None)])

2156
    def _handle_simple_function_tuple(self, node, function, pos_args):
2157
        """Replace tuple([...]) by PyList_AsTuple or PySequence_Tuple.
2158
        """
2159
        if len(pos_args) != 1:
2160
            return node
2161 2162 2163
        arg = pos_args[0]
        if arg.type is Builtin.tuple_type and not arg.may_be_none():
            return arg
2164 2165 2166
        if arg.type is Builtin.list_type:
            pos_args[0] = arg.as_none_safe_node(
                "'NoneType' object is not iterable")
2167

2168 2169 2170 2171 2172 2173 2174
            return ExprNodes.PythonCapiCallNode(
                node.pos, "PyList_AsTuple", self.PyList_AsTuple_func_type,
                args=pos_args, is_temp=node.is_temp)
        else:
            return ExprNodes.PythonCapiCallNode(
                node.pos, "PySequence_Tuple", self.PySequence_Tuple_func_type,
                args=pos_args, is_temp=node.is_temp)
2175

2176 2177 2178 2179 2180
    PySet_New_func_type = PyrexTypes.CFuncType(
        Builtin.set_type, [
            PyrexTypes.CFuncTypeArg("it", PyrexTypes.py_object_type, None)
        ])

2181
    def _handle_simple_function_set(self, node, function, pos_args):
2182 2183 2184 2185 2186 2187 2188 2189 2190 2191 2192 2193 2194 2195 2196 2197
        if len(pos_args) != 1:
            return node
        if pos_args[0].is_sequence_constructor:
            # We can optimise set([x,y,z]) safely into a set literal,
            # but only if we create all items before adding them -
            # adding an item may raise an exception if it is not
            # hashable, but creating the later items may have
            # side-effects.
            args = []
            temps = []
            for arg in pos_args[0].args:
                if not arg.is_simple():
                    arg = UtilNodes.LetRefNode(arg)
                    temps.append(arg)
                args.append(arg)
            result = ExprNodes.SetNode(node.pos, is_temp=1, args=args)
2198
            self.replace(node, result)
2199 2200 2201 2202 2203
            for temp in temps[::-1]:
                result = UtilNodes.EvalWithTempExprNode(temp, result)
            return result
        else:
            # PySet_New(it) is better than a generic Python call to set(it)
2204
            return self.replace(node, ExprNodes.PythonCapiCallNode(
2205 2206 2207 2208
                node.pos, "PySet_New",
                self.PySet_New_func_type,
                args=pos_args,
                is_temp=node.is_temp,
2209
                py_name="set"))
2210 2211 2212 2213 2214 2215 2216

    PyFrozenSet_New_func_type = PyrexTypes.CFuncType(
        Builtin.frozenset_type, [
            PyrexTypes.CFuncTypeArg("it", PyrexTypes.py_object_type, None)
        ])

    def _handle_simple_function_frozenset(self, node, function, pos_args):
2217 2218 2219
        if not pos_args:
            pos_args = [ExprNodes.NullNode(node.pos)]
        elif len(pos_args) > 1:
2220
            return node
2221 2222
        elif pos_args[0].type is Builtin.frozenset_type and not pos_args[0].may_be_none():
            return pos_args[0]
2223 2224
        # PyFrozenSet_New(it) is better than a generic Python call to frozenset(it)
        return ExprNodes.PythonCapiCallNode(
2225
            node.pos, "__Pyx_PyFrozenSet_New",
2226 2227 2228
            self.PyFrozenSet_New_func_type,
            args=pos_args,
            is_temp=node.is_temp,
2229
            utility_code=UtilityCode.load_cached('pyfrozenset_new', 'Builtins.c'),
2230 2231 2232 2233 2234 2235 2236 2237
            py_name="frozenset")

    PyObject_AsDouble_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_double_type, [
            PyrexTypes.CFuncTypeArg("obj", PyrexTypes.py_object_type, None),
            ],
        exception_value = "((double)-1)",
        exception_check = True)
2238

2239
    def _handle_simple_function_float(self, node, function, pos_args):
Stefan Behnel's avatar
Stefan Behnel committed
2240 2241 2242
        """Transform float() into either a C type cast or a faster C
        function call.
        """
2243 2244
        # Note: this requires the float() function to be typed as
        # returning a C 'double'
2245
        if len(pos_args) == 0:
Stefan Behnel's avatar
typo  
Stefan Behnel committed
2246
            return ExprNodes.FloatNode(
2247 2248 2249 2250
                node, value="0.0", constant_result=0.0
                ).coerce_to(Builtin.float_type, self.current_env())
        elif len(pos_args) != 1:
            self._error_wrong_arg_count('float', node, pos_args, '0 or 1')
2251 2252 2253 2254 2255 2256 2257
            return node
        func_arg = pos_args[0]
        if isinstance(func_arg, ExprNodes.CoerceToPyTypeNode):
            func_arg = func_arg.arg
        if func_arg.type is PyrexTypes.c_double_type:
            return func_arg
        elif node.type.assignable_from(func_arg.type) or func_arg.type.is_numeric:
2258 2259
            return ExprNodes.TypecastNode(
                node.pos, operand=func_arg, type=node.type)
2260 2261 2262 2263 2264
        return ExprNodes.PythonCapiCallNode(
            node.pos, "__Pyx_PyObject_AsDouble",
            self.PyObject_AsDouble_func_type,
            args = pos_args,
            is_temp = node.is_temp,
2265
            utility_code = load_c_utility('pyobject_as_double'),
2266 2267
            py_name = "float")

2268 2269 2270 2271 2272 2273 2274 2275 2276
    PyNumber_Int_func_type = PyrexTypes.CFuncType(
        PyrexTypes.py_object_type, [
            PyrexTypes.CFuncTypeArg("o", PyrexTypes.py_object_type, None)
            ])

    def _handle_simple_function_int(self, node, function, pos_args):
        """Transform int() into a faster C function call.
        """
        if len(pos_args) == 0:
2277
            return ExprNodes.IntNode(node.pos, value="0", constant_result=0,
Stefan Behnel's avatar
Stefan Behnel committed
2278
                                     type=PyrexTypes.py_object_type)
2279 2280 2281 2282 2283 2284 2285 2286 2287 2288 2289
        elif len(pos_args) != 1:
            return node  # int(x, base)
        func_arg = pos_args[0]
        if isinstance(func_arg, ExprNodes.CoerceToPyTypeNode):
            return node  # handled in visit_CoerceFromPyTypeNode()
        if func_arg.type.is_pyobject and node.type.is_pyobject:
            return ExprNodes.PythonCapiCallNode(
                node.pos, "PyNumber_Int", self.PyNumber_Int_func_type,
                args=pos_args, is_temp=True)
        return node

2290
    def _handle_simple_function_bool(self, node, function, pos_args):
2291 2292
        """Transform bool(x) into a type coercion to a boolean.
        """
2293 2294 2295 2296 2297 2298
        if len(pos_args) == 0:
            return ExprNodes.BoolNode(
                node.pos, value=False, constant_result=False
                ).coerce_to(Builtin.bool_type, self.current_env())
        elif len(pos_args) != 1:
            self._error_wrong_arg_count('bool', node, pos_args, '0 or 1')
2299
            return node
Craig Citro's avatar
Craig Citro committed
2300
        else:
2301 2302 2303 2304 2305 2306
            # => !!<bint>(x)  to make sure it's exactly 0 or 1
            operand = pos_args[0].coerce_to_boolean(self.current_env())
            operand = ExprNodes.NotNode(node.pos, operand = operand)
            operand = ExprNodes.NotNode(node.pos, operand = operand)
            # coerce back to Python object as that's the result we are expecting
            return operand.coerce_to_pyobject(self.current_env())
2307

2308 2309
    ### builtin functions

2310 2311 2312 2313 2314
    Pyx_strlen_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_size_t_type, [
            PyrexTypes.CFuncTypeArg("bytes", PyrexTypes.c_char_ptr_type, None)
            ])

2315 2316 2317 2318 2319
    Pyx_Py_UNICODE_strlen_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_size_t_type, [
            PyrexTypes.CFuncTypeArg("unicode", PyrexTypes.c_py_unicode_ptr_type, None)
            ])

2320 2321 2322
    PyObject_Size_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_py_ssize_t_type, [
            PyrexTypes.CFuncTypeArg("obj", PyrexTypes.py_object_type, None)
2323 2324
            ],
        exception_value="-1")
2325 2326

    _map_to_capi_len_function = {
2327
        Builtin.unicode_type   : "__Pyx_PyUnicode_GET_LENGTH",
2328
        Builtin.bytes_type     : "PyBytes_GET_SIZE",
2329 2330 2331 2332
        Builtin.list_type      : "PyList_GET_SIZE",
        Builtin.tuple_type     : "PyTuple_GET_SIZE",
        Builtin.dict_type      : "PyDict_Size",
        Builtin.set_type       : "PySet_Size",
2333
        Builtin.frozenset_type : "__Pyx_PyFrozenSet_Size",
2334 2335
        }.get

2336 2337
    _ext_types_with_pysize = set(["cpython.array.array"])

2338
    def _handle_simple_function_len(self, node, function, pos_args):
2339 2340
        """Replace len(char*) by the equivalent call to strlen(),
        len(Py_UNICODE) by the equivalent Py_UNICODE_strlen() and
Stefan Behnel's avatar
Stefan Behnel committed
2341
        len(known_builtin_type) by an equivalent C-API call.
Stefan Behnel's avatar
Stefan Behnel committed
2342
        """
2343 2344 2345 2346 2347 2348
        if len(pos_args) != 1:
            self._error_wrong_arg_count('len', node, pos_args, 1)
            return node
        arg = pos_args[0]
        if isinstance(arg, ExprNodes.CoerceToPyTypeNode):
            arg = arg.arg
2349 2350 2351 2352 2353
        if arg.type.is_string:
            new_node = ExprNodes.PythonCapiCallNode(
                node.pos, "strlen", self.Pyx_strlen_func_type,
                args = [arg],
                is_temp = node.is_temp,
2354
                utility_code = UtilityCode.load_cached("IncludeStringH", "StringTools.c"))
2355
        elif arg.type.is_pyunicode_ptr:
2356 2357 2358
            new_node = ExprNodes.PythonCapiCallNode(
                node.pos, "__Pyx_Py_UNICODE_strlen", self.Pyx_Py_UNICODE_strlen_func_type,
                args = [arg],
2359
                is_temp = node.is_temp)
2360 2361 2362
        elif arg.type.is_pyobject:
            cfunc_name = self._map_to_capi_len_function(arg.type)
            if cfunc_name is None:
2363 2364 2365 2366 2367 2368
                arg_type = arg.type
                if ((arg_type.is_extension_type or arg_type.is_builtin_type)
                    and arg_type.entry.qualified_name in self._ext_types_with_pysize):
                    cfunc_name = 'Py_SIZE'
                else:
                    return node
2369 2370
            arg = arg.as_none_safe_node(
                "object of type 'NoneType' has no len()")
2371 2372 2373 2374
            new_node = ExprNodes.PythonCapiCallNode(
                node.pos, cfunc_name, self.PyObject_Size_func_type,
                args = [arg],
                is_temp = node.is_temp)
Stefan Behnel's avatar
Stefan Behnel committed
2375
        elif arg.type.is_unicode_char:
2376 2377
            return ExprNodes.IntNode(node.pos, value='1', constant_result=1,
                                     type=node.type)
2378
        else:
Stefan Behnel's avatar
Stefan Behnel committed
2379
            return node
2380
        if node.type not in (PyrexTypes.c_size_t_type, PyrexTypes.c_py_ssize_t_type):
2381
            new_node = new_node.coerce_to(node.type, self.current_env())
2382
        return new_node
2383

2384 2385 2386 2387 2388
    Pyx_Type_func_type = PyrexTypes.CFuncType(
        Builtin.type_type, [
            PyrexTypes.CFuncTypeArg("object", PyrexTypes.py_object_type, None)
            ])

2389
    def _handle_simple_function_type(self, node, function, pos_args):
Stefan Behnel's avatar
Stefan Behnel committed
2390 2391
        """Replace type(o) by a macro call to Py_TYPE(o).
        """
2392
        if len(pos_args) != 1:
2393 2394
            return node
        node = ExprNodes.PythonCapiCallNode(
2395 2396 2397 2398
            node.pos, "Py_TYPE", self.Pyx_Type_func_type,
            args = pos_args,
            is_temp = False)
        return ExprNodes.CastNode(node, PyrexTypes.py_object_type)
2399

2400 2401 2402 2403 2404
    Py_type_check_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_bint_type, [
            PyrexTypes.CFuncTypeArg("arg", PyrexTypes.py_object_type, None)
            ])

2405
    def _handle_simple_function_isinstance(self, node, function, pos_args):
2406 2407 2408 2409 2410 2411
        """Replace isinstance() checks against builtin types by the
        corresponding C-API call.
        """
        if len(pos_args) != 2:
            return node
        arg, types = pos_args
2412
        temps = []
2413 2414
        if isinstance(types, ExprNodes.TupleNode):
            types = types.args
2415 2416
            if len(types) == 1 and not types[0].type is Builtin.type_type:
                return node  # nothing to improve here
2417
            if arg.is_attribute or not arg.is_simple():
2418 2419
                arg = UtilNodes.ResultRefNode(arg)
                temps.append(arg)
2420 2421 2422 2423 2424 2425 2426 2427 2428
        elif types.type is Builtin.type_type:
            types = [types]
        else:
            return node

        tests = []
        test_nodes = []
        env = self.current_env()
        for test_type_node in types:
Robert Bradshaw's avatar
Robert Bradshaw committed
2429
            builtin_type = None
Stefan Behnel's avatar
Stefan Behnel committed
2430
            if test_type_node.is_name:
Robert Bradshaw's avatar
Robert Bradshaw committed
2431 2432 2433 2434
                if test_type_node.entry:
                    entry = env.lookup(test_type_node.entry.name)
                    if entry and entry.type and entry.type.is_builtin_type:
                        builtin_type = entry.type
2435 2436 2437 2438 2439 2440
            if builtin_type is Builtin.type_type:
                # all types have type "type", but there's only one 'type'
                if entry.name != 'type' or not (
                        entry.scope and entry.scope.is_builtin_scope):
                    builtin_type = None
            if builtin_type is not None:
Robert Bradshaw's avatar
Robert Bradshaw committed
2441
                type_check_function = entry.type.type_check_function(exact=False)
2442 2443 2444
                if type_check_function in tests:
                    continue
                tests.append(type_check_function)
Robert Bradshaw's avatar
Robert Bradshaw committed
2445 2446 2447 2448 2449
                type_check_args = [arg]
            elif test_type_node.type is Builtin.type_type:
                type_check_function = '__Pyx_TypeCheck'
                type_check_args = [arg, test_type_node]
            else:
2450 2451 2452 2453 2454
                if not test_type_node.is_literal:
                    test_type_node = UtilNodes.ResultRefNode(test_type_node)
                    temps.append(test_type_node)
                type_check_function = 'PyObject_IsInstance'
                type_check_args = [arg, test_type_node]
2455 2456 2457
            test_nodes.append(
                ExprNodes.PythonCapiCallNode(
                    test_type_node.pos, type_check_function, self.Py_type_check_func_type,
2458 2459 2460
                    args=type_check_args,
                    is_temp=True,
                ))
2461

2462
        def join_with_or(a, b, make_binop_node=ExprNodes.binop_node):
2463 2464
            or_node = make_binop_node(node.pos, 'or', a, b)
            or_node.type = PyrexTypes.c_bint_type
2465
            or_node.wrap_operands(env)
2466 2467 2468
            return or_node

        test_node = reduce(join_with_or, test_nodes).coerce_to(node.type, env)
2469
        for temp in temps[::-1]:
2470 2471 2472
            test_node = UtilNodes.EvalWithTempExprNode(temp, test_node)
        return test_node

2473
    def _handle_simple_function_ord(self, node, function, pos_args):
2474
        """Unpack ord(Py_UNICODE) and ord('X').
2475 2476 2477 2478 2479
        """
        if len(pos_args) != 1:
            return node
        arg = pos_args[0]
        if isinstance(arg, ExprNodes.CoerceToPyTypeNode):
Stefan Behnel's avatar
Stefan Behnel committed
2480
            if arg.arg.type.is_unicode_char:
2481
                return ExprNodes.TypecastNode(
2482
                    arg.pos, operand=arg.arg, type=PyrexTypes.c_long_type
2483 2484 2485 2486
                    ).coerce_to(node.type, self.current_env())
        elif isinstance(arg, ExprNodes.UnicodeNode):
            if len(arg.value) == 1:
                return ExprNodes.IntNode(
2487
                    arg.pos, type=PyrexTypes.c_int_type,
2488 2489 2490 2491 2492
                    value=str(ord(arg.value)),
                    constant_result=ord(arg.value)
                    ).coerce_to(node.type, self.current_env())
        elif isinstance(arg, ExprNodes.StringNode):
            if arg.unicode_value and len(arg.unicode_value) == 1 \
2493
                    and ord(arg.unicode_value) <= 255:  # Py2/3 portability
2494
                return ExprNodes.IntNode(
2495
                    arg.pos, type=PyrexTypes.c_int_type,
2496 2497 2498
                    value=str(ord(arg.unicode_value)),
                    constant_result=ord(arg.unicode_value)
                    ).coerce_to(node.type, self.current_env())
2499 2500
        return node

2501 2502
    ### special methods

2503 2504
    Pyx_tp_new_func_type = PyrexTypes.CFuncType(
        PyrexTypes.py_object_type, [
2505
            PyrexTypes.CFuncTypeArg("type",   PyrexTypes.py_object_type, None),
2506
            PyrexTypes.CFuncTypeArg("args",   Builtin.tuple_type, None),
2507 2508
            ])

2509 2510
    Pyx_tp_new_kwargs_func_type = PyrexTypes.CFuncType(
        PyrexTypes.py_object_type, [
2511
            PyrexTypes.CFuncTypeArg("type",   PyrexTypes.py_object_type, None),
2512 2513 2514 2515
            PyrexTypes.CFuncTypeArg("args",   Builtin.tuple_type, None),
            PyrexTypes.CFuncTypeArg("kwargs", Builtin.dict_type, None),
        ])

2516 2517
    def _handle_any_slot__new__(self, node, function, args,
                                is_unbound_method, kwargs=None):
2518
        """Replace 'exttype.__new__(exttype, ...)' by a call to exttype->tp_new()
2519
        """
2520
        obj = function.obj
2521
        if not is_unbound_method or len(args) < 1:
2522 2523 2524 2525 2526 2527 2528 2529 2530 2531 2532 2533 2534 2535
            return node
        type_arg = args[0]
        if not obj.is_name or not type_arg.is_name:
            # play safe
            return node
        if obj.type != Builtin.type_type or type_arg.type != Builtin.type_type:
            # not a known type, play safe
            return node
        if not type_arg.type_entry or not obj.type_entry:
            if obj.name != type_arg.name:
                return node
            # otherwise, we know it's a type and we know it's the same
            # type for both - that should do
        elif type_arg.type_entry != obj.type_entry:
2536
            # different types - may or may not lead to an error at runtime
2537 2538
            return node

2539 2540 2541
        args_tuple = ExprNodes.TupleNode(node.pos, args=args[1:])
        args_tuple = args_tuple.analyse_types(
            self.current_env(), skip_children=True)
2542

2543 2544
        if type_arg.type_entry:
            ext_type = type_arg.type_entry.type
2545 2546 2547
            if (ext_type.is_extension_type and ext_type.typeobj_cname and
                    ext_type.scope.global_scope() == self.current_env().global_scope()):
                # known type in current module
2548 2549 2550 2551 2552 2553 2554 2555 2556 2557 2558 2559 2560 2561 2562 2563 2564 2565 2566 2567 2568
                tp_slot = TypeSlots.ConstructorSlot("tp_new", '__new__')
                slot_func_cname = TypeSlots.get_slot_function(ext_type.scope, tp_slot)
                if slot_func_cname:
                    cython_scope = self.context.cython_scope
                    PyTypeObjectPtr = PyrexTypes.CPtrType(
                        cython_scope.lookup('PyTypeObject').type)
                    pyx_tp_new_kwargs_func_type = PyrexTypes.CFuncType(
                        PyrexTypes.py_object_type, [
                            PyrexTypes.CFuncTypeArg("type",   PyTypeObjectPtr, None),
                            PyrexTypes.CFuncTypeArg("args",   PyrexTypes.py_object_type, None),
                            PyrexTypes.CFuncTypeArg("kwargs", PyrexTypes.py_object_type, None),
                            ])

                    type_arg = ExprNodes.CastNode(type_arg, PyTypeObjectPtr)
                    if not kwargs:
                        kwargs = ExprNodes.NullNode(node.pos, type=PyrexTypes.py_object_type)  # hack?
                    return ExprNodes.PythonCapiCallNode(
                        node.pos, slot_func_cname,
                        pyx_tp_new_kwargs_func_type,
                        args=[type_arg, args_tuple, kwargs],
                        is_temp=True)
2569
        else:
2570
            # arbitrary variable, needs a None check for safety
2571
            type_arg = type_arg.as_none_safe_node(
2572 2573
                "object.__new__(X): X is not a type object (NoneType)")

2574 2575 2576 2577 2578 2579 2580 2581 2582 2583 2584 2585 2586 2587
        utility_code = UtilityCode.load_cached('tp_new', 'ObjectHandling.c')
        if kwargs:
            return ExprNodes.PythonCapiCallNode(
                node.pos, "__Pyx_tp_new_kwargs", self.Pyx_tp_new_kwargs_func_type,
                args=[type_arg, args_tuple, kwargs],
                utility_code=utility_code,
                is_temp=node.is_temp
                )
        else:
            return ExprNodes.PythonCapiCallNode(
                node.pos, "__Pyx_tp_new", self.Pyx_tp_new_func_type,
                args=[type_arg, args_tuple],
                utility_code=utility_code,
                is_temp=node.is_temp
2588 2589
            )

2590 2591 2592
    ### methods of builtin types

    PyObject_Append_func_type = PyrexTypes.CFuncType(
2593
        PyrexTypes.c_returncode_type, [
2594 2595
            PyrexTypes.CFuncTypeArg("list", PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("item", PyrexTypes.py_object_type, None),
2596 2597
            ],
        exception_value="-1")
2598

2599
    def _handle_simple_method_object_append(self, node, function, args, is_unbound_method):
Stefan Behnel's avatar
Stefan Behnel committed
2600 2601 2602
        """Optimistic optimisation as X.append() is almost always
        referring to a list.
        """
2603
        if len(args) != 2 or node.result_is_used:
2604 2605
            return node

2606 2607
        return ExprNodes.PythonCapiCallNode(
            node.pos, "__Pyx_PyObject_Append", self.PyObject_Append_func_type,
2608 2609 2610 2611 2612 2613
            args=args,
            may_return_none=False,
            is_temp=node.is_temp,
            result_is_used=False,
            utility_code=load_c_utility('append')
        )
2614

2615 2616 2617 2618 2619 2620 2621 2622 2623 2624 2625 2626 2627 2628 2629 2630 2631 2632 2633 2634 2635
    PyByteArray_Append_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_returncode_type, [
            PyrexTypes.CFuncTypeArg("bytearray", PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("value", PyrexTypes.c_int_type, None),
            ],
        exception_value="-1")

    PyByteArray_AppendObject_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_returncode_type, [
            PyrexTypes.CFuncTypeArg("bytearray", PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("value", PyrexTypes.py_object_type, None),
            ],
        exception_value="-1")

    def _handle_simple_method_bytearray_append(self, node, function, args, is_unbound_method):
        if len(args) != 2:
            return node
        func_name = "__Pyx_PyByteArray_Append"
        func_type = self.PyByteArray_Append_func_type

        value = unwrap_coerced_node(args[1])
2636
        if value.type.is_int or isinstance(value, ExprNodes.IntNode):
2637 2638 2639 2640 2641 2642 2643 2644 2645 2646 2647 2648 2649 2650 2651 2652 2653 2654 2655 2656 2657 2658 2659 2660 2661
            value = value.coerce_to(PyrexTypes.c_int_type, self.current_env())
            utility_code = UtilityCode.load_cached("ByteArrayAppend", "StringTools.c")
        elif value.is_string_literal:
            if not value.can_coerce_to_char_literal():
                return node
            value = value.coerce_to(PyrexTypes.c_char_type, self.current_env())
            utility_code = UtilityCode.load_cached("ByteArrayAppend", "StringTools.c")
        elif value.type.is_pyobject:
            func_name = "__Pyx_PyByteArray_AppendObject"
            func_type = self.PyByteArray_AppendObject_func_type
            utility_code = UtilityCode.load_cached("ByteArrayAppendObject", "StringTools.c")
        else:
            return node

        new_node = ExprNodes.PythonCapiCallNode(
            node.pos, func_name, func_type,
            args=[args[0], value],
            may_return_none=False,
            is_temp=node.is_temp,
            utility_code=utility_code,
        )
        if node.result_is_used:
            new_node = new_node.coerce_to(node.type, self.current_env())
        return new_node

Robert Bradshaw's avatar
Robert Bradshaw committed
2662 2663 2664 2665 2666 2667 2668 2669
    PyObject_Pop_func_type = PyrexTypes.CFuncType(
        PyrexTypes.py_object_type, [
            PyrexTypes.CFuncTypeArg("list", PyrexTypes.py_object_type, None),
            ])

    PyObject_PopIndex_func_type = PyrexTypes.CFuncType(
        PyrexTypes.py_object_type, [
            PyrexTypes.CFuncTypeArg("list", PyrexTypes.py_object_type, None),
2670
            PyrexTypes.CFuncTypeArg("index", PyrexTypes.c_py_ssize_t_type, None),
2671 2672 2673
            PyrexTypes.CFuncTypeArg("is_signed", PyrexTypes.c_int_type, None),
        ],
        has_varargs=True)  # to fake the additional macro args that lack a proper C type
Robert Bradshaw's avatar
Robert Bradshaw committed
2674

2675 2676 2677 2678 2679
    def _handle_simple_method_list_pop(self, node, function, args, is_unbound_method):
        return self._handle_simple_method_object_pop(
            node, function, args, is_unbound_method, is_list=True)

    def _handle_simple_method_object_pop(self, node, function, args, is_unbound_method, is_list=False):
Stefan Behnel's avatar
Stefan Behnel committed
2680 2681 2682
        """Optimistic optimisation as X.pop([n]) is almost always
        referring to a list.
        """
2683 2684
        if not args:
            return node
2685
        obj = args[0]
2686 2687
        if is_list:
            type_name = 'List'
2688
            obj = obj.as_none_safe_node(
2689 2690 2691 2692 2693
                "'NoneType' object has no attribute '%s'",
                error="PyExc_AttributeError",
                format_args=['pop'])
        else:
            type_name = 'Object'
Robert Bradshaw's avatar
Robert Bradshaw committed
2694 2695
        if len(args) == 1:
            return ExprNodes.PythonCapiCallNode(
2696 2697
                node.pos, "__Pyx_Py%s_Pop" % type_name,
                self.PyObject_Pop_func_type,
2698
                args=[obj],
2699 2700 2701 2702
                may_return_none=True,
                is_temp=node.is_temp,
                utility_code=load_c_utility('pop'),
            )
Robert Bradshaw's avatar
Robert Bradshaw committed
2703
        elif len(args) == 2:
2704
            index = unwrap_coerced_node(args[1])
2705 2706 2707 2708 2709 2710 2711 2712 2713 2714 2715 2716 2717 2718 2719 2720 2721 2722 2723
            orig_index_type = index.type
            if not index.type.is_int:
                if is_list or isinstance(index, ExprNodes.IntNode):
                    index = index.coerce_to(PyrexTypes.c_py_ssize_t_type, self.current_env())
                else:
                    return node
            elif not PyrexTypes.numeric_type_fits(index.type, PyrexTypes.c_py_ssize_t_type):
                return node
            # real type might still be larger at runtime
            if not orig_index_type.is_int:
                orig_index_type = index.type
            if not orig_index_type.create_to_py_utility_code(self.current_env()):
                return node
            convert_func = orig_index_type.to_py_function
            conversion_type = PyrexTypes.CFuncType(
                PyrexTypes.py_object_type, [PyrexTypes.CFuncTypeArg("intval", orig_index_type, None)])
            return ExprNodes.PythonCapiCallNode(
                node.pos, "__Pyx_Py%s_PopIndex" % type_name,
                self.PyObject_PopIndex_func_type,
2724 2725 2726 2727 2728
                args=[obj, index,
                      ExprNodes.IntNode(index.pos, value=str(orig_index_type.signed and 1 or 0),
                                        constant_result=orig_index_type.signed and 1 or 0,
                                        type=PyrexTypes.c_int_type),
                      ExprNodes.RawCNameExprNode(index.pos, PyrexTypes.c_void_type,
2729
                                                 orig_index_type.empty_declaration_code()),
2730
                      ExprNodes.RawCNameExprNode(index.pos, conversion_type, convert_func)],
2731 2732 2733 2734
                may_return_none=True,
                is_temp=node.is_temp,
                utility_code=load_c_utility("pop_index"),
            )
2735

Robert Bradshaw's avatar
Robert Bradshaw committed
2736 2737
        return node

2738
    single_param_func_type = PyrexTypes.CFuncType(
2739
        PyrexTypes.c_returncode_type, [
2740 2741 2742
            PyrexTypes.CFuncTypeArg("obj", PyrexTypes.py_object_type, None),
            ],
        exception_value = "-1")
2743

2744
    def _handle_simple_method_list_sort(self, node, function, args, is_unbound_method):
Stefan Behnel's avatar
Stefan Behnel committed
2745 2746
        """Call PyList_Sort() instead of the 0-argument l.sort().
        """
2747
        if len(args) != 1:
2748
            return node
2749
        return self._substitute_method_call(
2750
            node, function, "PyList_Sort", self.single_param_func_type,
2751
            'sort', is_unbound_method, args).coerce_to(node.type, self.current_env)
2752

2753 2754 2755 2756 2757
    Pyx_PyDict_GetItem_func_type = PyrexTypes.CFuncType(
        PyrexTypes.py_object_type, [
            PyrexTypes.CFuncTypeArg("dict", PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("key", PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("default", PyrexTypes.py_object_type, None),
2758
            ])
2759

2760
    def _handle_simple_method_dict_get(self, node, function, args, is_unbound_method):
Stefan Behnel's avatar
Stefan Behnel committed
2761 2762
        """Replace dict.get() by a call to PyDict_GetItem().
        """
2763 2764 2765 2766 2767 2768 2769
        if len(args) == 2:
            args.append(ExprNodes.NoneNode(node.pos))
        elif len(args) != 3:
            self._error_wrong_arg_count('dict.get', node, args, "2 or 3")
            return node

        return self._substitute_method_call(
2770 2771
            node, function,
            "__Pyx_PyDict_GetItemDefault", self.Pyx_PyDict_GetItem_func_type,
2772
            'get', is_unbound_method, args,
Stefan Behnel's avatar
Stefan Behnel committed
2773
            may_return_none = True,
2774
            utility_code = load_c_utility("dict_getitem_default"))
2775

2776 2777 2778 2779 2780
    Pyx_PyDict_SetDefault_func_type = PyrexTypes.CFuncType(
        PyrexTypes.py_object_type, [
            PyrexTypes.CFuncTypeArg("dict", PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("key", PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("default", PyrexTypes.py_object_type, None),
2781
            PyrexTypes.CFuncTypeArg("is_safe_type", PyrexTypes.c_int_type, None),
2782 2783
            ])

2784
    def _handle_simple_method_dict_setdefault(self, node, function, args, is_unbound_method):
2785 2786 2787 2788 2789 2790 2791
        """Replace dict.setdefault() by calls to PyDict_GetItem() and PyDict_SetItem().
        """
        if len(args) == 2:
            args.append(ExprNodes.NoneNode(node.pos))
        elif len(args) != 3:
            self._error_wrong_arg_count('dict.setdefault', node, args, "2 or 3")
            return node
2792 2793 2794 2795 2796 2797 2798 2799 2800
        key_type = args[1].type
        if key_type.is_builtin_type:
            is_safe_type = int(key_type.name in
                               'str bytes unicode float int long bool')
        elif key_type is PyrexTypes.py_object_type:
            is_safe_type = -1  # don't know
        else:
            is_safe_type = 0   # definitely not
        args.append(ExprNodes.IntNode(
2801
            node.pos, value=str(is_safe_type), constant_result=is_safe_type))
2802 2803

        return self._substitute_method_call(
2804 2805
            node, function,
            "__Pyx_PyDict_SetDefault", self.Pyx_PyDict_SetDefault_func_type,
2806
            'setdefault', is_unbound_method, args,
2807 2808
            may_return_none=True,
            utility_code=load_c_utility('dict_setdefault'))
2809

2810
    Pyx_PyInt_BinopInt_func_type = PyrexTypes.CFuncType(
2811 2812 2813
        PyrexTypes.py_object_type, [
            PyrexTypes.CFuncTypeArg("op1", PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("op2", PyrexTypes.py_object_type, None),
2814 2815 2816 2817 2818 2819 2820 2821 2822 2823
            PyrexTypes.CFuncTypeArg("intval", PyrexTypes.c_long_type, None),
            PyrexTypes.CFuncTypeArg("inplace", PyrexTypes.c_bint_type, None),
        ])

    Pyx_PyFloat_BinopInt_func_type = PyrexTypes.CFuncType(
        PyrexTypes.py_object_type, [
            PyrexTypes.CFuncTypeArg("op1", PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("op2", PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("fval", PyrexTypes.c_double_type, None),
            PyrexTypes.CFuncTypeArg("inplace", PyrexTypes.c_bint_type, None),
2824 2825 2826
        ])

    def _handle_simple_method_object___add__(self, node, function, args, is_unbound_method):
2827
        return self._optimise_num_binop('Add', node, function, args, is_unbound_method)
2828 2829

    def _handle_simple_method_object___sub__(self, node, function, args, is_unbound_method):
2830 2831
        return self._optimise_num_binop('Subtract', node, function, args, is_unbound_method)

2832 2833 2834 2835 2836 2837
    def _handle_simple_method_object___eq__(self, node, function, args, is_unbound_method):
        return self._optimise_num_binop('Eq', node, function, args, is_unbound_method)

    def _handle_simple_method_object___neq__(self, node, function, args, is_unbound_method):
        return self._optimise_num_binop('Ne', node, function, args, is_unbound_method)

2838 2839 2840 2841 2842 2843 2844 2845 2846
    def _handle_simple_method_object___and__(self, node, function, args, is_unbound_method):
        return self._optimise_num_binop('And', node, function, args, is_unbound_method)

    def _handle_simple_method_object___or__(self, node, function, args, is_unbound_method):
        return self._optimise_num_binop('Or', node, function, args, is_unbound_method)

    def _handle_simple_method_object___xor__(self, node, function, args, is_unbound_method):
        return self._optimise_num_binop('Xor', node, function, args, is_unbound_method)

2847 2848 2849 2850 2851 2852 2853
    def _handle_simple_method_object___rshift__(self, node, function, args, is_unbound_method):
        if len(args) != 2 or not isinstance(args[1], ExprNodes.IntNode):
            return node
        if not args[1].has_constant_result() or not (1 <= args[1].constant_result <= 63):
            return node
        return self._optimise_num_binop('Rshift', node, function, args, is_unbound_method)

Stefan Behnel's avatar
Stefan Behnel committed
2854 2855 2856 2857 2858 2859 2860
    def _handle_simple_method_object___mod__(self, node, function, args, is_unbound_method):
        if len(args) != 2 or not isinstance(args[1], ExprNodes.IntNode):
            return node
        if not args[1].has_constant_result() or not (2 <= args[1].constant_result <= 2**30):
            return node
        return self._optimise_num_binop('Remainder', node, function, args, is_unbound_method)

2861 2862 2863 2864 2865 2866
    def _handle_simple_method_object___floordiv__(self, node, function, args, is_unbound_method):
        return self._optimise_num_div('FloorDivide', node, function, args, is_unbound_method)

    def _handle_simple_method_object___truediv__(self, node, function, args, is_unbound_method):
        return self._optimise_num_div('TrueDivide', node, function, args, is_unbound_method)

2867 2868 2869
    def _handle_simple_method_object___div__(self, node, function, args, is_unbound_method):
        return self._optimise_num_div('Divide', node, function, args, is_unbound_method)

2870 2871 2872 2873 2874 2875 2876 2877 2878 2879 2880 2881 2882
    def _optimise_num_div(self, operator, node, function, args, is_unbound_method):
        if len(args) != 2 or not args[1].has_constant_result() or args[1].constant_result == 0:
            return node
        if isinstance(args[1], ExprNodes.IntNode):
            if not (-2**30 <= args[1].constant_result <= 2**30):
                return node
        elif isinstance(args[1], ExprNodes.FloatNode):
            if not (-2**53 <= args[1].constant_result <= 2**53):
                return node
        else:
            return node
        return self._optimise_num_binop(operator, node, function, args, is_unbound_method)

2883 2884 2885 2886 2887
    def _handle_simple_method_float___add__(self, node, function, args, is_unbound_method):
        return self._optimise_num_binop('Add', node, function, args, is_unbound_method)

    def _handle_simple_method_float___sub__(self, node, function, args, is_unbound_method):
        return self._optimise_num_binop('Subtract', node, function, args, is_unbound_method)
2888

2889 2890 2891 2892 2893 2894
    def _handle_simple_method_float___truediv__(self, node, function, args, is_unbound_method):
        return self._optimise_num_binop('TrueDivide', node, function, args, is_unbound_method)

    def _handle_simple_method_float___div__(self, node, function, args, is_unbound_method):
        return self._optimise_num_binop('Divide', node, function, args, is_unbound_method)

2895 2896 2897 2898 2899 2900
    def _handle_simple_method_float___eq__(self, node, function, args, is_unbound_method):
        return self._optimise_num_binop('Eq', node, function, args, is_unbound_method)

    def _handle_simple_method_float___neq__(self, node, function, args, is_unbound_method):
        return self._optimise_num_binop('Ne', node, function, args, is_unbound_method)

2901
    def _optimise_num_binop(self, operator, node, function, args, is_unbound_method):
2902
        """
2903
        Optimise math operators for (likely) float or small integer operations.
2904 2905 2906 2907 2908 2909
        """
        if len(args) != 2:
            return node
        if not node.type.is_pyobject:
            return node

2910 2911
        # When adding IntNode/FloatNode to something else, assume other operand is also numeric.
        # Prefer constants on RHS as they allows better size control for some operators.
2912
        num_nodes = (ExprNodes.IntNode, ExprNodes.FloatNode)
2913
        if isinstance(args[1], num_nodes):
2914 2915
            if args[0].type is not PyrexTypes.py_object_type:
                return node
2916 2917
            numval = args[1]
            arg_order = 'ObjC'
2918 2919 2920 2921 2922
        elif isinstance(args[0], num_nodes):
            if args[1].type is not PyrexTypes.py_object_type:
                return node
            numval = args[0]
            arg_order = 'CObj'
2923 2924 2925
        else:
            return node

2926 2927 2928
        if not numval.has_constant_result():
            return node

2929
        is_float = isinstance(numval, ExprNodes.FloatNode)
2930
        if is_float:
2931
            if operator not in ('Add', 'Subtract', 'TrueDivide', 'Divide', 'Eq', 'Ne'):
2932
                return node
2933 2934 2935
        elif operator == 'Divide':
            # mixed old-/new-style division is not currently optimised for integers
            return node
2936
        elif abs(numval.constant_result) > 2**30:
2937
            return node
2938

2939
        args = list(args)
2940 2941 2942
        args.append((ExprNodes.FloatNode if is_float else ExprNodes.IntNode)(
            numval.pos, value=numval.value, constant_result=numval.constant_result,
            type=PyrexTypes.c_double_type if is_float else PyrexTypes.c_long_type))
2943
        inplace = node.inplace if isinstance(node, ExprNodes.NumBinopNode) else False
2944
        args.append(ExprNodes.BoolNode(node.pos, value=inplace, constant_result=inplace))
2945 2946

        utility_code = TempitaUtilityCode.load_cached(
2947
            "PyFloatBinop" if is_float else "PyIntBinop", "Optimize.c",
2948 2949 2950
            context=dict(op=operator, order=arg_order))

        return self._substitute_method_call(
2951 2952
            node, function, "__Pyx_Py%s_%s%s" % ('Float' if is_float else 'Int', operator, arg_order),
            self.Pyx_PyFloat_BinopInt_func_type if is_float else self.Pyx_PyInt_BinopInt_func_type,
2953 2954 2955 2956
            '__%s__' % operator[:3].lower(), is_unbound_method, args,
            may_return_none=True,
            with_none_check=False,
            utility_code=utility_code)
2957 2958 2959

    ### unicode type methods

2960 2961
    PyUnicode_uchar_predicate_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_bint_type, [
2962
            PyrexTypes.CFuncTypeArg("uchar", PyrexTypes.c_py_ucs4_type, None),
2963 2964
            ])

2965
    def _inject_unicode_predicate(self, node, function, args, is_unbound_method):
2966 2967 2968 2969
        if is_unbound_method or len(args) != 1:
            return node
        ustring = args[0]
        if not isinstance(ustring, ExprNodes.CoerceToPyTypeNode) or \
Stefan Behnel's avatar
Stefan Behnel committed
2970
               not ustring.arg.type.is_unicode_char:
2971 2972
            return node
        uchar = ustring.arg
2973
        method_name = function.attribute
2974 2975
        if method_name == 'istitle':
            # istitle() doesn't directly map to Py_UNICODE_ISTITLE()
2976 2977
            utility_code = UtilityCode.load_cached(
                "py_unicode_istitle", "StringTools.c")
2978 2979 2980 2981 2982
            function_name = '__Pyx_Py_UNICODE_ISTITLE'
        else:
            utility_code = None
            function_name = 'Py_UNICODE_%s' % method_name.upper()
        func_call = self._substitute_method_call(
2983 2984
            node, function,
            function_name, self.PyUnicode_uchar_predicate_func_type,
2985 2986 2987 2988 2989 2990 2991 2992 2993 2994 2995 2996 2997 2998 2999 3000 3001
            method_name, is_unbound_method, [uchar],
            utility_code = utility_code)
        if node.type.is_pyobject:
            func_call = func_call.coerce_to_pyobject(self.current_env)
        return func_call

    _handle_simple_method_unicode_isalnum   = _inject_unicode_predicate
    _handle_simple_method_unicode_isalpha   = _inject_unicode_predicate
    _handle_simple_method_unicode_isdecimal = _inject_unicode_predicate
    _handle_simple_method_unicode_isdigit   = _inject_unicode_predicate
    _handle_simple_method_unicode_islower   = _inject_unicode_predicate
    _handle_simple_method_unicode_isnumeric = _inject_unicode_predicate
    _handle_simple_method_unicode_isspace   = _inject_unicode_predicate
    _handle_simple_method_unicode_istitle   = _inject_unicode_predicate
    _handle_simple_method_unicode_isupper   = _inject_unicode_predicate

    PyUnicode_uchar_conversion_func_type = PyrexTypes.CFuncType(
3002 3003
        PyrexTypes.c_py_ucs4_type, [
            PyrexTypes.CFuncTypeArg("uchar", PyrexTypes.c_py_ucs4_type, None),
3004 3005
            ])

Stefan Behnel's avatar
Stefan Behnel committed
3006
    def _inject_unicode_character_conversion(self, node, function, args, is_unbound_method):
3007 3008 3009 3010
        if is_unbound_method or len(args) != 1:
            return node
        ustring = args[0]
        if not isinstance(ustring, ExprNodes.CoerceToPyTypeNode) or \
Stefan Behnel's avatar
Stefan Behnel committed
3011
               not ustring.arg.type.is_unicode_char:
3012 3013
            return node
        uchar = ustring.arg
3014
        method_name = function.attribute
3015 3016
        function_name = 'Py_UNICODE_TO%s' % method_name.upper()
        func_call = self._substitute_method_call(
3017 3018
            node, function,
            function_name, self.PyUnicode_uchar_conversion_func_type,
3019 3020 3021 3022 3023 3024 3025 3026 3027
            method_name, is_unbound_method, [uchar])
        if node.type.is_pyobject:
            func_call = func_call.coerce_to_pyobject(self.current_env)
        return func_call

    _handle_simple_method_unicode_lower = _inject_unicode_character_conversion
    _handle_simple_method_unicode_upper = _inject_unicode_character_conversion
    _handle_simple_method_unicode_title = _inject_unicode_character_conversion

3028 3029 3030 3031 3032 3033
    PyUnicode_Splitlines_func_type = PyrexTypes.CFuncType(
        Builtin.list_type, [
            PyrexTypes.CFuncTypeArg("str", Builtin.unicode_type, None),
            PyrexTypes.CFuncTypeArg("keepends", PyrexTypes.c_bint_type, None),
            ])

Stefan Behnel's avatar
Stefan Behnel committed
3034
    def _handle_simple_method_unicode_splitlines(self, node, function, args, is_unbound_method):
3035 3036 3037 3038 3039 3040
        """Replace unicode.splitlines(...) by a direct call to the
        corresponding C-API function.
        """
        if len(args) not in (1,2):
            self._error_wrong_arg_count('unicode.splitlines', node, args, "1 or 2")
            return node
3041
        self._inject_bint_default_argument(node, args, 1, False)
3042 3043

        return self._substitute_method_call(
3044 3045
            node, function,
            "PyUnicode_Splitlines", self.PyUnicode_Splitlines_func_type,
3046 3047 3048 3049 3050 3051 3052 3053 3054 3055
            'splitlines', is_unbound_method, args)

    PyUnicode_Split_func_type = PyrexTypes.CFuncType(
        Builtin.list_type, [
            PyrexTypes.CFuncTypeArg("str", Builtin.unicode_type, None),
            PyrexTypes.CFuncTypeArg("sep", PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("maxsplit", PyrexTypes.c_py_ssize_t_type, None),
            ]
        )

3056
    def _handle_simple_method_unicode_split(self, node, function, args, is_unbound_method):
3057 3058 3059 3060 3061 3062 3063 3064
        """Replace unicode.split(...) by a direct call to the
        corresponding C-API function.
        """
        if len(args) not in (1,2,3):
            self._error_wrong_arg_count('unicode.split', node, args, "1-3")
            return node
        if len(args) < 2:
            args.append(ExprNodes.NullNode(node.pos))
3065 3066
        self._inject_int_default_argument(
            node, args, 2, PyrexTypes.c_py_ssize_t_type, "-1")
3067 3068

        return self._substitute_method_call(
3069 3070
            node, function,
            "PyUnicode_Split", self.PyUnicode_Split_func_type,
3071 3072
            'split', is_unbound_method, args)

3073
    PyString_Tailmatch_func_type = PyrexTypes.CFuncType(
3074
        PyrexTypes.c_bint_type, [
3075
            PyrexTypes.CFuncTypeArg("str", PyrexTypes.py_object_type, None),  # bytes/str/unicode
3076 3077 3078 3079 3080 3081 3082
            PyrexTypes.CFuncTypeArg("substring", PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("start", PyrexTypes.c_py_ssize_t_type, None),
            PyrexTypes.CFuncTypeArg("end", PyrexTypes.c_py_ssize_t_type, None),
            PyrexTypes.CFuncTypeArg("direction", PyrexTypes.c_int_type, None),
            ],
        exception_value = '-1')

3083
    def _handle_simple_method_unicode_endswith(self, node, function, args, is_unbound_method):
3084
        return self._inject_tailmatch(
3085
            node, function, args, is_unbound_method, 'unicode', 'endswith',
3086
            unicode_tailmatch_utility_code, +1)
3087

3088
    def _handle_simple_method_unicode_startswith(self, node, function, args, is_unbound_method):
3089
        return self._inject_tailmatch(
3090
            node, function, args, is_unbound_method, 'unicode', 'startswith',
3091
            unicode_tailmatch_utility_code, -1)
3092

3093
    def _inject_tailmatch(self, node, function, args, is_unbound_method, type_name,
3094
                          method_name, utility_code, direction):
3095 3096 3097 3098
        """Replace unicode.startswith(...) and unicode.endswith(...)
        by a direct call to the corresponding C-API function.
        """
        if len(args) not in (2,3,4):
3099
            self._error_wrong_arg_count('%s.%s' % (type_name, method_name), node, args, "2-4")
3100
            return node
3101 3102 3103 3104
        self._inject_int_default_argument(
            node, args, 2, PyrexTypes.c_py_ssize_t_type, "0")
        self._inject_int_default_argument(
            node, args, 3, PyrexTypes.c_py_ssize_t_type, "PY_SSIZE_T_MAX")
3105 3106 3107 3108
        args.append(ExprNodes.IntNode(
            node.pos, value=str(direction), type=PyrexTypes.c_int_type))

        method_call = self._substitute_method_call(
3109 3110
            node, function,
            "__Pyx_Py%s_Tailmatch" % type_name.capitalize(),
3111
            self.PyString_Tailmatch_func_type,
3112
            method_name, is_unbound_method, args,
3113
            utility_code = utility_code)
Stefan Behnel's avatar
Stefan Behnel committed
3114
        return method_call.coerce_to(Builtin.bool_type, self.current_env())
3115

3116 3117 3118 3119 3120 3121 3122 3123 3124 3125
    PyUnicode_Find_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_py_ssize_t_type, [
            PyrexTypes.CFuncTypeArg("str", Builtin.unicode_type, None),
            PyrexTypes.CFuncTypeArg("substring", PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("start", PyrexTypes.c_py_ssize_t_type, None),
            PyrexTypes.CFuncTypeArg("end", PyrexTypes.c_py_ssize_t_type, None),
            PyrexTypes.CFuncTypeArg("direction", PyrexTypes.c_int_type, None),
            ],
        exception_value = '-2')

3126
    def _handle_simple_method_unicode_find(self, node, function, args, is_unbound_method):
3127
        return self._inject_unicode_find(
3128
            node, function, args, is_unbound_method, 'find', +1)
3129

3130
    def _handle_simple_method_unicode_rfind(self, node, function, args, is_unbound_method):
3131
        return self._inject_unicode_find(
3132
            node, function, args, is_unbound_method, 'rfind', -1)
3133

3134
    def _inject_unicode_find(self, node, function, args, is_unbound_method,
3135 3136 3137 3138 3139 3140 3141
                             method_name, direction):
        """Replace unicode.find(...) and unicode.rfind(...) by a
        direct call to the corresponding C-API function.
        """
        if len(args) not in (2,3,4):
            self._error_wrong_arg_count('unicode.%s' % method_name, node, args, "2-4")
            return node
3142 3143 3144 3145
        self._inject_int_default_argument(
            node, args, 2, PyrexTypes.c_py_ssize_t_type, "0")
        self._inject_int_default_argument(
            node, args, 3, PyrexTypes.c_py_ssize_t_type, "PY_SSIZE_T_MAX")
3146 3147 3148 3149
        args.append(ExprNodes.IntNode(
            node.pos, value=str(direction), type=PyrexTypes.c_int_type))

        method_call = self._substitute_method_call(
3150
            node, function, "PyUnicode_Find", self.PyUnicode_Find_func_type,
3151
            method_name, is_unbound_method, args)
Stefan Behnel's avatar
Stefan Behnel committed
3152
        return method_call.coerce_to_pyobject(self.current_env())
3153

Stefan Behnel's avatar
Stefan Behnel committed
3154 3155 3156 3157 3158 3159 3160 3161 3162
    PyUnicode_Count_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_py_ssize_t_type, [
            PyrexTypes.CFuncTypeArg("str", Builtin.unicode_type, None),
            PyrexTypes.CFuncTypeArg("substring", PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("start", PyrexTypes.c_py_ssize_t_type, None),
            PyrexTypes.CFuncTypeArg("end", PyrexTypes.c_py_ssize_t_type, None),
            ],
        exception_value = '-1')

3163
    def _handle_simple_method_unicode_count(self, node, function, args, is_unbound_method):
Stefan Behnel's avatar
Stefan Behnel committed
3164 3165 3166 3167 3168 3169
        """Replace unicode.count(...) by a direct call to the
        corresponding C-API function.
        """
        if len(args) not in (2,3,4):
            self._error_wrong_arg_count('unicode.count', node, args, "2-4")
            return node
3170 3171 3172 3173
        self._inject_int_default_argument(
            node, args, 2, PyrexTypes.c_py_ssize_t_type, "0")
        self._inject_int_default_argument(
            node, args, 3, PyrexTypes.c_py_ssize_t_type, "PY_SSIZE_T_MAX")
Stefan Behnel's avatar
Stefan Behnel committed
3174 3175

        method_call = self._substitute_method_call(
3176
            node, function, "PyUnicode_Count", self.PyUnicode_Count_func_type,
Stefan Behnel's avatar
Stefan Behnel committed
3177
            'count', is_unbound_method, args)
Stefan Behnel's avatar
Stefan Behnel committed
3178
        return method_call.coerce_to_pyobject(self.current_env())
Stefan Behnel's avatar
Stefan Behnel committed
3179

Stefan Behnel's avatar
Stefan Behnel committed
3180 3181 3182 3183 3184 3185 3186 3187
    PyUnicode_Replace_func_type = PyrexTypes.CFuncType(
        Builtin.unicode_type, [
            PyrexTypes.CFuncTypeArg("str", Builtin.unicode_type, None),
            PyrexTypes.CFuncTypeArg("substring", PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("replstr", PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("maxcount", PyrexTypes.c_py_ssize_t_type, None),
            ])

3188
    def _handle_simple_method_unicode_replace(self, node, function, args, is_unbound_method):
Stefan Behnel's avatar
Stefan Behnel committed
3189 3190 3191 3192 3193 3194
        """Replace unicode.replace(...) by a direct call to the
        corresponding C-API function.
        """
        if len(args) not in (3,4):
            self._error_wrong_arg_count('unicode.replace', node, args, "3-4")
            return node
3195 3196
        self._inject_int_default_argument(
            node, args, 3, PyrexTypes.c_py_ssize_t_type, "-1")
Stefan Behnel's avatar
Stefan Behnel committed
3197 3198

        return self._substitute_method_call(
3199
            node, function, "PyUnicode_Replace", self.PyUnicode_Replace_func_type,
Stefan Behnel's avatar
Stefan Behnel committed
3200 3201
            'replace', is_unbound_method, args)

3202 3203
    PyUnicode_AsEncodedString_func_type = PyrexTypes.CFuncType(
        Builtin.bytes_type, [
3204
            PyrexTypes.CFuncTypeArg("obj", Builtin.unicode_type, None),
3205 3206
            PyrexTypes.CFuncTypeArg("encoding", PyrexTypes.c_char_ptr_type, None),
            PyrexTypes.CFuncTypeArg("errors", PyrexTypes.c_char_ptr_type, None),
3207
            ])
3208 3209 3210

    PyUnicode_AsXyzString_func_type = PyrexTypes.CFuncType(
        Builtin.bytes_type, [
3211
            PyrexTypes.CFuncTypeArg("obj", Builtin.unicode_type, None),
3212
            ])
3213 3214 3215 3216

    _special_encodings = ['UTF8', 'UTF16', 'Latin1', 'ASCII',
                          'unicode_escape', 'raw_unicode_escape']

3217 3218
    _special_codecs = [ (name, codecs.getencoder(name))
                        for name in _special_encodings ]
3219

3220
    def _handle_simple_method_unicode_encode(self, node, function, args, is_unbound_method):
Stefan Behnel's avatar
Stefan Behnel committed
3221 3222 3223
        """Replace unicode.encode(...) by a direct C-API call to the
        corresponding codec.
        """
3224
        if len(args) < 1 or len(args) > 3:
3225
            self._error_wrong_arg_count('unicode.encode', node, args, '1-3')
3226 3227 3228 3229 3230
            return node

        string_node = args[0]

        if len(args) == 1:
3231
            null_node = ExprNodes.NullNode(node.pos)
3232
            return self._substitute_method_call(
3233
                node, function, "PyUnicode_AsEncodedString",
3234 3235 3236
                self.PyUnicode_AsEncodedString_func_type,
                'encode', is_unbound_method, [string_node, null_node, null_node])

3237 3238 3239 3240 3241
        parameters = self._unpack_encoding_and_error_mode(node.pos, args)
        if parameters is None:
            return node
        encoding, encoding_node, error_handling, error_handling_node = parameters

3242
        if encoding and isinstance(string_node, ExprNodes.UnicodeNode):
3243 3244 3245 3246 3247 3248 3249 3250 3251 3252 3253 3254
            # constant, so try to do the encoding at compile time
            try:
                value = string_node.value.encode(encoding, error_handling)
            except:
                # well, looks like we can't
                pass
            else:
                value = BytesLiteral(value)
                value.encoding = encoding
                return ExprNodes.BytesNode(
                    string_node.pos, value=value, type=Builtin.bytes_type)

3255
        if encoding and error_handling == 'strict':
3256 3257 3258 3259 3260
            # try to find a specific encoder function
            codec_name = self._find_special_codec_name(encoding)
            if codec_name is not None:
                encode_function = "PyUnicode_As%sString" % codec_name
                return self._substitute_method_call(
3261
                    node, function, encode_function,
3262 3263 3264 3265
                    self.PyUnicode_AsXyzString_func_type,
                    'encode', is_unbound_method, [string_node])

        return self._substitute_method_call(
3266
            node, function, "PyUnicode_AsEncodedString",
3267 3268 3269 3270
            self.PyUnicode_AsEncodedString_func_type,
            'encode', is_unbound_method,
            [string_node, encoding_node, error_handling_node])

3271
    PyUnicode_DecodeXyz_func_ptr_type = PyrexTypes.CPtrType(PyrexTypes.CFuncType(
3272 3273 3274 3275
        Builtin.unicode_type, [
            PyrexTypes.CFuncTypeArg("string", PyrexTypes.c_char_ptr_type, None),
            PyrexTypes.CFuncTypeArg("size", PyrexTypes.c_py_ssize_t_type, None),
            PyrexTypes.CFuncTypeArg("errors", PyrexTypes.c_char_ptr_type, None),
3276
            ]))
3277

3278
    _decode_c_string_func_type = PyrexTypes.CFuncType(
3279 3280
        Builtin.unicode_type, [
            PyrexTypes.CFuncTypeArg("string", PyrexTypes.c_char_ptr_type, None),
3281 3282
            PyrexTypes.CFuncTypeArg("start", PyrexTypes.c_py_ssize_t_type, None),
            PyrexTypes.CFuncTypeArg("stop", PyrexTypes.c_py_ssize_t_type, None),
3283 3284
            PyrexTypes.CFuncTypeArg("encoding", PyrexTypes.c_char_ptr_type, None),
            PyrexTypes.CFuncTypeArg("errors", PyrexTypes.c_char_ptr_type, None),
3285
            PyrexTypes.CFuncTypeArg("decode_func", PyUnicode_DecodeXyz_func_ptr_type, None),
3286
            ])
3287

Stefan Behnel's avatar
Stefan Behnel committed
3288 3289 3290 3291 3292 3293 3294 3295 3296 3297
    _decode_bytes_func_type = PyrexTypes.CFuncType(
        Builtin.unicode_type, [
            PyrexTypes.CFuncTypeArg("string", PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("start", PyrexTypes.c_py_ssize_t_type, None),
            PyrexTypes.CFuncTypeArg("stop", PyrexTypes.c_py_ssize_t_type, None),
            PyrexTypes.CFuncTypeArg("encoding", PyrexTypes.c_char_ptr_type, None),
            PyrexTypes.CFuncTypeArg("errors", PyrexTypes.c_char_ptr_type, None),
            PyrexTypes.CFuncTypeArg("decode_func", PyUnicode_DecodeXyz_func_ptr_type, None),
            ])

Stefan Behnel's avatar
Stefan Behnel committed
3298
    _decode_cpp_string_func_type = None  # lazy init
3299

3300
    def _handle_simple_method_bytes_decode(self, node, function, args, is_unbound_method):
Stefan Behnel's avatar
Stefan Behnel committed
3301
        """Replace char*.decode() by a direct C-API call to the
Stefan Behnel's avatar
Stefan Behnel committed
3302
        corresponding codec, possibly resolving a slice on the char*.
Stefan Behnel's avatar
Stefan Behnel committed
3303
        """
3304
        if not (1 <= len(args) <= 3):
3305 3306
            self._error_wrong_arg_count('bytes.decode', node, args, '1-3')
            return node
3307 3308

        # normalise input nodes
Stefan Behnel's avatar
Stefan Behnel committed
3309 3310 3311 3312
        string_node = args[0]
        start = stop = None
        if isinstance(string_node, ExprNodes.SliceIndexNode):
            index_node = string_node
3313 3314 3315 3316
            string_node = index_node.base
            start, stop = index_node.start, index_node.stop
            if not start or start.constant_result == 0:
                start = None
Stefan Behnel's avatar
Stefan Behnel committed
3317 3318 3319 3320
        if isinstance(string_node, ExprNodes.CoerceToPyTypeNode):
            string_node = string_node.arg

        string_type = string_node.type
Stefan Behnel's avatar
Stefan Behnel committed
3321
        if string_type in (Builtin.bytes_type, Builtin.bytearray_type):
Stefan Behnel's avatar
Stefan Behnel committed
3322 3323 3324
            if is_unbound_method:
                string_node = string_node.as_none_safe_node(
                    "descriptor '%s' requires a '%s' object but received a 'NoneType'",
Stefan Behnel's avatar
Stefan Behnel committed
3325
                    format_args=['decode', string_type.name])
Stefan Behnel's avatar
Stefan Behnel committed
3326 3327 3328
            else:
                string_node = string_node.as_none_safe_node(
                    "'NoneType' object has no attribute '%s'",
Stefan Behnel's avatar
Stefan Behnel committed
3329 3330
                    error="PyExc_AttributeError",
                    format_args=['decode'])
Stefan Behnel's avatar
Stefan Behnel committed
3331
        elif not string_type.is_string and not string_type.is_cpp_string:
3332 3333
            # nothing to optimise here
            return node
3334 3335 3336 3337 3338 3339

        parameters = self._unpack_encoding_and_error_mode(node.pos, args)
        if parameters is None:
            return node
        encoding, encoding_node, error_handling, error_handling_node = parameters

3340 3341 3342 3343 3344 3345 3346
        if not start:
            start = ExprNodes.IntNode(node.pos, value='0', constant_result=0)
        elif not start.type.is_int:
            start = start.coerce_to(PyrexTypes.c_py_ssize_t_type, self.current_env())
        if stop and not stop.type.is_int:
            stop = stop.coerce_to(PyrexTypes.c_py_ssize_t_type, self.current_env())

3347
        # try to find a specific encoder function
3348 3349 3350
        codec_name = None
        if encoding is not None:
            codec_name = self._find_special_codec_name(encoding)
3351
        if codec_name is not None:
3352 3353 3354 3355
            decode_function = ExprNodes.RawCNameExprNode(
                node.pos, type=self.PyUnicode_DecodeXyz_func_ptr_type,
                cname="PyUnicode_Decode%s" % codec_name)
            encoding_node = ExprNodes.NullNode(node.pos)
3356
        else:
3357 3358 3359 3360
            decode_function = ExprNodes.NullNode(node.pos)

        # build the helper function call
        temps = []
Stefan Behnel's avatar
Stefan Behnel committed
3361
        if string_type.is_string:
3362 3363 3364 3365 3366 3367 3368 3369
            # C string
            if not stop:
                # use strlen() to find the string length, just as CPython would
                if not string_node.is_name:
                    string_node = UtilNodes.LetRefNode(string_node) # used twice
                    temps.append(string_node)
                stop = ExprNodes.PythonCapiCallNode(
                    string_node.pos, "strlen", self.Pyx_strlen_func_type,
Stefan Behnel's avatar
Stefan Behnel committed
3370 3371 3372 3373
                    args=[string_node],
                    is_temp=False,
                    utility_code=UtilityCode.load_cached("IncludeStringH", "StringTools.c"),
                ).coerce_to(PyrexTypes.c_py_ssize_t_type, self.current_env())
3374 3375
            helper_func_type = self._decode_c_string_func_type
            utility_code_name = 'decode_c_string'
Stefan Behnel's avatar
Stefan Behnel committed
3376
        elif string_type.is_cpp_string:
3377 3378 3379 3380 3381 3382 3383 3384
            # C++ std::string
            if not stop:
                stop = ExprNodes.IntNode(node.pos, value='PY_SSIZE_T_MAX',
                                         constant_result=ExprNodes.not_a_constant)
            if self._decode_cpp_string_func_type is None:
                # lazy init to reuse the C++ string type
                self._decode_cpp_string_func_type = PyrexTypes.CFuncType(
                    Builtin.unicode_type, [
Stefan Behnel's avatar
Stefan Behnel committed
3385
                        PyrexTypes.CFuncTypeArg("string", string_type, None),
3386 3387 3388 3389 3390
                        PyrexTypes.CFuncTypeArg("start", PyrexTypes.c_py_ssize_t_type, None),
                        PyrexTypes.CFuncTypeArg("stop", PyrexTypes.c_py_ssize_t_type, None),
                        PyrexTypes.CFuncTypeArg("encoding", PyrexTypes.c_char_ptr_type, None),
                        PyrexTypes.CFuncTypeArg("errors", PyrexTypes.c_char_ptr_type, None),
                        PyrexTypes.CFuncTypeArg("decode_func", self.PyUnicode_DecodeXyz_func_ptr_type, None),
Stefan Behnel's avatar
Stefan Behnel committed
3391
                    ])
3392 3393
            helper_func_type = self._decode_cpp_string_func_type
            utility_code_name = 'decode_cpp_string'
Stefan Behnel's avatar
Stefan Behnel committed
3394
        else:
Stefan Behnel's avatar
Stefan Behnel committed
3395
            # Python bytes/bytearray object
Stefan Behnel's avatar
Stefan Behnel committed
3396 3397 3398 3399
            if not stop:
                stop = ExprNodes.IntNode(node.pos, value='PY_SSIZE_T_MAX',
                                         constant_result=ExprNodes.not_a_constant)
            helper_func_type = self._decode_bytes_func_type
Stefan Behnel's avatar
Stefan Behnel committed
3400 3401 3402 3403
            if string_type is Builtin.bytes_type:
                utility_code_name = 'decode_bytes'
            else:
                utility_code_name = 'decode_bytearray'
3404 3405 3406

        node = ExprNodes.PythonCapiCallNode(
            node.pos, '__Pyx_%s' % utility_code_name, helper_func_type,
Stefan Behnel's avatar
Stefan Behnel committed
3407 3408
            args=[string_node, start, stop, encoding_node, error_handling_node, decode_function],
            is_temp=node.is_temp,
3409
            utility_code=UtilityCode.load_cached(utility_code_name, 'StringTools.c'),
Stefan Behnel's avatar
Stefan Behnel committed
3410
        )
3411

3412 3413 3414
        for temp in temps[::-1]:
            node = UtilNodes.EvalWithTempExprNode(temp, node)
        return node
3415

Stefan Behnel's avatar
Stefan Behnel committed
3416 3417
    _handle_simple_method_bytearray_decode = _handle_simple_method_bytes_decode

3418 3419 3420
    def _find_special_codec_name(self, encoding):
        try:
            requested_codec = codecs.getencoder(encoding)
Stefan Behnel's avatar
Stefan Behnel committed
3421
        except LookupError:
3422 3423 3424 3425
            return None
        for name, codec in self._special_codecs:
            if codec == requested_codec:
                if '_' in name:
Stefan Behnel's avatar
Stefan Behnel committed
3426 3427
                    name = ''.join([s.capitalize()
                                    for s in name.split('_')])
3428 3429 3430 3431
                return name
        return None

    def _unpack_encoding_and_error_mode(self, pos, args):
3432 3433 3434
        null_node = ExprNodes.NullNode(pos)

        if len(args) >= 2:
3435 3436
            encoding, encoding_node = self._unpack_string_and_cstring_node(args[1])
            if encoding_node is None:
3437
                return None
3438
        else:
3439 3440
            encoding = None
            encoding_node = null_node
3441 3442

        if len(args) == 3:
3443 3444
            error_handling, error_handling_node = self._unpack_string_and_cstring_node(args[2])
            if error_handling_node is None:
3445
                return None
3446 3447
            if error_handling == 'strict':
                error_handling_node = null_node
3448 3449 3450 3451
        else:
            error_handling = 'strict'
            error_handling_node = null_node

3452
        return (encoding, encoding_node, error_handling, error_handling_node)
3453

3454 3455 3456 3457 3458 3459 3460 3461 3462 3463 3464 3465 3466 3467 3468 3469 3470 3471
    def _unpack_string_and_cstring_node(self, node):
        if isinstance(node, ExprNodes.CoerceToPyTypeNode):
            node = node.arg
        if isinstance(node, ExprNodes.UnicodeNode):
            encoding = node.value
            node = ExprNodes.BytesNode(
                node.pos, value=BytesLiteral(encoding.utf8encode()),
                type=PyrexTypes.c_char_ptr_type)
        elif isinstance(node, (ExprNodes.StringNode, ExprNodes.BytesNode)):
            encoding = node.value.decode('ISO-8859-1')
            node = ExprNodes.BytesNode(
                node.pos, value=node.value, type=PyrexTypes.c_char_ptr_type)
        elif node.type is Builtin.bytes_type:
            encoding = None
            node = node.coerce_to(PyrexTypes.c_char_ptr_type, self.current_env())
        elif node.type.is_string:
            encoding = None
        else:
3472
            encoding = node = None
3473 3474
        return encoding, node

3475
    def _handle_simple_method_str_endswith(self, node, function, args, is_unbound_method):
3476
        return self._inject_tailmatch(
3477
            node, function, args, is_unbound_method, 'str', 'endswith',
3478
            str_tailmatch_utility_code, +1)
3479

3480
    def _handle_simple_method_str_startswith(self, node, function, args, is_unbound_method):
3481
        return self._inject_tailmatch(
3482
            node, function, args, is_unbound_method, 'str', 'startswith',
3483
            str_tailmatch_utility_code, -1)
3484

3485
    def _handle_simple_method_bytes_endswith(self, node, function, args, is_unbound_method):
3486
        return self._inject_tailmatch(
3487
            node, function, args, is_unbound_method, 'bytes', 'endswith',
3488
            bytes_tailmatch_utility_code, +1)
3489

3490
    def _handle_simple_method_bytes_startswith(self, node, function, args, is_unbound_method):
3491
        return self._inject_tailmatch(
3492
            node, function, args, is_unbound_method, 'bytes', 'startswith',
3493
            bytes_tailmatch_utility_code, -1)
3494

Stefan Behnel's avatar
Stefan Behnel committed
3495 3496 3497 3498 3499 3500 3501 3502 3503 3504 3505 3506
    '''   # disabled for now, enable when we consider it worth it (see StringTools.c)
    def _handle_simple_method_bytearray_endswith(self, node, function, args, is_unbound_method):
        return self._inject_tailmatch(
            node, function, args, is_unbound_method, 'bytearray', 'endswith',
            bytes_tailmatch_utility_code, +1)

    def _handle_simple_method_bytearray_startswith(self, node, function, args, is_unbound_method):
        return self._inject_tailmatch(
            node, function, args, is_unbound_method, 'bytearray', 'startswith',
            bytes_tailmatch_utility_code, -1)
    '''

3507 3508
    ### helpers

3509
    def _substitute_method_call(self, node, function, name, func_type,
3510
                                attr_name, is_unbound_method, args=(),
3511
                                utility_code=None, is_temp=None,
3512 3513
                                may_return_none=ExprNodes.PythonCapiCallNode.may_return_none,
                                with_none_check=True):
3514
        args = list(args)
3515
        if with_none_check and args and not args[0].is_literal:
3516 3517
            self_arg = args[0]
            if is_unbound_method:
3518
                self_arg = self_arg.as_none_safe_node(
3519
                    "descriptor '%s' requires a '%s' object but received a 'NoneType'",
3520
                    format_args=[attr_name, function.obj.name])
3521
            else:
3522
                self_arg = self_arg.as_none_safe_node(
3523 3524 3525
                    "'NoneType' object has no attribute '%s'",
                    error = "PyExc_AttributeError",
                    format_args = [attr_name])
3526
            args[0] = self_arg
3527 3528
        if is_temp is None:
            is_temp = node.is_temp
3529
        return ExprNodes.PythonCapiCallNode(
3530
            node.pos, name, func_type,
3531
            args = args,
3532
            is_temp = is_temp,
Stefan Behnel's avatar
Stefan Behnel committed
3533 3534
            utility_code = utility_code,
            may_return_none = may_return_none,
3535
            result_is_used = node.result_is_used,
3536 3537
            )

3538 3539 3540
    def _inject_int_default_argument(self, node, args, arg_index, type, default_value):
        assert len(args) >= arg_index
        if len(args) == arg_index:
3541 3542
            args.append(ExprNodes.IntNode(node.pos, value=str(default_value),
                                          type=type, constant_result=default_value))
3543
        else:
3544
            args[arg_index] = args[arg_index].coerce_to(type, self.current_env())
3545 3546 3547 3548

    def _inject_bint_default_argument(self, node, args, arg_index, default_value):
        assert len(args) >= arg_index
        if len(args) == arg_index:
3549 3550 3551
            default_value = bool(default_value)
            args.append(ExprNodes.BoolNode(node.pos, value=default_value,
                                           constant_result=default_value))
3552
        else:
3553
            args[arg_index] = args[arg_index].coerce_to_boolean(self.current_env())
3554

3555

3556 3557 3558
unicode_tailmatch_utility_code = UtilityCode.load_cached('unicode_tailmatch', 'StringTools.c')
bytes_tailmatch_utility_code = UtilityCode.load_cached('bytes_tailmatch', 'StringTools.c')
str_tailmatch_utility_code = UtilityCode.load_cached('str_tailmatch', 'StringTools.c')
3559

3560

3561 3562 3563 3564
class ConstantFolding(Visitor.VisitorTransform, SkipDeclarations):
    """Calculate the result of constant expressions to store it in
    ``expr_node.constant_result``, and replace trivial cases by their
    constant result.
3565 3566 3567 3568 3569 3570 3571 3572 3573 3574 3575

    General rules:

    - We calculate float constants to make them available to the
      compiler, but we do not aggregate them into a single literal
      node to prevent any loss of precision.

    - We recursively calculate constants from non-literal nodes to
      make them available to the compiler, but we only aggregate
      literal nodes at each step.  Non-literal nodes are never merged
      into a single node.
3576
    """
3577

Mark Florisson's avatar
Mark Florisson committed
3578 3579 3580 3581 3582 3583 3584
    def __init__(self, reevaluate=False):
        """
        The reevaluate argument specifies whether constant values that were
        previously computed should be recomputed.
        """
        super(ConstantFolding, self).__init__()
        self.reevaluate = reevaluate
3585

3586
    def _calculate_const(self, node):
Mark Florisson's avatar
Mark Florisson committed
3587
        if (not self.reevaluate and
3588
                node.constant_result is not ExprNodes.constant_value_not_set):
3589 3590 3591 3592 3593 3594 3595 3596
            return

        # make sure we always set the value
        not_a_constant = ExprNodes.not_a_constant
        node.constant_result = not_a_constant

        # check if all children are constant
        children = self.visitchildren(node)
3597
        for child_result in children.values():
3598 3599
            if type(child_result) is list:
                for child in child_result:
Stefan Behnel's avatar
Stefan Behnel committed
3600
                    if getattr(child, 'constant_result', not_a_constant) is not_a_constant:
3601
                        return
Stefan Behnel's avatar
Stefan Behnel committed
3602
            elif getattr(child_result, 'constant_result', not_a_constant) is not_a_constant:
3603 3604 3605 3606 3607 3608 3609
                return

        # now try to calculate the real constant value
        try:
            node.calculate_constant_result()
#            if node.constant_result is not ExprNodes.not_a_constant:
#                print node.__class__.__name__, node.constant_result
3610
        except (ValueError, TypeError, KeyError, IndexError, AttributeError, ArithmeticError):
3611 3612 3613 3614 3615 3616 3617
            # ignore all 'normal' errors here => no constant result
            pass
        except Exception:
            # this looks like a real error
            import traceback, sys
            traceback.print_exc(file=sys.stdout)

3618 3619
    NODE_TYPE_ORDER = [ExprNodes.BoolNode, ExprNodes.CharNode,
                       ExprNodes.IntNode, ExprNodes.FloatNode]
3620 3621 3622 3623 3624 3625 3626 3627

    def _widest_node_class(self, *nodes):
        try:
            return self.NODE_TYPE_ORDER[
                max(map(self.NODE_TYPE_ORDER.index, map(type, nodes)))]
        except ValueError:
            return None

3628 3629 3630 3631
    def _bool_node(self, node, value):
        value = bool(value)
        return ExprNodes.BoolNode(node.pos, value=value, constant_result=value)

3632 3633 3634 3635
    def visit_ExprNode(self, node):
        self._calculate_const(node)
        return node

3636
    def visit_UnopNode(self, node):
3637
        self._calculate_const(node)
Stefan Behnel's avatar
Stefan Behnel committed
3638
        if not node.has_constant_result():
3639
            if node.operator == '!':
3640
                return self._handle_NotNode(node)
3641 3642 3643
            return node
        if not node.operand.is_literal:
            return node
Stefan Behnel's avatar
Stefan Behnel committed
3644
        if node.operator == '!':
3645
            return self._bool_node(node, node.constant_result)
3646
        elif isinstance(node.operand, ExprNodes.BoolNode):
Stefan Behnel's avatar
Stefan Behnel committed
3647 3648 3649
            return ExprNodes.IntNode(node.pos, value=str(int(node.constant_result)),
                                     type=PyrexTypes.c_int_type,
                                     constant_result=int(node.constant_result))
3650
        elif node.operator == '+':
3651 3652 3653 3654 3655
            return self._handle_UnaryPlusNode(node)
        elif node.operator == '-':
            return self._handle_UnaryMinusNode(node)
        return node

3656 3657 3658 3659 3660 3661 3662
    _negate_operator = {
        'in': 'not_in',
        'not_in': 'in',
        'is': 'is_not',
        'is_not': 'is'
    }.get

3663
    def _handle_NotNode(self, node):
3664 3665 3666 3667 3668 3669 3670
        operand = node.operand
        if isinstance(operand, ExprNodes.PrimaryCmpNode):
            operator = self._negate_operator(operand.operator)
            if operator:
                node = copy.copy(operand)
                node.operator = operator
                node = self.visit_PrimaryCmpNode(node)
3671 3672
        return node

3673
    def _handle_UnaryMinusNode(self, node):
3674 3675 3676 3677 3678 3679 3680
        def _negate(value):
            if value.startswith('-'):
                value = value[1:]
            else:
                value = '-' + value
            return value

3681
        node_type = node.operand.type
3682 3683
        if isinstance(node.operand, ExprNodes.FloatNode):
            # this is a safe operation
3684
            return ExprNodes.FloatNode(node.pos, value=_negate(node.operand.value),
3685
                                       type=node_type,
3686
                                       constant_result=node.constant_result)
3687
        if node_type.is_int and node_type.signed or \
3688 3689 3690 3691 3692
                isinstance(node.operand, ExprNodes.IntNode) and node_type.is_pyobject:
            return ExprNodes.IntNode(node.pos, value=_negate(node.operand.value),
                                     type=node_type,
                                     longness=node.operand.longness,
                                     constant_result=node.constant_result)
3693 3694
        return node

3695
    def _handle_UnaryPlusNode(self, node):
3696 3697
        if (node.operand.has_constant_result() and
                    node.constant_result == node.operand.constant_result):
3698 3699 3700
            return node.operand
        return node

3701 3702
    def visit_BoolBinopNode(self, node):
        self._calculate_const(node)
Stefan Behnel's avatar
Stefan Behnel committed
3703
        if not node.operand1.has_constant_result():
3704
            return node
Stefan Behnel's avatar
Stefan Behnel committed
3705
        if node.operand1.constant_result:
3706 3707 3708 3709
            if node.operator == 'and':
                return node.operand2
            else:
                return node.operand1
3710
        else:
3711 3712 3713 3714
            if node.operator == 'and':
                return node.operand1
            else:
                return node.operand2
3715

3716 3717 3718 3719
    def visit_BinopNode(self, node):
        self._calculate_const(node)
        if node.constant_result is ExprNodes.not_a_constant:
            return node
3720 3721
        if isinstance(node.constant_result, float):
            return node
3722 3723
        operand1, operand2 = node.operand1, node.operand2
        if not operand1.is_literal or not operand2.is_literal:
3724 3725 3726
            return node

        # now inject a new constant node with the calculated value
3727
        try:
3728
            type1, type2 = operand1.type, operand2.type
3729
            if type1 is None or type2 is None:
3730 3731 3732 3733
                return node
        except AttributeError:
            return node

3734
        if type1.is_numeric and type2.is_numeric:
3735
            widest_type = PyrexTypes.widest_numeric_type(type1, type2)
3736 3737
        else:
            widest_type = PyrexTypes.py_object_type
3738

3739
        target_class = self._widest_node_class(operand1, operand2)
3740 3741
        if target_class is None:
            return node
Stefan Behnel's avatar
Stefan Behnel committed
3742
        elif target_class is ExprNodes.BoolNode and node.operator in '+-//<<%**>>':
3743 3744
            # C arithmetic results in at least an int type
            target_class = ExprNodes.IntNode
Stefan Behnel's avatar
Stefan Behnel committed
3745
        elif target_class is ExprNodes.CharNode and node.operator in '+-//<<%**>>&|^':
3746 3747 3748 3749
            # C arithmetic results in at least an int type
            target_class = ExprNodes.IntNode

        if target_class is ExprNodes.IntNode:
3750 3751 3752 3753
            unsigned = getattr(operand1, 'unsigned', '') and \
                       getattr(operand2, 'unsigned', '')
            longness = "LL"[:max(len(getattr(operand1, 'longness', '')),
                                 len(getattr(operand2, 'longness', '')))]
3754
            new_node = ExprNodes.IntNode(pos=node.pos,
3755 3756 3757
                                         unsigned=unsigned, longness=longness,
                                         value=str(int(node.constant_result)),
                                         constant_result=int(node.constant_result))
3758 3759 3760 3761
            # IntNode is smart about the type it chooses, so we just
            # make sure we were not smarter this time
            if widest_type.is_pyobject or new_node.type.is_pyobject:
                new_node.type = PyrexTypes.py_object_type
3762
            else:
3763
                new_node.type = PyrexTypes.widest_numeric_type(widest_type, new_node.type)
3764
        else:
3765
            if target_class is ExprNodes.BoolNode:
3766 3767 3768 3769 3770 3771
                node_value = node.constant_result
            else:
                node_value = str(node.constant_result)
            new_node = target_class(pos=node.pos, type = widest_type,
                                    value = node_value,
                                    constant_result = node.constant_result)
3772 3773
        return new_node

3774 3775 3776 3777 3778 3779 3780 3781 3782 3783 3784 3785 3786 3787 3788 3789 3790 3791 3792 3793 3794 3795 3796 3797 3798
    def visit_AddNode(self, node):
        self._calculate_const(node)
        if node.constant_result is ExprNodes.not_a_constant:
            return node
        if node.operand1.is_string_literal and node.operand2.is_string_literal:
            # some people combine string literals with a '+'
            str1, str2 = node.operand1, node.operand2
            if isinstance(str1, ExprNodes.UnicodeNode) and isinstance(str2, ExprNodes.UnicodeNode):
                bytes_value = None
                if str1.bytes_value is not None and str2.bytes_value is not None:
                    if str1.bytes_value.encoding == str2.bytes_value.encoding:
                        bytes_value = BytesLiteral(str1.bytes_value + str2.bytes_value)
                        bytes_value.encoding = str1.bytes_value.encoding
                string_value = EncodedString(node.constant_result)
                return ExprNodes.UnicodeNode(
                    str1.pos, value=string_value, constant_result=node.constant_result, bytes_value=bytes_value)
            elif isinstance(str1, ExprNodes.BytesNode) and isinstance(str2, ExprNodes.BytesNode):
                if str1.value.encoding == str2.value.encoding:
                    bytes_value = BytesLiteral(node.constant_result)
                    bytes_value.encoding = str1.value.encoding
                    return ExprNodes.BytesNode(str1.pos, value=bytes_value, constant_result=node.constant_result)
            # all other combinations are rather complicated
            # to get right in Py2/3: encodings, unicode escapes, ...
        return self.visit_BinopNode(node)

3799
    def visit_MulNode(self, node):
3800
        self._calculate_const(node)
3801 3802
        if node.operand1.is_sequence_constructor:
            return self._calculate_constant_seq(node, node.operand1, node.operand2)
3803
        if isinstance(node.operand1, ExprNodes.IntNode) and \
3804 3805
                node.operand2.is_sequence_constructor:
            return self._calculate_constant_seq(node, node.operand2, node.operand1)
3806 3807
        return self.visit_BinopNode(node)

3808
    def _calculate_constant_seq(self, node, sequence_node, factor):
3809 3810 3811
        if factor.constant_result != 1 and sequence_node.args:
            if isinstance(factor.constant_result, (int, long)) and factor.constant_result <= 0:
                del sequence_node.args[:]
3812 3813 3814 3815 3816 3817 3818 3819 3820 3821 3822 3823 3824
                sequence_node.mult_factor = None
            elif sequence_node.mult_factor is not None:
                if (isinstance(factor.constant_result, (int, long)) and
                        isinstance(sequence_node.mult_factor.constant_result, (int, long))):
                    value = sequence_node.mult_factor.constant_result * factor.constant_result
                    sequence_node.mult_factor = ExprNodes.IntNode(
                        sequence_node.mult_factor.pos,
                        value=str(value), constant_result=value)
                else:
                    # don't know if we can combine the factors, so don't
                    return self.visit_BinopNode(node)
            else:
                sequence_node.mult_factor = factor
3825 3826
        return sequence_node

3827
    def visit_PrimaryCmpNode(self, node):
3828
        # calculate constant partial results in the comparison cascade
3829
        self.visitchildren(node, ['operand1'])
3830 3831 3832
        left_node = node.operand1
        cmp_node = node
        while cmp_node is not None:
3833
            self.visitchildren(cmp_node, ['operand2'])
3834 3835 3836 3837 3838 3839 3840 3841 3842 3843
            right_node = cmp_node.operand2
            cmp_node.constant_result = not_a_constant
            if left_node.has_constant_result() and right_node.has_constant_result():
                try:
                    cmp_node.calculate_cascaded_constant_result(left_node.constant_result)
                except (ValueError, TypeError, KeyError, IndexError, AttributeError, ArithmeticError):
                    pass  # ignore all 'normal' errors here => no constant result
            left_node = right_node
            cmp_node = cmp_node.cascade

3844 3845 3846 3847 3848
        if not node.cascade:
            if node.has_constant_result():
                return self._bool_node(node, node.constant_result)
            return node

3849 3850
        # collect partial cascades: [[value, CmpNode...], [value, CmpNode, ...], ...]
        cascades = [[node.operand1]]
3851
        final_false_result = []
3852 3853 3854 3855 3856

        def split_cascades(cmp_node):
            if cmp_node.has_constant_result():
                if not cmp_node.constant_result:
                    # False => short-circuit
3857
                    final_false_result.append(self._bool_node(cmp_node, False))
3858 3859 3860 3861 3862 3863 3864 3865 3866 3867 3868 3869 3870 3871 3872 3873 3874 3875 3876 3877 3878 3879 3880 3881 3882 3883 3884 3885 3886 3887 3888
                    return
                else:
                    # True => discard and start new cascade
                    cascades.append([cmp_node.operand2])
            else:
                # not constant => append to current cascade
                cascades[-1].append(cmp_node)
            if cmp_node.cascade:
                split_cascades(cmp_node.cascade)

        split_cascades(node)

        cmp_nodes = []
        for cascade in cascades:
            if len(cascade) < 2:
                continue
            cmp_node = cascade[1]
            pcmp_node = ExprNodes.PrimaryCmpNode(
                cmp_node.pos,
                operand1=cascade[0],
                operator=cmp_node.operator,
                operand2=cmp_node.operand2,
                constant_result=not_a_constant)
            cmp_nodes.append(pcmp_node)

            last_cmp_node = pcmp_node
            for cmp_node in cascade[2:]:
                last_cmp_node.cascade = cmp_node
                last_cmp_node = cmp_node
            last_cmp_node.cascade = None

3889
        if final_false_result:
3890
            # last cascade was constant False
3891
            cmp_nodes.append(final_false_result[0])
3892
        elif not cmp_nodes:
3893 3894 3895 3896 3897 3898 3899 3900 3901 3902 3903 3904 3905 3906
            # only constants, but no False result
            return self._bool_node(node, True)
        node = cmp_nodes[0]
        if len(cmp_nodes) == 1:
            if node.has_constant_result():
                return self._bool_node(node, node.constant_result)
        else:
            for cmp_node in cmp_nodes[1:]:
                node = ExprNodes.BoolBinopNode(
                    node.pos,
                    operand1=node,
                    operator='and',
                    operand2=cmp_node,
                    constant_result=not_a_constant)
3907
        return node
3908

3909 3910
    def visit_CondExprNode(self, node):
        self._calculate_const(node)
3911
        if not node.test.has_constant_result():
3912 3913 3914 3915 3916 3917
            return node
        if node.test.constant_result:
            return node.true_val
        else:
            return node.false_val

3918 3919 3920 3921 3922
    def visit_IfStatNode(self, node):
        self.visitchildren(node)
        # eliminate dead code based on constant condition results
        if_clauses = []
        for if_clause in node.if_clauses:
Stefan Behnel's avatar
Stefan Behnel committed
3923 3924 3925 3926 3927 3928 3929 3930
            condition = if_clause.condition
            if condition.has_constant_result():
                if condition.constant_result:
                    # always true => subsequent clauses can safely be dropped
                    node.else_clause = if_clause.body
                    break
                # else: false => drop clause
            else:
3931
                # unknown result => normal runtime evaluation
3932
                if_clauses.append(if_clause)
3933 3934 3935 3936
        if if_clauses:
            node.if_clauses = if_clauses
            return node
        elif node.else_clause:
3937
            return node.else_clause
3938 3939
        else:
            return Nodes.StatListNode(node.pos, stats=[])
3940

3941 3942 3943
    def visit_SliceIndexNode(self, node):
        self._calculate_const(node)
        # normalise start/stop values
3944 3945 3946 3947 3948 3949 3950 3951 3952 3953 3954
        if node.start is None or node.start.constant_result is None:
            start = node.start = None
        else:
            start = node.start.constant_result
        if node.stop is None or node.stop.constant_result is None:
            stop = node.stop = None
        else:
            stop = node.stop.constant_result
        # cut down sliced constant sequences
        if node.constant_result is not not_a_constant:
            base = node.base
3955
            if base.is_sequence_constructor and base.mult_factor is None:
3956 3957 3958
                base.args = base.args[start:stop]
                return base
            elif base.is_string_literal:
3959 3960 3961
                base = base.as_sliced_node(start, stop)
                if base is not None:
                    return base
3962 3963
        return node

3964 3965 3966 3967
    def visit_ComprehensionNode(self, node):
        self.visitchildren(node)
        if isinstance(node.loop, Nodes.StatListNode) and not node.loop.stats:
            # loop was pruned already => transform into literal
3968
            if node.type is Builtin.list_type:
3969 3970
                return ExprNodes.ListNode(
                    node.pos, args=[], constant_result=[])
3971
            elif node.type is Builtin.set_type:
3972 3973
                return ExprNodes.SetNode(
                    node.pos, args=[], constant_result=set())
3974
            elif node.type is Builtin.dict_type:
3975 3976
                return ExprNodes.DictNode(
                    node.pos, key_value_pairs=[], constant_result={})
3977 3978
        return node

3979 3980 3981
    def visit_ForInStatNode(self, node):
        self.visitchildren(node)
        sequence = node.iterator.sequence
3982 3983 3984 3985 3986 3987 3988
        if isinstance(sequence, ExprNodes.SequenceNode):
            if not sequence.args:
                if node.else_clause:
                    return node.else_clause
                else:
                    # don't break list comprehensions
                    return Nodes.StatListNode(node.pos, stats=[])
Stefan Behnel's avatar
Stefan Behnel committed
3989 3990 3991
            # iterating over a list literal? => tuples are more efficient
            if isinstance(sequence, ExprNodes.ListNode):
                node.iterator.sequence = sequence.as_tuple()
3992 3993
        return node

3994 3995
    def visit_WhileStatNode(self, node):
        self.visitchildren(node)
Stefan Behnel's avatar
Stefan Behnel committed
3996
        if node.condition and node.condition.has_constant_result():
3997
            if node.condition.constant_result:
3998
                node.condition = None
3999 4000 4001 4002 4003
                node.else_clause = None
            else:
                return node.else_clause
        return node

4004 4005
    def visit_ExprStatNode(self, node):
        self.visitchildren(node)
Stefan Behnel's avatar
Stefan Behnel committed
4006 4007 4008
        if not isinstance(node.expr, ExprNodes.ExprNode):
            # ParallelRangeTransform does this ...
            return node
4009 4010 4011 4012 4013
        # drop unused constant expressions
        if node.expr.has_constant_result():
            return None
        return node

4014 4015
    # in the future, other nodes can have their own handler method here
    # that can replace them with a constant result node
Stefan Behnel's avatar
Stefan Behnel committed
4016

4017
    visit_Node = Visitor.VisitorTransform.recurse_to_children
4018 4019


4020
class FinalOptimizePhase(Visitor.CythonTransform, Visitor.NodeRefCleanupMixin):
4021 4022
    """
    This visitor handles several commuting optimizations, and is run
4023 4024 4025 4026
    just before the C code generation phase.

    The optimizations currently implemented in this class are:
        - eliminate None assignment and refcounting for first assignment.
4027
        - isinstance -> typecheck for cdef types
Stefan Behnel's avatar
Stefan Behnel committed
4028
        - eliminate checks for None and/or types that became redundant after tree changes
4029
        - replace Python function calls that look like method calls by a faster PyMethodCallNode
4030
    """
4031
    def visit_SingleAssignmentNode(self, node):
4032 4033 4034 4035
        """Avoid redundant initialisation of local variables before their
        first assignment.
        """
        self.visitchildren(node)
4036 4037
        if node.first:
            lhs = node.lhs
4038
            lhs.lhs_of_first_assignment = True
4039
        return node
4040

4041
    def visit_SimpleCallNode(self, node):
4042 4043 4044
        """
        Replace generic calls to isinstance(x, type) by a more efficient type check.
        Replace likely Python method calls by a specialised PyMethodCallNode.
4045
        """
4046
        self.visitchildren(node)
4047 4048 4049
        function = node.function
        if function.type.is_cfunction and function.is_name:
            if function.name == 'isinstance' and len(node.args) == 2:
4050 4051
                type_arg = node.args[1]
                if type_arg.type.is_builtin_type and type_arg.type.name == 'type':
4052
                    cython_scope = self.context.cython_scope
4053 4054
                    function.entry = cython_scope.lookup('PyObject_TypeCheck')
                    function.type = function.entry.type
4055
                    PyTypeObjectPtr = PyrexTypes.CPtrType(cython_scope.lookup('PyTypeObject').type)
4056
                    node.args[1] = ExprNodes.CastNode(node.args[1], PyTypeObjectPtr)
4057 4058
        elif (self.current_directives.get("optimize.unpack_method_calls")
                and node.is_temp and function.type.is_pyobject):
4059 4060 4061 4062 4063
            # optimise simple Python methods calls
            if isinstance(node.arg_tuple, ExprNodes.TupleNode) and not (
                    node.arg_tuple.mult_factor or (node.arg_tuple.is_literal and node.arg_tuple.args)):
                # simple call, now exclude calls to objects that are definitely not methods
                may_be_a_method = True
4064 4065 4066
                if function.type is Builtin.type_type:
                    may_be_a_method = False
                elif function.is_name:
4067 4068 4069
                    if function.entry.is_builtin:
                        may_be_a_method = False
                    elif function.cf_state:
4070 4071
                        # local functions/classes are definitely not methods
                        non_method_nodes = (ExprNodes.PyCFunctionNode, ExprNodes.ClassNode, ExprNodes.Py3ClassNode)
4072
                        may_be_a_method = any(
4073
                            assignment.rhs and not isinstance(assignment.rhs, non_method_nodes)
4074 4075
                            for assignment in function.cf_state)
                if may_be_a_method:
4076 4077
                    node = self.replace(node, ExprNodes.PyMethodCallNode.from_node(
                        node, function=function, arg_tuple=node.arg_tuple, type=node.type))
4078
        return node
Stefan Behnel's avatar
Stefan Behnel committed
4079 4080 4081 4082 4083 4084 4085 4086 4087 4088 4089

    def visit_PyTypeTestNode(self, node):
        """Remove tests for alternatively allowed None values from
        type tests when we know that the argument cannot be None
        anyway.
        """
        self.visitchildren(node)
        if not node.notnone:
            if not node.arg.may_be_none():
                node.notnone = True
        return node
4090 4091 4092 4093 4094 4095 4096 4097 4098

    def visit_NoneCheckNode(self, node):
        """Remove None checks from expressions that definitely do not
        carry a None value.
        """
        self.visitchildren(node)
        if not node.arg.may_be_none():
            return node.arg
        return node
4099 4100 4101 4102 4103 4104 4105 4106 4107

class ConsolidateOverflowCheck(Visitor.CythonTransform):
    """
    This class facilitates the sharing of overflow checking among all nodes
    of a nested arithmetic expression.  For example, given the expression
    a*b + c, where a, b, and x are all possibly overflowing ints, the entire
    sequence will be evaluated and the overflow bit checked only at the end.
    """
    overflow_bit_node = None
4108

4109 4110 4111 4112 4113 4114 4115 4116 4117
    def visit_Node(self, node):
        if self.overflow_bit_node is not None:
            saved = self.overflow_bit_node
            self.overflow_bit_node = None
            self.visitchildren(node)
            self.overflow_bit_node = saved
        else:
            self.visitchildren(node)
        return node
4118

4119
    def visit_NumBinopNode(self, node):
4120
        if node.overflow_check and node.overflow_fold:
4121 4122 4123 4124 4125 4126 4127 4128 4129 4130 4131 4132
            top_level_overflow = self.overflow_bit_node is None
            if top_level_overflow:
                self.overflow_bit_node = node
            else:
                node.overflow_bit_node = self.overflow_bit_node
                node.overflow_check = False
            self.visitchildren(node)
            if top_level_overflow:
                self.overflow_bit_node = None
        else:
            self.visitchildren(node)
        return node