Optimize.py 140 KB
Newer Older
1
from Cython.Compiler import TypeSlots
2
from Cython.Compiler.ExprNodes import not_a_constant
3 4 5 6 7
import cython
cython.declare(UtilityCode=object, EncodedString=object, BytesLiteral=object,
               Nodes=object, ExprNodes=object, PyrexTypes=object, Builtin=object,
               UtilNodes=object, Naming=object)

8 9
import Nodes
import ExprNodes
10
import PyrexTypes
11
import Visitor
12 13
import Builtin
import UtilNodes
14
import Options
15
import Naming
16

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

22 23
import codecs

24
try:
25 26
    from __builtin__ import reduce
except ImportError:
27 28
    from functools import reduce

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

34
def load_c_utility(name):
35
    return UtilityCode.load_cached(name, "Optimize.c")
36

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

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

def is_common_value(a, b):
48 49
    a = unwrap_node(a)
    b = unwrap_node(b)
50 51 52
    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):
53
        return not a.is_py_attr and is_common_value(a.obj, b.obj) and a.attribute == b.attribute
54 55
    return False

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

61
class IterationTransform(Visitor.EnvTransform):
62 63 64
    """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
65
    - for-in-enumerate is replaced by an external counter variable
66
    - for-in-range loop becomes a plain C for loop
67
    """
68 69
    def visit_PrimaryCmpNode(self, node):
        if node.is_ptr_contains():
70

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

78 79 80 81 82 83 84 85 86 87 88
            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(
89
                pos,
90 91 92 93 94 95 96 97 98 99 100 101 102 103 104
                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))))
105
            for_loop = for_loop.analyse_expressions(self.current_env())
106
            for_loop = self.visit(for_loop)
107
            new_node = UtilNodes.TempResultFromStatNode(result_ref, for_loop)
108

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

        else:
            self.visitchildren(node)
            return node
116

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

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

130
        # C array (slice) iteration?
131
        if iterator.type.is_ptr or iterator.type.is_array:
132
            return self._transform_carray_iteration(node, iterator, reversed=reversed)
133 134 135 136
        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)
137 138 139

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

142 143 144 145 146 147 148
        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
149
        function = iterator.function
150
        # dict iteration?
151
        if function.is_attribute and not reversed and not arg_count:
152
            base_obj = iterator.self or function.obj
153
            method = function.attribute
154
            # in Py3, items() is equivalent to Py2's iteritems()
155
            is_safe_iter = self.global_scope().context.language_level >= 3
156 157 158 159 160 161 162 163 164 165

            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
166 167

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

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

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

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

Stefan Behnel's avatar
Stefan Behnel committed
200
        return node
201

202 203 204 205 206 207 208 209 210 211
    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
212 213 214 215 216 217 218 219 220
        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)
221

222 223 224 225 226 227 228 229 230 231
    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)
            ])

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

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

        slice_base_node = ExprNodes.PythonCapiCallNode(
243 244
            slice_node.pos, "PyBytes_AS_STRING",
            self.PyBytes_AS_STRING_func_type,
245 246 247 248
            args = [unpack_temp_node],
            is_temp = 0,
            )
        len_node = ExprNodes.PythonCapiCallNode(
249 250
            slice_node.pos, "PyBytes_GET_SIZE",
            self.PyBytes_GET_SIZE_func_type,
251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266
            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,
267 268
                    ),
                reversed = reversed))
269

270 271 272 273 274 275 276
    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)
        ])

277 278 279 280 281 282 283 284
    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')
285 286 287 288 289 290 291

    def _transform_unicode_iteration(self, node, slice_node, reversed=False):
        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)
292 293
        length_temp = UtilNodes.TempHandle(PyrexTypes.c_py_ssize_t_type)
        end_node = length_temp.ref(node.pos)
294 295 296 297 298 299
        if reversed:
            relation1, relation2 = '>', '>='
            start_node, end_node = end_node, start_node
        else:
            relation1, relation2 = '<=', '<'

300 301 302
        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)
303

304 305 306
        target_value = ExprNodes.PythonCapiCallNode(
            slice_node.pos, "__Pyx_PyUnicode_READ",
            self.PyUnicode_READ_func_type,
307 308 309
            args = [kind_temp.ref(slice_node.pos),
                    data_temp.ref(slice_node.pos),
                    counter_temp.ref(node.target.pos)],
310 311 312 313
            is_temp = False,
            )
        if target_value.type != node.target.type:
            target_value = target_value.coerce_to(node.target.type,
314
                                                  self.current_env())
315 316 317
        target_assign = Nodes.SingleAssignmentNode(
            pos = node.target.pos,
            lhs = node.target,
318
            rhs = target_value)
319 320 321 322 323 324 325
        body = Nodes.StatListNode(
            node.pos,
            stats = [target_assign, node.body])

        loop_node = Nodes.ForFromStatNode(
            node.pos,
            bound1=start_node, relation1=relation1,
326
            target=counter_temp.ref(node.target.pos),
327 328 329 330 331
            relation2=relation2, bound2=end_node,
            step=None, body=body,
            else_clause=node.else_clause,
            from_range=True)

332 333 334
        setup_node = Nodes.ExprStatNode(
            node.pos,
            expr = ExprNodes.PythonCapiCallNode(
335 336 337 338 339 340 341 342 343 344
                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),
                        ],
345 346
                is_temp = True,
                result_is_used = False,
347
                utility_code=UtilityCode.load_cached("unicode_iter", "Optimize.c"),
348 349 350
                ))
        return UtilNodes.LetNode(
            unpack_temp_node,
351 352 353
            UtilNodes.TempsBlockNode(
                node.pos, temps=[counter_temp, length_temp, data_temp, kind_temp],
                body=Nodes.StatListNode(node.pos, stats=[setup_node, loop_node])))
354

355
    def _transform_carray_iteration(self, node, slice_node, reversed=False):
356
        neg_step = False
357 358
        if isinstance(slice_node, ExprNodes.SliceIndexNode):
            slice_base = slice_node.base
359 360
            start = filter_none_node(slice_node.start)
            stop = filter_none_node(slice_node.stop)
361 362
            step = None
            if not stop:
363 364
                if not slice_base.type.is_pyobject:
                    error(slice_node.pos, "C array iteration requires known end index")
365
                return node
366

367
        elif isinstance(slice_node, ExprNodes.IndexNode):
368
            assert isinstance(slice_node.index, ExprNodes.SliceNode)
369 370
            slice_base = slice_node.base
            index = slice_node.index
371 372 373
            start = filter_none_node(index.start)
            stop = filter_none_node(index.stop)
            step = filter_none_node(index.step)
374
            if step:
375
                if not isinstance(step.constant_result, (int,long)) \
376 377 378
                       or step.constant_result == 0 \
                       or step.constant_result > 0 and not stop \
                       or step.constant_result < 0 and not start:
379 380
                    if not slice_base.type.is_pyobject:
                        error(step.pos, "C array iteration requires known step size and end index")
381 382 383
                    return node
                else:
                    # step sign is handled internally by ForFromStatNode
384 385 386 387
                    step_value = step.constant_result
                    if reversed:
                        step_value = -step_value
                    neg_step = step_value < 0
388
                    step = ExprNodes.IntNode(step.pos, type=PyrexTypes.c_py_ssize_t_type,
389 390
                                             value=str(abs(step_value)),
                                             constant_result=abs(step_value))
391

392 393
        elif slice_node.type.is_array:
            if slice_node.type.size is None:
Stefan Behnel's avatar
Stefan Behnel committed
394
                error(slice_node.pos, "C array iteration requires known end index")
395
                return node
396 397 398
            slice_base = slice_node
            start = None
            stop = ExprNodes.IntNode(
399 400
                slice_node.pos, value=str(slice_node.type.size),
                type=PyrexTypes.c_py_ssize_t_type, constant_result=slice_node.type.size)
401
            step = None
402

403
        else:
404
            if not slice_node.type.is_pyobject:
405
                error(slice_node.pos, "C array iteration requires known end index")
406 407
            return node

408
        if start:
409
            start = start.coerce_to(PyrexTypes.c_py_ssize_t_type, self.current_env())
410
        if stop:
411
            stop = stop.coerce_to(PyrexTypes.c_py_ssize_t_type, self.current_env())
412 413 414 415 416 417 418
        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
419

420 421 422 423 424 425 426
        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

427 428 429
        ptr_type = slice_base.type
        if ptr_type.is_array:
            ptr_type = ptr_type.element_ptr_type()
430
        carray_ptr = slice_base.coerce_to_simple(self.current_env())
431

432
        if start and start.constant_result != 0:
433 434 435 436 437
            start_ptr_node = ExprNodes.AddNode(
                start.pos,
                operand1=carray_ptr,
                operator='+',
                operand2=start,
438
                type=ptr_type)
439
        else:
440
            start_ptr_node = carray_ptr
441

442 443 444 445 446 447 448
        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
449
                ).coerce_to_simple(self.current_env())
450 451
        else:
            stop_ptr_node = ExprNodes.CloneNode(carray_ptr)
452

453
        counter = UtilNodes.TempHandle(ptr_type)
454 455
        counter_temp = counter.ref(node.target.pos)

456
        if slice_base.type.is_string and node.target.type.is_pyobject:
457
            # special case: char* -> bytes
458 459
            target_value = ExprNodes.SliceIndexNode(
                node.target.pos,
460 461 462 463 464 465 466
                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,
467 468
                type=Builtin.bytes_type,
                is_temp=1)
469 470 471
        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
472 473 474
        else:
            target_value = ExprNodes.IndexNode(
                node.target.pos,
475 476 477 478
                index=ExprNodes.IntNode(node.target.pos, value='0',
                                        constant_result=0,
                                        type=PyrexTypes.c_int_type),
                base=counter_temp,
479
                is_buffer_access=False,
480
                type=ptr_type.base_type)
481 482 483

        if target_value.type != node.target.type:
            target_value = target_value.coerce_to(node.target.type,
484
                                                  self.current_env())
485 486 487 488 489 490 491 492 493 494

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

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

495 496
        relation1, relation2 = self._find_for_from_node_relations(neg_step, reversed)

497 498
        for_node = Nodes.ForFromStatNode(
            node.pos,
499
            bound1=start_ptr_node, relation1=relation1,
500
            target=counter_temp,
501
            relation2=relation2, bound2=stop_ptr_node,
502 503 504 505 506 507 508 509
            step=step, body=body,
            else_clause=node.else_clause,
            from_range=True)

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

510 511 512 513 514 515
    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
516
        elif len(args) > 2:
517
            error(enumerate_function.pos,
518
                  "enumerate() takes at most 2 arguments")
519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535
            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

536
        if len(args) == 2:
537
            start = unwrap_coerced_node(args[1]).coerce_to(counter_type, self.current_env())
538 539 540 541 542 543 544
        else:
            start = ExprNodes.IntNode(enumerate_function.pos,
                                      value='0',
                                      type=counter_type,
                                      constant_result=0)
        temp = UtilNodes.LetRefNode(start)

545 546
        inc_expression = ExprNodes.AddNode(
            enumerate_function.pos,
547
            operand1 = temp,
548
            operand2 = ExprNodes.IntNode(node.pos, value='1',
549 550
                                         type=counter_type,
                                         constant_result=1),
551 552
            operator = '+',
            type = counter_type,
Stefan Behnel's avatar
Stefan Behnel committed
553
            #inplace = True,   # not worth using in-place operation for Py ints
554 555 556
            is_temp = counter_type.is_pyobject
            )

557 558 559 560
        loop_body = [
            Nodes.SingleAssignmentNode(
                pos = enumerate_target.pos,
                lhs = enumerate_target,
561
                rhs = temp),
562 563
            Nodes.SingleAssignmentNode(
                pos = enumerate_target.pos,
564
                lhs = temp,
565 566
                rhs = inc_expression)
            ]
