From 4904610dc605f7a8d2aa8e7bd4887e6181e45e8d Mon Sep 17 00:00:00 2001
From: Sebastien Bourdeauducq <sb@m-labs.hk>
Date: Tue, 4 Jan 2022 21:52:26 +0800
Subject: [PATCH] flake: provide mimalloc-enabled Python

The Linux linker and the libc are garbage, so there isn't much of an alternative to using the Nix wrapper and LD_PRELOAD.
---
 flake.nix | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/flake.nix b/flake.nix
index 1e37b47d4..5a51737b4 100644
--- a/flake.nix
+++ b/flake.nix
@@ -66,6 +66,10 @@
               '';
           }
         );
+        python3-mimalloc = pkgs.python3 // rec {
+          withMimalloc = pkgs.python3.buildEnv.override({ makeWrapperArgs = [ "--set LD_PRELOAD ${pkgs.mimalloc}/lib/libmimalloc.so" ]; });
+          withPackages = f: let packages = f pkgs.python3.pkgs; in withMimalloc.override { extraLibs = packages; };
+        };
 
         # LLVM PGO support
         llvm-nac3-instrumented = pkgs.callPackage "${self}/llvm" {
@@ -96,7 +100,7 @@
         nac3artiq-profile = pkgs.stdenvNoCC.mkDerivation {
           name = "nac3artiq-profile";
           src = self;
-          buildInputs = [ (pkgs.python3.withPackages(ps: [ ps.numpy nac3artiq-instrumented ])) pkgs.lld_13 pkgs.llvmPackages_13.libllvm ];
+          buildInputs = [ (python3-mimalloc.withPackages(ps: [ ps.numpy nac3artiq-instrumented ])) pkgs.lld_13 pkgs.llvmPackages_13.libllvm ];
           phases = [ "buildPhase" "installPhase" ];
           # TODO: get more representative code.
           buildPhase = "python $src/nac3artiq/demo/demo.py";
@@ -179,7 +183,7 @@
           cargo-insta
           rustc
           clippy
-          (python3.withPackages(ps: [ ps.numpy ]))
+          (packages.x86_64-linux.python3-mimalloc.withPackages(ps: [ ps.numpy ]))
         ];
       };