/** @file
  Platform VTd Info Sample PEI driver.

  Copyright (c) 2017, Intel Corporation. All rights reserved.<BR>
  This program and the accompanying materials
  are licensed and made available under the terms and conditions of the BSD License
  which accompanies this distribution.  The full text of the license may be found at
  http://opensource.org/licenses/bsd-license.php

  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.

**/

#include <PiPei.h>

#include <Ppi/VtdInfo.h>

#include <Library/PeiServicesLib.h>
#include <Library/DebugLib.h>
#include <Library/PciLib.h>
#include <Library/IoLib.h>
#include <Library/MemoryAllocationLib.h>

#define R_SA_MCHBAR               (0x48)
#define R_SA_GGC                  (0x50)
#define N_SKL_SA_GGC_GGMS_OFFSET  (0x6)
#define B_SKL_SA_GGC_GGMS_MASK    (0xc0)
#define N_SKL_SA_GGC_GMS_OFFSET   (0x8)
#define B_SKL_SA_GGC_GMS_MASK     (0xff00)
#define V_SKL_SA_GGC_GGMS_8MB     3
#define R_SA_TOLUD                (0xbc)

#define R_SA_MCHBAR_VTD1_OFFSET  0x5400  ///< HW UNIT for IGD
#define R_SA_MCHBAR_VTD2_OFFSET  0x5410  ///< HW UNIT for all other - PEG, USB, SATA etc

EFI_GUID gEdkiiSiliconInitializedPpiGuid = {0x82a72dc8, 0x61ec, 0x403e, {0xb1, 0x5a, 0x8d, 0x7a, 0x3a, 0x71, 0x84, 0x98}};

typedef struct {
  EFI_ACPI_DMAR_HEADER                         DmarHeader;
  //
  // VTd engine 1 - integrated graphic
  //
  EFI_ACPI_DMAR_DRHD_HEADER                    Drhd1;
  EFI_ACPI_DMAR_DEVICE_SCOPE_STRUCTURE_HEADER  Drhd11;
  EFI_ACPI_DMAR_PCI_PATH                       Drhd111;
  //
  // VTd engine 2 - all rest
  //
  EFI_ACPI_DMAR_DRHD_HEADER                    Drhd2;
  //
  // RMRR 1 - integrated graphic
  //
  EFI_ACPI_DMAR_RMRR_HEADER                    Rmrr1;
  EFI_ACPI_DMAR_DEVICE_SCOPE_STRUCTURE_HEADER  Rmrr11;
  EFI_ACPI_DMAR_PCI_PATH                       Rmrr111;
} MY_VTD_INFO_PPI;

MY_VTD_INFO_PPI  mPlatformVTdSample = {
  { // DmarHeader
    { // Header
      EFI_ACPI_4_0_DMA_REMAPPING_TABLE_SIGNATURE,
      sizeof(MY_VTD_INFO_PPI),
      EFI_ACPI_DMAR_REVISION,
    },
    0x26, // HostAddressWidth
  },

  { // Drhd1
    { // Header
      EFI_ACPI_DMAR_TYPE_DRHD,
      sizeof(EFI_ACPI_DMAR_DRHD_HEADER) +
        sizeof(EFI_ACPI_DMAR_DEVICE_SCOPE_STRUCTURE_HEADER) +
        sizeof(EFI_ACPI_DMAR_PCI_PATH)
    },
    0, // Flags
    0, // Reserved
    0, // SegmentNumber
    0xFED90000 // RegisterBaseAddress -- TO BE PATCHED
  },
  { // Drhd11
    EFI_ACPI_DEVICE_SCOPE_ENTRY_TYPE_PCI_ENDPOINT,
    sizeof(EFI_ACPI_DMAR_DEVICE_SCOPE_STRUCTURE_HEADER) +
      sizeof(EFI_ACPI_DMAR_PCI_PATH),
    0, // Reserved2
    0, // EnumerationId
    0  // StartBusNumber
  },
  { // Drhd111
    2,  // Device
    0   // Function
  },

  { // Drhd2
    { // Header
      EFI_ACPI_DMAR_TYPE_DRHD,
      sizeof(EFI_ACPI_DMAR_DRHD_HEADER)
    },
    EFI_ACPI_DMAR_DRHD_FLAGS_INCLUDE_PCI_ALL, // Flags
    0, // Reserved
    0, // SegmentNumber
    0xFED91000 // RegisterBaseAddress -- TO BE PATCHED
  },

  { // Rmrr1
    { // Header
      EFI_ACPI_DMAR_TYPE_RMRR,
      sizeof(EFI_ACPI_DMAR_RMRR_HEADER) +
        sizeof(EFI_ACPI_DMAR_DEVICE_SCOPE_STRUCTURE_HEADER) +
        sizeof(EFI_ACPI_DMAR_PCI_PATH)
    },
    {0}, // Reserved
    0, // SegmentNumber
    0x0, // ReservedMemoryRegionBaseAddress -- TO BE PATCHED
    0x0 // ReservedMemoryRegionLimitAddress -- TO BE PATCHED
  },
  { // Rmrr11
    EFI_ACPI_DEVICE_SCOPE_ENTRY_TYPE_PCI_ENDPOINT,
    sizeof(EFI_ACPI_DMAR_DEVICE_SCOPE_STRUCTURE_HEADER) +
      sizeof(EFI_ACPI_DMAR_PCI_PATH),
    0, // Reserved2
    0, // EnumerationId
    0  // StartBusNumber
  },
  { // Rmrr111
    2,  // Device
    0   // Function
  },
};

