From 8525059c0bfc593e77969722e31d537c179ef438 Mon Sep 17 00:00:00 2001
From: Jake Arkinstall <65358059+jake-arkinstall@users.noreply.github.com>
Date: Wed, 23 Oct 2024 12:58:46 +0100
Subject: [PATCH] Updated nix test to utilise nixgl for cuda library handling

---
 flake.lock                         | 54 ++++++++++++++++++++++++++++--
 flake.nix                          | 13 ++++---
 nix-support/pytket-cutensornet.nix |  8 +++--
 3 files changed, 66 insertions(+), 9 deletions(-)

diff --git a/flake.lock b/flake.lock
index eaae203b..78253789 100644
--- a/flake.lock
+++ b/flake.lock
@@ -19,6 +19,21 @@
       }
     },
     "flake-utils_2": {
+      "locked": {
+        "lastModified": 1659877975,
+        "narHash": "sha256-zllb8aq3YO3h8B/U0/J1WBgAL8EX5yWf5pMj3G0NAmc=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "c0e246b9b83f637f4681389ecabcb2681b4f3af0",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "flake-utils_3": {
       "inputs": {
         "systems": "systems_2"
       },
@@ -36,7 +51,41 @@
         "type": "github"
       }
     },
+    "nixgl": {
+      "inputs": {
+        "flake-utils": "flake-utils_2",
+        "nixpkgs": "nixpkgs"
+      },
+      "locked": {
+        "lastModified": 1713543440,
+        "narHash": "sha256-lnzZQYG0+EXl/6NkGpyIz+FEOc/DSEG57AP1VsdeNrM=",
+        "owner": "nix-community",
+        "repo": "nixGL",
+        "rev": "310f8e49a149e4c9ea52f1adf70cdc768ec53f8a",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-community",
+        "repo": "nixGL",
+        "type": "github"
+      }
+    },
     "nixpkgs": {
+      "locked": {
+        "lastModified": 1660551188,
+        "narHash": "sha256-a1LARMMYQ8DPx1BgoI/UN4bXe12hhZkCNqdxNi6uS0g=",
+        "owner": "nixos",
+        "repo": "nixpkgs",
+        "rev": "441dc5d512153039f19ef198e662e4f3dbb9fd65",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nixos",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "nixpkgs_2": {
       "locked": {
         "lastModified": 1722087241,
         "narHash": "sha256-2ShmEaFi0kJVOEEu5gmlykN5dwjWYWYUJmlRTvZQRpU=",
@@ -55,6 +104,7 @@
     "root": {
       "inputs": {
         "flake-utils": "flake-utils",
+        "nixgl": "nixgl",
         "nixpkgs": [
           "tket",
           "nixpkgs"
@@ -94,8 +144,8 @@
     },
     "tket": {
       "inputs": {
-        "flake-utils": "flake-utils_2",
-        "nixpkgs": "nixpkgs"
+        "flake-utils": "flake-utils_3",
+        "nixpkgs": "nixpkgs_2"
       },
       "locked": {
         "lastModified": 1722327570,
diff --git a/flake.nix b/flake.nix
index 05266460..04b67de9 100644
--- a/flake.nix
+++ b/flake.nix
@@ -7,10 +7,13 @@
     cuda-maintainers.cachix.org-1:0dq3bujKpuEPMCX6U4WylrUDZ9JyUG0VpVZa7CNfq5E=
     nix-community.cachix.org-1:mB9FSh9qf2dCimDSUo8Zy7bkq5CX+/rkCWyvRCYg3Fs=
   '';
-  inputs.flake-utils.url = "github:numtide/flake-utils";
-  inputs.tket.url = "github:CQCL/tket";
-  inputs.nixpkgs.follows = "tket/nixpkgs";
-  outputs = { self, nixpkgs, flake-utils, tket }:
+  inputs = {
+    flake-utils.url = "github:numtide/flake-utils";
+    tket.url = "github:CQCL/tket";
+    nixpkgs.follows = "tket/nixpkgs";
+    nixgl.url = "github:nix-community/nixGL";
+  };
+  outputs = { self, nixpkgs, flake-utils, tket, nixgl }:
     flake-utils.lib.eachDefaultSystem (system:
       let
         pkgs = import nixpkgs {
@@ -20,6 +23,8 @@
             cudaSupport = true;
           };
           overlays = [
+            (nixgl.overlay)
+
             (self: super: {
               inherit (tket.packages."${system}") tket pytket;
             })
diff --git a/nix-support/pytket-cutensornet.nix b/nix-support/pytket-cutensornet.nix
index d95840e7..4895c30d 100644
--- a/nix-support/pytket-cutensornet.nix
+++ b/nix-support/pytket-cutensornet.nix
@@ -55,12 +55,14 @@ EOF
       pytest
       pytest-lazy-fixture
     ]);
+    nixgl-bin = self.lib.getExe self.nixgl.auto.nixGLNvidia;
   in super.writeShellScriptBin "run-pytket-cutensornet-tests" ''
       HOME=$(mktemp -d);
       export HOME;
-      echo "---------------------------";
-      env;
-      echo "---------------------------";
+      NIXGL_PATH="$(${nixgl-bin} printenv LD_LIBRARY_PATH)";
+      WSL_PATH="/usr/lib/wsl/lib";
+      LD_LIBRARY_PATH="$NIXGL_PATH:$WSL_PATH:$LD_LIBRARY_PATH";
+      export LD_LIBRARY_PATH;
       ${test-env}/bin/pytest -s ${../tests};
   '';
 }