#ifndef LLVM_ANALYSIS_TBAATYPERECONSTRUCTION_H
#define LLVM_ANALYSIS_TBAATYPERECONSTRUCTION_H

#include "llvm/ADT/StringMap.h"
#include "llvm/IR/PassManager.h"
#include "llvm/Analysis/ClassHierarchyAnalysis.h"
#include "llvm/Pass.h"
#include <set>

/**
 * Structure of TBAA
 * Type 1 -> 4 node layout
 * {struct layout (0), type node (1), offset (2), size (3)}
 *
 * Type 2 -> 3 node layout
 * {tbaa identifier (0), size (1), type name (2)}
 *
 * Type 3 -> used to represent struct layout
 * {type 2 node (0), struct total size (1), type name (2), {type node (0),
 * offset (1), size (2)} (3), ... subnode for every element of the structure}
 *
 * Notes
 * * Type 2 is used to represent builtin datatypes
 * * In case of Type 1, when the variable is not part of a struct, (1) and (2)
 *   are same nodes.
 * * In all cases, the first node of tbaa is the parent node, so (0) in Type
 * 1/2/3
 */
namespace llvm {

class TBAATypeReconstruction;
class TBAATypeReconstruction {
public:
  TBAATypeReconstruction(Module &M, ClassHierarchyAnalysis *CHA) : M(M), CHA(CHA) {}
  class TypeInfo;
  Module &M;
  ClassHierarchyAnalysis *CHA = nullptr;

  // return whether the given value is expected to have TBAA Type Info
  bool hasTypeInfo(Value* V);
  bool hasTypeInfo(Value& V);

  TBAATypeReconstruction::TypeInfo* getPointerTypeFor(TBAATypeReconstruction::TypeInfo *TBAAType);
  
  /// return the TypeInfo object generated for the given instruction.
  /// Since only load and store instructions are tagged with TBAA MD Nodes,
  /// this function returns nullptr incase the instruction passed isn't load
  /// or store. This also return nullptr, if it cannot find an analysed TypeInfo
  /// node for the given instruction.
  TBAATypeReconstruction::TypeInfo *getTypeInfoForInstr(Instruction *I);
  TBAATypeReconstruction::TypeInfo *getTypeInfoExtended(Value* V);
  