567

568 569 570 571 572 573 574
        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)
575 576

        node.target = iterable_target
577
        node.item = node.item.coerce_to(iterable_target.type, self.current_env())
578
        node.iterator.sequence = args[0]
579 580

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

583 584 585 586 587 588 589 590 591 592 593
    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 '<=', '<'
594

595
    def _transform_range_iteration(self, node, range_function, reversed=False):
596 597 598 599
        args = range_function.arg_tuple.args
        if len(args) < 3:
            step_pos = range_function.pos
            step_value = 1
600 601
            step = ExprNodes.IntNode(step_pos, value='1',
                                     constant_result=1)
602 603 604
        else:
            step = args[2]
            step_pos = step.pos
605
            if not isinstance(step.constant_result, (int, long)):
606 607
                # cannot determine step direction
                return node
608 609 610
            step_value = step.constant_result
            if step_value == 0:
                # will lead to an error elsewhere
611
                return node
612 613 614
            if reversed and step_value not in (1, -1):
                # FIXME: currently broken - requires calculation of the correct bounds
                return node
615
            if not isinstance(step, ExprNodes.IntNode):
616 617
                step = ExprNodes.IntNode(step_pos, value=str(step_value),
                                         constant_result=step_value)
618 619

        if len(args) == 1:
620 621
            bound1 = ExprNodes.IntNode(range_function.pos, value='0',
                                       constant_result=0)
622
            bound2 = args[0].coerce_to_integer(self.current_env())
623
        else:
624 625
            bound1 = args[0].coerce_to_integer(self.current_env())
            bound2 = args[1].coerce_to_integer(self.current_env())
626

627 628
        relation1, relation2 = self._find_for_from_node_relations(step_value < 0, reversed)

629 630 631 632 633 634 635 636 637 638
        if reversed:
            bound1, bound2 = bound2, bound1
            if step_value < 0:
                step_value = -step_value
        else:
            if step_value < 0:
                step_value = -step_value

        step.value = str(step_value)
        step.constant_result = step_value
639
        step = step.coerce_to_integer(self.current_env())
640

641
        if not bound2.is_literal:
642 643 644 645 646 647
            # stop bound must be immutable => keep it in a temp var
            bound2_is_temp = True
            bound2 = UtilNodes.LetRefNode(bound2)
        else:
            bound2_is_temp = False

648 649 650 651 652 653 654
        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
655
            from_range=True)
656 657 658 659

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

660 661
        return for_node

662
    def _transform_dict_iteration(self, node, dict_obj, method, keys, values):
663
        temps = []
664 665 666
        temp = UtilNodes.TempHandle(PyrexTypes.py_object_type)
        temps.append(temp)
        dict_temp = temp.ref(dict_obj.pos)
667 668
        temp = UtilNodes.TempHandle(PyrexTypes.c_py_ssize_t_type)
        temps.append(temp)
669
        pos_temp = temp.ref(node.pos)
670

671
        key_target = value_target = tuple_target = None
672 673 674 675 676
        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
677
                    # unusual case that may or may not lead to an error
678 679 680
                    return node
            else:
                tuple_target = node.target
681 682
        elif keys:
            key_target = node.target
683
        else:
684
            value_target = node.target
685 686 687 688 689 690 691 692 693 694

        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)
695 696 697 698 699 700 701 702 703 704 705 706 707 708
        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)
709
        iter_next_node = iter_next_node.analyse_expressions(self.current_env())
710 711 712 713 714 715 716 717 718 719 720 721
        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")
722

723 724 725
        def flag_node(value):
            value = value and 1 or 0
            return ExprNodes.IntNode(node.pos, value=str(value), constant_result=value)
726 727

        result_code = [
728
            Nodes.SingleAssignmentNode(
Stefan Behnel's avatar
Stefan Behnel committed
729
                node.pos,
730
                lhs = pos_temp,
731 732
                rhs = ExprNodes.IntNode(node.pos, value='0',
                                        constant_result=0)),
733
            Nodes.SingleAssignmentNode(
Stefan Behnel's avatar
Stefan Behnel committed
734
                dict_obj.pos,
735 736 737 738 739 740
                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"),
741
                    args = [dict_obj, flag_node(dict_obj.type is Builtin.dict_type),
742 743 744
                            method_node, dict_len_temp_addr, is_dict_temp_addr,
                            ],
                    is_temp=True,
745 746
                )),
            Nodes.WhileStatNode(
Stefan Behnel's avatar
Stefan Behnel committed
747
                node.pos,
748
                condition = None,
749 750 751 752 753 754 755 756
                body = body,
                else_clause = node.else_clause
                )
            ]

        return UtilNodes.TempsBlockNode(
            node.pos, temps=temps,
            body=Nodes.StatListNode(
757
                node.pos,
758 759 760
                stats = result_code
                ))

761 762 763 764 765 766 767 768 769
    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),
            ])

770

771 772
class SwitchTransform(Visitor.VisitorTransform):
    """
773
    This transformation tries to turn long if statements into C switch statements.
774
    The requirement is that every clause be an (or of) var == value, where the var
775
    is common among all clauses and both var and value are ints.
776
    """
777 778 779
    NO_MATCH = (None, None, None)

    def extract_conditions(self, cond, allow_not_in):
780
        while True:
781 782
            if isinstance(cond, (ExprNodes.CoerceToTempNode,
                                 ExprNodes.CoerceToBooleanNode)):
783 784 785 786 787 788 789 790
                cond = cond.arg
            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
791

792
        if isinstance(cond, ExprNodes.PrimaryCmpNode):
793 794 795 796 797 798 799 800 801 802 803 804 805 806 807
            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():
808 809 810 811 812 813 814 815 816 817 818 819 820 821 822 823 824 825 826 827 828 829 830 831 832 833 834 835 836 837
                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]
        elif isinstance(cond, ExprNodes.BoolBinopNode):
            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

838 839
    def extract_in_string_conditions(self, string_literal):
        if isinstance(string_literal, ExprNodes.UnicodeNode):
840
            charvals = list(map(ord, set(string_literal.value)))
841 842 843 844 845 846 847 848 849
            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
850 851
            characters = list(set([ characters[i:i+1] for i in range(len(characters)) ]))
            characters.sort()
852 853 854 855
            return [ ExprNodes.CharNode(string_literal.pos, value=charval,
                                        constant_result=charval)
                     for charval in characters ]

856 857
    def extract_common_conditions(self, common_var, condition, allow_not_in):
        not_in, var, conditions = self.extract_conditions(condition, allow_not_in)
858
        if var is None:
859
            return self.NO_MATCH
860
        elif common_var is not None and not is_common_value(var, common_var):
861
            return self.NO_MATCH
862
        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]):
863 864 865 866 867 868 869
            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:
870
            if value.has_constant_result():
871 872 873 874 875 876
                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
877 878 879 880 881 882
                try:
                    if value.entry.cname in seen:
                        return True
                except AttributeError:
                    return True  # play safe
                seen.add(value.entry.cname)
883
        return False
884

885 886 887 888
    def visit_IfStatNode(self, node):
        common_var = None
        cases = []
        for if_clause in node.if_clauses:
889 890
            _, common_var, conditions = self.extract_common_conditions(
                common_var, if_clause.condition, False)
891
            if common_var is None:
892
                self.visitchildren(node)
893
                return node
894 895 896 897
            cases.append(Nodes.SwitchCaseNode(pos = if_clause.pos,
                                              conditions = conditions,
                                              body = if_clause.body))

898 899 900
        condition_values = [
            cond for case in cases for cond in case.conditions]
        if len(condition_values) < 2:
901 902
            self.visitchildren(node)
            return node
903
        if self.has_duplicate_values(condition_values):
904
            self.visitchildren(node)
905
            return node
906

Robert Bradshaw's avatar
Robert Bradshaw committed
907
        common_var = unwrap_node(common_var)
908 909 910 911 912 913 914
        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):
915 916 917 918 919 920
        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)
921
            return node
922 923 924
        return self.build_simple_switch_statement(
            node, common_var, conditions, not_in,
            node.true_val, node.false_val)
925 926

    def visit_BoolBinopNode(self, node):
927 928 929 930 931 932
        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)
933 934
            return node

935 936
        return self.build_simple_switch_statement(
            node, common_var, conditions, not_in,
937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952
            ExprNodes.BoolNode(node.pos, value=True, constant_result=True),
            ExprNodes.BoolNode(node.pos, value=False, constant_result=False))

    def visit_PrimaryCmpNode(self, node):
        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))
953 954 955

    def build_simple_switch_statement(self, node, common_var, conditions,
                                      not_in, true_val, false_val):
956 957 958 959
        result_ref = UtilNodes.ResultRefNode(node)
        true_body = Nodes.SingleAssignmentNode(
            node.pos,
            lhs = result_ref,
960
            rhs = true_val,
961 962 963 964
            first = True)
        false_body = Nodes.SingleAssignmentNode(
            node.pos,
            lhs = result_ref,
965
            rhs = false_val,
966 967
            first = True)

968 969 970
        if not_in:
            true_body, false_body = false_body, true_body

971 972 973 974 975 976 977 978 979
        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)
980 981 982 983 984 985 986 987 988 989 990 991 992
        replacement = UtilNodes.TempResultFromStatNode(result_ref, switch_node)
        return replacement

    def visit_EvalWithTempExprNode(self, node):
        # 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
993

994
    visit_Node = Visitor.VisitorTransform.recurse_to_children
995

996

997
class FlattenInListTransform(Visitor.VisitorTransform, SkipDeclarations):
998 999
    """
    This transformation flattens "x in [val1, ..., valn]" into a sequential list
1000
    of comparisons.
1001
    """
1002

1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014
    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
1015

1016 1017 1018
        if not isinstance(node.operand2, (ExprNodes.TupleNode,
                                          ExprNodes.ListNode,
                                          ExprNodes.SetNode)):
Stefan Behnel's avatar
Stefan Behnel committed
1019
            return node
1020

Stefan Behnel's avatar
Stefan Behnel committed
1021 1022
        args = node.operand2.args
        if len(args) == 0:
1023 1024 1025
            constant_result = node.operator == 'not_in'
            return ExprNodes.BoolNode(pos = node.pos, value = constant_result,
                                      constant_result = constant_result)
1026

1027
        lhs = UtilNodes.ResultRefNode(node.operand1)
Stefan Behnel's avatar
Stefan Behnel committed
1028 1029

        conds = []
1030
        temps = []
Stefan Behnel's avatar
Stefan Behnel committed
1031
        for arg in args:
1032 1033 1034 1035 1036 1037 1038 1039 1040
            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:
1041 1042 1043
                # must evaluate all non-simple RHS before doing the comparisons
                arg = UtilNodes.LetRefNode(arg)
                temps.append(arg)
Stefan Behnel's avatar
Stefan Behnel committed
1044 1045 1046 1047 1048 1049 1050
            cond = ExprNodes.PrimaryCmpNode(
                                pos = node.pos,
                                operand1 = lhs,
                                operator = eq_or_neq,
                                operand2 = arg,
                                cascade = None)
            conds.append(ExprNodes.TypecastNode(
1051
                                pos = node.pos,
Stefan Behnel's avatar
Stefan Behnel committed
1052 1053 1054 1055
                                operand = cond,
                                type = PyrexTypes.c_bint_type))
        def concat(left, right):
            return ExprNodes.BoolBinopNode(
1056
                                pos = node.pos,
Stefan Behnel's avatar
Stefan Behnel committed
1057 1058 1059 1060
                                operator = conjunction,
                                operand1 = left,
                                operand2 = right)

1061
        condition = reduce(concat, conds)
