SPIRV-Tools/tools/util/flags.cpp

// Copyright (c) 2023 Google LLC.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "flags.h"

#include <cstring>
#include <iostream>
#include <regex>
#include <string>
#include <unordered_set>
#include <variant>
#include <vector>

namespace flags {

std::vector<std::string> positional_arguments;

namespace {

using token_t = const char*;
using token_iterator_t = token_t*;

// Extracts the flag name from a potential token.
// This function only looks for a '=', to split the flag name from the value for
// long-form flags. Returns the name of the flag, prefixed with the hyphen(s).
inline std::string get_flag_name(const std::string& flag, bool is_short_flag) {
  if (is_short_flag) {
    return flag;
  }

  size_t equal_index = flag.find('=');
  if (equal_index == std::string::npos) {
    return flag;
  }
  return flag.substr(0, equal_index);
}

// Parse a boolean flag. Returns `true` if the parsing succeeded, `false`
// otherwise.
bool parse_flag(Flag<bool>& flag, bool is_short_flag,
                const std::string& token) {
  if (is_short_flag) {
    flag.value() = true;
    return true;
  }

  const std::string raw_flag(token);
  size_t equal_index = raw_flag.find('=');
  if (equal_index == std::string::npos) {
    flag.value() = true;
    return true;
  }

  const std::string value = raw_flag.substr(equal_index + 1);
  if (value == "true") {
    flag.value() = true;
    return true;
  }

  if (value == "false") {
    flag.value() = false;
    return true;
  }

  return false;
}

// Parse a string flag. Moved the iterator to the last flag's token if it's a
// multi-token flag. Returns `true` if the parsing succeeded.
// The iterator is moved to the last parsed token.
bool parse_flag(Flag<std::string>& flag, bool is_short_flag,
                token_iterator_t* iterator) {
  const std::string raw_flag(**iterator);
  const size_t equal_index = raw_flag.find('=');
  if (is_short_flag || equal_index == std::string::npos) {
    if ((*iterator)[1] == nullptr) {
      return false;
    }

    // This is a bi-token flag. Moving iterator to the last parsed token.
    flag.value() = (*iterator)[1];
    *iterator += 1;
    return true;
  }

  // This is a mono-token flag, no need to move the iterator.
  const std::string value = raw_flag.substr(equal_index + 1);
  flag.value() = value;
  return true;
}
}  // namespace

// This is the function to expand if you want to support a new type.
bool FlagList::parse_flag_info(FlagInfo& info, token_iterator_t* iterator) {
  bool success = false;

  std::visit(
      [&](auto&& item) {
        using T = std::decay_t<decltype(item.get())>;
        if constexpr (std::is_same_v<T, Flag<bool>>) {
          success = parse_flag(item.get(), info.is_short, **iterator);
        } else if constexpr (std::is_same_v<T, Flag<std::string>>) {
          success = parse_flag(item.get(), info.is_short, iterator);
        } else {
          static_assert(always_false_v<T>, "Unsupported flag type.");
        }
      },
      info.flag);

  return success;
}

bool FlagList::parse(token_t* argv) {
  flags::positional_arguments.clear();
  std::unordered_set<const FlagInfo*> parsed_flags;

  bool ignore_flags = false;
  for (const char** it = argv + 1; *it != nullptr; it++) {
    if (ignore_flags) {
      flags::positional_arguments.emplace_back(*it);
      continue;
    }

    // '--' alone is used to mark the end of the flags.
    if (std::strcmp(*it, "--") == 0) {
      ignore_flags = true;
      continue;
    }

    // '-' alone is not a flag, but often used to say 'stdin'.
    if (std::strcmp(*it, "-") == 0) {
      flags::positional_arguments.emplace_back(*it);
      continue;
    }

    const std::string raw_flag(*it);
    if (raw_flag.size() == 0) {
      continue;
    }

    if (raw_flag[0] != '-') {
      flags::positional_arguments.emplace_back(*it);
      continue;
    }

    // Only case left: flags (long and shorts).
    if (raw_flag.size() < 2) {
      std::cerr << "Unknown flag " << raw_flag << std::endl;
      return false;
    }
    const bool is_short_flag = std::strncmp(*it, "--", 2) != 0;
    const std::string flag_name = get_flag_name(raw_flag, is_short_flag);

    auto needle = std::find_if(
        get_flags().begin(), get_flags().end(),
        [&flag_name](const auto& item) { return item.name == flag_name; });
    if (needle == get_flags().end()) {
      std::cerr << "Unknown flag " << flag_name << std::endl;
      return false;
    }

    if (parsed_flags.count(&*needle) != 0) {
      std::cerr << "The flag " << flag_name << " was specified multiple times."
                << std::endl;
      return false;
    }
    parsed_flags.insert(&*needle);

    if (!parse_flag_info(*needle, &it)) {
      std::cerr << "Invalid usage for flag " << flag_name << std::endl;
      return false;
    }
  }

  // Check that we parsed all required flags.
  for (const auto& flag : get_flags()) {
    if (!flag.required) {
      continue;
    }

    if (parsed_flags.count(&flag) == 0) {
      std::cerr << "Missing required flag " << flag.name << std::endl;
      return false;
    }
  }

  return true;
}

// Just the public wrapper around the parse function.
bool Parse(const char** argv) { return FlagList::parse(argv); }

}  // namespace flags