Createing python bindings

BUILD.md

@@ -32,6 +32,7 @@ Pass these to the `cmake -S dependencies` step:
 | `USE_EXTERNAL_LLVM` | ON | Use system LLVM (OFF = build from source) |
 | `USE_EXTERNAL_ABSEIL` | OFF | Use system abseil (OFF = build from source) |
 | `USE_EXTERNAL_HIGHWAY` | OFF | Use system highway (OFF = build from source) |
+| `USE_EXTERNAL_NANOBIND` | OFF | Use system nanobind (OFF = build from source) |
 | `COBRA_BUILD_TESTS` | OFF | Build GoogleTest for tests |
 | `USE_EXTERNAL_GOOGLETEST` | OFF | Use system GoogleTest (OFF = build from source) |
 | `COBRA_ENABLE_Z3` | OFF | Enable Z3 support (requires lib/verify implementation) |
@@ -45,6 +46,7 @@ Pass these to the `cmake -S .` step:
 |--------|---------|-------------|
 | `COBRA_BUILD_LLVM_PASS` | OFF | Build the LLVM pass plugin (requires LLVM 19-22) |
 | `COBRA_BUILD_TESTS` | OFF | Build tests (requires GoogleTest in prefix) |
+| `COBRA_BUILD_PYTHON_BINDINGS` | OFF | Builds the Python module (requires nanobind) |
 
 ## With LLVM Pass Plugin
 

CMakeLists.txt

@@ -6,6 +6,27 @@ set(CMAKE_CXX_STANDARD_REQUIRED ON)
 set(CMAKE_CXX_EXTENSIONS OFF)
 set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 
+# Prefer local dependency prefix if present (helps wheel builds find absl/hwy)
+if(DEFINED COBRA_DEPENDENCY_PREFIX)
+    list(PREPEND CMAKE_PREFIX_PATH "${COBRA_DEPENDENCY_PREFIX}")
+else()
+    set(_cobra_default_dep_prefix "${CMAKE_SOURCE_DIR}/build-deps/install")
+    if(EXISTS "${_cobra_default_dep_prefix}")
+        list(PREPEND CMAKE_PREFIX_PATH "${_cobra_default_dep_prefix}")
+    endif()
+endif()
+
+if (CMAKE_EXPORT_COMPILE_COMMANDS AND UNIX)
+    set(_cobra_cc_src "${CMAKE_SOURCE_DIR}/compile_commands.json")
+    set(_cobra_cc_bin "${CMAKE_BINARY_DIR}/compile_commands.json")
+    if (NOT EXISTS "${_cobra_cc_src}")
+        execute_process(
+            COMMAND ${CMAKE_COMMAND} -E create_symlink "${_cobra_cc_bin}" "${_cobra_cc_src}"
+            OUTPUT_QUIET ERROR_QUIET
+        )
+    endif()
+endif()
+
 include(GNUInstallDirs)
 include(CMakePackageConfigHelpers)
 
