text: /*!
   *  Copyright (c) 2019 by Contributors
   * \file graph/nodeflow.cc
   * \brief DGL NodeFlow related functions.
   */
  
  #include <dgl/immutable_graph.h>
  #include <dgl/packed_func_ext.h>
  #include <dgl/nodeflow.h>
  
  #include <string>
  
  #include "../c_api_common.h"
  
  using dgl::runtime::DGLArgs;
  using dgl::runtime::DGLArgValue;
  using dgl::runtime::DGLRetValue;
  using dgl::runtime::PackedFunc;
  
  namespace dgl {
  
  std::vector<IdArray> GetNodeFlowSlice(const ImmutableGraph &graph, const std::string &fmt,
                                        size_t layer0_size, size_t layer1_start,
                                        size_t layer1_end, bool remap) {
    CHECK_GE(layer1_start, layer0_size);
    if (fmt == std::string("csr")) {
      dgl_id_t first_vid = layer1_start - layer0_size;
      auto csr = aten::CSRSliceRows(graph.GetInCSR()->ToCSRMatrix(), layer1_start, layer1_end);
      if (remap) {
        dgl_id_t *eid_data = static_cast<dgl_id_t*>(csr.data->data);
        const dgl_id_t first_eid = eid_data[0];
        IdArray new_indices = aten::Sub(csr.indices, first_vid);
        IdArray new_data = aten::Sub(csr.data, first_eid);
        return {csr.indptr, new_indices, new_data};
      } else {
        return {csr.indptr, csr.indices, csr.data};
      }
    } else if (fmt == std::string("coo")) {
      auto csr = graph.GetInCSR()->ToCSRMatrix();
      const dgl_id_t* indptr = static_cast<dgl_id_t*>(csr.indptr->data);
      const dgl_id_t* indices = static_cast<dgl_id_t*>(csr.indices->data);
      const dgl_id_t* edge_ids = static_cast<dgl_id_t*>(csr.data->data);
      int64_t nnz = indptr[layer1_end] - indptr[layer1_start];
      IdArray idx = aten::NewIdArray(2 * nnz);
      IdArray eid = aten::NewIdArray(nnz);
      int64_t *idx_data = static_cast<int64_t*>(idx->data);
      dgl_id_t *eid_data = static_cast<dgl_id_t*>(eid->data);
      size_t num_edges = 0;
      for (size_t i = layer1_start; i < layer1_end; i++) {
        for (dgl_id_t j = indptr[i]; j < indptr[i + 1]; j++) {
          // These nodes are all in a layer. We need to remap them to the node id
          // local to the layer.
          idx_data[num_edges] = remap ? i - layer1_start : i;
          num_edges++;
        }
      }
      CHECK_EQ(num_edges, nnz);
      if (remap) {
        size_t edge_start = indptr[layer1_start];
        dgl_id_t first_eid = edge_ids[edge_start];
        dgl_id_t first_vid = layer1_start - layer0_size;
        for (int64_t i = 0; i < nnz; i++) {
          CHECK_GE(indices[edge_start + i], first_vid);
          idx_data[nnz + i] = indices[edge_start + i] - first_vid;
          eid_data[i] = edge_ids[edge_start + i] - first_eid;
        }
      } else {
        std::copy(indices + indptr[layer1_start],
                  indices + indptr[layer1_end], idx_data + nnz);
        std::copy(edge_ids + indptr[layer1_start],
                  edge_ids + indptr[layer1_end], eid_data);
      }
      return std::vector<IdArray>{idx, eid};
    } else {
      LOG(FATAL) << "unsupported adjacency matrix format";
      return {};
    }
  }
  
  DGL_REGISTER_GLOBAL("_deprecate.nodeflow._CAPI_NodeFlowGetBlockAdj")
  .set_body([] (DGLArgs args, DGLRetValue* rv) {
      GraphRef g = args[0];
      std::string format = args[1];
      int64_t layer0_size = args[2];
      int64_t start = args[3];
      int64_t end = args[4];
      const bool remap = args[5];
      auto ig = CHECK_NOTNULL(std::dynamic_pointer_cast<ImmutableGraph>(g.sptr()));
      auto res = GetNodeFlowSlice(*ig, format, layer0_size, start, end, remap);
      *rv = ConvertNDArrayVectorToPackedFunc(res);
    });
  
  }  // namespace dgl
  