✨ Added pickle support to Aig (#145)

* 🚧 `pickle` support for AIGs

* 🚧 Initial tests for `pickle` support

* ✏️ Fixed a typo

* ✨ Add `write_dot` function for AIG network visualization

* 🔥 Remove `pybind11` dependency from build requirements

* 🔧 Add type stubs for `matplotlib` and `networkx` to pre-commit configuration

* 🔧 Ignore security warnings for `pickle` usage in `pyproject.toml`

* 🚧 Save commit

* 🐛 Fixed `pickle` support via `AigIndexList`

* ✅ Test cases demonstrating that `pickle` support works now

* 🏷️ Add parameter name and type hints for `pickle` support

* ✅ Add tests for exception handling in Aig unpickling

* 🎨 Move import statements for `pickle` into the test cases where explicitly needed

* ✅ Add test for pickling and unpickling multiple AIGs

* 📝 Add documentation for `pickle` support

* 🎨 Refactor AIG visualization code and remove unused dependencies

* 🎨 Update AIG documentation with improved section headings and additional details on AIG views

* 🔒 Update lock file

* 🔧 Update pre-commit configuration to skip mypy checks due to memory restrictions

* ✅ Add unit test for `write_dot` function to verify file creation and content

* 📝 Add documentation on `write_dot`

Author: marcelwa

.pre-commit-config.yaml

@@ -11,8 +11,9 @@ ci:
   autoupdate_commit_msg: "⬆️ Bump pre-commit hooks"
   autoupdate_schedule: quarterly
   autofix_commit_msg: "🎨 Incorporated pre-commit fixes"
+  skip: [mypy]
 
-exclude: "^(libs/|docs/examples/i10\\.aig$)"
+exclude: "^(libs/|docs/examples/)"
 
 repos:
   # Standard hooks
@@ -71,6 +72,8 @@ repos:
         additional_dependencies:
           - nox
           - pytest
+          - matplotlib-stubs
+          - networkx-stubs
 
   # clang-format the C++ part of the code base
   - repo: https://github.com/pre-commit/mirrors-clang-format

README.md

@@ -205,12 +205,14 @@ Verilog and PLA.
 #### ✏️ Writing
 
 ```python
-from aigverse import write_aiger, write_verilog
+from aigverse import write_aiger, write_verilog, write_dot
 
 # Write an AIG network to an AIGER file
 write_aiger(aig, "example.aig")
 # Write an AIG network to a Verilog file
 write_verilog(aig, "example.v")
+# Write an AIG network to a DOT file
+write_dot(aig, "example.dot")
 ```
 
 #### 👓 Parsing
@@ -235,6 +237,23 @@ aig4 = read_pla_into_aig("example.pla")
 Additionally, you can read AIGER files into sequential AIGs using `read_aiger_into_sequential_aig` and
 `read_ascii_aiger_into_sequential_aig`.
 
+### 🥒 `pickle` Support
+
+AIGs support Python's `pickle` protocol, allowing you to serialize and deserialize AIG objects for persistent storage or
+interface with data science or machine learning workflows.
+
+```python
+import pickle
+
+with open("aig.pkl", "wb") as f:
+    pickle.dump(aig, f)
+
+with open("aig.pkl", "rb") as f:
+    unpickled_aig = pickle.load(f)
+```
+
+You can also pickle multiple AIGs at once by storing them in a tuple or list.
+
 ### 🔌 Adapters
 
 Adapters provide alternative representations of AIGs for integration with other tools or workflows.

bindings/aigverse/__init__.py

@@ -34,6 +34,7 @@
     to_edge_list,
     to_index_list,
     write_aiger,
+    write_dot,
     write_verilog,
 )
 
@@ -69,5 +70,6 @@
     "to_edge_list",
     "to_index_list",
     "write_aiger",
+    "write_dot",
     "write_verilog",
 ]

bindings/aigverse/aigverse.cpp

@@ -13,6 +13,7 @@
 #include "aigverse/io/read_pla.hpp"
 #include "aigverse/io/read_verilog.hpp"
 #include "aigverse/io/write_aiger.hpp"
+#include "aigverse/io/write_dot.hpp"
 #include "aigverse/io/write_verilog.hpp"
 #include "aigverse/networks/logic_networks.hpp"
 #include "aigverse/truth_tables/operations.hpp"