1062 1063 1064 1065
        new_node = UtilNodes.EvalWithTempExprNode(lhs, condition)
        for temp in temps[::-1]:
            new_node = UtilNodes.EvalWithTempExprNode(temp, new_node)
        return new_node
1066

1067
    visit_Node = Visitor.VisitorTransform.recurse_to_children
1068 1069


1070 1071 1072 1073 1074 1075
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
1076 1077 1078
        """
        Parallel swap assignments like 'a,b = b,a' are safe.
        """
1079 1080 1081 1082
        left_names, right_names = [], []
        left_indices, right_indices = [], []
        temps = []

1083 1084
        for stat in node.stats:
            if isinstance(stat, Nodes.SingleAssignmentNode):
1085 1086
                if not self._extract_operand(stat.lhs, left_names,
                                             left_indices, temps):
1087
                    return node
1088 1089
                if not self._extract_operand(stat.rhs, right_names,
                                             right_indices, temps):
1090
                    return node
1091 1092 1093
            elif isinstance(stat, Nodes.CascadedAssignmentNode):
                # FIXME
                return node
1094 1095 1096
            else:
                return node

1097 1098
        if left_names or right_names:
            # lhs/rhs names must be a non-redundant permutation
1099 1100
            lnames = [ path for path, n in left_names ]
            rnames = [ path for path, n in right_names ]
1101 1102 1103
            if set(lnames) != set(rnames):
                return node
            if len(set(lnames)) != len(right_names):
1104 1105
                return node

1106 1107 1108 1109 1110 1111 1112 1113 1114 1115 1116 1117 1118 1119 1120
        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)
1121

1122 1123 1124 1125 1126 1127 1128
            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
1129 1130
            return node

1131 1132 1133 1134
        temp_args = [t.arg for t in temps]
        for temp in temps:
            temp.use_managed_ref = False

1135
        for _, name_node in left_names + right_names:
1136 1137 1138 1139 1140
            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
1141 1142 1143

        return node

1144 1145 1146 1147 1148 1149 1150
    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
1151 1152 1153 1154
        name_path = []
        obj_node = node
        while isinstance(obj_node, ExprNodes.AttributeNode):
            if obj_node.is_py_attr:
1155
                return False
1156 1157 1158 1159 1160
            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) )
1161 1162 1163 1164 1165 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179 1180 1181 1182 1183 1184
        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)

1185

1186 1187 1188 1189 1190 1191 1192
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
1193 1194 1195

    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
1196
    after type analysis.
1197
    """
1198 1199
    # only intercept on call nodes
    visit_Node = Visitor.VisitorTransform.recurse_to_children
1200

1201 1202 1203 1204 1205 1206 1207
    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)

1208
    def visit_GeneralCallNode(self, node):
1209
        self.visitchildren(node)
1210
        function = node.function
1211
        if not self._function_is_builtin_name(function):
1212 1213 1214 1215
            return node
        arg_tuple = node.positional_args
        if not isinstance(arg_tuple, ExprNodes.TupleNode):
            return node
1216
        args = arg_tuple.args
1217
        return self._dispatch_to_handler(
1218
            node, function, args, node.keyword_args)
1219

1220 1221 1222
    def _function_is_builtin_name(self, function):
        if not function.is_name:
            return False
1223
        env = self.current_env()
1224
        entry = env.lookup(function.name)
1225
        if entry is not env.builtin_scope().lookup_here(function.name):
1226
            return False
1227
        # if entry is None, it's at least an undeclared name, so likely builtin
1228 1229 1230 1231 1232 1233 1234 1235 1236 1237 1238 1239 1240 1241 1242 1243 1244 1245 1246 1247 1248 1249 1250 1251 1252 1253 1254 1255 1256 1257 1258 1259 1260 1261 1262 1263 1264 1265
        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

1266
    def _handle_simple_function_float(self, node, pos_args):
Stefan Behnel's avatar
Stefan Behnel committed
1267
        if not pos_args:
1268 1269 1270
            return ExprNodes.FloatNode(node.pos, value='0.0')
        if len(pos_args) > 1:
            self._error_wrong_arg_count('float', node, pos_args, 1)
1271 1272 1273
        arg_type = getattr(pos_args[0], 'type', None)
        if arg_type in (PyrexTypes.c_double_type, Builtin.float_type):
            return pos_args[0]
1274 1275
        return node

1276 1277 1278
    class YieldNodeCollector(Visitor.TreeVisitor):
        def __init__(self):
            Visitor.TreeVisitor.__init__(self)
1279
            self.yield_stat_nodes = {}
1280 1281 1282
            self.yield_nodes = []

        visit_Node = Visitor.TreeVisitor.visitchildren
1283 1284
        # XXX: disable inlining while it's not back supported
        def __visit_YieldExprNode(self, node):
1285 1286 1287
            self.yield_nodes.append(node)
            self.visitchildren(node)

1288
        def __visit_ExprStatNode(self, node):
1289 1290 1291 1292
            self.visitchildren(node)
            if node.expr in self.yield_nodes:
                self.yield_stat_nodes[node.expr] = node

1293 1294 1295 1296 1297 1298
        def __visit_GeneratorExpressionNode(self, node):
            # enable when we support generic generator expressions
            #
            # everything below this node is out of scope
            pass

1299
    def _find_single_yield_expression(self, node):
1300 1301 1302
        collector = self.YieldNodeCollector()
        collector.visitchildren(node)
        if len(collector.yield_nodes) != 1:
1303 1304
            return None, None
        yield_node = collector.yield_nodes[0]
1305 1306 1307 1308
        try:
            return (yield_node.arg, collector.yield_stat_nodes[yield_node])
        except KeyError:
            return None, None
1309

1310 1311 1312 1313 1314 1315 1316 1317 1318 1319 1320 1321 1322 1323 1324 1325 1326 1327 1328 1329 1330 1331 1332 1333 1334 1335 1336 1337 1338 1339 1340 1341 1342 1343 1344 1345 1346 1347 1348 1349 1350 1351 1352 1353 1354
    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
1355 1356
        gen_expr_node = pos_args[0]
        loop_node = gen_expr_node.loop
1357 1358
        yield_expression, yield_stat_node = self._find_single_yield_expression(loop_node)
        if yield_expression is None:
1359 1360 1361 1362 1363 1364 1365
            return node

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

1366
        result_ref = UtilNodes.ResultRefNode(pos=node.pos, type=PyrexTypes.c_bint_type)
1367
        test_node = Nodes.IfStatNode(
1368
            yield_expression.pos,
1369 1370
            else_clause = None,
            if_clauses = [ Nodes.IfClauseNode(
1371
                yield_expression.pos,
1372 1373 1374 1375 1376 1377 1378
                condition = condition,
                body = Nodes.StatListNode(
                    node.pos,
                    stats = [
                        Nodes.SingleAssignmentNode(
                            node.pos,
                            lhs = result_ref,
1379
                            rhs = ExprNodes.BoolNode(yield_expression.pos, value = is_any,
1380 1381 1382 1383 1384 1385 1386 1387 1388
                                                     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,
1389
                Nodes.BreakStatNode(yield_expression.pos)
1390
                ])
1391
            next_loop.else_clause = Nodes.ContinueStatNode(yield_expression.pos)
1392 1393 1394 1395
            loop = next_loop
        loop_node.else_clause = Nodes.SingleAssignmentNode(
            node.pos,
            lhs = result_ref,
1396
            rhs = ExprNodes.BoolNode(yield_expression.pos, value = not is_any,
1397 1398
                                     constant_result = not is_any))

1399
        Visitor.recursively_replace_node(loop_node, yield_stat_node, test_node)
1400

1401 1402
        return ExprNodes.InlinedGeneratorExpressionNode(
            gen_expr_node.pos, loop = loop_node, result_node = result_ref,
1403
            expr_scope = gen_expr_node.expr_scope, orig_func = is_any and 'any' or 'all')
1404

1405
    def _handle_simple_function_sorted(self, node, pos_args):
1406 1407 1408 1409 1410
        """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.
1411 1412 1413
        """
        if len(pos_args) != 1:
            return node
1414
        if isinstance(pos_args[0], ExprNodes.ComprehensionNode) \
1415
               and pos_args[0].type is Builtin.list_type:
1416 1417 1418 1419 1420 1421 1422 1423
            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
1424

1425
            append_node = ExprNodes.ComprehensionAppendNode(
1426
                yield_expression.pos, expr = yield_expression)
1427

1428
            Visitor.recursively_replace_node(loop_node, yield_stat_node, append_node)
1429

1430
            listcomp_node = ExprNodes.ComprehensionNode(
1431
                gen_expr_node.pos, loop = loop_node,
1432 1433 1434
                append = append_node, type = Builtin.list_type,
                expr_scope = gen_expr_node.expr_scope,
                has_local_scope = True)
1435
            append_node.target = listcomp_node
1436 1437
        else:
            return node
1438

1439 1440
        result_node = UtilNodes.ResultRefNode(
            pos = loop_node.pos, type = Builtin.list_type, may_hold_none=False)
1441 1442 1443 1444 1445 1446 1447 1448 1449 1450 1451 1452 1453 1454 1455 1456 1457
        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 ]))

1458
    def _handle_simple_function_sum(self, node, pos_args):
Stefan Behnel's avatar
Stefan Behnel committed
1459 1460
        """Transform sum(genexpr) into an equivalent inlined aggregation loop.
        """
1461 1462
        if len(pos_args) not in (1,2):
            return node
1463 1464
        if not isinstance(pos_args[0], (ExprNodes.GeneratorExpressionNode,
                                        ExprNodes.ComprehensionNode)):
1465 1466 1467 1468
            return node
        gen_expr_node = pos_args[0]
        loop_node = gen_expr_node.loop

1469 1470 1471 1472 1473 1474 1475 1476 1477 1478 1479 1480 1481 1482
        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
        else: # ComprehensionNode
            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
1483 1484 1485 1486 1487 1488 1489 1490

        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(
1491
            yield_expression.pos,
1492 1493 1494 1495
            lhs = result_ref,
            rhs = ExprNodes.binop_node(node.pos, '+', result_ref, yield_expression)
            )

1496
        Visitor.recursively_replace_node(loop_node, yield_stat_node, add_node)
1497 1498 1499 1500 1501 1502 1503 1504 1505 1506 1507 1508 1509 1510

        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,
1511 1512
            expr_scope = gen_expr_node.expr_scope, orig_func = 'sum',
            has_local_scope = gen_expr_node.has_local_scope)
1513

1514 1515 1516 1517 1518 1519 1520 1521 1522 1523 1524 1525 1526
    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:
            # leave this to Python
            return node

1527
        cascaded_nodes = list(map(UtilNodes.ResultRefNode, args[1:]))
1528 1529 1530 1531 1532 1533 1534 1535 1536 1537 1538 1539 1540 1541 1542 1543 1544 1545 1546 1547 1548 1549

        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

1550
    def _DISABLED_handle_simple_function_tuple(self, node, pos_args):
Stefan Behnel's avatar
Stefan Behnel committed
1551
        if not pos_args:
1552 1553 1554 1555 1556 1557 1558 1559
            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.
1560
        result = self._transform_list_set_genexpr(node, pos_args, Builtin.list_type)
1561 1562 1563 1564
        if result is not node:
            return ExprNodes.AsTupleNode(node.pos, arg=result)
        return node

1565
    def _handle_simple_function_list(self, node, pos_args):
Stefan Behnel's avatar
Stefan Behnel committed
1566
        if not pos_args:
1567
            return ExprNodes.ListNode(node.pos, args=[], constant_result=[])
1568
        return self._transform_list_set_genexpr(node, pos_args, Builtin.list_type)
1569 1570

    def _handle_simple_function_set(self, node, pos_args):
Stefan Behnel's avatar
Stefan Behnel committed
1571
        if not pos_args:
1572
            return ExprNodes.SetNode(node.pos, args=[], constant_result=set())
1573
        return self._transform_list_set_genexpr(node, pos_args, Builtin.set_type)
1574

1575
    def _transform_list_set_genexpr(self, node, pos_args, target_type):
1576 1577 1578 1579 1580 1581 1582 1583 1584
        """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

