1 files changed, 49 insertions, 15 deletions

diff --git a/src/main.cpp b/src/main.cpp
index cb1c611..1da5202 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -227,15 +227,32 @@ static bool isPathSubPathOf(const FileString &path, const FileString &subPathOf)
 }
 
 #if OS_FAMILY == OS_FAMILY_WINDOWS
-static char* join(const char *str1, const char *str2, const char separator)
+static char* join(const vector<const char *> &strs, const char separator)
 {
-    int str1len = strlen(str1);
-    int str2len = strlen(str2);
-    char *result = new char[str1len + 1 + str2len + 1];
-    memcpy(result, str1, str1len);
-    result[str1len] = separator;
-    memcpy(result + str1len + 1, str2, str2len);
-    result[str1len + 1 + str2len] = '\0';
+    vector<int> lengths;
+    lengths.reserve(strs.size());
+    int totalLength = strs.size() - 1;
+    for (const char *str : strs)
+    {
+        int length = strlen(str);
+        totalLength += length;
+        lengths.push_back(length);
+    }
+
+    char *result = new char[totalLength + 1];
+    result[totalLength] = '\0';
+    int offset = 0;
+    for (int i = 0; i < strs.size(); ++i)
+    {
+        if (i > 0)
+        {
+            result[offset] = separator;
+            ++offset;
+        }
+        memcpy(result + offset, strs[i], lengths[i]);
+        offset += lengths[i];
+    }
+
     return result;
 }
 
@@ -246,9 +263,17 @@ struct MicrosoftBuildTool
     // empty if not found
     char binPath[_TINYDIR_PATH_MAX];
     // empty if not found
+    char vsLibPath[_TINYDIR_PATH_MAX];
+    // empty if not found
     char umLibPath[_TINYDIR_PATH_MAX];
     // empty if not found
     char ucrtLibPath[_TINYDIR_PATH_MAX];
+    // empty if not found
+    char vsIncludePath[_TINYDIR_PATH_MAX];
+    // empty if not found
+    char umIncludePath[_TINYDIR_PATH_MAX];
+    // empty if not found
+    char ucrtIncludePath[_TINYDIR_PATH_MAX];
 
     bool found()
     {
@@ -259,32 +284,41 @@ struct MicrosoftBuildTool
 static MicrosoftBuildTool locateLatestMicrosoftBuildTool()
 {
     MicrosoftBuildTool result = { 0 };
-    Result<ExecResult> execResult = exec(TINYDIR_STRING("locate_windows_sdk"));
+    Result<ExecResult> execResult = exec(TINYDIR_STRING("locate_windows_sdk x64"));
     if (execResult && execResult.unwrap().exitCode == 0)
     {
         auto &str = execResult.unwrap().execStdout;
-        sscanf(execResult.unwrap().execStdout.c_str(), "%d %[^\r\n] %[^\r\n] %[^\r\n]", &result.version, result.binPath, result.umLibPath, result.ucrtLibPath);
+        sscanf(execResult.unwrap().execStdout.c_str(), "%d %[^\r\n] %[^\r\n] %[^\r\n] %[^\r\n] %[^\r\n] %[^\r\n] %[^\r\n]", 
+            &result.version, 
+            result.binPath, 
+            result.vsLibPath, 
+            result.umLibPath,
+            result.ucrtLibPath,
+            result.vsIncludePath,
+            result.umIncludePath,
+            result.ucrtIncludePath);
     }
     return result;
 }
 
+// We do not free allocated data here because they needs to live as long as they're used as env (in _putenv)
 static void appendMicrosoftBuildToolToPathEnv()
 {
     MicrosoftBuildTool msBuildTool = locateLatestMicrosoftBuildTool();
     if (msBuildTool.found())
     {
         fprintf(stderr, "Located microsoft build tools at %s\n", msBuildTool.binPath);
-        fprintf(stderr, "Located microsoft build libraries at %s;%s\n", msBuildTool.umLibPath, msBuildTool.ucrtLibPath);
 
         if (const char *pathEnv = getenv("PATH"))
         {
-            // We do not free this because it needs to live as long as it's used as env (in _putenv)
-            if (_putenv_s("PATH", join(pathEnv, msBuildTool.binPath, ';')) != 0)
+            if (_putenv_s("PATH", join({ pathEnv, msBuildTool.binPath }, ';')) != 0)
                 fprintf(stderr, "Warning: Failed to add microsoft build tools to PATH env\n");
         }
 
-        // We do not free this because it needs to live as long as it's used as env (in _putenv)
-        if (_putenv_s("LIB", join(msBuildTool.umLibPath, msBuildTool.ucrtLibPath, ';')) != 0)
+        if (_putenv_s("INCLUDE", join({ msBuildTool.vsIncludePath, msBuildTool.umIncludePath, msBuildTool.ucrtIncludePath }, ';')) != 0)
+            fprintf(stderr, "Warning: Failed to add microsoft build libraries to INCLUDE env\n");
+
+        if (_putenv_s("LIB", join({ msBuildTool.vsLibPath, msBuildTool.umLibPath, msBuildTool.ucrtLibPath }, ';')) != 0)
             fprintf(stderr, "Warning: Failed to add microsoft build libraries to LIB env\n");
     }
 }