registry.h

Go to the documentation of this file.

#ifndef DMLC_REGISTRY_H_
#define DMLC_REGISTRY_H_

#include <map>
#include <string>
#include <vector>
#include "./base.h"
#include "./logging.h"
#include "./parameter.h"
#include "./type_traits.h"

namespace dmlc {
template<typename EntryType>
class Registry {
 public:
  inline static const std::vector<const EntryType*>& List() {
    return Get()->const_list_;
  }
  inline static std::vector<std::string> ListAllNames() {
    const std::map<std::string, EntryType*> &fmap = Get()->fmap_;
    typename std::map<std::string, EntryType*>::const_iterator p;
    std::vector<std::string> names;
    for (p = fmap.begin(); p !=fmap.end(); ++p) {
      names.push_back(p->first);
    }
    return names;
  }
  inline static const EntryType *Find(const std::string &name) {
    const std::map<std::string, EntryType*> &fmap = Get()->fmap_;
    typename std::map<std::string, EntryType*>::const_iterator p = fmap.find(name);
    if (p != fmap.end()) {
      return p->second;
    } else {
      return NULL;
    }
  }
  inline void AddAlias(const std::string& key_name,
                       const std::string& alias) {
    EntryType* e = fmap_.at(key_name);
    if (fmap_.count(alias)) {
      CHECK_EQ(e, fmap_.at(alias))
          << "Trying to register alias " << alias << " for key " << key_name
          << " but " << alias << " is already taken";
    } else {
      fmap_[alias] = e;
    }
  }
  inline EntryType &__REGISTER__(const std::string& name) {
    std::lock_guard<std::mutex> guard(registering_mutex);
    if (fmap_.count(name) > 0) {
      return *fmap_[name];
    }
    EntryType *e = new EntryType();
    e->name = name;
    fmap_[name] = e;
    const_list_.push_back(e);
    entry_list_.push_back(e);
    return *e;
  }
  inline EntryType &__REGISTER_OR_GET__(const std::string& name) {
    if (fmap_.count(name) == 0) {
      return __REGISTER__(name);
    } else {
      return *fmap_.at(name);
    }
  }
  static Registry *Get();

 private:
  std::vector<EntryType*> entry_list_;
  std::vector<const EntryType*> const_list_;
  std::map<std::string, EntryType*> fmap_;
  std::mutex registering_mutex;
  Registry() {}
  ~Registry() {
    for (size_t i = 0; i < entry_list_.size(); ++i) {
      delete entry_list_[i];
    }
  }
};

template<typename EntryType, typename FunctionType>
class FunctionRegEntryBase {
 public:
  std::string name;
  std::string description;
  std::vector<ParamFieldInfo> arguments;
  FunctionType body;
  std::string return_type;

  inline EntryType &set_body(FunctionType body) {
    this->body = body;
    return this->self();
  }
  inline EntryType &describe(const std::string &description) {
    this->description = description;
    return this->self();
  }
  inline EntryType &add_argument(const std::string &name,
                                 const std::string &type,
                                 const std::string &description) {
    ParamFieldInfo info;
    info.name = name;
    info.type = type;
    info.type_info_str = info.type;
    info.description = description;
    arguments.push_back(info);
    return this->self();
  }
  inline EntryType &add_arguments(const std::vector<ParamFieldInfo> &args) {
    arguments.insert(arguments.end(), args.begin(), args.end());
    return this->self();
  }
  inline EntryType &set_return_type(const std::string &type) {
    return_type = type;
    return this->self();
  }

 protected:
  inline EntryType &self() {
    return *(static_cast<EntryType*>(this));
  }
};

#define DMLC_REGISTRY_ENABLE(EntryType)                                 \
  template<>                                                            \
  Registry<EntryType > *Registry<EntryType >::Get() {                   \
    static Registry<EntryType > inst;                                   \
    return &inst;                                                       \
  }                                                                     \


#define DMLC_REGISTRY_REGISTER(EntryType, EntryTypeName, Name)          \
  static DMLC_ATTRIBUTE_UNUSED EntryType & __make_ ## EntryTypeName ## _ ## Name ## __ = \
      ::dmlc::Registry<EntryType>::Get()->__REGISTER__(#Name)           \


#define DMLC_REGISTRY_FILE_TAG(UniqueTag)                                \
  int __dmlc_registry_file_tag_ ## UniqueTag ## __() { return 0; }

#define DMLC_REGISTRY_LINK_TAG(UniqueTag)                                \
  int __dmlc_registry_file_tag_ ## UniqueTag ## __();                   \
  static int DMLC_ATTRIBUTE_UNUSED __reg_file_tag_ ## UniqueTag ## __ = \
      __dmlc_registry_file_tag_ ## UniqueTag ## __();
}  // namespace dmlc
#endif  // DMLC_REGISTRY_H_