Support: Avoid SmallVector::set_size() in Windows code

[dexonsmith]

Summary

Replace a few reserve() / set_size() pairs with
resize_for_overwrite() / truncate() in the platform-specific
code for Windows.

Depends on https://reviews.llvm.org/D115383 or successor; blocker for https://reviews.llvm.org/D115380.

[dexonsmith]

Adding some reviewers that have touched llvm/lib/Support/Windows recently. Note that I don't have access to a Windows machine and CI hasn't returned so I cannot confirm yet that this builds.

[mstorsjo]

Looks good to me. The Windows CI seems to have failed on Other/test-changed.ll, but I applied this patch and ran all tests in a number of subprojects, so it seems safe in that sense at least.

[dexonsmith]

In D115390#3182349, @mstorsjo wrote:

Looks good to me. The Windows CI seems to have failed on Other/test-changed.ll, but I applied this patch and ran all tests in a number of subprojects, so it seems safe in that sense at least.

Thanks for the review and the testing! Landed in 4d4439470eaf015784b76fc84cd5373456ce3d5c.

llvm/lib/Support/Windows/Path.inc

 } // end namespace windows
 
 namespace fs {
 
 const file_t kInvalidFile = INVALID_HANDLE_VALUE;
 
 std::string getMainExecutable(const char *argv0, void *MainExecAddr) {
   SmallVector<wchar_t, MAX_PATH> PathName;
-  DWORD Size = ::GetModuleFileNameW(NULL, PathName.data(), PathName.capacity());
+  PathName.resize_for_overwrite(PathName.capacity());
+  DWORD Size = ::GetModuleFileNameW(NULL, PathName.data(), PathName.size());
 
   // A zero return value indicates a failure other than insufficient space.
   if (Size == 0)
     return "";
 
   // Insufficient space is determined by a return value equal to the size of
   // the buffer passed in.
   if (Size == PathName.capacity())
     return "";
 
   // On success, GetModuleFileNameW returns the number of characters written to
   // the buffer not including the NULL terminator.
-  PathName.set_size(Size);
+  PathName.truncate(Size);
 
   // Convert the result from UTF-16 to UTF-8.
   SmallVector<char, MAX_PATH> PathNameUTF8;
   if (UTF16ToUTF8(PathName.data(), PathName.size(), PathNameUTF8))
     return "";
 
   llvm::sys::path::make_preferred(PathNameUTF8);
   return std::string(PathNameUTF8.data());
uint32_t file_status::getLinkCount() const {
   return NumLinks;
 }
 
 std::error_code current_path(SmallVectorImpl<char> &result) {
   SmallVector<wchar_t, MAX_PATH> cur_path;
   DWORD len = MAX_PATH;
 
   do {
-    cur_path.reserve(len);
-    len = ::GetCurrentDirectoryW(cur_path.capacity(), cur_path.data());
+    cur_path.resize_for_overwrite(len);
+    len = ::GetCurrentDirectoryW(cur_path.size(), cur_path.data());
 
     // A zero return value indicates a failure other than insufficient space.
     if (len == 0)
       return mapWindowsError(::GetLastError());
 
     // If there's insufficient space, the len returned is larger than the len
     // given.
-  } while (len > cur_path.capacity());
+  } while (len > cur_path.size());
 
   // On success, GetCurrentDirectoryW returns the number of characters not
   // including the null-terminator.
-  cur_path.set_size(len);
+  cur_path.truncate(len);
 
   if (std::error_code EC =
           UTF16ToUTF8(cur_path.begin(), cur_path.size(), result))
     return EC;
 
   llvm::sys::path::make_preferred(result);
   return std::error_code();
 }
if (Err != ERROR_INSUFFICIENT_BUFFER)
       return mapWindowsError(Err);
 
     Len *= 2;
   }
   // If the output buffer has exactly enough space for the path name, but not
   // the null terminator, it will leave the output unterminated.  Push a null
   // terminator onto the end to ensure that this never happens.
   VolumePath.push_back(L'\0');