@@ -55,6 +56,7 @@ PYBIND11_MODULE(pyaigverse, m, pybind11::mod_gil_not_used())
     aigverse::read_pla(m);
     aigverse::read_verilog(m);
     aigverse::write_verilog(m);
+    aigverse::write_dot(m);
 
     /**
      * Adapters

bindings/aigverse/include/aigverse/adapters/index_list.hpp

@@ -37,7 +37,7 @@ void ntk_index_list(pybind11::module& m, const std::string& network_name)
         /**
          * Index list.
          */
-        using IndexList = mockturtle::xag_index_list<true>;
+        using IndexList = aigverse::aig_index_list;
         py::class_<IndexList>(m, fmt::format("{}IndexList", network_name).c_str())
             .def(py::init<const uint32_t>(), "num_pis"_a = 0)
             .def(py::init<const std::vector<uint32_t>&>(), "values"_a)
@@ -156,7 +156,7 @@ namespace fmt
 
 // make index_list compatible with fmt::format
 template <>
-struct formatter<mockturtle::xag_index_list<true>>
+struct formatter<aigverse::aig_index_list>
 {
     template <typename ParseContext>
     constexpr auto parse(ParseContext& ctx)
@@ -165,7 +165,7 @@ struct formatter<mockturtle::xag_index_list<true>>
     }
 
     template <typename FormatContext>
