/** @file
  Copyright (c) 2020, PMheart. All rights reserved.
  SPDX-License-Identifier: BSD-3-Clause
**/

#include <Uefi.h>
#include <Library/BaseMemoryLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/DebugLib.h>
#include <Library/UefiLib.h>
#include <Library/UefiApplicationEntryPoint.h>

#include <string.h>
#include <stdlib.h>

#ifdef WIN32
#include <malloc.h>
#endif // WIN32

UINTN mPoolAllocations;
UINTN mPageAllocations;

VOID *
EFIAPI
CopyMem (
  OUT VOID        *DestinationBuffer,
  IN  CONST VOID  *SourceBuffer,
  IN  UINTN       Length
  )
{
  return memmove (DestinationBuffer, SourceBuffer, Length);
}

VOID *
EFIAPI
SetMem (
  OUT VOID   *Buffer,
  IN  UINTN  Length,
  IN  UINT8  Value
  )
{
  return memset (Buffer, Value, Length);
}

VOID *
EFIAPI
ZeroMem (
  OUT VOID   *Buffer,
  IN  UINTN  Length
  )
{
  return memset (Buffer, 0, Length);
}

INTN
EFIAPI
CompareMem (
  IN  CONST VOID  *DestinationBuffer,
  IN  CONST VOID  *SourceBuffer,
  IN  UINTN       Length
  )
{
  return memcmp (DestinationBuffer, SourceBuffer, Length);
}

VOID *
EFIAPI
ScanMem16 (
  IN  CONST VOID  *Buffer,
  IN  UINTN       Length,
  IN  UINT16      Value
  )
{
  UINT16  *Pointer;
  UINTN   Index;

  Pointer = (UINT16 *) Buffer;

  for (Index = 0; Index < Length; ++Index) {
    if (Pointer[Index] == Value) {
      return Pointer;
    }
  }

  return NULL;
}

VOID *
EFIAPI
AllocatePool (
  IN  UINTN  AllocationSize
  )
{
  // UEFI guarantees 8-byte alignment.
  void *p = malloc ((AllocationSize + 7U) & ~7U);
  DEBUG ((
    DEBUG_POOL,
    "UMEM: Allocating pool %u at %p\n",
    (UINT32) AllocationSize,
    p
    ));
  ASSERT (((UINTN)p & 7U) == 0);
  if (p != NULL) {
    ++mPoolAllocations;
  }
  return p;
}

VOID *
EFIAPI
AllocateCopyPool (
  IN  UINTN       AllocationSize,
  IN  CONST VOID  *Buffer
  )
{
  VOID *Memory;

  ASSERT (Buffer != NULL);

  Memory = AllocatePool (AllocationSize);

  if (Memory != NULL) {
    Memory = CopyMem (Memory, Buffer, AllocationSize);
  }

  return Memory;
}

VOID *
EFIAPI
AllocateZeroPool (
  IN  UINTN  AllocationSize
  )
{
  VOID *Memory;

  Memory = AllocatePool (AllocationSize);

  if (Memory != NULL) {
    Memory = ZeroMem (Memory, AllocationSize);
  }

  return Memory;
}

VOID *
ReallocatePool (
  IN  UINTN  OldSize,
  IN  UINTN  NewSize,
  IN  VOID   *OldBuffer  OPTIONAL
  )
{
  VOID  *NewBuffer;

  NewBuffer = AllocateZeroPool (NewSize);

  if (NewBuffer != NULL && OldBuffer != NULL) {
    memcpy (NewBuffer, OldBuffer, MIN (OldSize, NewSize));
    FreePool (OldBuffer);
  }

  return NewBuffer;
}

VOID *
EFIAPI
AllocatePages (
  IN UINTN  Pages
  )
{
  VOID  *Memory;

  #ifdef WIN32
  Memory = _aligned_malloc (Pages * EFI_PAGE_SIZE, EFI_PAGE_SIZE);
  #else // !WIN32
  INTN  RetVal;

  Memory = NULL;

  RetVal = posix_memalign (&Memory, EFI_PAGE_SIZE, Pages * EFI_PAGE_SIZE);

  if (RetVal != 0) {
    DEBUG ((DEBUG_ERROR, "posix_memalign returns error %d\n", RetVal));
    Memory = NULL;
  }
  #endif // WIN32

  DEBUG ((
    DEBUG_PAGE,
    "UMEM: Allocating %u pages at %p\n",
    (UINT32) Pages,
    Memory
    ));

  if (Memory != NULL) {
    mPageAllocations += Pages;
  }

  return Memory;
}

VOID
EFIAPI
FreePool (
  IN VOID   *Buffer
  )
{
  ASSERT (Buffer != NULL);
  DEBUG ((
    DEBUG_POOL,
    "UMEM: Deallocating pool %p\n",
    Buffer
    ));

  --mPoolAllocations;

  free (Buffer);
}

VOID
EFIAPI
FreePages (
  IN VOID   *Buffer,
  IN UINTN  Pages
  )
{
  ASSERT (Buffer != NULL);

  DEBUG ((
    DEBUG_PAGE,
    "UMEM: Deallocating %u pages at %p\n",
    (UINT32) Pages,
    Buffer
    ));

  mPageAllocations -= Pages;

  free (Buffer);
}

GUID *
EFIAPI
CopyGuid (
  OUT GUID        *DestinationGuid,
  IN  CONST GUID  *SourceGuid
  )
{
  memmove (DestinationGuid, SourceGuid, sizeof (GUID));

  return DestinationGuid;
}

BOOLEAN
EFIAPI
CompareGuid (
  IN  CONST GUID  *Guid1,
  IN  CONST GUID  *Guid2
  )
{
  return memcmp (Guid1, Guid2, sizeof (GUID)) == 0;
}

UINT16
EFIAPI
ReadUnaligned16 (
  IN CONST UINT16  *Buffer
  )
{
  UINT16 Value;

  memmove (&Value, Buffer, sizeof (UINT16));

  return Value;
}

UINT16
EFIAPI
WriteUnaligned16 (
  OUT UINT16  *Buffer,
  IN  UINT16  Value
  )
{
  ASSERT (Buffer != NULL);

  memmove (Buffer, &Value, sizeof (UINT16));

  return Value;
}

UINT32
EFIAPI
ReadUnaligned24 (
  IN CONST UINT32  *Buffer
  )
{
  UINT32 Value;

  Value = ReadUnaligned32 (Buffer) & 0xFFFFFFU;

  return Value;
}

UINT32
EFIAPI
ReadUnaligned32 (
  IN CONST UINT32  *Buffer
  )
{
  UINT32 Value;

  memmove (&Value, Buffer, sizeof (UINT32));

  return Value;
}

UINT64
EFIAPI
ReadUnaligned64 (
  IN CONST UINT64  *Buffer
  )
{
  UINT64 Value;

  memmove (&Value, Buffer, sizeof (UINT64));

  return Value;
}

UINT32
EFIAPI
WriteUnaligned32 (
  OUT UINT32  *Buffer,
  IN  UINT32  Value
  )
{
  ASSERT (Buffer != NULL);

  memmove (Buffer, &Value, sizeof (UINT32));

  return Value;
}

UINT64
EFIAPI
WriteUnaligned64 (
  OUT UINT64  *Buffer,
  IN  UINT64  Value
  )
{
  ASSERT (Buffer != NULL);

  memmove (Buffer, &Value, sizeof (UINT64));

  return Value;
}