// Part of the Crubit project, under the Apache License v2.0 with LLVM
// Exceptions. See /LICENSE for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#ifndef CRUBIT_LIFETIME_ANNOTATIONS_LIFETIME_SYMBOL_TABLE_H_
#define CRUBIT_LIFETIME_ANNOTATIONS_LIFETIME_SYMBOL_TABLE_H_

#include <optional>
#include <string>

#include "lifetime_annotations/lifetime.h"
#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/DenseMap.h"
#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/StringMap.h"
#include "third_party/llvm/llvm-project/llvm/include/llvm/ADT/StringRef.h"

namespace devtools_rust {

// One-to-one mapping between lifetime names and the corresponding lifetimes.
class LifetimeSymbolTable {
 public:
  // Looks up a lifetime name in the symbol table.
  // Returns the corresponding lifetime if the name was present in the symbol
  // table, or nullopt if the lifetime name wasn't found.
  // If `name` is "static", returns `Lifetime::Static()`.
  std::optional<Lifetime> LookupName(llvm::StringRef name) const;

  // Looks up a lifetime name in the symbol table and inserts a new variable
  // lifetime in the table if the name was not found.
  // If `name` is "static", returns `Lifetime::Static()`.
  Lifetime LookupNameAndMaybeDeclare(llvm::StringRef name);

  // Looks up a lifetime in the symbol table.
  // Returns the corresponding name if the lifetime was present in the symbol
  // table, or nullopt if the lifetime wasn't found.
  // If `lifetime` is `Lifetime::Static()`, returns "static".
  std::optional<llvm::StringRef> LookupLifetime(Lifetime lifetime) const;

  // Looks up a lifetime in the symbol table and inserts a new autogenerated
  // name for it if the lifetime was not found.
  // If `lifetime` is `Lifetime::Static()`, returns "static".
  // The autogenerated name will be the first name from the sequence
  // "a", ..., "z", "aa", "ab", "ac", ... that is not yet contained in the
  // symbol table.
  llvm::StringRef LookupLifetimeAndMaybeDeclare(Lifetime lifetime);

  // Associates a lifetime to a given name.
  // The name must not be already used.
  // Calling this function with a lifetime that is already used will cause calls
  // to LookupLifetime() to have undefined behaviour.
  void Add(llvm::StringRef name, Lifetime lifetime);

  // Accessor for hashing/debugging purposes.
  const llvm::StringMap<Lifetime>& GetMapping() const {
    return name_to_lifetime_;
  }

 private:
  llvm::StringMap<Lifetime> name_to_lifetime_;
  llvm::DenseMap<Lifetime, std::string> lifetime_to_name_;
  int next_name_index_ = 0;
};

}  // namespace devtools_rust

#endif  // CRUBIT_LIFETIME_ANNOTATIONS_LIFETIME_SYMBOL_TABLE_H_