From ad8ca3cf02a414a1240566766116b108602b82a1 Mon Sep 17 00:00:00 2001
From: Hind-M <70631848+Hind-M@users.noreply.github.com>
Date: Fri, 20 Sep 2024 14:36:17 +0200
Subject: [PATCH] Compute `root prefix` as mamba install path (#3447)

* Compute root prefix as mamba install path

* Run test only on linux

* Test conflict between conda and mamba with same base env in CI

* Change strategy

* Use target_compile_definitions instead of add_definitions

* Add test logs

* Add ENABLE_MAMBA_ROOT_PREFIX_FALLBACK OPTION
---
 .github/workflows/unix_impl.yml    |  3 ++-
 .github/workflows/windows_impl.yml |  3 ++-
 libmamba/CMakeLists.txt            | 10 ++++++++++
 libmamba/src/api/configuration.cpp | 24 ++++++++++++++++++++++++
 4 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/unix_impl.yml b/.github/workflows/unix_impl.yml
index 7d04dde8e3..14614bfd07 100644
--- a/.github/workflows/unix_impl.yml
+++ b/.github/workflows/unix_impl.yml
@@ -40,7 +40,8 @@ jobs:
             --preset mamba-unix-shared-${{ inputs.build_type }}  \
             -D CMAKE_CXX_COMPILER_LAUNCHER=sccache \
             -D CMAKE_C_COMPILER_LAUNCHER=sccache \
-            -D BUILD_LIBMAMBAPY=OFF
+            -D BUILD_LIBMAMBAPY=OFF \
+            -D ENABLE_MAMBA_ROOT_PREFIX_FALLBACK=OFF
           cmake --build build/ --parallel
       - name: Show build cache statistics
         run: sccache --show-stats
diff --git a/.github/workflows/windows_impl.yml b/.github/workflows/windows_impl.yml
index ad33f1314d..480e20dae0 100644
--- a/.github/workflows/windows_impl.yml
+++ b/.github/workflows/windows_impl.yml
@@ -42,7 +42,8 @@ jobs:
             -D CMAKE_MSVC_RUNTIME_LIBRARY="MultiThreadedDLL" ^
             -D CMAKE_CXX_COMPILER_LAUNCHER=sccache ^
             -D CMAKE_C_COMPILER_LAUNCHER=sccache ^
-            -D BUILD_LIBMAMBAPY=OFF
+            -D BUILD_LIBMAMBAPY=OFF ^
+            -D ENABLE_MAMBA_ROOT_PREFIX_FALLBACK=OFF
           if %errorlevel% neq 0 exit /b %errorlevel%
           cmake --build build/ --parallel
           if %errorlevel% neq 0 exit /b %errorlevel%
diff --git a/libmamba/CMakeLists.txt b/libmamba/CMakeLists.txt
index 8bd51b591f..16716339ce 100644
--- a/libmamba/CMakeLists.txt
+++ b/libmamba/CMakeLists.txt
@@ -659,9 +659,19 @@ endmacro()
 
 set(libmamba_targets "")
 
+option(
+    ENABLE_MAMBA_ROOT_PREFIX_FALLBACK
+    "Enable mamba (shared) root prefix to be set to install prefix"
+    ON
+)
+
 if(BUILD_SHARED)
     message(STATUS "Adding shared libmamba target")
     libmamba_create_target(libmamba-dyn SHARED libmamba)
+    if(ENABLE_MAMBA_ROOT_PREFIX_FALLBACK)
+        # Use mamba installation prefix to set root prefix (base)
+        target_compile_definitions(libmamba-dyn PUBLIC MAMBA_USE_INSTALL_PREFIX_AS_BASE)
+    endif()
 endif()
 
 if(BUILD_STATIC)
diff --git a/libmamba/src/api/configuration.cpp b/libmamba/src/api/configuration.cpp
index c8c0654951..ce2c45320d 100644
--- a/libmamba/src/api/configuration.cpp
+++ b/libmamba/src/api/configuration.cpp
@@ -629,6 +629,21 @@ namespace mamba
             }
         }
 
+        auto
+        get_root_prefix_from_mamba_bin(const fs::u8path& mamba_bin_path) -> expected_t<fs::u8path>
+        {
+            if (mamba_bin_path.empty())
+            {
+                return make_unexpected(
+                    "`mamba` binary not found.\nPlease set `MAMBA_ROOT_PREFIX`.",
+                    mamba_error_code::incorrect_usage
+                );
+            }
+            // In linux and osx, the install path would be install_prefix/bin/mamba
+            // In windows, install_prefix/Scripts/mamba.exe
+            return { fs::weakly_canonical(mamba_bin_path.parent_path().parent_path()) };
+        }
+
         auto validate_existing_root_prefix(const fs::u8path& candidate) -> expected_t<fs::u8path>
         {
             auto prefix = fs::u8path(util::expand_home(candidate.string()));
@@ -724,11 +739,20 @@ namespace mamba
                 }
                 else
                 {
+#ifdef MAMBA_USE_INSTALL_PREFIX_AS_BASE
+                    // mamba case
+                    // set the root prefix as the mamba installation path
+                    get_root_prefix_from_mamba_bin(util::which("mamba"))
+                        .transform([&](fs::u8path&& p) { prefix = std::move(p); })
+                        .or_else([](mamba_error&& error) { throw std::move(error); });
+#else
+                    // micromamba case
                     validate_existing_root_prefix(default_root_prefix_v1())
                         .or_else([](const auto& /* error */)
                                  { return validate_root_prefix(default_root_prefix_v2()); })
                         .transform([&](fs::u8path&& p) { prefix = std::move(p); })
                         .or_else([](mamba_error&& error) { throw std::move(error); });
+#endif
                 }
 
                 if (env_name.configured())