1585 1586
        yield_expression, yield_stat_node = self._find_single_yield_expression(loop_node)
        if yield_expression is None:
1587 1588 1589
            return node

        append_node = ExprNodes.ComprehensionAppendNode(
1590
            yield_expression.pos,
1591
            expr = yield_expression)
1592

1593
        Visitor.recursively_replace_node(loop_node, yield_stat_node, append_node)
1594

1595
        comp = ExprNodes.ComprehensionNode(
1596 1597 1598 1599 1600
            node.pos,
            has_local_scope = True,
            expr_scope = gen_expr_node.expr_scope,
            loop = loop_node,
            append = append_node,
1601 1602 1603
            type = target_type)
        append_node.target = comp
        return comp
1604 1605 1606 1607 1608 1609 1610 1611 1612 1613 1614 1615 1616

    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

1617 1618
        yield_expression, yield_stat_node = self._find_single_yield_expression(loop_node)
        if yield_expression is None:
1619 1620 1621 1622 1623 1624 1625 1626
            return node

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

        append_node = ExprNodes.DictComprehensionAppendNode(
1627
            yield_expression.pos,
1628
            key_expr = yield_expression.args[0],
1629
            value_expr = yield_expression.args[1])
1630

1631
        Visitor.recursively_replace_node(loop_node, yield_stat_node, append_node)
1632 1633 1634 1635 1636 1637 1638

        dictcomp = ExprNodes.ComprehensionNode(
            node.pos,
            has_local_scope = True,
            expr_scope = gen_expr_node.expr_scope,
            loop = loop_node,
            append = append_node,
1639
            type = Builtin.dict_type)
1640 1641 1642
        append_node.target = dictcomp
        return dictcomp

1643 1644 1645 1646 1647 1648 1649 1650 1651 1652 1653 1654
    # 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

1655 1656

class InlineDefNodeCalls(Visitor.NodeRefCleanupMixin, Visitor.EnvTransform):
1657 1658
    visit_Node = Visitor.VisitorTransform.recurse_to_children

1659
    def get_constant_value_node(self, name_node):
1660 1661 1662
        if name_node.cf_state is None:
            return None
        if name_node.cf_state.cf_is_null:
1663 1664 1665 1666 1667 1668
            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
1669
        return entry.cf_assignments[0].rhs
1670

1671 1672 1673 1674 1675 1676 1677
    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
1678
        function = self.get_constant_value_node(function_name)
1679 1680 1681 1682 1683 1684
        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():
1685
            return self.replace(node, inlined)
1686 1687
        return node

1688

1689
class OptimizeBuiltinCalls(Visitor.MethodDispatcherTransform):
Stefan Behnel's avatar
Stefan Behnel committed
1690
    """Optimize some common methods calls and instantiation patterns
1691 1692 1693 1694 1695
    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.
1696
    """
1697 1698
    ### cleanup to avoid redundant coercions to/from Python types

1699 1700 1701
    def _visit_PyTypeTestNode(self, node):
        # disabled - appears to break assignments in some cases, and
        # also drops a None check, which might still be required
1702 1703 1704 1705 1706 1707 1708 1709
        """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

1710 1711 1712
    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
1713 1714 1715 1716 1717 1718 1719 1720
        """
        Drop redundant type casts.
        """
        self.visitchildren(node)
        if node.type == node.operand.type:
            return node.operand
        return node

1721 1722 1723 1724 1725 1726 1727 1728 1729
    def visit_ExprStatNode(self, node):
        """
        Drop useless coercions.
        """
        self.visitchildren(node)
        if isinstance(node.expr, ExprNodes.CoerceToPyTypeNode):
            node.expr = node.expr.arg
        return node

1730 1731 1732 1733 1734
    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
1735 1736
        if isinstance(arg, ExprNodes.PyTypeTestNode):
            arg = arg.arg
1737 1738
        if isinstance(arg, ExprNodes.CoerceToPyTypeNode):
            if arg.type in (PyrexTypes.py_object_type, Builtin.bool_type):
1739
                return arg.arg.coerce_to_boolean(self.current_env())
1740 1741
        return node

1742 1743 1744 1745 1746 1747 1748 1749 1750 1751 1752 1753 1754 1755
    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:
1756
                return arg.coerce_to(node.type, self.current_env())
1757 1758
        if isinstance(arg, ExprNodes.PyTypeTestNode):
            arg = arg.arg
1759 1760 1761 1762
        if isinstance(arg, ExprNodes.CoerceToPyTypeNode):
            if arg.type is PyrexTypes.py_object_type:
                if node.type.assignable_from(arg.arg.type):
                    # completely redundant C->Py->C coercion
1763
                    return arg.arg.coerce_to(node.type, self.current_env())
Stefan Behnel's avatar
Stefan Behnel committed
1764 1765 1766
        if isinstance(arg, ExprNodes.SimpleCallNode):
            if node.type.is_int or node.type.is_float:
                return self._optimise_numeric_cast_call(node, arg)
1767 1768 1769 1770 1771 1772
        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
1773 1774
        return node

1775 1776 1777 1778 1779 1780 1781 1782 1783 1784 1785 1786
    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
1787
        if arg.base.type is Builtin.bytes_type:
1788 1789 1790 1791 1792 1793 1794 1795
            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,
1796
                    args=[
Stefan Behnel's avatar
Stefan Behnel committed
1797
                        arg.base.as_none_safe_node("'NoneType' object is not subscriptable"),
1798 1799 1800
                        index_node.coerce_to(PyrexTypes.c_py_ssize_t_type, env),
                        bound_check_node,
                        ],
1801 1802 1803
                    is_temp=True,
                    utility_code=UtilityCode.load_cached(
                        'bytes_index', 'StringTools.c'))
1804 1805 1806 1807 1808
                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
1809
    def _optimise_numeric_cast_call(self, node, arg):
1810 1811 1812 1813 1814 1815 1816 1817 1818 1819 1820 1821 1822 1823 1824 1825 1826 1827 1828
        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:
1829 1830
                    return ExprNodes.TypecastNode(
                        node.pos, operand=func_arg, type=node.type)
1831 1832 1833 1834 1835
        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:
1836 1837
                    return ExprNodes.TypecastNode(
                        node.pos, operand=func_arg, type=node.type)
1838 1839
        return node

1840 1841 1842 1843 1844 1845 1846 1847 1848 1849 1850 1851 1852 1853 1854 1855
    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)))

1856 1857
    ### builtin types

1858 1859 1860 1861 1862
    PyDict_Copy_func_type = PyrexTypes.CFuncType(
        Builtin.dict_type, [
            PyrexTypes.CFuncTypeArg("dict", Builtin.dict_type, None)
            ])

1863
    def _handle_simple_function_dict(self, node, function, pos_args):
1864
        """Replace dict(some_dict) by PyDict_Copy(some_dict).
1865
        """
1866
        if len(pos_args) != 1:
1867
            return node
1868
        arg = pos_args[0]
1869
        if arg.type is Builtin.dict_type:
1870
            arg = arg.as_none_safe_node("'NoneType' is not iterable")
1871 1872
            return ExprNodes.PythonCapiCallNode(
                node.pos, "PyDict_Copy", self.PyDict_Copy_func_type,
1873
                args = [arg],
1874 1875 1876
                is_temp = node.is_temp
                )
        return node
1877

1878 1879 1880 1881 1882
    PyList_AsTuple_func_type = PyrexTypes.CFuncType(
        Builtin.tuple_type, [
            PyrexTypes.CFuncTypeArg("list", Builtin.list_type, None)
            ])

1883
    def _handle_simple_function_tuple(self, node, function, pos_args):
1884 1885
        """Replace tuple([...]) by a call to PyList_AsTuple.
        """
1886
        if len(pos_args) != 1:
1887
            return node
1888
        list_arg = pos_args[0]
1889 1890 1891 1892
        if list_arg.type is not Builtin.list_type:
            return node
        if not isinstance(list_arg, (ExprNodes.ComprehensionNode,
                                     ExprNodes.ListNode)):
1893
            pos_args[0] = list_arg.as_none_safe_node(
1894
                "'NoneType' object is not iterable")
1895

1896 1897
        return ExprNodes.PythonCapiCallNode(
            node.pos, "PyList_AsTuple", self.PyList_AsTuple_func_type,
1898
            args = pos_args,
1899 1900 1901
            is_temp = node.is_temp
            )

1902 1903 1904 1905 1906 1907 1908
    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)

1909
    def _handle_simple_function_set(self, node, function, pos_args):
1910 1911 1912 1913 1914 1915 1916 1917 1918 1919 1920 1921 1922 1923 1924 1925 1926 1927 1928 1929
        if len(pos_args) == 1 and isinstance(pos_args[0], (ExprNodes.ListNode,
                                                           ExprNodes.TupleNode)):
            # 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)
            for temp in temps[::-1]:
                result = UtilNodes.EvalWithTempExprNode(temp, result)
            return result
        return node

1930
    def _handle_simple_function_float(self, node, function, pos_args):
Stefan Behnel's avatar
Stefan Behnel committed
1931 1932 1933
        """Transform float() into either a C type cast or a faster C
        function call.
        """
1934 1935
        # Note: this requires the float() function to be typed as
        # returning a C 'double'
1936
        if len(pos_args) == 0:
Stefan Behnel's avatar
typo  
Stefan Behnel committed
1937
            return ExprNodes.FloatNode(
1938 1939 1940 1941
                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')
1942 1943 1944 1945 1946 1947 1948
            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:
1949 1950
            return ExprNodes.TypecastNode(
                node.pos, operand=func_arg, type=node.type)
1951 1952 1953 1954 1955
        return ExprNodes.PythonCapiCallNode(
            node.pos, "__Pyx_PyObject_AsDouble",
            self.PyObject_AsDouble_func_type,
            args = pos_args,
            is_temp = node.is_temp,
1956
            utility_code = load_c_utility('pyobject_as_double'),
1957 1958
            py_name = "float")

1959
    def _handle_simple_function_bool(self, node, function, pos_args):
1960 1961
        """Transform bool(x) into a type coercion to a boolean.
        """
1962 1963 1964 1965 1966 1967
        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')
1968
            return node
Craig Citro's avatar
Craig Citro committed
1969
        else:
1970 1971 1972 1973 1974 1975
            # => !!<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())
1976

1977 1978
    ### builtin functions

1979 1980 1981 1982 1983
    Pyx_strlen_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_size_t_type, [
            PyrexTypes.CFuncTypeArg("bytes", PyrexTypes.c_char_ptr_type, None)
            ])

1984 1985 1986 1987 1988
    Pyx_Py_UNICODE_strlen_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_size_t_type, [
            PyrexTypes.CFuncTypeArg("unicode", PyrexTypes.c_py_unicode_ptr_type, None)
            ])

1989 1990 1991
    PyObject_Size_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_py_ssize_t_type, [
            PyrexTypes.CFuncTypeArg("obj", PyrexTypes.py_object_type, None)
1992 1993
            ],
        exception_value="-1")
1994 1995

    _map_to_capi_len_function = {
1996
        Builtin.unicode_type   : "__Pyx_PyUnicode_GET_LENGTH",
1997
        Builtin.bytes_type     : "PyBytes_GET_SIZE",
1998 1999 2000 2001 2002 2003 2004
        Builtin.list_type      : "PyList_GET_SIZE",
        Builtin.tuple_type     : "PyTuple_GET_SIZE",
        Builtin.dict_type      : "PyDict_Size",
        Builtin.set_type       : "PySet_Size",
        Builtin.frozenset_type : "PySet_Size",
        }.get

2005 2006
    _ext_types_with_pysize = set(["cpython.array.array"])

2007
    def _handle_simple_function_len(self, node, function, pos_args):
2008 2009
        """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
