/**
 * This module provides OS specific helper function for threads support
 *
 * Copyright: Copyright Digital Mars 2010.
 * License:   <a href="http://www.boost.org/LICENSE_1_0.txt">Boost License 1.0</a>.
 * Authors:   Rainer Schuetze
 *
 *          Copyright Digital Mars 2010.
 * Distributed under the Boost Software License, Version 1.0.
 *    (See accompanying file LICENSE_1_0.txt or copy at
 *          http://www.boost.org/LICENSE_1_0.txt)
 */

module core.thread_helper;

version( Windows )
{
    import std.c.windows.windows;
    import std.c.stdlib;

    public import core.thread;

    extern(Windows)
    HANDLE OpenThread(DWORD dwDesiredAccess, BOOL bInheritHandle, DWORD dwThreadId);

    extern (C) extern __gshared int _tls_index;
    
private:
    ///////////////////////////////////////////////////////////////////
    struct thread_helper_aux
    {
	// don't let symbols leak into other modules
	
	const SystemProcessInformation = 5;
	const STATUS_INFO_LENGTH_MISMATCH = 0xc0000004;

	// abbreviated versions of these structs (full info can be found 
	//  here: http://undocumented.ntinternals.net )
	struct _SYSTEM_PROCESS_INFORMATION
	{
	    int NextEntryOffset; // When this entry is 0, there are no more processes to be read.
	    int NumberOfThreads;
	    int[15] fill1;
	    int ProcessId;
	    int[28] fill2;
	
	    // SYSTEM_THREAD_INFORMATION or SYSTEM_EXTENDED_THREAD_INFORMATION structures follow.
	}

	struct _SYSTEM_THREAD_INFORMATION
	{
	    int[8] fill1;
	    int ProcessId;
	    int ThreadId;
	    int[6] fill2;
	}

	alias extern(Windows)
	HRESULT fnNtQuerySystemInformation( uint SystemInformationClass, void* info, uint infoLength, uint* ReturnLength );

	const ThreadBasicInformation = 0;

	struct THREAD_BASIC_INFORMATION 
	{
	    int    ExitStatus;
	    void** TebBaseAddress;
	    int    ProcessId;
	    int    ThreadId;
	    int    AffinityMask;
	    int    Priority;
	    int    BasePriority;
        }

	alias extern(Windows)
	int fnNtQueryInformationThread( HANDLE ThreadHandle, uint ThreadInformationClass, void* buf, uint size, uint* ReturnLength );

	const SYNCHRONIZE = 0x00100000;
	const THREAD_GET_CONTEXT = 8;
	const THREAD_QUERY_INFORMATION = 0x40;
	const THREAD_SUSPEND_RESUME = 2;

	///////////////////////////////////////////////////////////////////
	// get the thread environment block (TEB) of the thread with the given handle
	static void** getTEB( HANDLE hnd )
	{
	    HANDLE nthnd = GetModuleHandleA( "NTDLL" );
	    assert( nthnd, "cannot get module handle for ntdll" );
	    fnNtQueryInformationThread* fn = cast(fnNtQueryInformationThread*) GetProcAddress( nthnd, "NtQueryInformationThread" );
	    assert( fn, "cannot find NtQueryInformationThread in ntdll" );

	    THREAD_BASIC_INFORMATION tbi;
	    int Status = (*fn)(hnd, ThreadBasicInformation, &tbi, tbi.sizeof, null);
	    assert(Status == 0);
	
	    return tbi.TebBaseAddress;
	}

	// get the thread environment block (TEB) of the thread with the given identifier
	static void** getTEB( uint id )
	{
	    HANDLE hnd = OpenThread( THREAD_QUERY_INFORMATION, FALSE, id );
	    assert( hnd, "OpenThread failed" );

	    void** teb = getTEB( hnd );
	    CloseHandle( hnd );
	    return teb;
	}

	// get linear address of TEB of current thread
	static void** getTEB()
	{
	    asm
	    {
		naked;
		mov EAX,FS:[0x18];
		ret;
	    }
	}

	// get the stack bottom (the top address) of the thread with the given handle
	static void* getThreadStackBottom( HANDLE hnd )
	{
	    void** teb = getTEB( hnd );
	    return teb[1];
	}

	// get the stack bottom (the top address) of the thread with the given identifier
	static void* getThreadStackBottom( uint id )
	{
	    void** teb = getTEB( id );
	    return teb[1];
	}

	// create a thread handle with full access to the thread with the given identifier
	static HANDLE OpenThreadHandle( uint id )
	{
	    return OpenThread( SYNCHRONIZE|THREAD_GET_CONTEXT|THREAD_QUERY_INFORMATION|THREAD_SUSPEND_RESUME, FALSE, id );
	}

	///////////////////////////////////////////////////////////////////
	// enumerate threads of the given process calling the passed function on each thread
	// using function instead of delegate here to avoid allocating closure
	static bool enumProcessThreads( uint procid, bool function( uint id, void* context ) dg, void* context )
	{
	    HANDLE hnd = GetModuleHandleA( "NTDLL" );
	    fnNtQuerySystemInformation* fn = cast(fnNtQuerySystemInformation*) GetProcAddress( hnd, "NtQuerySystemInformation" );
	    if( !fn )
		return false;

	    uint sz = 16384;
	    uint retLength;
	    HRESULT rc;
	    char* buf;
	    for( ; ; )
	    {
		buf = cast(char*) std.c.stdlib.malloc(sz);
		if(!buf)
	    	    return false;
		rc = (*fn)( SystemProcessInformation, buf, sz, &retLength );
		if( rc != STATUS_INFO_LENGTH_MISMATCH )
		    break;
		std.c.stdlib.free( buf );
		sz *= 2;
	    }
	    scope(exit) std.c.stdlib.free( buf );

	    if(rc != 0)
		return false;

	    auto pinfo = cast(_SYSTEM_PROCESS_INFORMATION*) buf;
	    auto pend  = cast(_SYSTEM_PROCESS_INFORMATION*) (buf + retLength);
	    for( ; pinfo < pend; )
	    {
		if( pinfo.ProcessId == procid )
		{
		    auto tinfo = cast(_SYSTEM_THREAD_INFORMATION*)(pinfo + 1);
		    for( int i = 0; i < pinfo.NumberOfThreads; i++, tinfo++ )
			if( tinfo.ProcessId == procid )
			    if( !dg( tinfo.ThreadId, context ) )
				return false;
		}
		if( pinfo.NextEntryOffset == 0 )
		    break;
		pinfo = cast(_SYSTEM_PROCESS_INFORMATION*) (cast(char*) pinfo + pinfo.NextEntryOffset);
	    }
	    return true;
	}

	static bool enumProcessThreads( bool function( uint id, void* context ) dg, void* context )
	{
	    return enumProcessThreads( GetCurrentProcessId(), dg, context );
	}

	// execute function on the TLS for the given thread
	static void impersonate_thread( uint id, void function() fn )
	{
	    if( id == GetCurrentThreadId() )
	    {
		fn();
		return;
	    }

	    // temporarily set current TLS data pointer to the data pointer of the referenced thread
	    void** curteb = getTEB();
	    void** teb    = getTEB( id );
	    assert( teb && curteb );

	    void** curtlsarray = cast(void**) curteb[11];
	    void** tlsarray    = cast(void**) teb[11];
	    if( !curtlsarray || !tlsarray )
		return;

	    void** curtlsdata = cast(void**) curtlsarray[_tls_index];
	    void** tlsdata    = cast(void**) tlsarray[_tls_index];
	    if( !curtlsdata || !tlsdata )
		return;

	    curtlsarray[_tls_index] = tlsdata;
	    fn();
	    curtlsarray[_tls_index] = curtlsdata;
        }
    }

public:
    // forward as few symbols as possible into the "global" name space
    alias thread_helper_aux.getTEB getTEB;
    alias thread_helper_aux.getThreadStackBottom getThreadStackBottom;
    alias thread_helper_aux.OpenThreadHandle OpenThreadHandle;
    alias thread_helper_aux.enumProcessThreads enumProcessThreads;

    void* GetTlsDataAddress( HANDLE hnd )
    {
	if( void** teb = getTEB( hnd ) )
	    if( void** tlsarray = cast(void**) teb[11] )
		return tlsarray[_tls_index];
	return null;
    }
    
    ///////////////////////////////////////////////////////////////////
    // run _moduleTlsCtor in the context of the given thread
    void thread_moduleTlsCtor( uint id )
    {
	thread_helper_aux.impersonate_thread(id, &_moduleTlsCtor);
    }

    ///////////////////////////////////////////////////////////////////
    // run _moduleTlsDtor in the context of the given thread
    void thread_moduleTlsDtor( uint id )
    {
	thread_helper_aux.impersonate_thread(id, &_moduleTlsDtor);
    }
}

