// ptxasvs v0.001
// Written by David Shust, Summer 2009
// This program modifies the redirected stderr of ptxas such that its
// error messages are more suitable for Visual Studio.
// For some reason, stdout of the child process has to be redirected as well.
// To compile for UNICODE, #define UNICODE below.

#define REDIR_STDOUT 1
#ifndef UNICODE
#	define  UNICODE
#endif
#if defined(UNICODE) && !defined(_UNICODE)
#	define _UNICODE
#endif
#if defined(_UNICODE) && !defined(UNICODE)
#	define   UNICODE
#endif
#include <windows.h>
#include <stdio.h>
#include <tchar.h>

#ifndef _CRT_INSECURE_DEPRECATE
#define _tcscpy_s( strDest, numberOfElements, strSource ) \
		_tcscpy( strDest, strSource )
#define _tcscat_s( strDest, numberOfElements, strSource ) \
		_tcscat( strDest, strSource )
#define strcat_s( strDest, numberOfElements, strSource ) \
		strcat( strDest, strSource )
#define strncat_s( strDest, numberOfElements, strSource, count ) \
		strncat( strDest, strSource, count )
#define strncpy_s( strDest, numberOfElements, strSource, count ) \
		strncpy( strDest, strSource, count )
#endif

#define BUFSIZE 4096

HANDLE hProcess = NULL;
HANDLE hStdout, hStdoutRd, hStdoutRdTmp, hStdoutWr;
HANDLE hStderr, hStderrRd, hStderrRdTmp, hStderrWr;

LPTSTR CreateCmdLine( void );
BOOL CreateChildProcess( LPTSTR pCmdLine );
void ReadFromPipes( void );
void ErrorExit( LPSTR msg );

int main( int argc, char *argv[] )
{
	LPTSTR pCmdLine = CreateCmdLine();
	SECURITY_ATTRIBUTES sa;
	DWORD exitCode;

	memset( &sa, 0, sizeof sa );
	sa.nLength = sizeof sa;
	sa.bInheritHandle = TRUE; 

	hProcess = GetCurrentProcess();
#if REDIR_STDOUT
	if (!CreatePipe( &hStdoutRdTmp, &hStdoutWr, &sa, 0 ))
		ErrorExit( "Stdout pipe creation failed" );
	if (!DuplicateHandle( hProcess, hStdoutRdTmp, hProcess,
			&hStdoutRd, 0, FALSE, DUPLICATE_SAME_ACCESS ))
		ErrorExit( "DuplicateHandle failed" );
	CloseHandle( hStdoutRdTmp );
	hStdout = GetStdHandle( STD_OUTPUT_HANDLE );
	if (!SetStdHandle( STD_OUTPUT_HANDLE, hStdoutWr ))
		ErrorExit( "Redirecting STDOUT failed" );
#endif
	if (!CreatePipe( &hStderrRdTmp, &hStderrWr, &sa, 0 ))
		ErrorExit( "Stderr pipe creation failed" );
	if (!DuplicateHandle( hProcess, hStderrRdTmp, hProcess,
			&hStderrRd, 0, FALSE, DUPLICATE_SAME_ACCESS ))
		ErrorExit( "DuplicateHandle failed" );
	CloseHandle( hStderrRdTmp );
	hStderr = GetStdHandle( STD_ERROR_HANDLE );
	if (!SetStdHandle( STD_ERROR_HANDLE, hStderrWr ))
		ErrorExit( "Redirecting STDERR failed" );

	if (!CreateChildProcess( pCmdLine ))
		ErrorExit( "Create process failed" );

#if REDIR_STDOUT
	if (!SetStdHandle( STD_OUTPUT_HANDLE, hStdout ))
		ErrorExit( "Resetting STDOUT failed" );
	if (!CloseHandle( hStdoutWr ))
		ErrorExit( "Closing handle failed" );
#endif
	if (!SetStdHandle( STD_ERROR_HANDLE, hStderr ))
		ErrorExit( "Resetting STDERR failed" );
	if (!CloseHandle( hStderrWr ))
		ErrorExit( "Closing handle failed ");

	ReadFromPipes();

	if (GetExitCodeProcess( hProcess, &exitCode ))
		return (exitCode == 0) ? 0 : 1;
	else
		return 0;
}