2010
        len(known_builtin_type) by an equivalent C-API call.
Stefan Behnel's avatar
Stefan Behnel committed
2011
        """
2012 2013 2014 2015 2016 2017
        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
2018 2019 2020 2021 2022
        if arg.type.is_string:
            new_node = ExprNodes.PythonCapiCallNode(
                node.pos, "strlen", self.Pyx_strlen_func_type,
                args = [arg],
                is_temp = node.is_temp,
2023
                utility_code = UtilityCode.load_cached("IncludeStringH", "StringTools.c"))
2024
        elif arg.type.is_pyunicode_ptr:
2025 2026 2027
            new_node = ExprNodes.PythonCapiCallNode(
                node.pos, "__Pyx_Py_UNICODE_strlen", self.Pyx_Py_UNICODE_strlen_func_type,
                args = [arg],
2028
                is_temp = node.is_temp)
2029 2030 2031
        elif arg.type.is_pyobject:
            cfunc_name = self._map_to_capi_len_function(arg.type)
            if cfunc_name is None:
2032 2033 2034 2035 2036 2037
                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
2038 2039
            arg = arg.as_none_safe_node(
                "object of type 'NoneType' has no len()")
2040 2041 2042 2043
            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
2044
        elif arg.type.is_unicode_char:
2045 2046
            return ExprNodes.IntNode(node.pos, value='1', constant_result=1,
                                     type=node.type)
2047
        else:
Stefan Behnel's avatar
Stefan Behnel committed
2048
            return node
2049
        if node.type not in (PyrexTypes.c_size_t_type, PyrexTypes.c_py_ssize_t_type):
2050
            new_node = new_node.coerce_to(node.type, self.current_env())
2051
        return new_node
2052

2053 2054 2055 2056 2057
    Pyx_Type_func_type = PyrexTypes.CFuncType(
        Builtin.type_type, [
            PyrexTypes.CFuncTypeArg("object", PyrexTypes.py_object_type, None)
            ])

2058
    def _handle_simple_function_type(self, node, function, pos_args):
Stefan Behnel's avatar
Stefan Behnel committed
2059 2060
        """Replace type(o) by a macro call to Py_TYPE(o).
        """
2061
        if len(pos_args) != 1:
2062 2063
            return node
        node = ExprNodes.PythonCapiCallNode(
2064 2065 2066 2067
            node.pos, "Py_TYPE", self.Pyx_Type_func_type,
            args = pos_args,
            is_temp = False)
        return ExprNodes.CastNode(node, PyrexTypes.py_object_type)
2068

2069 2070 2071 2072 2073
    Py_type_check_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_bint_type, [
            PyrexTypes.CFuncTypeArg("arg", PyrexTypes.py_object_type, None)
            ])

2074
    def _handle_simple_function_isinstance(self, node, function, pos_args):
2075 2076 2077 2078 2079 2080 2081 2082 2083 2084 2085 2086 2087 2088 2089 2090 2091 2092 2093
        """Replace isinstance() checks against builtin types by the
        corresponding C-API call.
        """
        if len(pos_args) != 2:
            return node
        arg, types = pos_args
        temp = None
        if isinstance(types, ExprNodes.TupleNode):
            types = types.args
            arg = temp = UtilNodes.ResultRefNode(arg)
        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
2094
            builtin_type = None
Stefan Behnel's avatar
Stefan Behnel committed
2095
            if test_type_node.is_name:
Robert Bradshaw's avatar
Robert Bradshaw committed
2096 2097 2098 2099
                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
2100 2101 2102 2103 2104 2105
            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
2106
                type_check_function = entry.type.type_check_function(exact=False)
2107 2108 2109
                if type_check_function in tests:
                    continue
                tests.append(type_check_function)
Robert Bradshaw's avatar
Robert Bradshaw committed
2110 2111 2112 2113 2114
                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:
2115
                return node
2116 2117 2118 2119 2120 2121
            test_nodes.append(
                ExprNodes.PythonCapiCallNode(
                    test_type_node.pos, type_check_function, self.Py_type_check_func_type,
                    args = type_check_args,
                    is_temp = True,
                    ))
2122 2123 2124 2125 2126 2127 2128 2129 2130 2131 2132 2133

        def join_with_or(a,b, make_binop_node=ExprNodes.binop_node):
            or_node = make_binop_node(node.pos, 'or', a, b)
            or_node.type = PyrexTypes.c_bint_type
            or_node.is_temp = True
            return or_node

        test_node = reduce(join_with_or, test_nodes).coerce_to(node.type, env)
        if temp is not None:
            test_node = UtilNodes.EvalWithTempExprNode(temp, test_node)
        return test_node

2134
    def _handle_simple_function_ord(self, node, function, pos_args):
2135
        """Unpack ord(Py_UNICODE) and ord('X').
2136 2137 2138 2139 2140
        """
        if len(pos_args) != 1:
            return node
        arg = pos_args[0]
        if isinstance(arg, ExprNodes.CoerceToPyTypeNode):
Stefan Behnel's avatar
Stefan Behnel committed
2141
            if arg.arg.type.is_unicode_char:
2142 2143 2144 2145 2146 2147 2148 2149 2150 2151 2152 2153 2154 2155 2156 2157 2158 2159
                return ExprNodes.TypecastNode(
                    arg.pos, operand=arg.arg, type=PyrexTypes.c_int_type
                    ).coerce_to(node.type, self.current_env())
        elif isinstance(arg, ExprNodes.UnicodeNode):
            if len(arg.value) == 1:
                return ExprNodes.IntNode(
                    ord(arg.value), type=PyrexTypes.c_int_type,
                    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 \
                   and ord(arg.unicode_value) <= 255: # Py2/3 portability
                return ExprNodes.IntNode(
                    ord(arg.unicode_value), type=PyrexTypes.c_int_type,
                    value=str(ord(arg.unicode_value)),
                    constant_result=ord(arg.unicode_value)
                    ).coerce_to(node.type, self.current_env())
2160 2161
        return node

2162 2163
    ### special methods

2164 2165
    Pyx_tp_new_func_type = PyrexTypes.CFuncType(
        PyrexTypes.py_object_type, [
2166
            PyrexTypes.CFuncTypeArg("type",   PyrexTypes.py_object_type, None),
2167
            PyrexTypes.CFuncTypeArg("args",   Builtin.tuple_type, None),
2168 2169
            ])

2170 2171
    Pyx_tp_new_kwargs_func_type = PyrexTypes.CFuncType(
        PyrexTypes.py_object_type, [
2172
            PyrexTypes.CFuncTypeArg("type",   PyrexTypes.py_object_type, None),
2173 2174 2175 2176
            PyrexTypes.CFuncTypeArg("args",   Builtin.tuple_type, None),
            PyrexTypes.CFuncTypeArg("kwargs", Builtin.dict_type, None),
        ])

2177 2178
    def _handle_any_slot__new__(self, node, function, args,
                                is_unbound_method, kwargs=None):
2179
        """Replace 'exttype.__new__(exttype, ...)' by a call to exttype->tp_new()
2180
        """
2181
        obj = function.obj
2182
        if not is_unbound_method or len(args) < 1:
2183 2184 2185 2186 2187 2188 2189 2190 2191 2192 2193 2194 2195 2196
            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:
2197
            # different types - may or may not lead to an error at runtime
2198 2199
            return node

2200 2201 2202
        args_tuple = ExprNodes.TupleNode(node.pos, args=args[1:])
        args_tuple = args_tuple.analyse_types(
            self.current_env(), skip_children=True)
2203

2204 2205
        if type_arg.type_entry:
            ext_type = type_arg.type_entry.type
2206 2207 2208
            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
2209 2210 2211 2212 2213 2214 2215 2216 2217 2218 2219 2220 2221 2222 2223 2224 2225 2226 2227 2228 2229
                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)
2230
        else:
2231
            # arbitrary variable, needs a None check for safety
2232
            type_arg = type_arg.as_none_safe_node(
2233 2234
                "object.__new__(X): X is not a type object (NoneType)")

2235 2236 2237 2238 2239 2240 2241 2242 2243 2244 2245 2246 2247 2248
        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
2249 2250
            )

2251 2252 2253 2254 2255 2256 2257 2258
    ### methods of builtin types

    PyObject_Append_func_type = PyrexTypes.CFuncType(
        PyrexTypes.py_object_type, [
            PyrexTypes.CFuncTypeArg("list", PyrexTypes.py_object_type, None),
            PyrexTypes.CFuncTypeArg("item", PyrexTypes.py_object_type, None),
            ])

2259
    def _handle_simple_method_object_append(self, node, function, args, is_unbound_method):
Stefan Behnel's avatar
Stefan Behnel committed
2260 2261 2262
        """Optimistic optimisation as X.append() is almost always
        referring to a list.
        """
2263
        if len(args) != 2:
2264 2265
            return node

2266 2267 2268
        return ExprNodes.PythonCapiCallNode(
            node.pos, "__Pyx_PyObject_Append", self.PyObject_Append_func_type,
            args = args,
Stefan Behnel's avatar
Stefan Behnel committed
2269
            may_return_none = True,
2270
            is_temp = node.is_temp,
2271
            utility_code = load_c_utility('append')
2272 2273
            )

Robert Bradshaw's avatar
Robert Bradshaw committed
2274 2275 2276 2277 2278 2279 2280 2281 2282 2283 2284
    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),
            PyrexTypes.CFuncTypeArg("index", PyrexTypes.c_long_type, None),
            ])

2285
    def _handle_simple_method_object_pop(self, node, function, args, is_unbound_method):
Stefan Behnel's avatar
Stefan Behnel committed
2286 2287 2288
        """Optimistic optimisation as X.pop([n]) is almost always
        referring to a list.
        """
Robert Bradshaw's avatar
Robert Bradshaw committed
2289 2290 2291 2292
        if len(args) == 1:
            return ExprNodes.PythonCapiCallNode(
                node.pos, "__Pyx_PyObject_Pop", self.PyObject_Pop_func_type,
                args = args,
Stefan Behnel's avatar
Stefan Behnel committed
2293
                may_return_none = True,
Robert Bradshaw's avatar
Robert Bradshaw committed
2294
                is_temp = node.is_temp,
2295
                utility_code = load_c_utility('pop')
Robert Bradshaw's avatar
Robert Bradshaw committed
2296 2297
                )
        elif len(args) == 2:
2298 2299 2300 2301 2302 2303 2304 2305 2306 2307
            index = args[1]
            if isinstance(index, ExprNodes.CoerceToPyTypeNode):
                index = index.arg
            if isinstance(index, ExprNodes.IntNode):
                index = index.coerce_to(PyrexTypes.c_py_ssize_t_type, None)
            if index.type.is_int:
                widest = PyrexTypes.widest_numeric_type(
                    index.type, PyrexTypes.c_py_ssize_t_type)
                if widest == PyrexTypes.c_py_ssize_t_type:
                    args[1] = index
Robert Bradshaw's avatar
Robert Bradshaw committed
2308
                    return ExprNodes.PythonCapiCallNode(
2309 2310
                        node.pos, "__Pyx_PyObject_PopIndex",
                        self.PyObject_PopIndex_func_type,
Robert Bradshaw's avatar
Robert Bradshaw committed
2311
                        args = args,
Stefan Behnel's avatar
Stefan Behnel committed
2312
                        may_return_none = True,
Robert Bradshaw's avatar
Robert Bradshaw committed
2313
                        is_temp = node.is_temp,
2314
                        utility_code = load_c_utility("pop_index")
Robert Bradshaw's avatar
Robert Bradshaw committed
2315
                        )
2316

Robert Bradshaw's avatar
Robert Bradshaw committed
2317 2318
        return node

2319 2320
    _handle_simple_method_list_pop = _handle_simple_method_object_pop

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

2327
    def _handle_simple_method_list_sort(self, node, function, args, is_unbound_method):
Stefan Behnel's avatar
Stefan Behnel committed
2328 2329
        """Call PyList_Sort() instead of the 0-argument l.sort().
        """
2330
        if len(args) != 1:
2331
            return node
2332
        return self._substitute_method_call(
2333
            node, function, "PyList_Sort", self.single_param_func_type,
2334
            'sort', is_unbound_method, args).coerce_to(node.type, self.current_env)
2335

2336 2337 2338 2339 2340
    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),
2341
            ])
2342

2343
    def _handle_simple_method_dict_get(self, node, function, args, is_unbound_method):
Stefan Behnel's avatar
Stefan Behnel committed
2344 2345
        """Replace dict.get() by a call to PyDict_GetItem().
        """
2346 2347 2348 2349 2350 2351 2352
        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(
2353 2354
            node, function,
            "__Pyx_PyDict_GetItemDefault", self.Pyx_PyDict_GetItem_func_type,
2355
            'get', is_unbound_method, args,
Stefan Behnel's avatar
Stefan Behnel committed
2356
            may_return_none = True,
2357
            utility_code = load_c_utility("dict_getitem_default"))
2358

2359 2360 2361 2362 2363
    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),
2364
            PyrexTypes.CFuncTypeArg("is_safe_type", PyrexTypes.c_int_type, None),
2365 2366
            ])

2367
    def _handle_simple_method_dict_setdefault(self, node, function, args, is_unbound_method):
2368 2369 2370 2371 2372 2373 2374
        """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
