/* Copyright 2019 The OpenXLA Authors.

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.

// SPDX-License-Identifier: Apache-2.0

// Copyright (C) 2024 Intel Corporation

==============================================================================*/

#ifndef XLA_SERVICE_GRAPHCYCLES_ORDERED_SET_H_
#define XLA_SERVICE_GRAPHCYCLES_ORDERED_SET_H_

#include <vector>
#include <iostream>
#include <unordered_map>
#include <cassert>

// #include "absl/container/flat_hash_map.h"
// #include "absl/types/span.h" 

//#include "tsl/platform/logging.h"

namespace xla {
// This is a set data structure that provides a deterministic iteration order.
// The iteration order of elements only depends on the sequence of
// inserts/deletes, so as long as the inserts/deletes happen in the same
// sequence, the set will have the same iteration order.
//
// Assumes that T can be cheaply copied for simplicity.
template <typename T>
class OrderedSet {
 public:
  // Inserts `value` into the ordered set.  Returns true if the value was not
  // present in the set before the insertion.
  bool Insert(T value) {
    bool new_insertion =
      value_to_index_.insert({value, static_cast<int32_t>(value_sequence_.size())}).second;
    if (new_insertion) {
      value_sequence_.push_back(value);
    }
    return new_insertion;
  }

  // Removes `value` from the set.  Assumes `value` is already present in the
  // set.
  void Erase(T value) {
    auto it = value_to_index_.find(value);
    if (it == value_to_index_.end()) {
      std::cerr << "Value not found in OrderedSet" << std::endl;
      exit(0);
    }

    auto index = it->second;
    // Since we don't want to move values around in `value_sequence_` we swap
    // the value in the last position and with value to be deleted and then
    // pop_back.
    value_to_index_[value_sequence_.back()] = index;
    std::swap(value_sequence_[it->second], value_sequence_.back());
    value_sequence_.pop_back();
    value_to_index_.erase(it);
  }

  void Reserve(size_t new_size) {
    value_to_index_.reserve(new_size);
    value_sequence_.reserve(new_size);
  }

  void Clear() {
    value_to_index_.clear();
    value_sequence_.clear();
  }

  bool Contains(T value) const { return value_to_index_.find(value) != value_to_index_.end(); }
  size_t Size() const { return value_sequence_.size(); }

  std::vector<T> const& GetSequence() const { return value_sequence_; }

 private:
  // The stable order that we maintain through insertions and deletions.
  std::vector<T> value_sequence_;

  // Maps values to their indices in `value_sequence_`.
  std::unordered_map<T, int> value_to_index_;
};;
}  // namespace xla

#endif  // XLA_SERVICE_GRAPHCYCLES_ORDERED_SET_H_
