IndexMap.h

/* Copyright (C) 2012-2020 IBM Corp.
 * This program is 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. See accompanying LICENSE file.
 */
#ifndef HELIB_INDEXMAP_H
#define HELIB_INDEXMAP_H
 
#include <unordered_map>
#include <helib/IndexSet.h>
#include <helib/clonedPtr.h>
 
namespace helib {
 
template <typename T>
class IndexMapInit
{
public:
  virtual void init(T&) = 0;
 
  virtual IndexMapInit<T>* clone() const = 0;
  virtual ~IndexMapInit() {} // ensure that derived destructor is called
};
 
template <typename T>
class IndexMap
{
 
  std::unordered_map<long, T> map;
 
  IndexSet indexSet;
  cloned_ptr<IndexMapInit<T>> init;
 
public:
  IndexMap();
 
  explicit IndexMap(IndexMapInit<T>* _init) : init(_init) {}
 
  const IndexSet& getIndexSet() const { return indexSet; }
 
  T& operator[](long j)
  {
    assertTrue(indexSet.contains(j), "Key not found");
    return map[j];
  }
  const T& operator[](long j) const
  {
    assertTrue(indexSet.contains(j), "Key not found");
    // unordered_map does not support a const [] operator,
    // so we have to artificially strip away the const-ness here
    std::unordered_map<long, T>& map1 =
        const_cast<std::unordered_map<long, T>&>(map);
    return map1[j];
  }
 
  void insert(long j)
  {
    if (!indexSet.contains(j)) {
      indexSet.insert(j);
      if (!init.null())
        init->init(map[j]);
    }
  }
  void insert(const IndexSet& s)
  {
    for (long i = s.first(); i <= s.last(); i = s.next(i))
      insert(i);
  }
 
  void remove(long j)
  {
    indexSet.remove(j);
    map.erase(j);
  }
  void remove(const IndexSet& s)
  {
    for (long i = s.first(); i <= s.last(); i = s.next(i))
      map.erase(i);
    indexSet.remove(s);
  }
 
  void clear()
  {
    map.clear();
    indexSet.clear();
  }
};
 
template <typename T>
bool operator==(const IndexMap<T>& map1, const IndexMap<T>& map2)
{
  if (map1.getIndexSet() != map2.getIndexSet())
    return false;
  const IndexSet& s = map1.getIndexSet();
  for (long i = s.first(); i <= s.last(); i = s.next(i)) {
    if (map1[i] == map2[i])
      continue;
    return false;
  }
  return true;
}
 
template <typename T>
bool operator!=(const IndexMap<T>& map1, const IndexMap<T>& map2)
{
  return !(map1 == map2);
}
 
} // namespace helib
 
#endif // ifndef HELIB_INDEXMAP_H