-  VolumePath.set_size(wcslen(VolumePath.data()));
+  VolumePath.truncate(wcslen(VolumePath.data()));
   const wchar_t *P = VolumePath.data();
 
   UINT Type = ::GetDriveTypeW(P);
   switch (Type) {
   case DRIVE_FIXED:
     Result = true;
     return std::error_code();
   case DRIVE_REMOTE:
std::error_code is_local(const Twine &path, bool &result) {
   SmallVector<wchar_t, 128> WidePath;
   if (std::error_code ec = widenPath(P, WidePath))
     return ec;
   return is_local_internal(WidePath, result);
 }
 
 static std::error_code realPathFromHandle(HANDLE H,
                                           SmallVectorImpl<wchar_t> &Buffer) {
+  Buffer.resize_for_overwrite(Buffer.capacity());
   DWORD CountChars = ::GetFinalPathNameByHandleW(
       H, Buffer.begin(), Buffer.capacity(), FILE_NAME_NORMALIZED);
   if (CountChars && CountChars >= Buffer.capacity()) {
     // The buffer wasn't big enough, try again.  In this case the return value
     // *does* indicate the size of the null terminator.
-    Buffer.reserve(CountChars);
+    Buffer.resize_for_overwrite(CountChars);
     CountChars = ::GetFinalPathNameByHandleW(
-        H, Buffer.begin(), Buffer.capacity(), FILE_NAME_NORMALIZED);
+        H, Buffer.begin(), Buffer.size(), FILE_NAME_NORMALIZED);
   }
+  Buffer.truncate(CountChars);
   if (CountChars == 0)
     return mapWindowsError(GetLastError());
-  Buffer.set_size(CountChars);
   return std::error_code();
 }
 
 static std::error_code realPathFromHandle(HANDLE H,
                                           SmallVectorImpl<char> &RealPath) {
   RealPath.clear();
   SmallVector<wchar_t, MAX_PATH> Buffer;
   if (std::error_code EC = realPathFromHandle(H, Buffer))

 bool cache_directory(SmallVectorImpl<char> &result) {
   return getKnownFolderPath(FOLDERID_LocalAppData, result);
 }
 
 static bool getTempDirEnvVar(const wchar_t *Var, SmallVectorImpl<char> &Res) {
   SmallVector<wchar_t, 1024> Buf;
   size_t Size = 1024;
   do {
-    Buf.reserve(Size);
-    Size = GetEnvironmentVariableW(Var, Buf.data(), Buf.capacity());
+    Buf.resize_for_overwrite(Size);
+    Size = GetEnvironmentVariableW(Var, Buf.data(), Buf.size());
     if (Size == 0)
       return false;
 
     // Try again with larger buffer.
-  } while (Size > Buf.capacity());
-  Buf.set_size(Size);
+  } while (Size > Buf.size());
+  Buf.truncate(Size);
 
   return !windows::UTF16ToUTF8(Buf.data(), Size, Res);
 }
 
 static bool getTempDirEnvVar(SmallVectorImpl<char> &Res) {
   const wchar_t *EnvironmentVariables[] = {L"TMP", L"TEMP", L"USERPROFILE"};
   for (auto *Env : EnvironmentVariables) {
     if (getTempDirEnvVar(Env, Res))
if (!original.empty()) {
     int len = ::MultiByteToWideChar(codepage, MB_ERR_INVALID_CHARS, original.begin(),
                                     original.size(), utf16.begin(), 0);
 
     if (len == 0) {
       return mapWindowsError(::GetLastError());
     }
 
     utf16.reserve(len + 1);
-    utf16.set_size(len);
+    utf16.resize_for_overwrite(len);
 
     len = ::MultiByteToWideChar(codepage, MB_ERR_INVALID_CHARS, original.begin(),
                                 original.size(), utf16.begin(), utf16.size());
 
     if (len == 0) {
       return mapWindowsError(::GetLastError());
     }
   }
if (utf16_len) {
     // Get length.
     int len = ::WideCharToMultiByte(codepage, 0, utf16, utf16_len, converted.begin(),
                                     0, NULL, NULL);
 
     if (len == 0) {
       return mapWindowsError(::GetLastError());
     }
 
-    converted.reserve(len);
-    converted.set_size(len);
+    converted.reserve(len + 1);
+    converted.resize_for_overwrite(len);
 
     // Now do the actual conversion.
     len = ::WideCharToMultiByte(codepage, 0, utf16, utf16_len, converted.data(),
                                 converted.size(), NULL, NULL);
 
     if (len == 0) {
       return mapWindowsError(::GetLastError());
     }

llvm/lib/Support/Windows/Process.inc

if (windows::UTF8ToUTF16(Name, NameUTF16))
     return None;
 
   // Environment variable can be encoded in non-UTF8 encoding, and there's no
   // way to know what the encoding is. The only reliable way to look up
   // multibyte environment variable is to use GetEnvironmentVariableW().
   SmallVector<wchar_t, MAX_PATH> Buf;
   size_t Size = MAX_PATH;
   do {
-    Buf.reserve(Size);
+    Buf.resize_for_overwrite(Size);
     SetLastError(NO_ERROR);
     Size =
-      GetEnvironmentVariableW(NameUTF16.data(), Buf.data(), Buf.capacity());
+      GetEnvironmentVariableW(NameUTF16.data(), Buf.data(), Buf.size());
     if (Size == 0 && GetLastError() == ERROR_ENVVAR_NOT_FOUND)
       return None;
 
     // Try again with larger buffer.
-  } while (Size > Buf.capacity());
-  Buf.set_size(Size);
+  } while (Size > Buf.size());
+  Buf.truncate(Size);
 
   // Convert the result from UTF-16 to UTF-8.
   SmallVector<char, MAX_PATH> Res;
   if (windows::UTF16ToUTF8(Buf.data(), Size, Res))
     return None;
   return std::string(Res.data());
 }
 

llvm/lib/Support/Windows/Program.inc

ErrorOr<std::string> sys::findProgramByName(StringRef Name,
   if (const char *PathExtEnv = std::getenv("PATHEXT"))
     SplitString(PathExtEnv, PathExts, ";");
 
   SmallVector<char, MAX_PATH> U8Result;
   for (StringRef Ext : PathExts) {
     SmallVector<wchar_t, MAX_PATH> U16Result;
     DWORD Len = MAX_PATH;
     do {
-      U16Result.reserve(Len);
+      U16Result.resize_for_overwrite(Len);
       // Lets attach the extension manually. That is needed for files
       // with a point in name like aaa.bbb. SearchPathW will not add extension
       // from its argument to such files because it thinks they already had one.
       SmallVector<wchar_t, MAX_PATH> U16NameExt;
       if (std::error_code EC =
               windows::UTF8ToUTF16(Twine(Name + Ext).str(), U16NameExt))
         return EC;
 
       Len = ::SearchPathW(Path, c_str(U16NameExt), nullptr,
-                          U16Result.capacity(), U16Result.data(), nullptr);
-    } while (Len > U16Result.capacity());
+                          U16Result.size(), U16Result.data(), nullptr);
+    } while (Len > U16Result.size());
 
     if (Len == 0)
       continue;
 
-    U16Result.set_size(Len);
+    U16Result.truncate(Len);
 
     if (std::error_code EC =
         windows::UTF16ToUTF8(U16Result.data(), U16Result.size(), U8Result))
       return EC;
 
     if (sys::fs::can_execute(U8Result))
       break; // Found it.