From 5db36a47e11539316f22a0689f1c4ff845d4266c Mon Sep 17 00:00:00 2001
From: Simeon Armenchev <monyarm@gmail.com>
Date: Thu, 13 Jun 2024 10:46:21 +0300
Subject: [PATCH] feat(rapidsnark-gpu): Add rapidsnark-gpu package

---
 flake.nix                           | 20 +++++++++---
 packages/all-packages.nix           |  5 ++-
 packages/rapidsnark-gpu/default.nix | 50 +++++++++++++++++++++++++++++
 3 files changed, 69 insertions(+), 6 deletions(-)
 create mode 100644 packages/rapidsnark-gpu/default.nix

diff --git a/flake.nix b/flake.nix
index c517a1e..d541d77 100644
--- a/flake.nix
+++ b/flake.nix
@@ -36,20 +36,30 @@
     };
   };
 
-  outputs = inputs @ {flake-parts, ...}:
+  outputs = inputs @ {
+    flake-parts,
+    nixpkgs,
+    ...
+  }:
     flake-parts.lib.mkFlake {inherit inputs;} {
       systems = ["x86_64-linux" "aarch64-linux" "x86_64-darwin" "aarch64-darwin"];
       imports = [./packages];
       perSystem = {
-        final,
+        pkgs,
         self',
+        system,
         ...
       }: {
+        _module.args = {
+          pkgs = import nixpkgs {
+            inherit system;
+            config.allowUnfree = true;
+          };
+        };
         devShells.default = import ./shells/all.nix {
-          pkgs = final;
-          inherit self';
+          inherit pkgs self';
         };
-        devShells.ci = import ./shells/ci.nix {pkgs = final;};
+        devShells.ci = import ./shells/ci.nix {inherit pkgs;};
       };
     };
 }
diff --git a/packages/all-packages.nix b/packages/all-packages.nix
index b579d97..1e11bc2 100644
--- a/packages/all-packages.nix
+++ b/packages/all-packages.nix
@@ -59,6 +59,9 @@
     rapidsnark = callPackage ./rapidsnark/default.nix {
       inherit ffiasm zqfield-bn254;
     };
+    rapidsnark-gpu = callPackage ./rapidsnark-gpu/default.nix {
+      inherit ffiasm zqfield-bn254;
+    };
 
     # Elrond / MultiversX
     # copied from https://github.com/NixOS/nixpkgs/blob/8df7949791250b580220eb266e72e77211bedad9/pkgs/development/python-modules/cryptography/default.nix
@@ -157,7 +160,7 @@
         cdt = callPackage ./cdt/default.nix {};
       }
       // lib.optionalAttrs hostPlatform.isx86 rec {
-        inherit zqfield-bn254 ffiasm ffiasm-src rapidsnark;
+        inherit zqfield-bn254 ffiasm ffiasm-src rapidsnark rapidsnark-gpu;
 
         inherit cardano graphql;
       }
diff --git a/packages/rapidsnark-gpu/default.nix b/packages/rapidsnark-gpu/default.nix
new file mode 100644
index 0000000..47c7e42
--- /dev/null
+++ b/packages/rapidsnark-gpu/default.nix
@@ -0,0 +1,50 @@
+{
+  lib,
+  cudaPackages,
+  stdenv,
+  ffiasm,
+  zqfield-bn254,
+  nlohmann_json,
+  gmp,
+  libsodium,
+  cmake,
+  fetchFromGitHub,
+  pkg-config,
+}: let
+  ffiasm-c = "${ffiasm}/lib/node_modules/ffiasm/c";
+in
+  stdenv.mkDerivation rec {
+    pname = "rapidsnark-gpu";
+    version = "2023-04-08";
+
+    src = fetchFromGitHub {
+      owner = "Orbiter-Finance";
+      repo = "rapidsnark";
+      rev = "77016322808ac58a3acd25a6235510b55172f967";
+      hash = "sha256-8vy+iXkSINFregve+rej1rXyXdWxm0n1wvYfoy/0idk=";
+    };
+
+    nativeBuildInputs = [pkg-config];
+    buildInputs = [nlohmann_json gmp libsodium cudaPackages.cudatoolkit] ++ ffiasm.passthru.openmp;
+
+    buildPhase = ''
+      mkdir -p $out/bin
+      c++ \
+        -I{${ffiasm-c},${zqfield-bn254}/lib} \
+        ./src/{main_prover,binfile_utils,zkey_utils,wtns_utils,logger}.cpp \
+        ${ffiasm-c}/{alt_bn128,misc,naf,splitparstr}.cpp \
+        ${zqfield-bn254}/lib/{fq,fr}.{cpp,o} \
+        $(pkg-config --cflags --libs libsodium gmp nlohmann_json) \
+        -std=c++17 -pthread -O3 -fopenmp \
+        -o $out/bin/prover
+    '';
+
+    installPhase = ''
+      # Already done in buildPhase
+    '';
+
+    meta = {
+      homepage = "https://github.com/iden3/rapidsnark";
+      platforms = with lib.platforms; linux ++ darwin;
+    };
+  }