//===- LocationDetail.h - MLIR Location storage details ---------*- C++ -*-===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This holds implementation details of the location attributes.
//
//===----------------------------------------------------------------------===//
#ifndef MLIR_IR_LOCATIONDETAIL_H_
#define MLIR_IR_LOCATIONDETAIL_H_

#include "mlir/IR/Attributes.h"
#include "mlir/IR/Identifier.h"
#include "mlir/IR/Location.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/TrailingObjects.h"

namespace mlir {

namespace detail {

struct CallSiteLocationStorage : public AttributeStorage {
  CallSiteLocationStorage(Location callee, Location caller)
      : callee(callee), caller(caller) {}

  /// The hash key used for uniquing.
  using KeyTy = std::pair<Location, Location>;
  bool operator==(const KeyTy &key) const {
    return key == KeyTy(callee, caller);
  }

  /// Construct a new storage instance.
  static CallSiteLocationStorage *
  construct(AttributeStorageAllocator &allocator, const KeyTy &key) {
    return new (allocator.allocate<CallSiteLocationStorage>())
        CallSiteLocationStorage(key.first, key.second);
  }

  Location callee, caller;
};

struct FileLineColLocationStorage : public AttributeStorage {
  FileLineColLocationStorage(Identifier filename, unsigned line,
                             unsigned column)
      : filename(filename), line(line), column(column) {}

  /// The hash key used for uniquing.
  using KeyTy = std::tuple<Identifier, unsigned, unsigned>;
  bool operator==(const KeyTy &key) const {
    return key == KeyTy(filename, line, column);
  }

  /// Construct a new storage instance.
  static FileLineColLocationStorage *
  construct(AttributeStorageAllocator &allocator, const KeyTy &key) {
    return new (allocator.allocate<FileLineColLocationStorage>())
        FileLineColLocationStorage(std::get<0>(key), std::get<1>(key),
                                   std::get<2>(key));
  }

  Identifier filename;
  unsigned line, column;
};

struct FusedLocationStorage final
    : public AttributeStorage,
      public llvm::TrailingObjects<FusedLocationStorage, Location> {
  FusedLocationStorage(unsigned numLocs, Attribute metadata)
      : numLocs(numLocs), metadata(metadata) {}

  ArrayRef<Location> getLocations() const {
    return ArrayRef<Location>(getTrailingObjects<Location>(), numLocs);
  }

  /// The hash key used for uniquing.
  using KeyTy = std::pair<ArrayRef<Location>, Attribute>;
  bool operator==(const KeyTy &key) const {
    return key == KeyTy(getLocations(), metadata);
  }

  /// Construct a new storage instance.
  static FusedLocationStorage *construct(AttributeStorageAllocator &allocator,
                                         const KeyTy &key) {
    ArrayRef<Location> locs = key.first;

    auto byteSize = totalSizeToAlloc<Location>(locs.size());
    auto rawMem = allocator.allocate(byteSize, alignof(FusedLocationStorage));
    auto result = new (rawMem) FusedLocationStorage(locs.size(), key.second);

    std::uninitialized_copy(locs.begin(), locs.end(),
                            result->getTrailingObjects<Location>());
    return result;
  }

  // This stuff is used by the TrailingObjects template.
  friend llvm::TrailingObjects<FusedLocationStorage, Location>;
  size_t numTrailingObjects(OverloadToken<Location>) const { return numLocs; }

  /// Number of trailing location objects.
  unsigned numLocs;

  /// Metadata used to reason about the generation of this fused location.
  Attribute metadata;
};

struct NameLocationStorage : public AttributeStorage {
  NameLocationStorage(Identifier name, Location child)
      : name(name), child(child) {}

  /// The hash key used for uniquing.
  using KeyTy = std::pair<Identifier, Location>;
  bool operator==(const KeyTy &key) const { return key == KeyTy(name, child); }

  /// Construct a new storage instance.
  static NameLocationStorage *construct(AttributeStorageAllocator &allocator,
                                        const KeyTy &key) {
    return new (allocator.allocate<NameLocationStorage>())
        NameLocationStorage(key.first, key.second);
  }

  Identifier name;
  Location child;
};

struct OpaqueLocationStorage : public AttributeStorage {
  OpaqueLocationStorage(uintptr_t underlyingLocation, TypeID typeID,
                        Location fallbackLocation)
      : underlyingLocation(underlyingLocation), typeID(typeID),
        fallbackLocation(fallbackLocation) {}

  /// The hash key used for uniquing.
  using KeyTy = std::tuple<uintptr_t, TypeID, Location>;
  bool operator==(const KeyTy &key) const {
    return key == KeyTy(underlyingLocation, typeID, fallbackLocation);
  }

  /// Construct a new storage instance.
  static OpaqueLocationStorage *construct(AttributeStorageAllocator &allocator,
                                          const KeyTy &key) {
    return new (allocator.allocate<OpaqueLocationStorage>())
        OpaqueLocationStorage(std::get<0>(key), std::get<1>(key),
                              std::get<2>(key));
  }

  /// Pointer to the corresponding object.
  uintptr_t underlyingLocation;

  /// A unique pointer for each type of underlyingLocation.
  TypeID typeID;

  /// An additional location that can be used if the external one is not
  /// suitable.
  Location fallbackLocation;
};

} // end namespace detail
} // end namespace mlir

#endif // MLIR_IR_LOCATIONDETAIL_H_