EFI_PEI_PPI_DESCRIPTOR mPlatformVTdInfoSampleDesc = {
  (EFI_PEI_PPI_DESCRIPTOR_PPI | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST),
  &gEdkiiVTdInfoPpiGuid,
  &mPlatformVTdSample
};

typedef struct {
  EFI_ACPI_DMAR_HEADER                         DmarHeader;
  //
  // VTd engine 2 - all rest
  //
  EFI_ACPI_DMAR_DRHD_HEADER                    Drhd2;
} MY_VTD_INFO_NO_IGD_PPI;

MY_VTD_INFO_NO_IGD_PPI  mPlatformVTdNoIgdSample = {
  { // DmarHeader
    { // Header
      EFI_ACPI_4_0_DMA_REMAPPING_TABLE_SIGNATURE,
      sizeof(MY_VTD_INFO_NO_IGD_PPI),
      EFI_ACPI_DMAR_REVISION,
    },
    0x26, // HostAddressWidth
  },

  { // Drhd2
    { // Header
      EFI_ACPI_DMAR_TYPE_DRHD,
      sizeof(EFI_ACPI_DMAR_DRHD_HEADER)
    },
    EFI_ACPI_DMAR_DRHD_FLAGS_INCLUDE_PCI_ALL, // Flags
    0, // Reserved
    0, // SegmentNumber
    0xFED91000 // RegisterBaseAddress -- TO BE PATCHED
  },
};

EFI_PEI_PPI_DESCRIPTOR mPlatformVTdNoIgdInfoSampleDesc = {
  (EFI_PEI_PPI_DESCRIPTOR_PPI | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST),
  &gEdkiiVTdInfoPpiGuid,
  &mPlatformVTdNoIgdSample
};

/**
  Initialize VTd register.
**/
VOID
InitDmar (
  VOID
  )
{
  UINT32              MchBar;

  DEBUG ((DEBUG_INFO, "InitDmar\n"));

  MchBar = PciRead32 (PCI_LIB_ADDRESS(0, 0, 0, R_SA_MCHBAR)) & ~BIT0;
  PciWrite32 (PCI_LIB_ADDRESS(0, 0, 0, R_SA_MCHBAR), 0xFED10000 | BIT0);
  DEBUG ((DEBUG_INFO, "MchBar - %x\n", MchBar));

  MmioWrite32 (MchBar + R_SA_MCHBAR_VTD2_OFFSET, (UINT32)mPlatformVTdSample.Drhd2.RegisterBaseAddress | 1);
  DEBUG ((DEBUG_INFO, "VTd2 - %x\n", (MmioRead32 (MchBar + R_SA_MCHBAR_VTD2_OFFSET))));
}