2375 2376 2377 2378 2379 2380 2381 2382 2383 2384
        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(
            node.pos, value=is_safe_type, constant_result=is_safe_type))
2385 2386

        return self._substitute_method_call(
2387 2388
            node, function,
            "__Pyx_PyDict_SetDefault", self.Pyx_PyDict_SetDefault_func_type,
2389
            'setdefault', is_unbound_method, args,
2390 2391
            may_return_none=True,
            utility_code=load_c_utility('dict_setdefault'))
2392

2393 2394 2395

    ### unicode type methods

2396 2397
    PyUnicode_uchar_predicate_func_type = PyrexTypes.CFuncType(
        PyrexTypes.c_bint_type, [
2398
            PyrexTypes.CFuncTypeArg("uchar", PyrexTypes.c_py_ucs4_type, None),
2399 2400
            ])

2401
    def _inject_unicode_predicate(self, node, function, args, is_unbound_method):
2402 2403 2404 2405
        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
2406
               not ustring.arg.type.is_unicode_char:
2407 2408
            return node
        uchar = ustring.arg
2409
        method_name = function.attribute
2410 2411
        if method_name == 'istitle':
            # istitle() doesn't directly map to Py_UNICODE_ISTITLE()
2412 2413
            utility_code = UtilityCode.load_cached(
                "py_unicode_istitle", "StringTools.c")
2414 2415 2416 2417 2418
            function_name = '__Pyx_Py_UNICODE_ISTITLE'
        else:
            utility_code = None
            function_name = 'Py_UNICODE_%s' % method_name.upper()
        func_call = self._substitute_method_call(
2419 2420
            node, function,
            function_name, self.PyUnicode_uchar_predicate_func_type,
2421 2422 2423 2424 2425 2426 2427 2428 2429 2430 2431 2432 2433 2434 2435 2436 2437
            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(
2438 2439
        PyrexTypes.c_py_ucs4_type, [
            PyrexTypes.CFuncTypeArg("uchar", PyrexTypes.c_py_ucs4_type, None),
2440 2441
            ])

Stefan Behnel's avatar
Stefan Behnel committed
2442
    def _inject_unicode_character_conversion(self, node, function, args, is_unbound_method):
2443 2444 2445 2446
        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
2447
               not ustring.arg.type.is_unicode_char:
2448 2449
            return node
        uchar = ustring.arg
2450
        method_name = function.attribute
2451 2452
        function_name = 'Py_UNICODE_TO%s' % method_name.upper()
        func_call = self._substitute_method_call(
2453 2454
            node, function,
            function_name, self.PyUnicode_uchar_conversion_func_type,
2455 2456 2457 2458 2459 2460 2461 2462 2463
            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

2464 2465 2466 2467 2468 2469
    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
2470
    def _handle_simple_method_unicode_splitlines(self, node, function, args, is_unbound_method):
2471 2472 2473 2474 2475 2476
        """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
2477
        self._inject_bint_default_argument(node, args, 1, False)
2478 2479

        return self._substitute_method_call(
2480 2481
            node, function,
            "PyUnicode_Splitlines", self.PyUnicode_Splitlines_func_type,
2482 2483 2484 2485 2486 2487 2488 2489 2490 2491
            '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),
            ]
        )

2492
    def _handle_simple_method_unicode_split(self, node, function, args, is_unbound_method):
2493 2494 2495 2496 2497 2498 2499 2500
        """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))
2501 2502
        self._inject_int_default_argument(
            node, args, 2, PyrexTypes.c_py_ssize_t_type, "-1")
2503 2504

        return self._substitute_method_call(
2505 2506
            node, function,
            "PyUnicode_Split", self.PyUnicode_Split_func_type,
2507 2508
            'split', is_unbound_method, args)

2509
    PyString_Tailmatch_func_type = PyrexTypes.CFuncType(
2510
        PyrexTypes.c_bint_type, [
2511
            PyrexTypes.CFuncTypeArg("str", PyrexTypes.py_object_type, None),  # bytes/str/unicode
2512 2513 2514 2515 2516 2517 2518
            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')

2519
    def _handle_simple_method_unicode_endswith(self, node, function, args, is_unbound_method):
2520
        return self._inject_tailmatch(
2521
            node, function, args, is_unbound_method, 'unicode', 'endswith',
2522
            unicode_tailmatch_utility_code, +1)
2523

2524
    def _handle_simple_method_unicode_startswith(self, node, function, args, is_unbound_method):
2525
        return self._inject_tailmatch(
2526
            node, function, args, is_unbound_method, 'unicode', 'startswith',
2527
            unicode_tailmatch_utility_code, -1)
2528

2529
    def _inject_tailmatch(self, node, function, args, is_unbound_method, type_name,
2530
                          method_name, utility_code, direction):
2531 2532 2533 2534
        """Replace unicode.startswith(...) and unicode.endswith(...)
        by a direct call to the corresponding C-API function.
        """
        if len(args) not in (2,3,4):
2535
            self._error_wrong_arg_count('%s.%s' % (type_name, method_name), node, args, "2-4")
2536
            return node
2537 2538 2539 2540
        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")
2541 2542 2543 2544
        args.append(ExprNodes.IntNode(
            node.pos, value=str(direction), type=PyrexTypes.c_int_type))

        method_call = self._substitute_method_call(
2545 2546
            node, function,
            "__Pyx_Py%s_Tailmatch" % type_name.capitalize(),
2547
            self.PyString_Tailmatch_func_type,
2548
            method_name, is_unbound_method, args,
2549
            utility_code = utility_code)
Stefan Behnel's avatar
Stefan Behnel committed
2550
        return method_call.coerce_to(Builtin.bool_type, self.current_env())
2551

2552 2553 2554 2555 2556 2557 2558 2559 2560 2561
    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')

2562
    def _handle_simple_method_unicode_find(self, node, function, args, is_unbound_method):
2563
        return self._inject_unicode_find(
2564
            node, function, args, is_unbound_method, 'find', +1)
2565

2566
    def _handle_simple_method_unicode_rfind(self, node, function, args, is_unbound_method):
2567
        return self._inject_unicode_find(
2568
            node, function, args, is_unbound_method, 'rfind', -1)
2569

2570
    def _inject_unicode_find(self, node, function, args, is_unbound_method,
2571 2572 2573 2574 2575 2576 2577
                             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
2578 2579 2580 2581
        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")
2582 2583 2584 2585
        args.append(ExprNodes.IntNode(
            node.pos, value=str(direction), type=PyrexTypes.c_int_type))

        method_call = self._substitute_method_call(
2586
            node, function, "PyUnicode_Find", self.PyUnicode_Find_func_type,
2587
            method_name, is_unbound_method, args)
Stefan Behnel's avatar
Stefan Behnel committed
2588
        return method_call.coerce_to_pyobject(self.current_env())
2589

Stefan Behnel's avatar
Stefan Behnel committed
2590 2591 2592 2593 2594 2595 2596 2597 2598
    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')

2599
    def _handle_simple_method_unicode_count(self, node, function, args, is_unbound_method):
Stefan Behnel's avatar
Stefan Behnel committed
2600 2601 2602 2603 2604 2605
        """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
2606 2607 2608 2609
        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
2610 2611

        method_call = self._substitute_method_call(
2612
            node, function, "PyUnicode_Count", self.PyUnicode_Count_func_type,
Stefan Behnel's avatar
Stefan Behnel committed
2613
            'count', is_unbound_method, args)
Stefan Behnel's avatar
Stefan Behnel committed
2614
        return method_call.coerce_to_pyobject(self.current_env())
Stefan Behnel's avatar
Stefan Behnel committed
2615

Stefan Behnel's avatar
Stefan Behnel committed
2616 2617 2618 2619 2620 2621 2622 2623
    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),
            ])

2624
    def _handle_simple_method_unicode_replace(self, node, function, args, is_unbound_method):
Stefan Behnel's avatar
Stefan Behnel committed
2625 2626 2627 2628 2629 2630
        """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
2631 2632
        self._inject_int_default_argument(
            node, args, 3, PyrexTypes.c_py_ssize_t_type, "-1")
Stefan Behnel's avatar
Stefan Behnel committed
2633 2634

        return self._substitute_method_call(
2635
            node, function, "PyUnicode_Replace", self.PyUnicode_Replace_func_type,
Stefan Behnel's avatar
Stefan Behnel committed
2636 2637
            'replace', is_unbound_method, args)

2638 2639
    PyUnicode_AsEncodedString_func_type = PyrexTypes.CFuncType(
        Builtin.bytes_type, [
2640
            PyrexTypes.CFuncTypeArg("obj", Builtin.unicode_type, None),
2641 2642
            PyrexTypes.CFuncTypeArg("encoding", PyrexTypes.c_char_ptr_type, None),
            PyrexTypes.CFuncTypeArg("errors", PyrexTypes.c_char_ptr_type, None),
2643
            ])
2644 2645 2646

    PyUnicode_AsXyzString_func_type = PyrexTypes.CFuncType(
        Builtin.bytes_type, [
2647
            PyrexTypes.CFuncTypeArg("obj", Builtin.unicode_type, None),
2648
            ])
2649 2650 2651 2652

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

2653 2654
    _special_codecs = [ (name, codecs.getencoder(name))
                        for name in _special_encodings ]
2655

2656
    def _handle_simple_method_unicode_encode(self, node, function, args, is_unbound_method):
Stefan Behnel's avatar
Stefan Behnel committed
2657 2658 2659
        """Replace unicode.encode(...) by a direct C-API call to the
        corresponding codec.
        """
2660
        if len(args) < 1 or len(args) > 3:
2661
            self._error_wrong_arg_count('unicode.encode', node, args, '1-3')
2662 2663 2664 2665 2666
            return node

        string_node = args[0]

        if len(args) == 1:
2667
            null_node = ExprNodes.NullNode(node.pos)
2668
            return self._substitute_method_call(
2669
                node, function, "PyUnicode_AsEncodedString",
2670 2671 2672
                self.PyUnicode_AsEncodedString_func_type,
                'encode', is_unbound_method, [string_node, null_node, null_node])

2673 2674 2675 2676 2677 2678 2679 2680 2681 2682 2683 2684 2685 2686 2687 2688 2689 2690 2691 2692 2693 2694 2695 2696
        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

        if isinstance(string_node, ExprNodes.UnicodeNode):
            # 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)

        if error_handling == 'strict':
            # 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(
2697
                    node, function, encode_function,
2698 2699 2700 2701
                    self.PyUnicode_AsXyzString_func_type,
                    'encode', is_unbound_method, [string_node])

        return self._substitute_method_call(
2702
            node, function, "PyUnicode_AsEncodedString",
2703 2704 2705 2706
            self.PyUnicode_AsEncodedString_func_type,
            'encode', is_unbound_method,
            [string_node, encoding_node, error_handling_node])

2707
    PyUnicode_DecodeXyz_func_ptr_type = PyrexTypes.CPtrType(PyrexTypes.CFuncType(
2708 2709 2710 2711
        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),
2712
            ]))
