/** @file
  Provides cache info for each package, core type, cache level and cache type.

  Copyright (c) 2020 - 2021, Intel Corporation. All rights reserved.<BR>
  SPDX-License-Identifier: BSD-2-Clause-Patent

**/

#include "InternalCpuCacheInfoLib.h"

/**
  Print CpuCacheInfo array.

  @param[in]  CpuCacheInfo        Pointer to the CpuCacheInfo array.
  @param[in]  CpuCacheInfoCount   The length of CpuCacheInfo array.

**/
VOID
CpuCacheInfoPrintCpuCacheInfoTable (
  IN CPU_CACHE_INFO         *CpuCacheInfo,
  IN UINTN                  CpuCacheInfoCount
  )
{
  UINTN                     Index;

  DEBUG ((DEBUG_INFO, "+-------+--------------------------------------------------------------------------------------+\n"));
  DEBUG ((DEBUG_INFO, "| Index | Packge  CoreType  CacheLevel  CacheType  CacheWays (FA|DM) CacheSizeinKB  CacheCount |\n"));
  DEBUG ((DEBUG_INFO, "+-------+--------------------------------------------------------------------------------------+\n"));

  for (Index = 0; Index < CpuCacheInfoCount; Index++) {
    DEBUG ((DEBUG_INFO, "| %4x  | %4x       %2x        %2x          %2x       %4x     ( %x| %x) %8x         %4x     |\n",
        Index, CpuCacheInfo[Index].Package, CpuCacheInfo[Index].CoreType, CpuCacheInfo[Index].CacheLevel,
        CpuCacheInfo[Index].CacheType, CpuCacheInfo[Index].CacheWays, CpuCacheInfo[Index].FullyAssociativeCache,
        CpuCacheInfo[Index].DirectMappedCache, CpuCacheInfo[Index].CacheSizeinKB, CpuCacheInfo[Index].CacheCount));
  }

  DEBUG ((DEBUG_INFO, "+-------+--------------------------------------------------------------------------------------+\n"));
}

/**
  Get the total number of package and package ID in the platform.

  @param[in]      ProcessorInfo       Pointer to the ProcessorInfo array.
  @param[in]      NumberOfProcessors  Total number of logical processors in the platform.
  @param[in, out] Package             Pointer to the Package array.

  @retval  Return the total number of package and package ID in the platform.
**/
UINT32
CpuCacheInfoGetNumberOfPackages (
  IN CPUID_PROCESSOR_INFO   *ProcessorInfo,
  IN UINTN                  NumberOfProcessors,
  IN OUT UINT32             *Package
  )
{
  UINTN                     ProcessorIndex;
  UINT32                    PackageIndex;
  UINT32                    PackageCount;
  UINT32                    CurrentPackage;

  PackageCount = 0;

  for (ProcessorIndex = 0; ProcessorIndex < NumberOfProcessors; ProcessorIndex++) {
    CurrentPackage = ProcessorInfo[ProcessorIndex].Package;

    //
    // For the package that already exists in Package array, break out the loop.
    //
    for (PackageIndex = 0; PackageIndex < PackageCount; PackageIndex++) {
      if (CurrentPackage == Package[PackageIndex]) {
        break;
      }
    }

    //
    // For the new package, save it in Package array.
    //
    if (PackageIndex == PackageCount) {
      ASSERT (PackageCount < MAX_NUM_OF_PACKAGE);
      Package[PackageCount++] = CurrentPackage;
    }
  }

  return PackageCount;
}