  //update the type info from TBAA Metadata.
  void updateTypeInfoForInstr(Instruction &I);
  /// return the TypeInfo object with the given Type Name. If it cannot find
  /// such an object it returns nullptr. Format of the Type Name is as follows
  /// struct/class/array/pointer@<MangledTypeName>
  TBAATypeReconstruction::TypeInfo *getTypeInfoFromTypeName(StringRef TypeName);
  /// return the TypeInfo object with the given Raw Type Name. If it cannot find
  /// such an object it returns nullptr. Format of the Raw Type Name is as
  /// follows <MangledTypeName>
  TBAATypeReconstruction::TypeInfo *
  getTypeInfoFromRawTypeName(StringRef RawtypeName);
  // return a SmallVector of instructions with the given Access TypeInfo node.
  // If it cannot find instructions with given TypeInfo node, it returns a
  // std::nullopt. It returns the SmallVector by reference, as thus has been
  // wrapped in a reference_wrapper.
  //
  // A short example of the usage:
  //
  // auto instrVector = TypeRecon.getInstructionsOfTypeInfo(*typeInfo);
  // if (instrVector.has_value())
  //   dbgs() << instrVector.value().get().size() << "\n";
  //
  // We need to use value() to get the underlying reference_wrapper and then
  // get() to get the underlying reference object (SmallVector<Instruction*>).
  std::optional<std::reference_wrapper<SmallVector<Instruction *>>>
  getInstructionsOfAccessTypeInfo(
      TBAATypeReconstruction::TypeInfo &InstrAccessType);

  
  // return a SmallVector of instructions with the given Base TypeInfo node. If
  // it cannot find instructions with given TypeInfo node, it returns a
  // std::nullopt. It returns the SmallVector by reference, as thus has been
  // wrapped in a reference_wrapper.
  //
  // A short example of the usage:
  //
  // auto instrVector =
  // TypeRecon.getInstructionsOfTypeInfo(typeInfo->getBaseStructLayout()); if
  // (instrVector.has_value())
  //   dbgs() << instrVector.value().get().size() << "\n";
  //
  // We need to use value() to get the underlying reference_wrapper and then
  // get() to get the underlying reference object (SmallVector<Instruction*>).
  std::optional<std::reference_wrapper<SmallVector<Instruction *>>>
  getInstructionsOfBaseTypeInfo(
      TBAATypeReconstruction::TypeInfo &InstrBaseType);
  void analyze();
  bool invalidate(Module &M, const PreservedAnalyses &PA,
                  ModuleAnalysisManager::Invalidator &Inv);
  void getDerivedPointers(TypeInfo *TI, std::set<StringRef> &DNames);

private:
  // stores the global struct types defined in the module
  StringMap<StructType *> GlobalStructTypes;
  // StringMap between Type Name and the TypeInfo object
  StringMap<std::shared_ptr<TypeInfo>> NameTypeInfoMap;
  // StringMap between Raw Type Name and the TypeInfo object
  StringMap<std::shared_ptr<TypeInfo>> RawNameTypeInfoMap;
  // DenseMap between valid instructions and their TypeInfo object
  DenseMap<Instruction *, std::shared_ptr<TypeInfo>>
      InstructionTypeInfoMap;
  std::unordered_map<const MDNode *,
                     std::shared_ptr<TBAATypeReconstruction::TypeInfo>>
      TypeInfoMap;
  std::set<std::shared_ptr<TBAATypeReconstruction::TypeInfo>> TypeInfoFields;
  // StringMap between Access Type Name and a SmallVector of instructions
  // associated with this Type Name (or TypeInfo node).
  StringMap<SmallVector<Instruction *>> AccessTypeInfoInstructionsMap;
  // StringMap between Base Type Name and a SmallVector of instructions
  // associated with this Type Name (or TypeInfo node).
  StringMap<SmallVector<Instruction *>> BaseTypeInfoInstructionsMap;
  /// Internal function which runs after analyzeTypeInfo, and finds out the
  /// type of pointee, and searches for the pointee type in the analyzed
  /// TypeInfo nodes. If it finds it, it will link it to the PointerType member.
  /// Please note that in the mangled string we need to remove a P after _ZTS to
  /// get the pointee type, removePointerFromMangledName static function can be
  /// used for the same. Please note that this will only work for C++ programs
  /// as in C, mangled names are not generated for TBAA Type Nodes.
  void findPointerTypes();
  /// Internal function which runs after analyzeTypeInfo and findPointerTypes,
  /// this function should only be called at the end, as it populates the
  /// InstructionTypeInfoMap map with the correct mapping of Instruction ->
  /// TypeInfo.
  void finalizeResultMap();
  /// Internal function which goes over all the instructions in the passed
  /// function and builds the TypeInfo for the instructions and adds them to the
  /// TypeInfoMap map. It is a mapping between MDNode -> TypeInfo.
  void analyzeTypeInfo(Function &F);