2713

2714
    _decode_c_string_func_type = PyrexTypes.CFuncType(
2715 2716
        Builtin.unicode_type, [
            PyrexTypes.CFuncTypeArg("string", PyrexTypes.c_char_ptr_type, None),
2717 2718
            PyrexTypes.CFuncTypeArg("start", PyrexTypes.c_py_ssize_t_type, None),
            PyrexTypes.CFuncTypeArg("stop", PyrexTypes.c_py_ssize_t_type, None),
2719 2720
            PyrexTypes.CFuncTypeArg("encoding", PyrexTypes.c_char_ptr_type, None),
            PyrexTypes.CFuncTypeArg("errors", PyrexTypes.c_char_ptr_type, None),
2721
            PyrexTypes.CFuncTypeArg("decode_func", PyUnicode_DecodeXyz_func_ptr_type, None),
2722
            ])
2723

Stefan Behnel's avatar
Stefan Behnel committed
2724 2725 2726 2727 2728 2729 2730 2731 2732 2733
    _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),
            ])

2734 2735
    _decode_cpp_string_func_type = None # lazy init

2736
    def _handle_simple_method_bytes_decode(self, node, function, args, is_unbound_method):
Stefan Behnel's avatar
Stefan Behnel committed
2737
        """Replace char*.decode() by a direct C-API call to the
Stefan Behnel's avatar
Stefan Behnel committed
2738
        corresponding codec, possibly resolving a slice on the char*.
Stefan Behnel's avatar
Stefan Behnel committed
2739
        """
2740
        if not (1 <= len(args) <= 3):
2741 2742
            self._error_wrong_arg_count('bytes.decode', node, args, '1-3')
            return node
2743 2744

        # normalise input nodes
Stefan Behnel's avatar
Stefan Behnel committed
2745 2746 2747 2748
        string_node = args[0]
        start = stop = None
        if isinstance(string_node, ExprNodes.SliceIndexNode):
            index_node = string_node
2749 2750 2751 2752
            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
2753 2754 2755 2756 2757 2758 2759 2760 2761 2762 2763 2764 2765 2766 2767
        if isinstance(string_node, ExprNodes.CoerceToPyTypeNode):
            string_node = string_node.arg

        string_type = string_node.type
        if string_type is Builtin.bytes_type:
            if is_unbound_method:
                string_node = string_node.as_none_safe_node(
                    "descriptor '%s' requires a '%s' object but received a 'NoneType'",
                    format_args = ['decode', 'bytes'])
            else:
                string_node = string_node.as_none_safe_node(
                    "'NoneType' object has no attribute '%s'",
                    error = "PyExc_AttributeError",
                    format_args = ['decode'])
        elif not string_type.is_string and not string_type.is_cpp_string:
2768 2769
            # nothing to optimise here
            return node
2770 2771 2772 2773 2774 2775

        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

2776 2777 2778 2779 2780 2781 2782
        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())

2783
        # try to find a specific encoder function
2784 2785 2786
        codec_name = None
        if encoding is not None:
            codec_name = self._find_special_codec_name(encoding)
2787
        if codec_name is not None:
2788 2789 2790 2791
            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)
2792
        else:
2793 2794 2795 2796
            decode_function = ExprNodes.NullNode(node.pos)

        # build the helper function call
        temps = []
Stefan Behnel's avatar
Stefan Behnel committed
2797
        if string_type.is_string:
2798 2799 2800 2801 2802 2803 2804 2805 2806 2807 2808 2809 2810 2811
            # 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,
                    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())
            helper_func_type = self._decode_c_string_func_type
            utility_code_name = 'decode_c_string'
Stefan Behnel's avatar
Stefan Behnel committed
2812
        elif string_type.is_cpp_string:
2813 2814 2815 2816 2817 2818 2819 2820
            # 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
2821
                        PyrexTypes.CFuncTypeArg("string", string_type, None),
2822 2823 2824 2825 2826 2827 2828 2829
                        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),
                        ])
            helper_func_type = self._decode_cpp_string_func_type
            utility_code_name = 'decode_cpp_string'
Stefan Behnel's avatar
Stefan Behnel committed
2830 2831 2832 2833 2834 2835 2836
        else:
            # Python bytes object
            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
            utility_code_name = 'decode_bytes'
2837 2838 2839 2840 2841

        node = ExprNodes.PythonCapiCallNode(
            node.pos, '__Pyx_%s' % utility_code_name, helper_func_type,
            args = [string_node, start, stop, encoding_node, error_handling_node, decode_function],
            is_temp = node.is_temp,
2842
            utility_code=UtilityCode.load_cached(utility_code_name, 'StringTools.c'),
2843
            )
2844

2845 2846 2847
        for temp in temps[::-1]:
            node = UtilNodes.EvalWithTempExprNode(temp, node)
        return node
2848 2849 2850 2851 2852 2853 2854 2855 2856 2857 2858 2859 2860 2861 2862

    def _find_special_codec_name(self, encoding):
        try:
            requested_codec = codecs.getencoder(encoding)
        except:
            return None
        for name, codec in self._special_codecs:
            if codec == requested_codec:
                if '_' in name:
                    name = ''.join([ s.capitalize()
                                     for s in name.split('_')])
                return name
        return None

    def _unpack_encoding_and_error_mode(self, pos, args):
2863 2864 2865
        null_node = ExprNodes.NullNode(pos)

        if len(args) >= 2:
2866 2867
            encoding, encoding_node = self._unpack_string_and_cstring_node(args[1])
            if encoding_node is None:
2868
                return None
2869
        else:
2870 2871
            encoding = None
            encoding_node = null_node
2872 2873

        if len(args) == 3:
2874 2875
            error_handling, error_handling_node = self._unpack_string_and_cstring_node(args[2])
            if error_handling_node is None:
2876
                return None
2877 2878
            if error_handling == 'strict':
                error_handling_node = null_node
2879 2880 2881 2882
        else:
            error_handling = 'strict'
            error_handling_node = null_node

2883
        return (encoding, encoding_node, error_handling, error_handling_node)
2884

2885 2886 2887 2888 2889 2890 2891 2892 2893 2894 2895 2896 2897 2898 2899 2900 2901 2902
    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:
2903
            encoding = node = None
2904 2905
        return encoding, node

2906
    def _handle_simple_method_str_endswith(self, node, function, args, is_unbound_method):
2907
        return self._inject_tailmatch(
2908
            node, function, args, is_unbound_method, 'str', 'endswith',
2909
            str_tailmatch_utility_code, +1)
2910

2911
    def _handle_simple_method_str_startswith(self, node, function, args, is_unbound_method):
2912
        return self._inject_tailmatch(
2913
            node, function, args, is_unbound_method, 'str', 'startswith',
2914
            str_tailmatch_utility_code, -1)
2915

2916
    def _handle_simple_method_bytes_endswith(self, node, function, args, is_unbound_method):
2917
        return self._inject_tailmatch(
2918
            node, function, args, is_unbound_method, 'bytes', 'endswith',
2919
            bytes_tailmatch_utility_code, +1)
2920

2921
    def _handle_simple_method_bytes_startswith(self, node, function, args, is_unbound_method):
2922
        return self._inject_tailmatch(
2923
            node, function, args, is_unbound_method, 'bytes', 'startswith',
2924
            bytes_tailmatch_utility_code, -1)
2925 2926 2927

    ### helpers

2928
    def _substitute_method_call(self, node, function, name, func_type,
2929
                                attr_name, is_unbound_method, args=(),
2930
                                utility_code=None, is_temp=None,
Stefan Behnel's avatar
Stefan Behnel committed
2931
                                may_return_none=ExprNodes.PythonCapiCallNode.may_return_none):
2932
        args = list(args)
2933
        if args and not args[0].is_literal:
2934 2935
            self_arg = args[0]
            if is_unbound_method:
2936
                self_arg = self_arg.as_none_safe_node(
2937
                    "descriptor '%s' requires a '%s' object but received a 'NoneType'",
2938
                    format_args=[attr_name, function.obj.name])
2939
            else:
2940
                self_arg = self_arg.as_none_safe_node(
2941 2942 2943
                    "'NoneType' object has no attribute '%s'",
                    error = "PyExc_AttributeError",
                    format_args = [attr_name])
2944
            args[0] = self_arg
2945 2946
        if is_temp is None:
            is_temp = node.is_temp
2947
        return ExprNodes.PythonCapiCallNode(
2948
            node.pos, name, func_type,
2949
            args = args,
2950
            is_temp = is_temp,
Stefan Behnel's avatar
Stefan Behnel committed
2951 2952
            utility_code = utility_code,
            may_return_none = may_return_none,
2953
            result_is_used = node.result_is_used,
2954 2955
            )

2956 2957 2958
    def _inject_int_default_argument(self, node, args, arg_index, type, default_value):
        assert len(args) >= arg_index
        if len(args) == arg_index:
2959 2960
            args.append(ExprNodes.IntNode(node.pos, value=str(default_value),
                                          type=type, constant_result=default_value))
2961
        else:
2962
            args[arg_index] = args[arg_index].coerce_to(type, self.current_env())
2963 2964 2965 2966

    def _inject_bint_default_argument(self, node, args, arg_index, default_value):
        assert len(args) >= arg_index
        if len(args) == arg_index:
2967 2968 2969
            default_value = bool(default_value)
            args.append(ExprNodes.BoolNode(node.pos, value=default_value,
                                           constant_result=default_value))
2970
        else:
2971
            args[arg_index] = args[arg_index].coerce_to_boolean(self.current_env())
2972

2973

2974 2975 2976
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')
2977

2978

2979 2980 2981 2982
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.
2983 2984 2985 2986 2987 2988 2989 2990 2991 2992 2993

    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.