/**
  Get the number of CoreType of requested package.

  @param[in]  ProcessorInfo       Pointer to the ProcessorInfo array.
  @param[in]  NumberOfProcessors  Total number of logical processors in the platform.
  @param[in]  Package             The requested package number.

  @retval  Return the number of CoreType of requested package.
**/
UINTN
CpuCacheInfoGetNumberOfCoreTypePerPackage(
  IN CPUID_PROCESSOR_INFO   *ProcessorInfo,
  IN UINTN                  NumberOfProcessors,
  IN UINTN                  Package
  )
{
  UINTN                     ProcessorIndex;
  //
  // Core Type value comes from CPUID.1Ah.EAX[31:24].
  // So max number of core types should be MAX_UINT8.
  //
  UINT8                     CoreType[MAX_UINT8];
  UINTN                     CoreTypeIndex;
  UINTN                     CoreTypeCount;
  UINT8                     CurrentCoreType;

  //
  // CoreType array is empty.
  //
  CoreTypeCount = 0;

  for (ProcessorIndex = 0; ProcessorIndex < NumberOfProcessors; ProcessorIndex++) {
    CurrentCoreType = ProcessorInfo[ProcessorIndex].CoreType;

    if (ProcessorInfo[ProcessorIndex].Package != Package) {
      continue;
    }

    //
    // For the type that already exists in CoreType array, break out the loop.
    //
    for (CoreTypeIndex = 0; CoreTypeIndex < CoreTypeCount; CoreTypeIndex++) {
      if (CurrentCoreType == CoreType[CoreTypeIndex]) {
        break;
      }
    }

    //
    // For the new type, save it in CoreType array.
    //
    if (CoreTypeIndex == CoreTypeCount) {
      ASSERT (CoreTypeCount < MAX_UINT8);
      CoreType[CoreTypeCount++] = CurrentCoreType;
    }
  }

  return CoreTypeCount;
}

/**
  Collect core and cache information of calling processor via CPUID instructions.

  @param[in, out] Buffer              The pointer to private data buffer.
**/
VOID
EFIAPI
CpuCacheInfoCollectCoreAndCacheData (
  IN OUT VOID               *Buffer
  )
{
  UINTN                     ProcessorIndex;
  UINT32                    CpuidMaxInput;
  UINT8                     CacheParamLeafIndex;
  CPUID_CACHE_PARAMS_EAX    CacheParamEax;
  CPUID_CACHE_PARAMS_EBX    CacheParamEbx;
  UINT32                    CacheParamEcx;
  CPUID_CACHE_PARAMS_EDX    CacheParamEdx;
  CPUID_NATIVE_MODEL_ID_AND_CORE_TYPE_EAX   NativeModelIdAndCoreTypeEax;
  COLLECT_CPUID_CACHE_DATA_CONTEXT  *Context;
  CPUID_CACHE_DATA          *CacheData;

  Context = (COLLECT_CPUID_CACHE_DATA_CONTEXT *)Buffer;
  ProcessorIndex = CpuCacheInfoWhoAmI (Context->MpServices);
  CacheData = &Context->CacheData[MAX_NUM_OF_CACHE_PARAMS_LEAF * ProcessorIndex];

  AsmCpuid (CPUID_SIGNATURE, &CpuidMaxInput, NULL, NULL, NULL);

  //
  // get CoreType if CPUID_HYBRID_INFORMATION leaf is supported.
  //
  Context->ProcessorInfo[ProcessorIndex].CoreType = 0;
  if (CpuidMaxInput >= CPUID_HYBRID_INFORMATION) {
    AsmCpuidEx (CPUID_HYBRID_INFORMATION, CPUID_HYBRID_INFORMATION_MAIN_LEAF, &NativeModelIdAndCoreTypeEax.Uint32, NULL, NULL, NULL);
    Context->ProcessorInfo[ProcessorIndex].CoreType = (UINT8) NativeModelIdAndCoreTypeEax.Bits.CoreType;
  }

  //
  // cache hierarchy starts with an index value of 0.
  //
  CacheParamLeafIndex = 0;

  while (CacheParamLeafIndex < MAX_NUM_OF_CACHE_PARAMS_LEAF) {
    AsmCpuidEx (CPUID_CACHE_PARAMS, CacheParamLeafIndex, &CacheParamEax.Uint32, &CacheParamEbx.Uint32, &CacheParamEcx, &CacheParamEdx.Uint32);

    if (CacheParamEax.Bits.CacheType == 0) {
      break;
    }

    CacheData[CacheParamLeafIndex].CacheLevel            = (UINT8)CacheParamEax.Bits.CacheLevel;
    CacheData[CacheParamLeafIndex].CacheType             = (UINT8)CacheParamEax.Bits.CacheType;
    CacheData[CacheParamLeafIndex].CacheWays             = (UINT16)CacheParamEbx.Bits.Ways;
    CacheData[CacheParamLeafIndex].FullyAssociativeCache = (UINT8)CacheParamEax.Bits.FullyAssociativeCache;
    CacheData[CacheParamLeafIndex].DirectMappedCache     = (UINT8)CacheParamEdx.Bits.ComplexCacheIndexing;
    CacheData[CacheParamLeafIndex].CacheShareBits        = (UINT16)CacheParamEax.Bits.MaximumAddressableIdsForLogicalProcessors;
    CacheData[CacheParamLeafIndex].CacheSizeinKB         = (CacheParamEbx.Bits.Ways + 1) *
        (CacheParamEbx.Bits.LinePartitions + 1) * (CacheParamEbx.Bits.LineSize + 1) * (CacheParamEcx + 1) / SIZE_1KB;

    CacheParamLeafIndex++;
  }
}