/**
  Patch Graphic UMA address in RMRR and base address.
**/
EFI_PEI_PPI_DESCRIPTOR *
PatchDmar (
  VOID
  )
{
  UINT32                  MchBar;
  UINT16                  IgdMode;
  UINT16                  GttMode;
  UINT32                  IgdMemSize;
  UINT32                  GttMemSize;
  MY_VTD_INFO_PPI         *PlatformVTdSample;
  EFI_PEI_PPI_DESCRIPTOR  *PlatformVTdInfoSampleDesc;
  MY_VTD_INFO_NO_IGD_PPI  *PlatformVTdNoIgdSample;
  EFI_PEI_PPI_DESCRIPTOR  *PlatformVTdNoIgdInfoSampleDesc;

  DEBUG ((DEBUG_INFO, "PatchDmar\n"));

  if (PciRead16 (PCI_LIB_ADDRESS(0, 2, 0, 0)) != 0xFFFF) {
    PlatformVTdSample = AllocateCopyPool (sizeof(MY_VTD_INFO_PPI), &mPlatformVTdSample);
    ASSERT(PlatformVTdSample != NULL);
    PlatformVTdInfoSampleDesc = AllocateCopyPool (sizeof(EFI_PEI_PPI_DESCRIPTOR), &mPlatformVTdInfoSampleDesc);
    ASSERT(PlatformVTdInfoSampleDesc != NULL);
    PlatformVTdInfoSampleDesc->Ppi = PlatformVTdSample;

    ///
    /// Calculate IGD memsize
    ///
    IgdMode = ((PciRead16 (PCI_LIB_ADDRESS(0, 0, 0, R_SA_GGC)) & B_SKL_SA_GGC_GMS_MASK) >> N_SKL_SA_GGC_GMS_OFFSET) & 0xFF;
    if (IgdMode < 0xF0) {
      IgdMemSize = IgdMode * 32 * (1024) * (1024);
    } else {
      IgdMemSize = 4 * (IgdMode - 0xF0 + 1) * (1024) * (1024);
    }

    ///
    /// Calculate GTT mem size
    ///
    GttMemSize = 0;
    GttMode = (PciRead16 (PCI_LIB_ADDRESS(0, 0, 0, R_SA_GGC)) & B_SKL_SA_GGC_GGMS_MASK) >> N_SKL_SA_GGC_GGMS_OFFSET;
    if (GttMode <= V_SKL_SA_GGC_GGMS_8MB) {
      GttMemSize = (1 << GttMode) * (1024) * (1024);
    }

    PlatformVTdSample->Rmrr1.ReservedMemoryRegionBaseAddress  = (PciRead32 (PCI_LIB_ADDRESS(0, 0, 0, R_SA_TOLUD)) & ~(0x01)) - IgdMemSize - GttMemSize;
    PlatformVTdSample->Rmrr1.ReservedMemoryRegionLimitAddress = PlatformVTdSample->Rmrr1.ReservedMemoryRegionBaseAddress + IgdMemSize + GttMemSize - 1;

    ///
    /// Update DRHD structures of DmarTable
    ///
    MchBar = PciRead32 (PCI_LIB_ADDRESS(0, 0, 0, R_SA_MCHBAR)) & ~BIT0;

    if ((MmioRead32 (MchBar + R_SA_MCHBAR_VTD1_OFFSET) &~1) != 0) {
      PlatformVTdSample->Drhd1.RegisterBaseAddress = (MmioRead32 (MchBar + R_SA_MCHBAR_VTD1_OFFSET) &~1);
    } else {
      MmioWrite32 (MchBar + R_SA_MCHBAR_VTD1_OFFSET, (UINT32)PlatformVTdSample->Drhd1.RegisterBaseAddress | 1);
    }
    DEBUG ((DEBUG_INFO, "VTd1 - %x\n", (MmioRead32 (MchBar + R_SA_MCHBAR_VTD1_OFFSET))));

    if ((MmioRead32 (MchBar + R_SA_MCHBAR_VTD2_OFFSET) &~1) != 0) {
      PlatformVTdSample->Drhd2.RegisterBaseAddress = (MmioRead32 (MchBar + R_SA_MCHBAR_VTD2_OFFSET) &~1);
    } else {
      MmioWrite32 (MchBar + R_SA_MCHBAR_VTD2_OFFSET, (UINT32)PlatformVTdSample->Drhd2.RegisterBaseAddress | 1);
    }
    DEBUG ((DEBUG_INFO, "VTd2 - %x\n", (MmioRead32 (MchBar + R_SA_MCHBAR_VTD2_OFFSET))));

    return PlatformVTdInfoSampleDesc;
  } else {
    PlatformVTdNoIgdSample = AllocateCopyPool (sizeof(MY_VTD_INFO_NO_IGD_PPI), &mPlatformVTdNoIgdSample);
    ASSERT(PlatformVTdNoIgdSample != NULL);
    PlatformVTdNoIgdInfoSampleDesc = AllocateCopyPool (sizeof(EFI_PEI_PPI_DESCRIPTOR), &mPlatformVTdNoIgdInfoSampleDesc);
    ASSERT(PlatformVTdNoIgdInfoSampleDesc != NULL);
    PlatformVTdNoIgdInfoSampleDesc->Ppi = PlatformVTdNoIgdSample;

    ///
    /// Update DRHD structures of DmarTable
    ///
    MchBar = PciRead32 (PCI_LIB_ADDRESS(0, 0, 0, R_SA_MCHBAR)) & ~BIT0;

    if ((MmioRead32 (MchBar + R_SA_MCHBAR_VTD2_OFFSET) &~1) != 0) {
      PlatformVTdNoIgdSample->Drhd2.RegisterBaseAddress = (MmioRead32 (MchBar + R_SA_MCHBAR_VTD2_OFFSET) &~1);
    } else {
      MmioWrite32 (MchBar + R_SA_MCHBAR_VTD2_OFFSET, (UINT32)PlatformVTdNoIgdSample->Drhd2.RegisterBaseAddress | 1);
    }
    DEBUG ((DEBUG_INFO, "VTd2 - %x\n", (MmioRead32 (MchBar + R_SA_MCHBAR_VTD2_OFFSET))));

    return PlatformVTdNoIgdInfoSampleDesc;
  }
}