@@ -57,6 +78,7 @@ add_link_options(
 
 option(COBRA_BUILD_LLVM_PASS "Build the LLVM pass plugin (requires LLVM 19-22)" OFF)
 option(COBRA_BUILD_TESTS "Build tests (requires GoogleTest in prefix)" OFF)
+option(COBRA_BUILD_PYTHON_BINDINGS "Build Python bindings (nanobind)" OFF)
 option(COBRA_ENABLE_TRACE "Enable detailed pipeline tracing to stderr (debug builds)" OFF)
 option(COBRA_ENABLE_SIG_STATS "Enable signature evaluation counters for profiling builds" OFF)
 option(COBRA_ENABLE_TRACY "Enable Tracy profiler instrumentation" OFF)
@@ -98,6 +120,10 @@ if(COBRA_BUILD_LLVM_PASS)
     add_subdirectory(lib/llvm)
 endif()
 
+if(COBRA_BUILD_PYTHON_BINDINGS)
+    add_subdirectory(lib/bindings/python)
+endif()
+
 if(COBRA_BUILD_TESTS)
     enable_testing()
     add_subdirectory(test)
@@ -113,7 +139,11 @@ install(DIRECTORY include/cobra
 
 # Collect installed targets for the package config export set
 set(_cobra_install_targets cobra-core cobra-cli)
-if(Z3_FOUND)
+if(DEFINED SKBUILD_PLATLIB_DIR)
+    # Wheel builds only need the extension module; skip installing the CLI.
+    set(_cobra_install_targets cobra-core)
+endif()
+if(Z3_FOUND AND NOT DEFINED SKBUILD_PLATLIB_DIR)
     list(APPEND _cobra_install_targets cobra-verify)
 endif()
 

dependencies/CMakeLists.txt

@@ -6,19 +6,26 @@
 #   cmake --build build-deps
 #
 # Options:
-#   USE_EXTERNAL_LLVM       (default ON)  - Use system LLVM vs build from source
-#   USE_EXTERNAL_HIGHWAY    (default OFF) - Use system highway vs build from source
-#   COBRA_BUILD_TESTS       (default OFF) - Build GoogleTest for tests
-#   USE_EXTERNAL_GOOGLETEST (default OFF) - Use system GoogleTest vs build from source
-#   COBRA_ENABLE_Z3         (default OFF) - Enable Z3 dependency
-#   USE_EXTERNAL_Z3         (default OFF) - Use system Z3 vs build from source
+#   USE_EXTERNAL_LLVM         (default ON)  - Use system LLVM vs build from source
+#   COBRA_BUILD_LLVM_PASS      (default ON)  - Prepare LLVM dependency for pass plugin
+#   USE_EXTERNAL_ABSEIL       (default OFF) - Use system abseil vs build from source
+#   USE_EXTERNAL_HIGHWAY      (default OFF) - Use system highway vs build from source
+#   COBRA_BUILD_PYTHON_BINDINGS (default OFF) - Build nanobind for Python bindings
+#   USE_EXTERNAL_NANOBIND     (default OFF) - Use system nanobind vs build from source
+#   COBRA_BUILD_TESTS         (default OFF) - Build GoogleTest for tests
+#   USE_EXTERNAL_GOOGLETEST   (default OFF) - Use system GoogleTest vs build from source
+#   COBRA_ENABLE_Z3           (default OFF) - Enable Z3 dependency
+#   USE_EXTERNAL_Z3           (default OFF) - Use system Z3 vs build from source
 
 cmake_minimum_required(VERSION 3.20)
 project(cobra-dependencies LANGUAGES C CXX)
 
 option(USE_EXTERNAL_LLVM "Use system LLVM instead of building from source" ON)
+option(COBRA_BUILD_LLVM_PASS "Prepare LLVM dependency for the pass plugin" ON)
 option(USE_EXTERNAL_ABSEIL "Use system abseil instead of building from source" OFF)
 option(USE_EXTERNAL_HIGHWAY "Use system highway instead of building from source" OFF)
+option(COBRA_BUILD_PYTHON_BINDINGS "Build nanobind for Python bindings" OFF)
+option(USE_EXTERNAL_NANOBIND "Use system nanobind instead of building from source" OFF)
 option(COBRA_BUILD_TESTS "Build GoogleTest for CoBRA tests" OFF)
 option(USE_EXTERNAL_GOOGLETEST "Use system GoogleTest instead of building from source" OFF)
 option(COBRA_ENABLE_Z3 "Enable Z3 dependency (requires lib/verify to be implemented)" OFF)
@@ -28,7 +35,13 @@ include(superbuild.cmake)
 
 include(abseil.cmake)
 include(highway.cmake)
-include(llvm.cmake)
+if(COBRA_BUILD_LLVM_PASS)
+    include(llvm.cmake)
+endif()
+
+if(COBRA_BUILD_PYTHON_BINDINGS)
+    include(nanobind.cmake)
+endif()
 
 if(COBRA_BUILD_TESTS)
     include(googletest.cmake)

dependencies/abseil.cmake

@@ -22,6 +22,7 @@ else()
         CMAKE_ARGS
             ${COBRA_COMMON_CMAKE_ARGS}
             -DABSL_BUILD_TESTING=OFF
+            -DABSL_ENABLE_INSTALL=ON
             -DABSL_PROPAGATE_CXX_STD=ON
             -DCMAKE_POSITION_INDEPENDENT_CODE=ON
     )

dependencies/nanobind.cmake

@@ -0,0 +1,37 @@
+# dependencies/nanobind.cmake
+# nanobind dependency: forwarding config (external) or source build.
+
+set(_nanobind_config_dir "${COBRA_INSTALL_PREFIX}/lib/cmake/nanobind")
+file(MAKE_DIRECTORY "${_nanobind_config_dir}")
+
+if(USE_EXTERNAL_NANOBIND)
+    find_package(nanobind CONFIG REQUIRED)
+    message(STATUS "Using external nanobind")
+
+    file(WRITE "${_nanobind_config_dir}/nanobind-config.cmake"
+        "# Forwarding config - delegates to system nanobind\n"
+        "include(\"${nanobind_DIR}/nanobind-config.cmake\")\n"
+    )
+
+    cobra_mark_satisfied(nanobind)
+else()
+    message(STATUS "Building nanobind from source (v2.12.0)")
+    cobra_add_dependency(nanobind
+        GIT_REPOSITORY https://github.com/wjakob/nanobind.git
+        GIT_TAG 2a61ad2494d09fecb2e13322c1383342c299900d # v2.12.0
+        GIT_SHALLOW ON
+        GIT_PROGRESS ON
+        GIT_SUBMODULES_RECURSE ON
+        CMAKE_ARGS
+            ${COBRA_COMMON_CMAKE_ARGS}
+            -DNB_TEST=OFF
+            -DNB_CREATE_INSTALL_RULES=ON
+            -DNB_USE_SUBMODULE_DEPS=ON
+            -DCMAKE_POSITION_INDEPENDENT_CODE=ON
+    )
+
+    file(WRITE "${_nanobind_config_dir}/nanobind-config.cmake"
+        "# Forwarding config - delegates to installed nanobind\n"
+        "include(\"${COBRA_INSTALL_PREFIX}/nanobind/cmake/nanobind-config.cmake\")\n"
+    )
+endif()

lib/bindings/python/CMakeLists.txt

@@ -0,0 +1,39 @@
+find_package(Python COMPONENTS Interpreter Development.Module REQUIRED)
+if(Python_VERSION VERSION_LESS 3.10)
+  message(FATAL_ERROR "Python 3.10+ is required for cobra_mba")
+endif()
+
+if (NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES)
+  set(CMAKE_BUILD_TYPE Release CACHE STRING "Choose the type of build." FORCE)
+  set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS "Debug" "Release" "MinSizeRel" "RelWithDebInfo")
+endif()
+
+if (NOT nanobind_DIR AND NOT nanobind_ROOT)
+  execute_process(
+    COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
+    OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE nanobind_ROOT
+  )
+endif()
+if (nanobind_ROOT AND NOT nanobind_DIR)
+  set(nanobind_DIR "${nanobind_ROOT}")
+endif()
+find_package(nanobind CONFIG REQUIRED)
+
+nanobind_add_module(
+  cobra_mba
+  PythonInterface.cpp
+  CobraPython.cpp
+  CobraPython.hpp
+  ${PROJECT_SOURCE_DIR}/tools/cobra-cli/ExprParser.cpp
+)
+target_link_libraries(cobra_mba PRIVATE cobra-core)
+target_include_directories(cobra_mba PRIVATE ${PROJECT_SOURCE_DIR}/tools/cobra-cli)
+
+if (DEFINED SKBUILD_PLATLIB_DIR)
+  install(TARGETS cobra_mba
+    LIBRARY DESTINATION "${SKBUILD_PLATLIB_DIR}"
+    RUNTIME DESTINATION "${SKBUILD_PLATLIB_DIR}"
+  )
+endif()
+
+

