From 292892d1aeeb5a7b0c1f68342874d49a93559dd3 Mon Sep 17 00:00:00 2001
From: Matthias Koefferlein <matthias@klayout.de>
Date: Sat, 15 Jul 2023 22:50:29 +0200
Subject: [PATCH] Some performance enhancement for Python binding

---
 src/pya/pya/pyaInternal.cc | 24 +++++++++++++++++++++++-
 src/pya/pya/pyaInternal.h  |  1 +
 src/pya/pya/pyaModule.cc   | 23 +----------------------
 3 files changed, 25 insertions(+), 23 deletions(-)

diff --git a/src/pya/pya/pyaInternal.cc b/src/pya/pya/pyaInternal.cc
index 37ad0395f..433c03220 100644
--- a/src/pya/pya/pyaInternal.cc
+++ b/src/pya/pya/pyaInternal.cc
@@ -869,10 +869,17 @@ MethodTable::method_table_by_class (const gsi::ClassBase *cls_decl)
 // -------------------------------------------------------------------
 //  PythonClassClientData implementation
 
+static std::map<PyTypeObject *, const gsi::ClassBase *> s_type2cls;
+
 PythonClassClientData::PythonClassClientData (const gsi::ClassBase *_cls, PyTypeObject *_py_type, PyTypeObject *_py_type_static, PythonModule *module)
   : py_type_object ((PyObject *) _py_type), py_type_object_static ((PyObject *) _py_type_static), method_table (_cls, module)
 {
-  //  .. nothing yet ..
+  if (_py_type != NULL) {
+    s_type2cls.insert (std::make_pair (_py_type, _cls));
+  }
+  if (_py_type_static != NULL) {
+    s_type2cls.insert (std::make_pair (_py_type_static, _cls));
+  }
 }
 
 PythonClassClientData::~PythonClassClientData ()
@@ -890,11 +897,26 @@ PythonClassClientData::py_type (const gsi::ClassBase &cls_decl, bool as_static)
   return (PyTypeObject *) (cd ? (as_static ? cd->py_type_object_static.get () : cd->py_type_object.get ()) : 0);
 }
 
+const gsi::ClassBase *
+PythonClassClientData::cls_for_type (PyTypeObject *type)
+{
+  while (type && type != &PyBaseObject_Type) {
+    auto t2c = s_type2cls.find (type);
+    if (t2c != s_type2cls.end ()) {
+      return t2c->second;
+    }
+    type = type->tp_base;
+  }
+
+  return 0;
+}
+
 void
 PythonClassClientData::initialize (const gsi::ClassBase &cls_decl, PyTypeObject *py_type, bool as_static, PythonModule *module)
 {
   PythonClassClientData *cd = dynamic_cast<PythonClassClientData *>(cls_decl.data (gsi::ClientIndex::Python));
   if (cd) {
+    s_type2cls.insert (std::make_pair (py_type, &cls_decl));
     if (as_static) {
       cd->py_type_object_static = (PyObject *) py_type;
     } else {
diff --git a/src/pya/pya/pyaInternal.h b/src/pya/pya/pyaInternal.h
index 92b227252..8ecffd25d 100644
--- a/src/pya/pya/pyaInternal.h
+++ b/src/pya/pya/pyaInternal.h
@@ -309,6 +309,7 @@ struct PythonClassClientData
   MethodTable method_table;
 
   static PyTypeObject *py_type (const gsi::ClassBase &cls_decl, bool as_static);
+  static const gsi::ClassBase *cls_for_type (PyTypeObject *);
   static void initialize (const gsi::ClassBase &cls_decl, PyTypeObject *py_type, bool as_static, PythonModule *module);
 };
 
diff --git a/src/pya/pya/pyaModule.cc b/src/pya/pya/pyaModule.cc
index d06e811a3..ad77f0ee0 100644
--- a/src/pya/pya/pyaModule.cc
+++ b/src/pya/pya/pyaModule.cc
@@ -60,7 +60,6 @@ set_type_attr (PyTypeObject *type, const std::string &name, PythonRef &attr)
 
 std::map<const gsi::MethodBase *, std::string> PythonModule::m_python_doc;
 std::vector<const gsi::ClassBase *> PythonModule::m_classes;
-std::map<PyTypeObject *, const gsi::ClassBase *> PythonModule::m_class_by_type;
 
 const std::string pymod_name ("klayout");
 
@@ -738,27 +737,7 @@ PythonModule::make_classes (const char *mod_name)
 
 const gsi::ClassBase *PythonModule::cls_for_type (PyTypeObject *type)
 {
-  auto cls = m_class_by_type.find (type);
-  if (cls != m_class_by_type.end ()) {
-    return cls->second;
-  }
-
-  //  GSI classes store their class index inside the __gsi_id__ attribute
-  if (PyObject_HasAttrString ((PyObject *) type, "__gsi_id__")) {
-
-    PyObject *cls_id = PyObject_GetAttrString ((PyObject *) type, "__gsi_id__");
-    if (cls_id != NULL && pya::test_type<size_t> (cls_id)) {
-      size_t i = pya::python2c<size_t> (cls_id);
-      if (i < m_classes.size ()) {
-        const gsi::ClassBase *gsi_cls = m_classes [i];
-        m_class_by_type.insert (std::make_pair (type, gsi_cls));
-        return gsi_cls;
-      }
-    }
-
-  }
-
-  return 0;
+  return PythonClassClientData::cls_for_type (type);
 }
 
 PyTypeObject *PythonModule::type_for_cls (const gsi::ClassBase *cls)