#ifndef CAFFE2_CORE_NET_SIMPLE_REFCOUNT_H_
#define CAFFE2_CORE_NET_SIMPLE_REFCOUNT_H_

#include <vector>

#include "c10/util/Registry.h"
#include "caffe2/core/common.h"
#include "caffe2/core/logging.h"
#include "caffe2/core/net.h"
#include "caffe2/core/net_simple.h"
#include "caffe2/core/tensor.h"
#include "caffe2/core/workspace.h"
#include "caffe2/proto/caffe2_pb.h"

namespace caffe2 {

// SimpleRefcountNet is an implementation that adds an additional abstraction
// on top of SimpleRefCountNet: it tracks all the tensors and for those that are
// considered internal/temporary, delete them once their refcount go to zero.
// In the context of a simple static run, this can be carried out during
// construction time: we will do a pass through the network and track what
// blobs we need to do reset on, after the execution of every op.
//
// To identify which blob is considered temporary, we employ the following
// strategy: any blob that is
// (1) consumed but not produced by ops in the net, or
// (2) produced but not consumed by ops in the net, or
// (3) is marked as external_output in the protobuf
// will NOT be considered temporary.
//
// In the long run, we should design proper functional interfaces so that
// nets are less imperative and more functional.
//
// Also, for now, SimpleRefCountNet should only be used for benchmarking
// purposes and not product use, since it is not going to provide better
// performance gain, and is implicitly incompatible with the contract that
// earlier Nets expose - that all intermediate blobs are visible to the users.
class SimpleRefCountNet final : public SimpleNet {
 public:
  SimpleRefCountNet(
      const std::shared_ptr<const NetDef>& net_def,
      Workspace* ws);

 protected:
  bool Run() override;

  using SimpleNet::operators_;

 private:
  // The list of blobs to delete when each operator finishes its run.
  // This will be populated during construction time.
  vector<vector<Blob*>> delete_list_;

  C10_DISABLE_COPY_AND_ASSIGN(SimpleRefCountNet);
};

} // namespace caffe2

#endif // CAFFE2_CORE_NET_SIMPLE_REFCOUNT_H_