//===-- InstrumentorConfigFile.cpp ----------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// The implementation of the utilities for the Instrumentor JSON configuration
// file.
//
//===----------------------------------------------------------------------===//

#include "llvm/Transforms/IPO/Instrumentor.h"

#include "llvm/ADT/StringMap.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/StringSet.h"
#include "llvm/IR/DiagnosticInfo.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/JSON.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/Path.h"
#include "llvm/Support/StringSaver.h"
#include "llvm/Support/VirtualFileSystem.h"

#include <string>

using namespace llvm;

static Expected<std::unique_ptr<MemoryBuffer>>
setupMemoryBuffer(const Twine &Filename, vfs::FileSystem &FS) {
  auto BufferOrErr = Filename.str() == "-" ? MemoryBuffer::getSTDIN()
                                           : FS.getBufferForFile(Filename);
  if (std::error_code EC = BufferOrErr.getError())
    return errorCodeToError(EC);
  return std::move(BufferOrErr.get());
}

namespace llvm {
namespace instrumentor {

void writeConfigToJSON(InstrumentationConfig &IConf, StringRef OutputFile,
                       LLVMContext &Ctx) {
  if (OutputFile.empty())
    return;

  std::error_code EC;
  raw_fd_stream OS(OutputFile, EC);
  if (EC) {
    Ctx.diagnose(DiagnosticInfoInstrumentation(
        Twine("failed to open instrumentor configuration file for writing: ") +
            EC.message(),
        DS_Warning));
    return;
  }

  json::OStream J(OS, 2);
  J.objectBegin();

  J.attributeBegin("configuration");
  J.objectBegin();
  for (auto *BaseCO : IConf.BaseConfigurationOptions) {
    switch (BaseCO->Kind) {
    case BaseConfigurationOption::STRING:
      J.attribute(BaseCO->Name, BaseCO->getString());
      break;
    case BaseConfigurationOption::BOOLEAN:
      J.attribute(BaseCO->Name, BaseCO->getBool());
      break;
    }
    if (!BaseCO->Description.empty())
      J.attribute(std::string(BaseCO->Name) + ".description",
                  BaseCO->Description);
  }
  J.objectEnd();
  J.attributeEnd();

  for (unsigned KindVal = 0; KindVal <= InstrumentationLocation::Last;
       ++KindVal) {
    auto Kind = InstrumentationLocation::KindTy(KindVal);

    auto &KindChoices = IConf.IChoices[Kind];
    if (KindChoices.empty())
      continue;

    J.attributeBegin(InstrumentationLocation::getKindStr(Kind));
    J.objectBegin();
    for (auto &[Name, Choice] : KindChoices) {
      J.attributeBegin(Name);
      J.objectBegin();
      J.attribute("enabled", Choice->Enabled);
      J.attribute("filter", Choice->Filter);
      J.attribute("filter.description",
                  "Static property filter to exclude instrumentation.");
      for (auto &ArgIt : Choice->IRTArgs) {
        J.attribute(ArgIt.Name, ArgIt.Enabled);
        if ((ArgIt.Flags & IRTArg::REPLACABLE) ||
            (ArgIt.Flags & IRTArg::REPLACABLE_CUSTOM))
          J.attribute(std::string(ArgIt.Name) + ".replace", true);
        if (!ArgIt.Description.empty())
          J.attribute(std::string(ArgIt.Name) + ".description",
                      ArgIt.Description);
      }
      J.objectEnd();
      J.attributeEnd();
    }
    J.objectEnd();
    J.attributeEnd();
  }

  J.objectEnd();
}

template <typename Map>
static StringRef closestOption(const Map &Options, StringRef Missing) {
  uint32_t MaxEdit = 5;
  StringRef Closest;
  for (const auto &Key : Options.keys()) {
    auto EditDist = Missing.edit_distance_insensitive(Key, true, MaxEdit);
    if (EditDist < MaxEdit) {
      Closest = Key;
      MaxEdit = EditDist;
    }
  }
  return Closest;
}

bool readConfigFromJSON(InstrumentationConfig &IConf, StringRef InputFile,
                        LLVMContext &Ctx, vfs::FileSystem &FS) {
  if (InputFile.empty())
    return true;

  auto BufferOrErr = setupMemoryBuffer(InputFile, FS);
  if (Error E = BufferOrErr.takeError()) {
    Ctx.diagnose(DiagnosticInfoInstrumentation(
        Twine("failed to open instrumentor configuration file for reading: ") +
            toString(std::move(E)),
        DS_Warning));
    return false;
  }
  auto Buffer = std::move(BufferOrErr.get());
  json::Path::Root NullRoot;
  auto Parsed = json::parse(Buffer->getBuffer());
  if (!Parsed) {
    Ctx.diagnose(DiagnosticInfoInstrumentation(
        Twine("failed to parse instrumentor configuration file: ") +
            toString(Parsed.takeError()),
        DS_Warning));
    return false;
  }
  auto *Config = Parsed->getAsObject();
  if (!Config) {
    Ctx.diagnose(DiagnosticInfoInstrumentation(
        "failed to parse instrumentor configuration file, expected an object "
        "'{ ... }'",
        DS_Warning));
    return false;
  }

  StringMap<BaseConfigurationOption *> BCOMap;
  for (auto *BO : IConf.BaseConfigurationOptions)
    BCOMap[BO->Name] = BO;

  SmallPtrSet<InstrumentationOpportunity *, 32> SeenIOs;
  for (auto &It : *Config) {
    auto *Obj = It.second.getAsObject();
    if (!Obj) {
      Ctx.diagnose(DiagnosticInfoInstrumentation(
          "malformed JSON configuration, expected an object", DS_Warning));
      continue;
    }
    if (It.first == "configuration") {
      for (auto &ObjIt : *Obj) {
        if (auto *BO = BCOMap.lookup(ObjIt.first)) {
          switch (BO->Kind) {
          case BaseConfigurationOption::STRING:
            if (auto V = ObjIt.second.getAsString()) {
              BO->setString(IConf.SS.save(*V));
            } else {
              Ctx.diagnose(DiagnosticInfoInstrumentation(
                  Twine("configuration key '") + StringRef(ObjIt.first) +
                      Twine("' expects a string, value ignored"),
                  DS_Warning));
            }
            break;
          case BaseConfigurationOption::BOOLEAN:
            if (auto V = ObjIt.second.getAsBoolean())
              BO->setBool(*V);
            else {
              Ctx.diagnose(DiagnosticInfoInstrumentation(
                  Twine("configuration key '") + StringRef(ObjIt.first) +
                      Twine("' expects a boolean, value ignored"),
                  DS_Warning));
            }
            break;
          }
        } else if (!StringRef(ObjIt.first).ends_with(".description")) {
          std::string Diag = "configuration key '" + ObjIt.first.str() +
                             "' not found and ignored";
          StringRef Closest = closestOption(BCOMap, ObjIt.first);
          if (!Closest.empty())
            Diag += "; did you mean '" + Closest.str() + "'?";
          Ctx.diagnose(DiagnosticInfoInstrumentation(Diag, DS_Warning));
        }
      }
      continue;
    }

    auto &IChoiceMap =
        IConf.IChoices[InstrumentationLocation::getKindFromStr(It.first)];
    for (auto &ObjIt : *Obj) {
      auto *InnerObj = ObjIt.second.getAsObject();
      if (!InnerObj) {
        Ctx.diagnose(DiagnosticInfoInstrumentation(
            "malformed JSON configuration, expected an object", DS_Warning));
        continue;
      }
      auto *IO = IChoiceMap.lookup(ObjIt.first);
      if (!IO) {
        std::string Diag =
            "malformed JSON configuration, expected an object matching "
            "an instrumentor choice, got '" +
            ObjIt.first.str() + "'";
        StringRef Closest = closestOption(IChoiceMap, ObjIt.first);
        if (!Closest.empty())
          Diag += "; did you mean '" + Closest.str() + "'?";
        Ctx.diagnose(DiagnosticInfoInstrumentation(Diag, DS_Warning));
        continue;
      }
      SeenIOs.insert(IO);
      StringMap<bool> ValueMap, ReplaceMap;
      StringRef FilterStr;
      StringSet<> IOOpts;
      IOOpts.insert("enabled");
      IOOpts.insert("filter");
      for (auto &IRArg : IO->IRTArgs)
        IOOpts.insert(IRArg.Name);
      for (auto &InnerObjIt : *InnerObj) {
        auto Name = StringRef(InnerObjIt.first);
        if (Name == "filter") {
          if (auto V = InnerObjIt.second.getAsString())
            FilterStr = IConf.SS.save(*V);
        } else if (Name.consume_back(".replace")) {
          ReplaceMap[Name] = InnerObjIt.second.getAsBoolean().value_or(false);
        } else {
          ValueMap[Name] = InnerObjIt.second.getAsBoolean().value_or(false);
        }
        if (!IOOpts.contains(Name)) {
          std::string Diag = "unrecognized JSON property '" + Name.str() +
                             "' in configuration for '" + IO->getName().str() +
                             "'";
          StringRef Closest = closestOption(IOOpts, Name);
          if (!Closest.empty())
            Diag += "; did you mean '" + Closest.str() + "'?";
          Ctx.diagnose(DiagnosticInfoInstrumentation(Diag, DS_Warning));
        }
      }
      IO->Enabled = ValueMap["enabled"];
      IO->Filter = FilterStr;
      for (auto &IRArg : IO->IRTArgs) {
        IRArg.Enabled = ValueMap[IRArg.Name];
        if (!ReplaceMap.lookup(IRArg.Name)) {
          IRArg.Flags &= ~IRTArg::REPLACABLE;
          IRArg.Flags &= ~IRTArg::REPLACABLE_CUSTOM;
        }
      }
    }
  }

  for (auto &IChoiceMap : IConf.IChoices)
    for (auto &It : IChoiceMap)
      if (!SeenIOs.count(It.second))
        It.second->Enabled = false;

  return true;
}

bool readConfigPathsFile(StringRef InputFile, cl::list<std::string> &Configs,
                         LLVMContext &Ctx, vfs::FileSystem &FS) {
  if (InputFile.empty())
    return true;

  auto BufferOrErr = setupMemoryBuffer(InputFile, FS);
  if (Error E = BufferOrErr.takeError()) {
    Ctx.diagnose(DiagnosticInfoInstrumentation(
        Twine("failed to open instrumentor configuration paths file for "
              "reading: ") +
            toString(std::move(E)),
        DS_Warning));
    return false;
  }

  StringRef InputFilePath(sys::path::parent_path(InputFile));

  auto Buffer = std::move(BufferOrErr.get());
  StringRef Content = Buffer->getBuffer();
  StringRef EOL = Content.detectEOL();
  do {
    auto [LHS, RHS] = Content.split(EOL);
    std::string ConfigPath = LHS.trim().str();
    if (!sys::path::is_absolute(ConfigPath)) {
      SmallString<128> InputFilePathStringVec(InputFilePath);
      sys::path::append(InputFilePathStringVec, ConfigPath);
      ConfigPath = InputFilePathStringVec.c_str();
    }
    Configs.push_back(ConfigPath);
    Content = RHS.trim();
  } while (!Content.empty());

  return true;
}

} // end namespace instrumentor
} // end namespace llvm