LPTSTR CreateCmdLine( void )
{
	LPTSTR pCmdLine = NULL, pNewCmdLine = NULL;
	LPTSTR pModName = TEXT( "ptxas.exe " );
	size_t elts;

	if ((pCmdLine = GetCommandLine()) == NULL)
		ErrorExit( "GetCommandLine" );
	if (*pCmdLine == TEXT( '"' ))
	{
		pCmdLine++;
		while (*pCmdLine != TEXT( '\0' ) && *pCmdLine != TEXT( '"' ))
			pCmdLine++;
		if (*pCmdLine == TEXT( '"' ))
			pCmdLine++;
	}
	else
	{
		while (*pCmdLine != TEXT( '\0' ) && *pCmdLine != TEXT( ' ' ))
			pCmdLine++;
	}
	while (*pCmdLine != TEXT( '\0' ) && *pCmdLine == TEXT( ' ' ))
		pCmdLine++;
	elts = _tcslen( pModName ) + _tcslen( pCmdLine ) + 1;
	if ((pNewCmdLine = (LPTSTR) malloc( elts * sizeof (TCHAR) )) == NULL)
		ErrorExit( "malloc" );
	_tcscpy_s( pNewCmdLine, elts, pModName );
	_tcscat_s( pNewCmdLine, elts, pCmdLine );
	return pNewCmdLine;
}

BOOL CreateChildProcess( LPTSTR pCmdLine )
{ 
	PROCESS_INFORMATION pi;
	STARTUPINFO si;

	memset( &pi, 0, sizeof pi );
	memset( &si, 0, sizeof si );
	si.cb = sizeof si;
	if (!CreateProcess( NULL, pCmdLine,
			NULL, NULL, TRUE, 0, NULL, NULL, &si, &pi ))
		return FALSE;
	hProcess = pi.hProcess;
	if (!CloseHandle( pi.hThread ))
		ErrorExit( "CloseHandle" );
	return TRUE;
}

void ReadFromPipes( void )
{
	DWORD dwToRead, dwRead, dwToWrite, dwWritten;
	CHAR chSrc[BUFSIZE + 1], chDst[BUFSIZE + 1];
	LPSTR cmdstr = "ptxas ", linstr = ", line ";
	LPSTR errstr = "; error ", wrnstr = "; warning ";
	LPSTR endstr, substr, srcstr = chSrc, dststr = srcstr;
	size_t avail = 0;

#if REDIR_STDOUT
	for (;;) 
	{ 
		if (!ReadFile( hStdoutRd, chSrc, BUFSIZE, &dwRead, NULL ))
			break;
		if (dwRead == 0)
			break;
		if (!WriteFile( hStdout, chSrc, dwRead, &dwWritten, NULL ))
			break;
	}
#endif
	for (;;)
	{
		if (avail == 0 || (endstr = strstr( srcstr, "\n" )) == NULL)
		{
			dwToRead = (DWORD) (BUFSIZE - avail);
			memmove( chSrc, srcstr, avail );
			srcstr = chSrc;
			if (!ReadFile( hStderrRd, srcstr + avail, dwToRead, &dwRead, NULL ))
				break;
			avail += dwRead;
			srcstr[avail] = '\0';
			if (dwRead == 0 || (endstr = strstr( srcstr, "\n" )) == NULL)
			{
				dwToWrite = (DWORD) avail;
				if (!WriteFile( hStderr, srcstr, dwToWrite, &dwWritten, NULL ))
					break;
				avail = 0;
				continue;
			}
		}
		if ((substr = strstr( srcstr, linstr )) != NULL && substr < endstr)
		{
			size_t offset, length = strlen( cmdstr );
			if (strncmp( srcstr, cmdstr, length ) == 0)
			{
				srcstr += length;
				avail -= length;
			}
			offset = substr - srcstr;
			dststr = chDst;
			strncpy_s( dststr, sizeof chDst, srcstr, offset );
			dststr[offset] = '\0';
			length = offset + strlen( linstr );
			srcstr += length;
			avail -= length;
			if (*srcstr >= '0' && *srcstr <= '9')
			{
				dststr += offset;
				*dststr++ = '(';
				while (*srcstr >= '0' && *srcstr <= '9')
				{
					*dststr++ = *srcstr++;
					avail--;
				}
				*dststr = '\0';
				dststr = chDst;
				strcat_s( dststr, sizeof chDst, ") : " );
				if (strncmp( srcstr, errstr, strlen( errstr ) ) == 0 ||
					strncmp( srcstr, wrnstr, strlen( wrnstr ) ) == 0)
				{
					length = endstr - srcstr + strlen( "\n" );
					strncat_s( dststr, sizeof chDst, srcstr + 2, length - 2 );
					srcstr += length;
					avail -= length;
				}
			}
		}
		else
		{
			size_t offset = endstr - srcstr;
			size_t length = offset + strlen( "\n" );
			dststr = chDst;
			strncpy_s( dststr, sizeof chDst, srcstr, length );
			dststr[length] = '\0';
			srcstr += length;
			avail -= length;
		}
		dwToWrite = (DWORD) strlen( dststr );
		if (!WriteFile( hStderr, dststr, dwToWrite, &dwWritten, NULL ))
			break;
		dststr = srcstr;
	}
}

void ErrorExit( LPSTR msg )
{
	_ftprintf( stderr, TEXT( "%hs\n" ), msg );
	ExitProcess( 0 );
}