/**
  Collect CacheInfo data from the CacheData.

  @param[in]      CacheData           Pointer to the CacheData array.
  @param[in]      ProcessorInfo       Pointer to the ProcessorInfo array.
  @param[in]      NumberOfProcessors  Total number of logical processors in the platform.
  @param[in, out] CacheInfo           Pointer to the CacheInfo array.
  @param[in, out] CacheInfoCount      As input, point to the length of response CacheInfo array.
                                      As output, point to the actual length of response CacheInfo array.

  @retval         EFI_SUCCESS             Function completed successfully.
  @retval         EFI_OUT_OF_RESOURCES    Required resources could not be allocated.
  @retval         EFI_BUFFER_TOO_SMALL    CacheInfoCount is too small to hold the response CacheInfo
                                          array. CacheInfoCount has been updated with the length needed
                                          to complete the request.
**/
EFI_STATUS
CpuCacheInfoCollectCpuCacheInfoData (
  IN CPUID_CACHE_DATA       *CacheData,
  IN CPUID_PROCESSOR_INFO   *ProcessorInfo,
  IN UINTN                  NumberOfProcessors,
  IN OUT CPU_CACHE_INFO     *CacheInfo,
  IN OUT UINTN              *CacheInfoCount
  )
{
  EFI_STATUS                Status;
  UINT32                    NumberOfPackage;
  UINT32                    Package[MAX_NUM_OF_PACKAGE];
  UINTN                     PackageIndex;
  UINTN                     TotalNumberOfCoreType;
  UINTN                     MaxCacheInfoCount;
  CPU_CACHE_INFO            *LocalCacheInfo;
  UINTN                     CacheInfoIndex;
  UINTN                     LocalCacheInfoCount;
  UINTN                     Index;
  UINTN                     NextIndex;

  //
  // Get number of Packages and Package ID.
  //
  NumberOfPackage = CpuCacheInfoGetNumberOfPackages (ProcessorInfo, NumberOfProcessors, Package);

  //
  // Get number of core types for each package and count the total number.
  // E.g. If Package1 and Package2 both have 2 core types, the total number is 4.
  //
  TotalNumberOfCoreType = 0;
  for (PackageIndex = 0; PackageIndex < NumberOfPackage; PackageIndex++) {
    TotalNumberOfCoreType += CpuCacheInfoGetNumberOfCoreTypePerPackage (ProcessorInfo, NumberOfProcessors, Package[PackageIndex]);
  }

  MaxCacheInfoCount = TotalNumberOfCoreType * MAX_NUM_OF_CACHE_PARAMS_LEAF;
  LocalCacheInfo = AllocatePages (EFI_SIZE_TO_PAGES (MaxCacheInfoCount * sizeof (*LocalCacheInfo)));
  ASSERT (LocalCacheInfo != NULL);
  if (LocalCacheInfo == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  LocalCacheInfoCount = 0;

  for (Index = 0; Index < NumberOfProcessors * MAX_NUM_OF_CACHE_PARAMS_LEAF; Index++) {
    if (CacheData[Index].CacheSizeinKB == 0) {
      continue;
    }

    //
    // For the sharing caches, clear their CacheSize.
    //
    for (NextIndex = Index + 1; NextIndex < NumberOfProcessors * MAX_NUM_OF_CACHE_PARAMS_LEAF; NextIndex++) {
      if (CacheData[NextIndex].CacheSizeinKB == 0) {
        continue;
      }

      if (CacheData[Index].CacheLevel == CacheData[NextIndex].CacheLevel &&
          CacheData[Index].CacheType == CacheData[NextIndex].CacheType &&
          ProcessorInfo[Index / MAX_NUM_OF_CACHE_PARAMS_LEAF].Package == ProcessorInfo[NextIndex / MAX_NUM_OF_CACHE_PARAMS_LEAF].Package &&
          ProcessorInfo[Index / MAX_NUM_OF_CACHE_PARAMS_LEAF].CoreType == ProcessorInfo[NextIndex / MAX_NUM_OF_CACHE_PARAMS_LEAF].CoreType &&
          (ProcessorInfo[Index / MAX_NUM_OF_CACHE_PARAMS_LEAF].ApicId & ~CacheData[Index].CacheShareBits) ==
          (ProcessorInfo[NextIndex / MAX_NUM_OF_CACHE_PARAMS_LEAF].ApicId & ~CacheData[NextIndex].CacheShareBits)) {
        CacheData[NextIndex].CacheSizeinKB = 0; // uses the sharing cache
      }
    }

    //
    // For the cache that already exists in LocalCacheInfo, increase its CacheCount.
    //
    for (CacheInfoIndex = 0; CacheInfoIndex < LocalCacheInfoCount; CacheInfoIndex++) {
      if (LocalCacheInfo[CacheInfoIndex].Package    == ProcessorInfo[Index / MAX_NUM_OF_CACHE_PARAMS_LEAF].Package &&
          LocalCacheInfo[CacheInfoIndex].CoreType   == ProcessorInfo[Index / MAX_NUM_OF_CACHE_PARAMS_LEAF].CoreType &&
          LocalCacheInfo[CacheInfoIndex].CacheLevel == CacheData[Index].CacheLevel &&
          LocalCacheInfo[CacheInfoIndex].CacheType  == CacheData[Index].CacheType) {
        LocalCacheInfo[CacheInfoIndex].CacheCount++;
        break;
      }
    }

    //
    // For the new cache with different Package, CoreType, CacheLevel or CacheType, copy its
    // data into LocalCacheInfo buffer.
    //
    if (CacheInfoIndex == LocalCacheInfoCount) {
      ASSERT (LocalCacheInfoCount < MaxCacheInfoCount);

      LocalCacheInfo[LocalCacheInfoCount].Package               = ProcessorInfo[Index / MAX_NUM_OF_CACHE_PARAMS_LEAF].Package;
      LocalCacheInfo[LocalCacheInfoCount].CoreType              = ProcessorInfo[Index / MAX_NUM_OF_CACHE_PARAMS_LEAF].CoreType;
      LocalCacheInfo[LocalCacheInfoCount].CacheLevel            = CacheData[Index].CacheLevel;
      LocalCacheInfo[LocalCacheInfoCount].CacheType             = CacheData[Index].CacheType;
      LocalCacheInfo[LocalCacheInfoCount].CacheWays             = CacheData[Index].CacheWays;
      LocalCacheInfo[LocalCacheInfoCount].FullyAssociativeCache = CacheData[Index].FullyAssociativeCache;
      LocalCacheInfo[LocalCacheInfoCount].DirectMappedCache     = CacheData[Index].DirectMappedCache;
      LocalCacheInfo[LocalCacheInfoCount].CacheSizeinKB         = CacheData[Index].CacheSizeinKB;
      LocalCacheInfo[LocalCacheInfoCount].CacheCount            = 1;

      LocalCacheInfoCount++;
    }
  }

  if (*CacheInfoCount < LocalCacheInfoCount) {
    Status = EFI_BUFFER_TOO_SMALL;
  } else {
    CopyMem (CacheInfo, LocalCacheInfo, sizeof (*CacheInfo) * LocalCacheInfoCount);
    DEBUG_CODE (
      CpuCacheInfoPrintCpuCacheInfoTable (CacheInfo, LocalCacheInfoCount);
    );
    Status = EFI_SUCCESS;
  }

  *CacheInfoCount = LocalCacheInfoCount;

  FreePages (LocalCacheInfo, EFI_SIZE_TO_PAGES (MaxCacheInfoCount * sizeof (*LocalCacheInfo)));

  return Status;
}

/**
  Get CpuCacheInfo data array.

  @param[in, out] CpuCacheInfo        Pointer to the CpuCacheInfo array.
  @param[in, out] CpuCacheInfoCount   As input, point to the length of response CpuCacheInfo array.
                                      As output, point to the actual length of response CpuCacheInfo array.

  @retval         EFI_SUCCESS             Function completed successfully.
  @retval         EFI_INVALID_PARAMETER   CpuCacheInfoCount is NULL.
  @retval         EFI_INVALID_PARAMETER   CpuCacheInfo is NULL while CpuCacheInfoCount contains the value
                                          greater than zero.
  @retval         EFI_UNSUPPORTED         Processor does not support CPUID_CACHE_PARAMS Leaf.
  @retval         EFI_OUT_OF_RESOURCES    Required resources could not be allocated.
  @retval         EFI_BUFFER_TOO_SMALL    CpuCacheInfoCount is too small to hold the response CpuCacheInfo
                                          array. CpuCacheInfoCount has been updated with the length needed
                                          to complete the request.
**/
EFI_STATUS
EFIAPI
GetCpuCacheInfo (
  IN OUT CPU_CACHE_INFO     *CpuCacheInfo,
  IN OUT UINTN              *CpuCacheInfoCount
  )
{
  EFI_STATUS                Status;
  UINT32                    CpuidMaxInput;
  UINT32                    NumberOfProcessors;
  UINTN                     CacheDataCount;
  UINTN                     ProcessorIndex;
  EFI_PROCESSOR_INFORMATION ProcessorInfo;
  COLLECT_CPUID_CACHE_DATA_CONTEXT  Context;

  if (CpuCacheInfoCount == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  if (*CpuCacheInfoCount != 0 && CpuCacheInfo == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  AsmCpuid (CPUID_SIGNATURE, &CpuidMaxInput, NULL, NULL, NULL);
  if (CpuidMaxInput < CPUID_CACHE_PARAMS) {
    return EFI_UNSUPPORTED;
  }

  //
  // Initialize COLLECT_CPUID_CACHE_DATA_CONTEXT.MpServices.
  //
  CpuCacheInfoGetMpServices (&Context.MpServices);

  NumberOfProcessors = CpuCacheInfoGetNumberOfProcessors (Context.MpServices);

  //
  // Initialize COLLECT_CPUID_CACHE_DATA_CONTEXT.ProcessorInfo.
  //
  Context.ProcessorInfo = AllocatePages (EFI_SIZE_TO_PAGES (NumberOfProcessors * sizeof (*Context.ProcessorInfo)));
  ASSERT (Context.ProcessorInfo != NULL);
  if (Context.ProcessorInfo == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }
  //
  // Initialize COLLECT_CPUID_CACHE_DATA_CONTEXT.CacheData.
  // CacheData array consists of CPUID_CACHE_DATA data structure for each Cpuid Cache Parameter Leaf
  // per logical processor. The array begin with data of each Cache Parameter Leaf of processor 0, followed
  // by data of each Cache Parameter Leaf of processor 1 ...
  //
  CacheDataCount = NumberOfProcessors * MAX_NUM_OF_CACHE_PARAMS_LEAF;
  Context.CacheData = AllocatePages (EFI_SIZE_TO_PAGES (CacheDataCount * sizeof (*Context.CacheData)));
  ASSERT (Context.CacheData != NULL);
  if (Context.CacheData == NULL) {
    FreePages (Context.ProcessorInfo, EFI_SIZE_TO_PAGES (NumberOfProcessors * sizeof (*Context.ProcessorInfo)));
    return EFI_OUT_OF_RESOURCES;
  }

  ZeroMem (Context.CacheData, CacheDataCount * sizeof (*Context.CacheData));

  //
  // Collect Package ID and APIC ID of all processors.
  //
  for (ProcessorIndex = 0; ProcessorIndex < NumberOfProcessors; ProcessorIndex++) {
    CpuCacheInfoGetProcessorInfo (Context.MpServices, ProcessorIndex, &ProcessorInfo);
    Context.ProcessorInfo[ProcessorIndex].Package = ProcessorInfo.Location.Package;
    Context.ProcessorInfo[ProcessorIndex].ApicId  = (UINT32) ProcessorInfo.ProcessorId;
  }

  //
  // Wakeup all processors for CacheData(core type and cache data) collection.
  //
  CpuCacheInfoStartupAllCPUs (Context.MpServices, CpuCacheInfoCollectCoreAndCacheData, &Context);

  //
  // Collect CpuCacheInfo data from CacheData.
  //
  Status = CpuCacheInfoCollectCpuCacheInfoData (Context.CacheData, Context.ProcessorInfo, NumberOfProcessors, CpuCacheInfo, CpuCacheInfoCount);

  FreePages (Context.CacheData, EFI_SIZE_TO_PAGES (CacheDataCount * sizeof (*Context.CacheData)));
  FreePages (Context.ProcessorInfo, EFI_SIZE_TO_PAGES (NumberOfProcessors * sizeof (*Context.ProcessorInfo)));

  return Status;
}