  /// This function processes Type 1 nodes and returns the constructed TypeInfo
  /// node. Type 1 nodes are root nodes that are attached to an instruction. The
  /// format is as follows -> Type 1 has 4 nodes.
  ///
  /// {struct layout (0), type node (1), offset (2), size (3)}
  /// for a non-struct node, element 0 and 1 are the same. Below is an example
  /// snippet. !9 = !{!10, !11, i64 0, i64 4}
  std::shared_ptr<TBAATypeReconstruction::TypeInfo>
  processType1Node(const MDNode *MD);
  /// This function processes Type 2 nodes and returns the constructed TypeInfo
  /// node. Type 2 nodes are access nodes that represent the type of the access.
  /// The format is as follows -> Type 2 has 3 nodes.
  ///
  /// {tbaa identifier (0), size (1), type name (2)}
  /// !11 = !{!7, i64 4, !"int"}
  std::shared_ptr<TBAATypeReconstruction::TypeInfo>
  processType2Node(const MDNode *MD);
  /// This function processes Type 3 nodes and return the constructed TypeInfo
  /// node. Type 3 nodes are base access nodes that represent the type of the
  /// base type. In a simple way, it represents the type and layout of the
  /// underlying struct. The format is as follows -> Type 3 has 3 or more nodes,
  /// and a multiple of 3.
  ///
  /// {type 2 node (0), struct total size (1), type name (2), {type node (0),
  /// offset (1), size (2)} (3), ... subnode for every element of the structure}
  /// !10 = !{!7, i64 8, !"struct@_ZTS2TT", !11, i64 0, i64 4, !12, i64 4, i64
  /// 4}
  ///
  /// In case the node represents a pointer or an array, the element 0 will not
  /// point to the TBAA identifier node, but rather to the Pointee or Array
  /// Element type. In case of pointers, only scalar pointee types are handled.
  /// RecordType is too buggy to support at this moment.
  std::shared_ptr<TBAATypeReconstruction::TypeInfo>
  processType3Node(const MDNode *MD);
  /// This function calls the above three function based on the type of the
  /// passed root MD node.
  std::shared_ptr<TBAATypeReconstruction::TypeInfo>
  processTBAANodeTree(const MDNode *Root);
  /// The front end change associated with this analysis tries to preserve types
  /// after type merging that happens during LTO. For every llvm::StructType
  /// defined globally, it will emit a global variable with a specific name and
  /// of the the said llvm::StructType. An example of this is noted below.
  ///
  /// %struct.TT = type { i32, float }
  /// @"struct#_ZTS2TT" = weak_odr constant %struct.TT zeroinitializer
  ///
  /// The global variable name is defined as follows class/struct#<mangled type
  /// name> and the type of the global variable is set to the associated
  /// llvm::StructType. This functionality was added in commit
  /// (a34251d31a232b66d5edab372c13bc6b867f7478). Please refer to the commit
  /// message for more details.
  ///
  /// Coming to the usage of this function, the said GV helps us to map a type 3
  /// node to its llvm::StructType defined globally in the module. We get the
  /// mangled name from the TypeInfo node and search for a GV with the name
  /// <class/struct> + # + <mangled_type_name>, and once found look for the type
  /// of the GV.
  StructType *matchStructTypes(
      std::shared_ptr<TBAATypeReconstruction::TypeInfo> ConstructedStruct);
};

class TBAATypeReconstruction::TypeInfo {
private:
  /// Type Name of a TypeInfo object is as follows:
  /// class/struct/array/pointer@<mangled_type_name>
  /// So, raw Type Name is just the mangled type name
  /// without the type qualifiers.
  StringRef RawTypeName;
  /// Type Name of a TypeInfo object is as follows:
  /// class/struct/array/pointer@<mangled_type_name>
  StringRef TypeName;
  /// Memory offset of the member in the class
  unsigned int Offset;
  /// Size of the type
  unsigned int Size;
  /// If the TypeInfo object represents a pointer type
  /// the name will be as follows pointer@<mangled_type_name>
  /// This variable points to the TypeInfo object of the pointee
  /// type. This is not guaranteed to be filled. In case of
  /// scalar types we can say that this must be filled.
  std::shared_ptr<TypeInfo> PointerType;
  /// If the TypeInfo object represents a array type
  /// the name will be as follows array@<mangled_type_name>
  /// This variable points to the TypeInfo object of the array
  /// element type.
  std::shared_ptr<TypeInfo> ArrayType;
  /// This points to the parent struct layout node this TypeInfo
  /// object is a member of.
  std::shared_ptr<TypeInfo> StructLayout;
  /// This points to the llvm::StructType associated with this
  /// struct layout node.
  StructType *SType;
  /// Enumerates the member elements of the struct node.
  std::vector<std::shared_ptr<TypeInfo>> StructElements;
  /// Enumerates the memory offsets of the member elements.
  std::vector<unsigned int> StructOffsets;
  /// Enumerates the sizes of the member elements.
  std::vector<unsigned int> StructSizes;

public:
  friend TBAATypeReconstruction;
  /// Type Name of a TypeInfo object is as follows:
  /// class/struct/array/pointer@<mangled_type_name>
  /// So, raw Type Name is just the mangled type name
  /// without the type qualifiers.
  ///
  /// Returns the raw type name of the TypeInfo node.
  StringRef getRawTypeName() { return RawTypeName; }
  /// Type Name of a TypeInfo object is as follows:
  /// class/struct/array/pointer@<mangled_type_name>
  ///
  /// Returns the type name of the TypeInfo node.
  StringRef getTypeName() { return TypeName; }
  /// Returns the memory offset of the given type.
  unsigned int getOffset() { return Offset; }
  /// Returns the size of the given type.
  unsigned int getSize() { return Size; }

