fix: Adapt root prefix determination

libmamba/include/mamba/core/util_os.hpp

@@ -22,6 +22,7 @@ namespace mamba
 
     bool is_admin();
     fs::u8path get_self_exe_path();
+    fs::u8path get_libmamba_path();
 
     using PID =
 #ifdef _WIN32

libmamba/src/api/configuration.cpp

@@ -629,13 +629,21 @@ namespace mamba
             }
         }
 
+        bool is_root_prefix(const fs::u8path& prefix)
+        {
+            // TODO: consider the conjunction (i.e. &&-chaining) of the following conditions
+            // instead.
+            return fs::exists(prefix / "pkgs") || fs::exists(prefix / "conda-meta")
+                   || fs::exists(prefix / "envs");
+        }
+
         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`.",
+                    "The root prefix of your installation cannot be found.\nPlease set `MAMBA_ROOT_PREFIX`.",
                     mamba_error_code::incorrect_usage
                 );
             }
@@ -653,9 +661,7 @@ namespace mamba
                 return make_unexpected("Empty root prefix.", mamba_error_code::incorrect_usage);
             }
 
-            if (!fs::exists(prefix / "pkgs")           //
-                && !fs::exists(prefix / "conda-meta")  //
-                && !fs::exists(prefix / "envs"))
+            if (!is_root_prefix(prefix))
             {
                 return make_unexpected(
                     fmt::format(
@@ -713,51 +719,77 @@ namespace mamba
             return { fs::weakly_canonical(std::move(prefix)) };
         }
 
-        auto get_default_root_prefix(fs::u8path& prefix) -> void
+        auto get_root_prefix() -> fs::u8path
         {
-            if (util::get_env("MAMBA_DEFAULT_ROOT_PREFIX"))
+            fs::u8path root_prefix = util::get_env("MAMBA_ROOT_PREFIX").value_or("");
+
+            if (!root_prefix.empty())
+            {
+                LOG_TRACE << "Using root prefix set in `MAMBA_ROOT_PREFIX`: " << root_prefix;
+                return root_prefix;
+            }
+
+            root_prefix = util::get_env("MAMBA_DEFAULT_ROOT_PREFIX").value_or("");
+
+            if (!root_prefix.empty())
             {
-                prefix = util::get_env("MAMBA_DEFAULT_ROOT_PREFIX").value();
                 LOG_WARNING << unindent(R"(
-                                    'MAMBA_DEFAULT_ROOT_PREFIX' is meant for testing purpose.
-                                    Consider using 'MAMBA_ROOT_PREFIX' instead)");
+                                'MAMBA_DEFAULT_ROOT_PREFIX' is meant for testing purpose.
+                                Consider using 'MAMBA_ROOT_PREFIX' instead)");
+                LOG_TRACE << "Using root prefix set in `MAMBA_DEFAULT_ROOT_PREFIX`: " << root_prefix;
+                return root_prefix;
             }
-            else
+
+            // Find the location of libmamba
+            fs::u8path libmamba_library_path = get_libmamba_path();
+
+            // Find the environment directory of the executable
+            fs::u8path env_prefix = fs::weakly_canonical(
+                libmamba_library_path.parent_path().parent_path()
+            );
+
+            if (is_root_prefix(env_prefix))
             {
-#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
-
-                // In 1.0, only micromamba was using this location.
-                const fs::u8path default_root_prefix_v1 = fs::u8path(util::user_home_dir())
-                                                          / "micromamba";
-
-                // In 2.0, we change the default location.
-                // We unconditionally name the subfolder "mamba" for compatibility between ``mamba``
-                // and ``micromamba``, as well as consistency with ``MAMBA_`` environment variables.
-                const fs::u8path default_root_prefix_v2 = fs::u8path(util::user_data_dir()) / "mamba";
-
-                validate_existing_root_prefix(default_root_prefix_v1)
-                    .or_else([&default_root_prefix_v2](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
+                LOG_TRACE << "Using `libmamba`'s current environment as the root prefix: "
+                          << env_prefix;
+                return env_prefix;
             }
-        }
 
-        auto get_root_prefix() -> fs::u8path
-        {
-            fs::u8path root_prefix = util::get_env("MAMBA_ROOT_PREFIX").value_or("");
-            if (root_prefix.empty())
+            // From the environment directory, we might infer the root prefix.
+            fs::u8path supposed_root_prefix = fs::weakly_canonical(
+                env_prefix.parent_path().parent_path()
+            );
+
+            if (is_root_prefix(supposed_root_prefix))
             {
-                get_default_root_prefix(root_prefix);
+                LOG_TRACE << "Inferring and using the root prefix from `libmamba`'s current environment' as: "
+                          << supposed_root_prefix;
+                return supposed_root_prefix;
             }
+
+#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) { root_prefix = std::move(p); })
+                .or_else([](mamba_error&& error) { throw std::move(error); });
+#else
+            // micromamba case
+            // In 1.0, only micromamba was using this location.
+            const fs::u8path default_root_prefix_v1 = fs::u8path(util::user_home_dir())
+                                                      / "micromamba";
+
+            // In 2.0, we change the default location.
+            // We unconditionally name the subfolder "mamba" for compatibility between ``mamba``
+            // and ``micromamba``, as well as consistency with ``MAMBA_`` environment variables.
+            const fs::u8path default_root_prefix_v2 = fs::u8path(util::user_data_dir()) / "mamba";
+
+            validate_existing_root_prefix(default_root_prefix_v1)
+                .or_else([&default_root_prefix_v2](const auto& /* error */)
+                         { return validate_root_prefix(default_root_prefix_v2); })
+                .transform([&](fs::u8path&& p) { root_prefix = std::move(p); })
+                .or_else([](mamba_error&& error) { throw std::move(error); });
+#endif
             return root_prefix;
         }
 

libmamba/src/core/util_os.cpp

@@ -3,7 +3,8 @@
 
 #ifndef _WIN32
 #include <clocale>
-
+// To find the path of `libmamba`'s library.
+#include <dlfcn.h>
 #include <sys/ioctl.h>
 #include <sys/utsname.h>
 #include <unistd.h>
@@ -89,6 +90,29 @@ namespace mamba
 #endif
     }
 
+    fs::u8path get_libmamba_path()
+    {
+        fs::u8path libmamba_library_path;
+#if defined(__linux__) || __APPLE__ || __MACH__
+        Dl_info dl_info;
+        if (dladdr(reinterpret_cast<void*>(get_self_exe_path), &dl_info))
+        {
+            libmamba_library_path = dl_info.dli_fname;
+        }
+#else
+        HMODULE hModule = NULL;
+        CHAR path[MAX_PATH];
+        GetModuleHandleEx(
+            GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
+            (LPCTSTR) get_self_exe_path,
+            &hModule
+        );
+        GetModuleFileName(hModule, path, MAX_PATH);
+        libmamba_library_path = fs::u8path(std::string(path));
+#endif
+        return libmamba_library_path;
+    }
+
     bool is_admin()
     {
 #ifdef _WIN32

micromamba/etc/profile.d/mamba.sh.in

@@ -1,4 +1,8 @@
-export MAMBA_ROOT_PREFIX="@CMAKE_INSTALL_PREFIX@"
+# Do not reset `MAMBA_ROOT_PREFIX` if it is already set
+if [ -z "${MAMBA_ROOT_PREFIX}" ]; then
+    export MAMBA_ROOT_PREFIX="@CMAKE_INSTALL_PREFIX@"
+fi
+
 __mamba_setup="$("@CMAKE_INSTALL_PREFIX@/@CMAKE_INSTALL_BINDIR@/mamba" shell hook --shell posix 2> /dev/null)"
 if [ $? -eq 0 ]; then
     eval "$__mamba_setup"