/**
  The callback function for SiliconInitializedPpi.
  It reinstalls VTD_INFO_PPI.

  @param[in]  PeiServices       General purpose services available to every PEIM.
  @param[in]  NotifyDescriptor  Notify that this module published.
  @param[in]  Ppi               PPI that was installed.

  @retval     EFI_SUCCESS       The function completed successfully.
**/
EFI_STATUS
EFIAPI
SiliconInitializedPpiNotifyCallback (
  IN CONST EFI_PEI_SERVICES     **PeiServices,
  IN EFI_PEI_NOTIFY_DESCRIPTOR  *NotifyDescriptor,
  IN VOID                       *Ppi
  )
{
  EFI_STATUS               Status;
  EFI_PEI_PPI_DESCRIPTOR   *PpiDesc;

  PpiDesc = PatchDmar ();

  Status = PeiServicesReInstallPpi (&mPlatformVTdNoIgdInfoSampleDesc, PpiDesc);
  ASSERT_EFI_ERROR (Status);
  return EFI_SUCCESS;
}

EFI_PEI_NOTIFY_DESCRIPTOR mSiliconInitializedNotifyList = {
  (EFI_PEI_PPI_DESCRIPTOR_NOTIFY_CALLBACK | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST),
  &gEdkiiSiliconInitializedPpiGuid,
  (EFI_PEIM_NOTIFY_ENTRY_POINT) SiliconInitializedPpiNotifyCallback
};

/**
  Platform VTd Info sample driver.

  @param[in] FileHandle  Handle of the file being invoked.
  @param[in] PeiServices Describes the list of possible PEI Services.

  @retval EFI_SUCCESS if it completed successfully.
**/
EFI_STATUS
EFIAPI
PlatformVTdInfoSampleInitialize (
  IN       EFI_PEI_FILE_HANDLE  FileHandle,
  IN CONST EFI_PEI_SERVICES     **PeiServices
  )
{
  EFI_STATUS               Status;
  BOOLEAN                  SiliconInitialized;
  VOID                     *SiliconInitializedPpi;
  EFI_PEI_PPI_DESCRIPTOR   *PpiDesc;

  SiliconInitialized = FALSE;
  //
  // Check if silicon is initialized.
  //
  Status = PeiServicesLocatePpi (
             &gEdkiiSiliconInitializedPpiGuid,
             0,
             NULL,
             &SiliconInitializedPpi
             );
  if (!EFI_ERROR(Status)) {
    SiliconInitialized = TRUE;
  }
  DEBUG ((DEBUG_INFO, "SiliconInitialized - %x\n", SiliconInitialized));
  if (!SiliconInitialized) {
    Status = PeiServicesNotifyPpi (&mSiliconInitializedNotifyList);
    InitDmar ();

    Status = PeiServicesInstallPpi (&mPlatformVTdNoIgdInfoSampleDesc);
    ASSERT_EFI_ERROR (Status);
  } else {
    PpiDesc = PatchDmar ();

    Status = PeiServicesInstallPpi (PpiDesc);
    ASSERT_EFI_ERROR (Status);
  }

  return Status;
}