impl using pure rust

flake.nix

@@ -0,0 +1,94 @@
+{
+  # Build Pyo3 package
+  inputs = {
+    nixpkgs.url = "github:nixos/nixpkgs/nixpkgs-unstable";
+    flake-utils.url = "github:numtide/flake-utils";
+    rust-overlay = {
+      url = "github:oxalica/rust-overlay";
+      inputs.nixpkgs.follows = "nixpkgs";
+    };
+    crane.url = "github:ipetkov/crane";
+    fenix = {
+      url = "github:nix-community/fenix";
+      inputs.nixpkgs.follows = "nixpkgs";
+    };
+  };
+
+  outputs =
+    inputs:
+    inputs.flake-utils.lib.eachDefaultSystem (
+      system:
+      let
+        pkgs = import inputs.nixpkgs {
+          inherit system;
+          overlays = [ inputs.rust-overlay.overlays.default ];
+          config = {
+            allowUnfree = true;
+            cudaSupport = false;
+          };
+        };
+        lib = pkgs.lib;
+
+        python_version = pkgs.python313;
+        wheel_tail = "cp313-cp313-linux_x86_64"; # Change if python_version changes
+
+        # Get a custom rust toolchain
+        craneLib = (inputs.crane.mkLib pkgs).overrideToolchain (
+          p: inputs.fenix.packages.${system}.complete.toolchain
+        );
+
+        project_name = (craneLib.crateNameFromCargoToml { cargoToml = ./shuffles/Cargo.toml; }).pname;
+        project_version = (craneLib.crateNameFromCargoToml { cargoToml = ./shuffles/Cargo.toml; }).version;
+
+        crate_cfg = {
+          src =
+            let
+              fs = lib.fileset;
+            in
+            fs.toSource {
+              root = ./shuffles;
+              fileset = fs.unions [
+                ./shuffles/src/lib.rs
+                ./shuffles/Cargo.lock
+                ./shuffles/Cargo.toml
+                ./shuffles/pyproject.toml
+              ];
+            };
+          nativeBuildInputs = [ python_version ];
+          # doCheck = true;
+          # buildInputs = [];
+        };
+
+        crate_artifacts = craneLib.buildDepsOnly (
+          crate_cfg
+          // {
+            pname = "${project_name}-artifacts";
+            version = project_version;
+          }
+        );
+
+        # Build the library, then re-use the target dir to generate the wheel file with maturin
+        crate_pkg = (
+          craneLib.buildPackage (
+            crate_cfg
+            // {
+              pname = project_name;
+              version = project_version;
+              cargoArtifacts = crate_artifacts;
+            }
+          )
+        );
+      in
+      rec {
+        devShells.default = craneLib.devShell {
+          packages = [
+            pkgs.cargo
+            pkgs.cargo-flamegraph
+          ];
+          shellHook = ''
+            export CUDA_PATH=${pkgs.cudatoolkit}
+          '';
+        };
+      }
+    );
+}