From 350f55b4af726722a09e288f35f92a9b4a18fe27 Mon Sep 17 00:00:00 2001
From: Stefan Behnel <scoder@users.berlios.de>
Date: Mon, 22 Mar 2010 16:11:19 +0100
Subject: [PATCH] optimise len() for known builtin types

---
 Cython/Compiler/Builtin.py    |  6 +++--
 Cython/Compiler/Optimize.py   | 45 +++++++++++++++++++++++++++++------
 Cython/Compiler/PyrexTypes.py |  2 ++
 3 files changed, 44 insertions(+), 9 deletions(-)

diff --git a/Cython/Compiler/Builtin.py b/Cython/Compiler/Builtin.py
index 23f5ec71..2a4a4186 100644
--- a/Cython/Compiler/Builtin.py
+++ b/Cython/Compiler/Builtin.py
@@ -413,13 +413,15 @@ def init_builtins():
     init_builtin_funcs()
     init_builtin_types()
     init_builtin_structs()
-    global list_type, tuple_type, dict_type, set_type, type_type
-    global bytes_type, str_type, unicode_type, float_type, bool_type
+    global list_type, tuple_type, dict_type, set_type, frozenset_type
+    global bytes_type, str_type, unicode_type
+    global float_type, bool_type, type_type
     type_type  = builtin_scope.lookup('type').type
     list_type  = builtin_scope.lookup('list').type
     tuple_type = builtin_scope.lookup('tuple').type
     dict_type  = builtin_scope.lookup('dict').type
     set_type   = builtin_scope.lookup('set').type
+    frozenset_type = builtin_scope.lookup('frozenset').type
     bytes_type = builtin_scope.lookup('bytes').type
     str_type   = builtin_scope.lookup('str').type
     unicode_type = builtin_scope.lookup('unicode').type
diff --git a/Cython/Compiler/Optimize.py b/Cython/Compiler/Optimize.py
index 35e99528..42e23635 100644
--- a/Cython/Compiler/Optimize.py
+++ b/Cython/Compiler/Optimize.py
@@ -1230,6 +1230,22 @@ class OptimizeBuiltinCalls(Visitor.EnvTransform):
             PyrexTypes.CFuncTypeArg("bytes", PyrexTypes.c_char_ptr_type, None)
             ])
 
+    PyObject_Size_func_type = PyrexTypes.CFuncType(
+        PyrexTypes.c_py_ssize_t_type, [
+            PyrexTypes.CFuncTypeArg("obj", PyrexTypes.py_object_type, None)
+            ])
+
+    _map_to_capi_len_function = {
+        Builtin.unicode_type   : "PyUnicode_GET_SIZE",
+        Builtin.str_type       : "Py_SIZE",
+        Builtin.bytes_type     : "__Pyx_PyBytes_GET_SIZE",
+        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
+
     def _handle_simple_function_len(self, node, pos_args):
         """Replace len(char*) by the equivalent call to strlen().
         """
@@ -1239,14 +1255,29 @@ class OptimizeBuiltinCalls(Visitor.EnvTransform):
         arg = pos_args[0]
         if isinstance(arg, ExprNodes.CoerceToPyTypeNode):
             arg = arg.arg
-        if not arg.type.is_string:
+        if arg.type.is_string:
+            new_node = ExprNodes.PythonCapiCallNode(
+                node.pos, "strlen", self.Pyx_strlen_func_type,
+                args = [arg],
+                is_temp = node.is_temp,
+                utility_code = include_string_h_utility_code)
+        elif arg.type.is_pyobject:
+            if isinstance(arg, ExprNodes.NoneNode):
+                error(node.pos, "object of type 'NoneType' has no len()")
+                return node
+            cfunc_name = self._map_to_capi_len_function(arg.type)
+            if cfunc_name is None:
+                return node
+            if not arg.is_literal:
+                arg = ExprNodes.NoneCheckNode(
+                    arg, "PyExc_TypeError",
+                    "object of type 'NoneType' has no len()")
+            new_node = ExprNodes.PythonCapiCallNode(
+                node.pos, cfunc_name, self.PyObject_Size_func_type,
+                args = [arg],
+                is_temp = node.is_temp)
+        else:
             return node
-        new_node = ExprNodes.PythonCapiCallNode(
-            node.pos, "strlen", self.Pyx_strlen_func_type,
-            args = [arg],
-            is_temp = node.is_temp,
-            utility_code = include_string_h_utility_code
-            )
         if node.type not in (PyrexTypes.c_size_t_type, PyrexTypes.c_py_ssize_t_type):
             new_node = new_node.coerce_to(node.type, self.env_stack[-1])
         return new_node
diff --git a/Cython/Compiler/PyrexTypes.py b/Cython/Compiler/PyrexTypes.py
index 19dd80dc..b3589219 100755
--- a/Cython/Compiler/PyrexTypes.py
+++ b/Cython/Compiler/PyrexTypes.py
@@ -2464,10 +2464,12 @@ type_conversion_predeclarations = """
 #define __Pyx_PyBytes_FromString          PyString_FromString
 #define __Pyx_PyBytes_FromStringAndSize   PyString_FromStringAndSize
 #define __Pyx_PyBytes_AsString            PyString_AsString
+#define __Pyx_PyBytes_GET_SIZE            PyString_GET_SIZE
 #else
 #define __Pyx_PyBytes_FromString          PyBytes_FromString
 #define __Pyx_PyBytes_FromStringAndSize   PyBytes_FromStringAndSize
 #define __Pyx_PyBytes_AsString            PyBytes_AsString
+#define __Pyx_PyBytes_GET_SIZE            PyBytes_GET_SIZE
 #endif
 
 #define __Pyx_PyBytes_FromUString(s)      __Pyx_PyBytes_FromString((char*)s)
-- 
2.26.2