2994
    """
2995

Mark Florisson's avatar
Mark Florisson committed
2996 2997 2998 2999 3000 3001 3002
    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
3003

3004
    def _calculate_const(self, node):
Mark Florisson's avatar
Mark Florisson committed
3005
        if (not self.reevaluate and
3006
                node.constant_result is not ExprNodes.constant_value_not_set):
3007 3008 3009 3010 3011 3012 3013 3014
            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)
3015
        for child_result in children.values():
3016 3017
            if type(child_result) is list:
                for child in child_result:
Stefan Behnel's avatar
Stefan Behnel committed
3018
                    if getattr(child, 'constant_result', not_a_constant) is not_a_constant:
3019
                        return
Stefan Behnel's avatar
Stefan Behnel committed
3020
            elif getattr(child_result, 'constant_result', not_a_constant) is not_a_constant:
3021 3022 3023 3024 3025 3026 3027
                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
3028
        except (ValueError, TypeError, KeyError, IndexError, AttributeError, ArithmeticError):
3029 3030 3031 3032 3033 3034 3035
            # 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)

3036 3037
    NODE_TYPE_ORDER = [ExprNodes.BoolNode, ExprNodes.CharNode,
                       ExprNodes.IntNode, ExprNodes.FloatNode]
3038 3039 3040 3041 3042 3043 3044 3045

    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

3046 3047 3048 3049
    def visit_ExprNode(self, node):
        self._calculate_const(node)
        return node

3050
    def visit_UnopNode(self, node):
3051 3052 3053 3054 3055
        self._calculate_const(node)
        if node.constant_result is ExprNodes.not_a_constant:
            return node
        if not node.operand.is_literal:
            return node
Stefan Behnel's avatar
Stefan Behnel committed
3056 3057 3058
        if node.operator == '!':
            return ExprNodes.BoolNode(node.pos, value=bool(node.constant_result),
                                      constant_result=bool(node.constant_result))
3059
        elif isinstance(node.operand, ExprNodes.BoolNode):
Stefan Behnel's avatar
Stefan Behnel committed
3060 3061 3062
            return ExprNodes.IntNode(node.pos, value=str(int(node.constant_result)),
                                     type=PyrexTypes.c_int_type,
                                     constant_result=int(node.constant_result))
3063
        elif node.operator == '+':
3064 3065 3066 3067 3068 3069
            return self._handle_UnaryPlusNode(node)
        elif node.operator == '-':
            return self._handle_UnaryMinusNode(node)
        return node

    def _handle_UnaryMinusNode(self, node):
3070 3071 3072 3073 3074 3075 3076
        def _negate(value):
            if value.startswith('-'):
                value = value[1:]
            else:
                value = '-' + value
            return value

3077 3078
        if isinstance(node.operand, ExprNodes.FloatNode):
            # this is a safe operation
3079 3080
            return ExprNodes.FloatNode(node.pos, value=_negate(node.operand.value),
                                       constant_result=node.constant_result)
3081 3082
        node_type = node.operand.type
        if node_type.is_int and node_type.signed or \
3083 3084 3085 3086 3087
                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)
3088 3089
        return node

3090
    def _handle_UnaryPlusNode(self, node):
3091 3092 3093 3094
        if node.constant_result == node.operand.constant_result:
            return node.operand
        return node

3095 3096
    def visit_BoolBinopNode(self, node):
        self._calculate_const(node)
3097
        if node.operand1.constant_result is ExprNodes.not_a_constant:
3098
            return node
3099 3100 3101 3102 3103
        elif node.operand1.constant_result:
            if node.operator == 'and':
                return node.operand2
            else:
                return node.operand1
3104
        else:
3105 3106 3107 3108
            if node.operator == 'and':
                return node.operand1
            else:
                return node.operand2
3109

3110 3111 3112 3113
    def visit_BinopNode(self, node):
        self._calculate_const(node)
        if node.constant_result is ExprNodes.not_a_constant:
            return node
3114 3115
        if isinstance(node.constant_result, float):
            return node
3116 3117
        operand1, operand2 = node.operand1, node.operand2
        if not operand1.is_literal or not operand2.is_literal:
3118 3119 3120
            return node

        # now inject a new constant node with the calculated value
3121
        try:
3122
            type1, type2 = operand1.type, operand2.type
3123
            if type1 is None or type2 is None:
3124 3125 3126 3127
                return node
        except AttributeError:
            return node

3128
        if type1.is_numeric and type2.is_numeric:
3129
            widest_type = PyrexTypes.widest_numeric_type(type1, type2)
3130 3131
        else:
            widest_type = PyrexTypes.py_object_type
3132

3133
        target_class = self._widest_node_class(operand1, operand2)
3134 3135
        if target_class is None:
            return node
3136 3137 3138 3139 3140 3141 3142 3143
        if target_class is ExprNodes.BoolNode and node.operator in '+-//<<%**>>':
            # C arithmetic results in at least an int type
            target_class = ExprNodes.IntNode
        if target_class is ExprNodes.CharNode and node.operator in '+-//<<%**>>&|^':
            # C arithmetic results in at least an int type
            target_class = ExprNodes.IntNode

        if target_class is ExprNodes.IntNode:
3144 3145 3146 3147
            unsigned = getattr(operand1, 'unsigned', '') and \
                       getattr(operand2, 'unsigned', '')
            longness = "LL"[:max(len(getattr(operand1, 'longness', '')),
                                 len(getattr(operand2, 'longness', '')))]
3148
            new_node = ExprNodes.IntNode(pos=node.pos,
3149 3150 3151
                                         unsigned=unsigned, longness=longness,
                                         value=str(int(node.constant_result)),
                                         constant_result=int(node.constant_result))
3152 3153 3154 3155
            # 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
3156
            else:
3157
                new_node.type = PyrexTypes.widest_numeric_type(widest_type, new_node.type)
3158
        else:
3159
            if target_class is ExprNodes.BoolNode:
3160 3161 3162 3163 3164 3165
                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)
3166 3167
        return new_node

3168 3169 3170 3171 3172
    def visit_MulNode(self, node):
        if isinstance(node.operand1, (ExprNodes.ListNode, ExprNodes.TupleNode)):
            sequence_node = node.operand1
            factor = node.operand2
            self._calculate_const(factor)
3173 3174 3175 3176 3177 3178 3179 3180 3181
            if factor.constant_result != 1:
                sequence_node.mult_factor = factor
            self.visitchildren(sequence_node)
            return sequence_node
        if isinstance(node.operand1, ExprNodes.IntNode) and \
               isinstance(node.operand2, (ExprNodes.ListNode, ExprNodes.TupleNode)):
            sequence_node = node.operand2
            factor = node.operand1
            self._calculate_const(factor)
3182 3183 3184 3185 3186 3187
            if factor.constant_result != 1:
                sequence_node.mult_factor = factor
            self.visitchildren(sequence_node)
            return sequence_node
        return self.visit_BinopNode(node)

3188 3189
    def visit_PrimaryCmpNode(self, node):
        self._calculate_const(node)
3190 3191 3192 3193 3194 3195 3196 3197
        if node.constant_result is not ExprNodes.not_a_constant:
            bool_result = bool(node.constant_result)
            return ExprNodes.BoolNode(node.pos, value=bool_result,
                                      constant_result=bool_result)
        if node.operator in ('in', 'not_in') and not node.cascade:
            if isinstance(node.operand2, ExprNodes.ListNode):
                node.operand2 = node.operand2.as_tuple()
        return node
3198

3199 3200 3201 3202 3203 3204 3205 3206 3207
    def visit_CondExprNode(self, node):
        self._calculate_const(node)
        if node.test.constant_result is ExprNodes.not_a_constant:
            return node
        if node.test.constant_result:
            return node.true_val
        else:
            return node.false_val

3208 3209 3210 3211 3212 3213
    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:
            condition_result = if_clause.get_constant_condition_result()
3214 3215
            if condition_result is None:
                # unknown result => normal runtime evaluation
3216
                if_clauses.append(if_clause)
3217 3218 3219 3220 3221
            elif condition_result == True:
                # subsequent clauses can safely be dropped
                node.else_clause = if_clause.body
                break
            else:
3222
                # False clauses can safely be deleted
3223
                assert condition_result == False
3224 3225 3226 3227
        if if_clauses:
            node.if_clauses = if_clauses
            return node
        elif node.else_clause:
3228
            return node.else_clause
3229 3230
        else:
            return Nodes.StatListNode(node.pos, stats=[])
3231

3232 3233 3234
    def visit_SliceIndexNode(self, node):
        self._calculate_const(node)
        # normalise start/stop values
3235 3236 3237 3238 3239 3240 3241 3242 3243 3244 3245 3246 3247 3248 3249
        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
            if base.is_sequence_constructor:
                base.args = base.args[start:stop]
                return base
            elif base.is_string_literal:
3250 3251 3252
                base = base.as_sliced_node(start, stop)
                if base is not None:
                    return base
3253 3254
        return node

3255 3256 3257 3258
    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
3259 3260 3261 3262 3263 3264
            if node.type is Builtin.list_type:
                return ExprNodes.ListNode(node.pos, args=[])
            elif node.type is Builtin.set_type:
                return ExprNodes.SetNode(node.pos, args=[])
            elif node.type is Builtin.dict_type:
                return ExprNodes.DictNode(node.pos, key_value_pairs=[])
3265 3266
        return node

3267 3268 3269
    def visit_ForInStatNode(self, node):
        self.visitchildren(node)
        sequence = node.iterator.sequence
3270 3271 3272 3273 3274 3275 3276
        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
3277 3278 3279
            # iterating over a list literal? => tuples are more efficient
            if isinstance(sequence, ExprNodes.ListNode):
                node.iterator.sequence = sequence.as_tuple()
3280 3281
        return node

3282 3283 3284 3285 3286 3287 3288 3289 3290
    def visit_WhileStatNode(self, node):
        self.visitchildren(node)
        if node.condition.has_constant_result():
            if node.condition.constant_result:
                node.else_clause = None
            else:
                return node.else_clause
        return node

3291
    def _find_genexpr_yield_stat(self, node):
3292 3293 3294 3295
        body_node_types = (Nodes.ForInStatNode, Nodes.IfStatNode)
        while isinstance(node, body_node_types):
            node = node.body
        if isinstance(node, Nodes.ExprStatNode):
3296 3297
            expr = node.expr
            if isinstance(expr, ExprNodes.YieldExprNode):
3298 3299 3300
                return node
        return None

3301 3302
    # 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
3303

3304
    visit_Node = Visitor.VisitorTransform.recurse_to_children
3305 3306


3307 3308 3309
class FinalOptimizePhase(Visitor.CythonTransform):
    """
    This visitor handles several commuting optimizations, and is run
3310 3311 3312 3313
    just before the C code generation phase.

    The optimizations currently implemented in this class are:
        - eliminate None assignment and refcounting for first assignment.
3314
        - isinstance -> typecheck for cdef types
Stefan Behnel's avatar
Stefan Behnel committed
3315
        - eliminate checks for None and/or types that became redundant after tree changes
3316
    """
3317
    def visit_SingleAssignmentNode(self, node):
3318 3319 3320 3321
        """Avoid redundant initialisation of local variables before their
        first assignment.
        """
        self.visitchildren(node)
3322 3323
        if node.first:
            lhs = node.lhs
3324
            lhs.lhs_of_first_assignment = True
3325
        return node
3326

3327
    def visit_SimpleCallNode(self, node):
3328 3329 3330
        """Replace generic calls to isinstance(x, type) by a more efficient
        type check.
        """
3331
        self.visitchildren(node)
Robert Bradshaw's avatar
Robert Bradshaw committed
3332
        if node.function.type.is_cfunction and isinstance(node.function, ExprNodes.NameNode):
3333 3334 3335
            if node.function.name == 'isinstance':
                type_arg = node.args[1]
                if type_arg.type.is_builtin_type and type_arg.type.name == 'type':
3336 3337
                    cython_scope = self.context.cython_scope
                    node.function.entry = cython_scope.lookup('PyObject_TypeCheck')
3338
                    node.function.type = node.function.entry.type
3339
                    PyTypeObjectPtr = PyrexTypes.CPtrType(cython_scope.lookup('PyTypeObject').type)
3340 3341
                    node.args[1] = ExprNodes.CastNode(node.args[1], PyTypeObjectPtr)
        return node
Stefan Behnel's avatar
Stefan Behnel committed
3342 3343 3344 3345 3346 3347 3348 3349 3350 3351 3352

    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
3353 3354 3355 3356 3357 3358 3359 3360 3361

    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
3362 3363 3364 3365 3366 3367 3368 3369 3370 3371 3372 3373 3374 3375 3376 3377 3378 3379 3380 3381 3382

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
    
    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
    
    def visit_NumBinopNode(self, node):
3383
        if node.overflow_check and node.overflow_fold:
3384 3385 3386 3387 3388 3389 3390 3391 3392 3393 3394 3395
            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