lib/bindings/python/CobraPython.cpp

@@ -0,0 +1,166 @@
+#include "CobraPython.hpp"
+#include "ExprParser.h"
+#include "cobra/core/Classifier.h"
+#include "cobra/core/ExprUtils.h"
+#include "cobra/core/SignatureChecker.h"
+#include "cobra/core/Simplifier.h"
+
+#include <stdexcept>
+#include <string>
+
+namespace {
+std::vector<uint64_t> EvaluateToSignature(const cobra::Expr &ast,
+                                          uint32_t num_vars,
+                                          uint32_t bitwidth) {
+  const size_t kLen = size_t{1} << num_vars;
+  std::vector<uint64_t> sig(kLen);
+  for (size_t i = 0; i < kLen; ++i) {
+    std::vector<uint64_t> var_values(num_vars);
+    for (uint32_t v = 0; v < num_vars; ++v)
+      var_values[v] = (i >> v) & 1;
+    sig[i] = cobra::EvalExpr(ast, var_values, bitwidth);
+  }
+  return sig;
+}
+} // namespace
+
+namespace cobra::py {
+
+// converts a core Expr into a python facing PyExpr
+// copys the entier tree to avoid lifetime issues due to the mutability of
+// PyExpr
+PyExpr PyExpr::FromExprNode(const Expr &expr) {
+  PyExpr out;
+  out.kind = expr.kind;
+  out.constant_val = expr.constant_val;
+  out.var_index = expr.var_index;
+  out.children.reserve(expr.children.size());
+  for (const auto &child : expr.children)
+    out.children.push_back(FromExprNode(*child));
+  return out;
+}
+
+// require given no. subnexpressions
+void PyExpr::RequireArity(const PyExpr &node, size_t expected,
+                          const char *label) {
+  if (node.children.size() != expected) {
+    throw std::runtime_error(std::string("PyExpr ") + label + " expects " +
+                             std::to_string(expected) + " child(ren)");
+  }
+}
+
+// Converts a PyExpr back into a core Expr
+// TODO: Maybe force tailcall to prevent stack overflows on large linear
+// expressions
+std::unique_ptr<Expr> PyExpr::ToExprNode(const PyExpr &node) {
+  switch (node.kind) {
+  case Expr::Kind::kConstant:
+    return Expr::Constant(node.constant_val);
+
+  case Expr::Kind::kVariable:
+    return Expr::Variable(node.var_index);
+
+  case Expr::Kind::kNot:
+    RequireArity(node, 1, "Not");
+    return Expr::BitwiseNot(ToExprNode(node.children[0]));
+
+  case Expr::Kind::kNeg:
+    RequireArity(node, 1, "Neg");
+    return Expr::Negate(ToExprNode(node.children[0]));
+
+  case Expr::Kind::kShr:
+    RequireArity(node, 1, "Shr");
+    return Expr::LogicalShr(ToExprNode(node.children[0]), node.constant_val);
+
+  case Expr::Kind::kAdd:
+    RequireArity(node, 2, "Add");
+    return Expr::Add(ToExprNode(node.children[0]),
+                     ToExprNode(node.children[1]));
+
+  case Expr::Kind::kMul:
+    RequireArity(node, 2, "Mul");
+    return Expr::Mul(ToExprNode(node.children[0]),
+                     ToExprNode(node.children[1]));
+
+  case Expr::Kind::kAnd:
+    RequireArity(node, 2, "And");
+    return Expr::BitwiseAnd(ToExprNode(node.children[0]),
+                            ToExprNode(node.children[1]));
+
+  case Expr::Kind::kOr:
+    RequireArity(node, 2, "Or");
+    return Expr::BitwiseOr(ToExprNode(node.children[0]),
+                           ToExprNode(node.children[1]));
+
+  case Expr::Kind::kXor:
+    RequireArity(node, 2, "Xor");
+    return Expr::BitwiseXor(ToExprNode(node.children[0]),
+                            ToExprNode(node.children[1]));
+  }
+  throw std::runtime_error("PyExpr has unknown kind");
+}
+
+// Parses an expression from a string. Uses the same parser as the cobra-cli
+// tool. default params declared in header
+PyExprTree::PyExprTree(const std::string &s, uint32_t max_vars = 16,
+                       uint32_t bitwidth = 64) {
+  auto parsed = cobra::ParseToAst(s, bitwidth);
+  if (!parsed.has_value())
+    throw std::runtime_error(parsed.error().message);
+
+  auto &ast = parsed.value();
+  if (ast.vars.size() > max_vars) {
+    throw std::runtime_error(
+        "expression has " + std::to_string(ast.vars.size()) +
+        " variables (max " + std::to_string(max_vars) + ")");
+  }
+
+  this->bitwidth = bitwidth;
+
+  auto folded = cobra::FoldConstantBitwise(std::move(ast.expr), bitwidth);
+  this->root = PyExpr::FromExprNode(*folded);
+  this->vars = std::move(ast.vars);
+}
+
+std::string PyExprTree::ToString() const {
+  auto expr = root.ToExpr();
+  return cobra::Render(*expr, vars, bitwidth);
+}
+
+void PyExprTree::Simplify(bool validate = false) {
+  auto expr = root.ToExpr();
+  auto num_vars = static_cast<uint32_t>(vars.size());
+  auto sig = EvaluateToSignature(*expr, num_vars, bitwidth);
+
+  cobra::Options opts{
+      .bitwidth = bitwidth, .max_vars = num_vars, .spot_check = true};
+  opts.evaluator = cobra::Evaluator::FromExpr(
+      *expr, bitwidth, cobra::EvaluatorTraceKind::kCliOriginalAst);
+
+  auto result = cobra::Simplify(sig, vars, expr.get(), opts);
+  if (!result.has_value())
+    throw std::runtime_error(result.error().message);
+
+  auto &outcome = result.value();
+  if (outcome.kind == cobra::SimplifyOutcome::Kind::kError)
+    throw std::runtime_error(outcome.diag.reason);
+  if (outcome.kind == cobra::SimplifyOutcome::Kind::kUnchangedUnsupported)
+    return;
+
+  if (validate) {
+    std::vector<uint32_t> var_map;
+    if (outcome.real_vars.size() < vars.size())
+      var_map = cobra::BuildVarSupport(vars, outcome.real_vars);
+    auto fw = cobra::FullWidthCheck(*expr, num_vars, *outcome.expr, var_map,
+                                    bitwidth);
+    if (!fw.passed) {
+      throw std::runtime_error(
+          "CoB result is only correct on {0,1} inputs (polynomial target)");
+    }
+  }
+
+  this->root = PyExpr::FromExprNode(*outcome.expr);
+  this->vars = std::move(outcome.real_vars);
+}
+
+} // namespace cobra::py