1 files changed, 20 insertions, 15 deletions
diff --git a/runtimes/neurun/core/include/model/operation/SubgraphContext.h b/runtimes/neurun/core/include/model/operation/SubgraphContext.h
index 100f7b267..4a0e0e26b 100644
--- a/runtimes/neurun/core/include/model/operation/SubgraphContext.h
+++ b/runtimes/neurun/core/include/model/operation/SubgraphContext.h
@@ -23,6 +23,8 @@
 
 #include <unordered_map>
 
+using neurun::model::OperationIndex;
+
 namespace neurun
 {
 namespace model
@@ -36,24 +38,27 @@ public:
   SubgraphContext() : _index_count(0) {}
 
 public:
-  Index append(std::unique_ptr<Subgraph> &&subg);
-  Index append(const Index &index, const Node &node);
-  void remove(const Index &index) { _subgs.erase(index); };
+  OperationIndex append(std::unique_ptr<Subgraph> &&subg);
+  OperationIndex append(const OperationIndex &index, const Node &node);
+  void remove(const OperationIndex &index) { _subgs.erase(index); };
 
 public:
-  const Subgraph &at(const Index &) const;
-  Subgraph &at(const Index &);
-  std::unique_ptr<Subgraph> releaseAt(const Index &index) { return std::move(_subgs.at(index)); }
-  Subgraph *getAt(const Index &index) { return _subgs.at(index).get(); }
-  bool exist(const Index &) const;
-
-  bool hasNode(const Index &node_index) const;
-  Index findNode(const Index &node_index) const;
+  const Subgraph &at(const OperationIndex &) const;
+  Subgraph &at(const OperationIndex &);
+  std::unique_ptr<Subgraph> releaseAt(const OperationIndex &index)
+  {
+    return std::move(_subgs.at(index));
+  }
+  Subgraph *getAt(const OperationIndex &index) { return _subgs.at(index).get(); }
+  bool exist(const OperationIndex &) const;
+
+  bool hasNode(const OperationIndex &node_index) const;
+  OperationIndex findNode(const OperationIndex &node_index) const;
 
   uint32_t size() const { return _subgs.size(); }
 
-  void iterate(const std::function<void(const Index &, const Subgraph &)> &fn) const;
-  void iterate(const std::function<void(const Index &, Subgraph &)> &fn);
+  void iterate(const std::function<void(const OperationIndex &, const Subgraph &)> &fn) const;
+  void iterate(const std::function<void(const OperationIndex &, Subgraph &)> &fn);
 
   // TODO: Extract this into external helper function
   void dump(const std::string &msg) const;
@@ -61,10 +66,10 @@ public:
   void clear() { _subgs.clear(); }
 
 private:
-  Index generateIndex();
+  OperationIndex generateIndex();
 
 private:
-  std::unordered_map<Index, std::unique_ptr<Subgraph>> _subgs;
+  std::unordered_map<OperationIndex, std::unique_ptr<Subgraph>> _subgs;
   uint32_t _index_count;
 };