api_container.cc

/*!
 *  Copyright (c) 2019 by Contributors
 * \file api/api_container.cc
 * \brief Runtime container APIs. (reference: tvm/src/api/api_lang.cc)
 */
#include <dgl/runtime/ndarray.h>
#include <dgl/runtime/container.h>
#include <dgl/runtime/registry.h>
#include <dgl/packed_func_ext.h>

namespace dgl {
namespace runtime {

DGL_REGISTER_GLOBAL("_List")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    auto ret_obj = std::make_shared<runtime::ListObject>();
    for (int i = 0; i < args.size(); ++i) {
      ret_obj->data.push_back(args[i].obj_sptr());
    }
    *rv = ret_obj;
  });

DGL_REGISTER_GLOBAL("_ListGetItem")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    auto& sptr = args[0].obj_sptr();
    CHECK(sptr->is_type<ListObject>());
    auto* o = static_cast<const ListObject*>(sptr.get());
    int64_t i = args[1];
    CHECK_LT(i, o->data.size()) << "list out of bound";
    *rv = o->data[i];
  });

DGL_REGISTER_GLOBAL("_ListSize")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    auto& sptr = args[0].obj_sptr();
    CHECK(sptr->is_type<ListObject>());
    auto* o = static_cast<const ListObject*>(sptr.get());
    *rv = static_cast<int64_t>(o->data.size());
  });

DGL_REGISTER_GLOBAL("_Map")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    CHECK_EQ(args.size() % 2, 0);
    if (args.size() != 0 && args[0].type_code() == kStr) {
      // StrMap
      StrMapObject::ContainerType data;
      for (int i = 0; i < args.size(); i += 2) {
        CHECK(args[i].type_code() == kStr)
          << "The key of the map must be string";
        CHECK(args[i + 1].type_code() == kObjectHandle)
          << "The value of the map must be an object type";
        data.emplace(std::make_pair(args[i].operator std::string(),
                                    args[i + 1].obj_sptr()));
      }
      auto obj = std::make_shared<StrMapObject>();
      obj->data = std::move(data);
      *rv = obj;
    } else {
      // object container
      MapObject::ContainerType data;
      for (int i = 0; i < args.size(); i += 2) {
        CHECK(args[i].type_code() == kObjectHandle)
          << "The key of the map must be an object type";
        CHECK(args[i + 1].type_code() == kObjectHandle)
          << "The value of the map must be an object type";
        data.emplace(std::make_pair(args[i].obj_sptr(), args[i + 1].obj_sptr()));
      }
      auto obj = std::make_shared<MapObject>();
      obj->data = std::move(data);
      *rv = obj;
    }
  });

DGL_REGISTER_GLOBAL("_MapSize")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    auto& sptr = args[0].obj_sptr();
    if (sptr->is_type<MapObject>()) {
      auto* o = static_cast<const MapObject*>(sptr.get());
      *rv = static_cast<int64_t>(o->data.size());
    } else {
      CHECK(sptr->is_type<StrMapObject>());
      auto* o = static_cast<const StrMapObject*>(sptr.get());
      *rv = static_cast<int64_t>(o->data.size());
    }
  });

DGL_REGISTER_GLOBAL("_MapGetItem")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    auto& sptr = args[0].obj_sptr();
    if (sptr->is_type<MapObject>()) {
      auto* o = static_cast<const MapObject*>(sptr.get());
      auto it = o->data.find(args[1].obj_sptr());
      CHECK(it != o->data.end()) << "cannot find the key in the map";
      *rv = (*it).second;
    } else {
      CHECK(sptr->is_type<StrMapObject>());
      auto* o = static_cast<const StrMapObject*>(sptr.get());
      auto it = o->data.find(args[1].operator std::string());
      CHECK(it != o->data.end()) << "cannot find the key in the map";
      *rv = (*it).second;
    }
  });

DGL_REGISTER_GLOBAL("_MapItems")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    auto& sptr = args[0].obj_sptr();
    if (sptr->is_type<MapObject>()) {
      auto* o = static_cast<const MapObject*>(sptr.get());
      auto rkvs = std::make_shared<ListObject>();
      for (const auto& kv : o->data) {
        rkvs->data.push_back(kv.first);
        rkvs->data.push_back(kv.second);
      }
      *rv = rkvs;
    } else {
      CHECK(sptr->is_type<StrMapObject>());
      auto* o = static_cast<const StrMapObject*>(sptr.get());
      auto rkvs = std::make_shared<ListObject>();
      for (const auto& kv : o->data) {
        rkvs->data.push_back(MakeValue(kv.first));
        rkvs->data.push_back(kv.second);
      }
      *rv = rkvs;
    }
  });

DGL_REGISTER_GLOBAL("_MapCount")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    auto& sptr = args[0].obj_sptr();
    if (sptr->is_type<MapObject>()) {
      auto* o = static_cast<const MapObject*>(sptr.get());
      *rv = static_cast<int64_t>(o->data.count(args[1].obj_sptr()));
    } else {
      CHECK(sptr->is_type<StrMapObject>());
      auto* o = static_cast<const StrMapObject*>(sptr.get());
      *rv = static_cast<int64_t>(o->data.count(args[1].operator std::string()));
    }
  });

DGL_REGISTER_GLOBAL("_Value")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    *rv = MakeValue(args[0]);
  });

DGL_REGISTER_GLOBAL("_ValueGet")
.set_body([] (DGLArgs args, DGLRetValue* rv) {
    auto& sptr = args[0].obj_sptr();
    CHECK(sptr->is_type<ValueObject>());
    auto* o = static_cast<const ValueObject*>(sptr.get());
    *rv = o->data;
  });

}  // namespace runtime
}  // namespace dgl