  /// Returns true if the Type is an array type, or
  /// else false.
  bool isArrayType() { return TypeName.starts_with("array@"); }
  /// Returns true if the Type is a pointer type, or
  /// else false.
  bool isPointerType() { return TypeName.starts_with("pointer@"); }
  /// Returns true if the Type is struct or class type, or
  /// else false.
  bool isStructOrClassType() {
    return TypeName.starts_with("struct@") || TypeName.starts_with("class@");
  }
  /// Returns true if the Type is a struct type or else false.
  bool isStructType() { return TypeName.starts_with("struct@"); }
  /// Returns true if the Type is a class type or else false.
  bool isClassType() { return TypeName.starts_with("class@"); }
  /// Returns true if this TypeInfo node is a member of a struct node,
  /// and has a valid associated struct layout node.
  bool isStructMember() { return (StructLayout ? true : false); }

  /// Returns the array element type, if it's not there it will
  /// just return nullptr.
  std::shared_ptr<TypeInfo> getArrayElementType() { return ArrayType; }
  /// Returns the pointee type, if it's not there it will
  /// just return nullptr.
  std::shared_ptr<TypeInfo> getPointeeType() { return PointerType; }
  /// Returns the llvm::StructType associated with the struct layout node
  StructType *getStructType() { return SType; }

  /// Returns the Base Struct layout of the TypeInfo node. This is only
  /// valid in case that the node is a member of a struct node.
  std::shared_ptr<TypeInfo> getBaseStructLayout() { return StructLayout; }
  /// Returns the struct layout members as a vector.
  std::vector<std::shared_ptr<TypeInfo>> &getStructMembers() {
    return StructElements;
  }
  /// Returns the memory offsets of struct layout members as a vector.
  std::vector<unsigned int> &getStructMemberOffsets() { return StructOffsets; }
  /// Returns the size of struct layout members as a vector.
  std::vector<unsigned int> &getStructMemberSizes() { return StructSizes; }

  void dump();

  bool isSameAs(TypeInfo &Other, bool strict=true);
  bool isSameAs(TypeInfo *Other, bool strict=true);
};

class TBAATypeReconstructionAnalysis
    : public AnalysisInfoMixin<TBAATypeReconstructionAnalysis> {
  friend AnalysisInfoMixin<TBAATypeReconstructionAnalysis>;
  static AnalysisKey Key;

public:
  typedef TBAATypeReconstruction Result;
  Result run(Module &M, ModuleAnalysisManager &AM);
};

class TBAATypeReconstructionAnalysisPrinterPass
    : public PassInfoMixin<TBAATypeReconstructionAnalysisPrinterPass> {
  raw_ostream &OS;

public:
  explicit TBAATypeReconstructionAnalysisPrinterPass(raw_ostream &OS)
      : OS(OS) {}
  explicit TBAATypeReconstructionAnalysisPrinterPass() : OS(dbgs()) {}
  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
};

class TBAATypeReconstructionAnalysisWrapperPass
    : public PassInfoMixin<TBAATypeReconstructionAnalysisWrapperPass> {
  raw_ostream &OS;

public:
  explicit TBAATypeReconstructionAnalysisWrapperPass(raw_ostream &OS)
      : OS(OS) {}
  explicit TBAATypeReconstructionAnalysisWrapperPass() : OS(dbgs()) {}
  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
};
} // namespace llvm
#endif
