From 09d04901bb0b4101ebea2f64ffea7616abf735ad Mon Sep 17 00:00:00 2001
From: Nicolas Cornu <nicolas.cornu@epfl.ch>
Date: Wed, 14 Aug 2024 11:01:09 +0200
Subject: [PATCH] Use nanobind for python BBS implementation (#3034)

---
 src/nrnpython/nrnpy_p2h.cpp | 61 +++++++++++++++++++------------------
 1 file changed, 32 insertions(+), 29 deletions(-)

diff --git a/src/nrnpython/nrnpy_p2h.cpp b/src/nrnpython/nrnpy_p2h.cpp
index 0d765bb6a5..fddf59ebbe 100644
--- a/src/nrnpython/nrnpy_p2h.cpp
+++ b/src/nrnpython/nrnpy_p2h.cpp
@@ -1,6 +1,7 @@
 #include <../../nrnconf.h>
 
 #include <cstdio>
+
 #include <InterViews/resource.h>
 #include <nrnoc2iv.h>
 #include <classreg.h>
@@ -11,6 +12,10 @@
 #include "oc_ansi.h"
 
 #include "parse.hpp"
+
+#include <nanobind/nanobind.h>
+namespace nb = nanobind;
+
 static void nrnpy_decref_defer(PyObject*);
 static char* nrnpyerr_str();
 static PyObject* nrnpy_pyCallObject(PyObject*, PyObject*);
@@ -607,32 +612,31 @@ static int guigetstr(Object* ho, char** cpp) {
     return 1;
 }
 
-static PyObject* loads;
-static PyObject* dumps;
+static nb::callable loads;
+static nb::callable dumps;
 
 static void setpickle() {
     if (dumps) {
         return;
     }
-    PyObject* pickle = PyImport_ImportModule("pickle");
-    if (pickle) {
-        Py_INCREF(pickle);
-        dumps = PyObject_GetAttrString(pickle, "dumps");
-        loads = PyObject_GetAttrString(pickle, "loads");
-        if (dumps) {
-            Py_INCREF(dumps);
-            Py_INCREF(loads);
-        }
-    }
+    nb::module_ pickle = nb::module_::import_("pickle");
+    dumps = pickle.attr("dumps");
+    loads = pickle.attr("loads");
     if (!dumps || !loads) {
         hoc_execerror("Neither Python cPickle nor pickle are available", 0);
     }
+    // We intentionally leak these, because if we don't
+    // we observe SEGFAULTS during application shutdown.
+    // Likely because "~nb::callable" runs after the Python
+    // library cleaned up.
+    dumps.inc_ref();
+    loads.inc_ref();
 }
 
 // note that *size includes the null terminating character if it exists
 static std::vector<char> pickle(PyObject* p) {
     PyObject* arg = PyTuple_Pack(1, p);
-    PyObject* r = nrnpy_pyCallObject(dumps, arg);
+    PyObject* r = nrnpy_pyCallObject(dumps.ptr(), arg);
     Py_XDECREF(arg);
     if (!r && PyErr_Occurred()) {
         PyErr_Print();
@@ -656,25 +660,21 @@ static std::vector<char> po2pickle(Object* ho) {
     }
 }
 
-static PyObject* unpickle(const char* s, std::size_t len) {
-    PyObject* ps = PyBytes_FromStringAndSize(s, len);
-    PyObject* arg = PyTuple_Pack(1, ps);
-    PyObject* po = nrnpy_pyCallObject(loads, arg);
-    assert(po);
-    Py_XDECREF(arg);
-    Py_XDECREF(ps);
-    return po;
+static nb::object unpickle(const char* s, std::size_t len) {
+    nb::bytes string(s, len);
+    nb::list args;
+    args.append(string);
+    return loads(*args);
 }
 
-static PyObject* unpickle(const std::vector<char>& s) {
+static nb::object unpickle(const std::vector<char>& s) {
     return unpickle(s.data(), s.size());
 }
 
 static Object* pickle2po(const std::vector<char>& s) {
     setpickle();
-    PyObject* po = unpickle(s);
-    Object* ho = nrnpy_pyobject_in_obj(po);
-    Py_XDECREF(po);
+    nb::object po = unpickle(s);
+    Object* ho = nrnpy_pyobject_in_obj(po.ptr());
     return ho;
 }
 
@@ -750,7 +750,7 @@ std::vector<char> call_picklef(const std::vector<char>& fname, int narg) {
     setpickle();
     PyObject* ps = PyBytes_FromStringAndSize(fname.data(), fname.size());
     args = PyTuple_Pack(1, ps);
-    callable = nrnpy_pyCallObject(loads, args);
+    callable = nrnpy_pyCallObject(loads.ptr(), args);
     assert(callable);
     Py_XDECREF(args);
     Py_XDECREF(ps);
@@ -801,7 +801,8 @@ static PyObject* char2pylist(char* buf, int np, int* cnt, int* displ) {
             Py_INCREF(Py_None);  // 'Fatal Python error: deallocating None' eventually
             PyList_SetItem(plist, i, Py_None);
         } else {
-            PyObject* p = unpickle(buf + displ[i], cnt[i]);
+            nb::object po = unpickle(buf + displ[i], cnt[i]);
+            PyObject* p = po.release().ptr();
             PyList_SetItem(plist, i, p);
         }
     }
@@ -874,7 +875,8 @@ static PyObject* py_broadcast(PyObject* psrc, int root) {
     nrnmpi_char_broadcast(buf.data(), cnt, root);
     PyObject* pdest = psrc;
     if (root != nrnmpi_myid) {
-        pdest = unpickle(buf);
+        nb::object po = unpickle(buf);
+        pdest = po.release().ptr();
     } else {
         Py_INCREF(pdest);
     }
@@ -1082,7 +1084,8 @@ static Object* py_alltoall_type(int size, int type) {
                 delete[] sdispl;
 
             if (rcnt[0]) {
-                pdest = unpickle(r);
+                nb::object po = unpickle(r);
+                pdest = po.release().ptr();
             } else {
                 pdest = Py_None;
                 Py_INCREF(pdest);