// -*- C++ -*-

/*

  The Hoard Multiprocessor Memory Allocator
  www.hoard.org

  Author: Emery Berger, http://www.emeryberger.com
  Copyright (c) 1998-2020 Emery Berger

  See the LICENSE file at the top-level directory of this
  distribution and at http://github.com/emeryberger/Hoard.

*/

/*
 * This file leverages compiler support for thread-local variables for
 * access to thread-local heaps, when available. It also intercepts
 * thread completions to flush these local heaps, returning any unused
 * memory to the global Hoard heap. On Windows, this happens in
 * DllMain. On Unix platforms, we interpose our own versions of
 * pthread_create and pthread_exit.
 */
#if defined(_WIN32)

#include <new>

#pragma warning(disable: 4447) // Disable weird warning about threading model.

// Windows TLS functions.

#include "VERSION.h"

#define versionMessage "Using the Hoard memory allocator (http://www.hoard.org), version " HOARD_VERSION_STRING "\n"

#include "hoardheap.h"
#include "hoardtlab.h"

using namespace Hoard;

#define USE_DECLSPEC_THREADLOCAL 0

#if USE_DECLSPEC_THREADLOCAL
__declspec(thread) TheCustomHeapType * threadLocalHeap = NULL;
#else
DWORD LocalTLABIndex = TLS_OUT_OF_INDEXES;
#endif

extern HoardHeapType * getMainHoardHeap();

static TheCustomHeapType * initializeCustomHeap()
{
  // Allocate a per-thread heap.
  auto * mainHeap = getMainHoardHeap();
  auto * customHeapBuf = mainHeap->malloc(sizeof(TheCustomHeapType));
  auto * perThreadHeap = new (customHeapBuf) TheCustomHeapType (mainHeap);

  // Store it in the appropriate thread-local area.
#if USE_DECLSPEC_THREADLOCAL
  threadLocalHeap = perThreadHeap;
#else
  TlsSetValue (LocalTLABIndex, perThreadHeap);
#endif

  return perThreadHeap;
}

bool isCustomHeapInitialized() {
#if USE_DECLSPEC_THREADLOCAL
  return (threadLocalHeap != NULL);
#else
  return (TlsGetValue(LocalTLABIndex) != NULL);
#endif
}

// Track if TLS is ready (set in DLL_PROCESS_ATTACH after TlsAlloc)
static bool g_tlsReady = false;

TheCustomHeapType * getCustomHeap() {
#if USE_DECLSPEC_THREADLOCAL
  if (threadLocalHeap != NULL)
    return threadLocalHeap;
  initializeCustomHeap();
  return threadLocalHeap;
#else
  // Fast path: check if TLS is ready
  if (!g_tlsReady) {
    return nullptr;
  }
  auto p = TlsGetValue(LocalTLABIndex);
  if (p == NULL) {
    initializeCustomHeap();
    p = TlsGetValue(LocalTLABIndex);
  }
  return (TheCustomHeapType *) p;
#endif
}

extern "C" void InitializeWinWrapper();
extern "C" void FinalizeWinWrapper();


//
// Intercept thread creation and destruction to flush the TLABs.
//

#include <iostream>
using namespace std;

// example usage: debugMessage(L"Hello, world.\n");
void debugMessage(const wchar_t * lpBuff) {
  static bool initialized = false;
  if (!initialized) {
    AllocConsole();
  }
  DWORD dwSize = 0;
  WriteConsoleW(GetStdHandle(STD_OUTPUT_HANDLE), lpBuff, lstrlenW(lpBuff), &dwSize, NULL);
}

extern "C" void _CRT_INIT();

extern "C" {

  BOOL APIENTRY DllMain (HANDLE hinstDLL,
			 DWORD fdwReason,
			 LPVOID lpreserved)
  {
    static auto np = HL::CPUInfo::computeNumProcessors();
    switch (fdwReason) {
      
    case DLL_PROCESS_ATTACH:
      {
	// Pin this DLL in memory to prevent unloading. This is critical because
	// we patch CRT DLLs with detours - if our DLL unloads before the CRT,
	// the trampolines become invalid and cause crashes during process exit.
	HMODULE hSelf = nullptr;
	GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_PIN,
			   (LPCWSTR)DllMain, &hSelf);

	// Before we do anything, force initialization of the C++
	// library. Without this pre-initialization, the Windows heap
	// and the Hoard heaps get mixed up, and then nothing
	// works. This is quite the hack but seems to do the trick.
	// -- Emery Berger, 24/1/2019
	cout << "";

#if !USE_DECLSPEC_THREADLOCAL
	// Allocate TLS index for per-thread heap pointer
	LocalTLABIndex = TlsAlloc();
	if (LocalTLABIndex == TLS_OUT_OF_INDEXES) {
	  return FALSE;
	}
	g_tlsReady = true;
#endif

	// Now we are good to go.
	InitializeWinWrapper();
	// Force creation of the heap.
	volatile auto * ch = getCustomHeap();
      }
      break;
      
    case DLL_THREAD_ATTACH:
      {
	if (np == 1) {
	  // We have exactly one processor - just assign the thread to
	  // heap 0.
	  getMainHoardHeap()->chooseZero();
	} else {
	  getMainHoardHeap()->findUnusedHeap();
	}
	// Force creation of the heap.
	volatile auto * ch = getCustomHeap();
      }
      break;
      
    case DLL_THREAD_DETACH:
      {
	// Dump the memory from the TLAB.
	getCustomHeap()->clear();
	
#if USE_DECLSPEC_THREADLOCAL
	auto * heap = threadLocalHeap;
#else
	auto * heap = (TheCustomHeapType *) TlsGetValue(LocalTLABIndex);
#endif
	
	if (np != 1) {
	  // If we're on a multiprocessor box, relinquish the heap
	  // assigned to this thread.
	  getMainHoardHeap()->releaseHeap();
	}

      }
      break;
      
    case DLL_PROCESS_DETACH:
      if (lpreserved == NULL) {
	// Dynamic unload (FreeLibrary) - shouldn't happen since we pinned the DLL
	FinalizeWinWrapper();
#if !USE_DECLSPEC_THREADLOCAL
	if (LocalTLABIndex != TLS_OUT_OF_INDEXES) {
	  TlsFree(LocalTLABIndex);
	  LocalTLABIndex = TLS_OUT_OF_INDEXES;
	}
	g_tlsReady = false;
#endif
      } else {
	// Process exit (lpreserved != NULL)
	// Try to ensure output is flushed before terminating.
	// Use WriteFile with INVALID_HANDLE_VALUE check for safety.
	HANDLE hOut = GetStdHandle(STD_OUTPUT_HANDLE);
	if (hOut != INVALID_HANDLE_VALUE && hOut != NULL) {
	  // Write an empty string to force buffer flush
	  DWORD written;
	  WriteFile(hOut, "", 0, &written, NULL);
	}
	// On ARM64 with Detours, allowing other DLLs to run their cleanup causes
	// crashes because they try to use detoured functions that point to
	// invalid memory. Force immediate termination to avoid this.
	TerminateProcess(GetCurrentProcess(), 0);
      }
      break;
      
    default:
      return TRUE;
    }

    return TRUE;
  }
}

#endif // _WIN32
