1 #pragma once 2 3 #include <functional> 4 #include <memory> 5 6 #include <ATen/core/ivalue.h> 7 #include <c10/macros/Export.h> 8 #include <torch/csrc/jit/frontend/source_range.h> 9 10 namespace torch::jit { 11 12 /** 13 * SourceRef does two things: 14 * 1. Owns a Source object. 15 * 2. Serves as lookup key to the owned Source in associative containers, for 16 * runtime data aggregation. 17 * We don't want to use std::shared_ptr<Source> directly because we want to 18 * support heteogeneous lookup, and also shared_ptr is an implementation detail 19 * which should be encapsulated. 20 */ 21 class TORCH_API SourceRef : public CustomClassHolder { 22 public: SourceRef(std::shared_ptr<Source> source_view)23 explicit SourceRef(std::shared_ptr<Source> source_view) 24 : source_view_(std::move(source_view)) {} 25 bool operator==(const SourceRef& other) const { 26 return source_view_ == other.source_view_; 27 } 28 bool operator<(const Source& other) const { 29 return source_view_.get() < &other; 30 } 31 friend bool operator<(const Source& other, const SourceRef& self) { 32 return &other < self.source_view_.get(); 33 } 34 bool operator<(const SourceRef& other) const { 35 return *this < *other.source_view_; 36 } 37 const Source* operator->() const { 38 return source_view_.get(); 39 } 40 41 private: 42 std::shared_ptr<Source> source_view_; 43 }; 44 45 } // namespace torch::jit 46