-    auto format(const mockturtle::xag_index_list<true>& il, FormatContext& ctx) const
+    auto format(const aigverse::aig_index_list& il, FormatContext& ctx) const
     {
         std::vector<std::tuple<uint32_t, uint32_t>> gates{};
         gates.reserve(il.num_gates());

bindings/aigverse/include/aigverse/io/write_dot.hpp

@@ -0,0 +1,40 @@
+//
+// Created by marcel on 22.04.24.
+//
+
+#ifndef AIGVERSE_WRITE_DOT_HPP
+#define AIGVERSE_WRITE_DOT_HPP
+
+#include <mockturtle/io/write_dot.hpp>
+#include <mockturtle/networks/aig.hpp>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include <string>
+
+namespace aigverse
+{
+
+namespace detail
+{
+
+template <typename Ntk>
+void write_dot(pybind11::module& m)
+{
+    using namespace pybind11::literals;
+
+    m.def(
+        "write_dot", [](const Ntk& ntk, const std::string& filename) { mockturtle::write_dot(ntk, filename); },
+        "network"_a, "filename"_a);
+}
+
+}  // namespace detail
+
+inline void write_dot(pybind11::module& m)
+{
+    detail::write_dot<mockturtle::aig_network>(m);
+}
+
+}  // namespace aigverse
+
+#endif  // AIGVERSE_WRITE_DOT_HPP

bindings/aigverse/include/aigverse/networks/logic_networks.hpp

@@ -232,7 +232,45 @@ void network(pybind11::module& m, const std::string& network_name)
         .def(
             "is_nary_or", [](const Ntk& ntk, const Node& n) { return ntk.is_nary_or(n); }, "n"_a)
 
-        .def("cleanup_dangling", [](Ntk& ntk) { ntk = mockturtle::cleanup_dangling(ntk); });
+        // pickle support
+        .def(py::pickle(
+                 [](const Ntk& ntk)  // __getstate__
+                 {
+                     aigverse::aig_index_list il{};
+                     mockturtle::encode(il, ntk);
+                     return py::make_tuple(py::cast(il.raw()));
+                 },
+                 [](const py::tuple& state)  // __setstate__
+                 {
+                     if (state.size() != 1)
+                     {
+                         throw pybind11::value_error(
+                             "Invalid state: expected a tuple of size 1 containing an index list");
+                     }
+                     try
+                     {
+                         aigverse::aig_index_list il{state[0].cast<std::vector<uint32_t>>()};
+
+                         Ntk ntk{};
+                         mockturtle::decode(ntk, il);
+                         return ntk;
+                     }
+                     catch (const pybind11::cast_error& e)
+                     {
+                         throw pybind11::value_error(
+                             fmt::format("Invalid state: expected an index list. {}", e.what()));
+                     }
+                     catch (const std::exception& e)
+                     {
+                         throw pybind11::value_error(fmt::format("Failed to restore network state: {}", e.what()));
+                     }
+                 }),
+             "state"_a)
+
+        // clean up dangling nodes (after optimization)
+        .def("cleanup_dangling", [](Ntk& ntk) { ntk = mockturtle::cleanup_dangling(ntk); })
+
+        ;
 
     /**
      * Depth network.

bindings/aigverse/include/aigverse/types.hpp

@@ -8,6 +8,7 @@
 #include <kitty/dynamic_truth_table.hpp>
 #include <mockturtle/networks/aig.hpp>
 #include <mockturtle/networks/sequential.hpp>
+#include <mockturtle/utils/index_list.hpp>
 #include <mockturtle/views/depth_view.hpp>
 #include <mockturtle/views/fanout_view.hpp>
 
@@ -30,6 +31,10 @@ using fanouts_aig = mockturtle::fanout_view<aig>;
  * Alias for the sequential AIG.
  */
 using sequential_aig = mockturtle::sequential<aig>;
+/**
+ * Alias for the AIG index list.
+ */
+using aig_index_list = mockturtle::xag_index_list<true>;
 /**
  * Alias for the truth table.
  */

bindings/aigverse/pyaigverse.pyi

@@ -77,6 +77,8 @@ class Aig:
     def is_xor3(self, n: AigNode | int) -> bool: ...
     def is_nary_and(self, n: AigNode | int) -> bool: ...
     def is_nary_or(self, n: AigNode | int) -> bool: ...
+    def __getstate__(self) -> tuple[list[int]]: ...
+    def __setstate__(self, state: tuple[list[int]]) -> Aig: ...
     def cleanup_dangling(self) -> None: ...
 
 class DepthAig(Aig):
@@ -139,10 +141,11 @@ def read_aiger_into_aig(filename: str) -> Aig: ...
 def read_ascii_aiger_into_aig(filename: str) -> Aig: ...
 def read_aiger_into_sequential_aig(filename: str) -> SequentialAig: ...
 def read_ascii_aiger_into_sequential_aig(filename: str) -> SequentialAig: ...
-def write_aiger(ntk: Aig, filename: str) -> None: ...
 def read_pla_into_aig(filename: str) -> Aig: ...
 def read_verilog_into_aig(filename: str) -> Aig: ...
+def write_aiger(ntk: Aig, filename: str) -> None: ...
 def write_verilog(ntk: Aig, filename: str) -> None: ...
+def write_dot(ntk: Aig, filename: str) -> None: ...
 
 class AigEdge:
     def __init__(

bindings/aigverse/test/inout/test_write_dot.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+import tempfile
+from pathlib import Path
+
+from aigverse import Aig, write_dot
+
+# Get the temporary directory as a Path object
+temp_dir = Path(tempfile.gettempdir())
+
+
+def test_write_dot_creates_file() -> None:
+    aig = Aig()
+    x1 = aig.create_pi()
+    x2 = aig.create_pi()
+    f = aig.create_and(x1, x2)
+    aig.create_po(f)
+
+    dot_path = temp_dir / "aig.dot"
+    if dot_path.exists():
+        dot_path.unlink()
+
+    write_dot(aig, str(dot_path))
+    assert dot_path.exists()
+
+    content = dot_path.read_text()
+
+    expected = """
+    digraph {
+        rankdir=BT;
+
+        0 [label="0",shape=box,style=filled,fillcolor=snow2]
+        1 [label="1",shape=triangle,style=filled,fillcolor=snow2]
+        2 [label="2",shape=triangle,style=filled,fillcolor=snow2]
+        3 [label="3",shape=ellipse,style=filled,fillcolor=white]
+
+        po0 [shape=invtriangle,style=filled,fillcolor=snow2]
+
+        1 -> 3 [style=solid]
+        2 -> 3 [style=solid]
+        3 -> po0 [style=solid]
+
+        {rank = same; 0; 1; 2; }
+        {rank = same; 3; }
+        {rank = same; po0; }
+    }
+    """
+
+    def normalize(s: str) -> str:
+        return "".join(s.split())
+
+    assert normalize(content) == normalize(expected)