From 984192dff797bc6c686754c97f5038d8245587a9 Mon Sep 17 00:00:00 2001
From: Dnomd343 <dnomd343@gmail.com>
Date: Sun, 2 Feb 2025 00:23:15 +0800
Subject: [PATCH] feat: add pickle support of PyLayout

---
 src/core_ffi/python_ffi/binder/klsk_code.cc     |  5 +++++
 src/core_ffi/python_ffi/include/py_ffi/layout.h | 10 ++++++++++
 2 files changed, 15 insertions(+)

diff --git a/src/core_ffi/python_ffi/binder/klsk_code.cc b/src/core_ffi/python_ffi/binder/klsk_code.cc
index 1e6065f..55a0718 100644
--- a/src/core_ffi/python_ffi/binder/klsk_code.cc
+++ b/src/core_ffi/python_ffi/binder/klsk_code.cc
@@ -61,6 +61,11 @@ static void bind_layout(const py::module_ &mod) {
             return PyLayout(self);
         }, py::arg("memo"))
 
+        .def("__getstate__", [](const PyLayout &self) {
+            return self.value();
+        })
+        .def("__setstate__", &PyLayout::set_state)
+
         .def("next_cases", &PyLayout::next_cases)
         .def("dump_seq", &PyLayout::dump_seq)
         // TODO: add fast_cal / fast_cal_multi / ...
diff --git a/src/core_ffi/python_ffi/include/py_ffi/layout.h b/src/core_ffi/python_ffi/include/py_ffi/layout.h
index 91c34e4..33b793f 100644
--- a/src/core_ffi/python_ffi/include/py_ffi/layout.h
+++ b/src/core_ffi/python_ffi/include/py_ffi/layout.h
@@ -11,6 +11,11 @@
 
 // TODO: add `copy` and `pickle` support
 
+#include <iostream>
+#include <pybind11/pybind11.h>
+
+namespace py = pybind11;
+
 namespace klotski::ffi {
 
 enum class PyBlock : uint8_t {
@@ -71,6 +76,11 @@ public:
     /// Wrapper of `__repr__` method in Python.
     static std::string repr(PyLayout code) noexcept;
 
+    void set_state(uint64_t state) {
+        // std::cout << "set_state called" << std::endl;
+        code_ = codec::CommonCode::unsafe_create(state);
+    }
+
     // ------------------------------------------------------------------------------------- //
 
     [[nodiscard]] std::vector<PyLayout> next_cases() const noexcept;