onnx/common/visitor.h

0001 // Copyright (c) ONNX Project Contributors
0002
0003 /*
0004  * SPDX-License-Identifier: Apache-2.0
0005  */
0006
0007 #pragma once
0008 #include "onnx/common/common.h"
0009 #include "onnx/onnx_pb.h"
0010
0011 namespace ONNX_NAMESPACE {
0012 namespace internal {
0013
0014 // Visitor: A readonly visitor class for ONNX Proto objects.
0015 // This class is restricted to Nodes, Graphs, Attributes, and Functions.
0016 // The VisitX methods invoke ProcessX, and if that returns true, will
0017 // continue to visit all children of the X.
0018
0019 struct Visitor {
0020   virtual void VisitGraph(const GraphProto& graph) {
0021     if (ProcessGraph(graph))
0022       for (auto& node : graph.node())
0023         VisitNode(node);
0024   }
0025
0026   virtual void VisitFunction(const FunctionProto& function) {
0027     if (ProcessFunction(function))
0028       for (auto& node : function.node())
0029         VisitNode(node);
0030   }
0031
0032   virtual void VisitNode(const NodeProto& node) {
0033     if (ProcessNode(node)) {
0034       for (auto& attr : node.attribute()) {
0035         VisitAttribute(attr);
0036       }
0037     }
0038   }
0039
0040   virtual void VisitAttribute(const AttributeProto& attr) {
0041     if (ProcessAttribute(attr)) {
0042       if (attr.has_g()) {
0043         VisitGraph(attr.g());
0044       }
0045       for (auto& graph : attr.graphs())
0046         VisitGraph(graph);
0047     }
0048   }
0049
0050   virtual bool ProcessGraph(const GraphProto& graph) {
0051     ONNX_UNUSED_PARAMETER(graph);
0052     return true;
0053   }
0054
0055   virtual bool ProcessFunction(const FunctionProto& function) {
0056     ONNX_UNUSED_PARAMETER(function);
0057     return true;
0058   }
0059
0060   virtual bool ProcessNode(const NodeProto& node) {
0061     ONNX_UNUSED_PARAMETER(node);
0062     return true;
0063   }
0064
0065   virtual bool ProcessAttribute(const AttributeProto& attr) {
0066     ONNX_UNUSED_PARAMETER(attr);
0067     return true;
0068   }
0069
0070   virtual ~Visitor() {}
0071 };
0072
0073 // MutableVisitor: A version of Visitor that allows mutation of the visited objects.
0074 struct MutableVisitor {
0075   virtual void VisitGraph(GraphProto* graph) {
0076     if (ProcessGraph(graph))
0077       for (auto& node : *(graph->mutable_node()))
0078         VisitNode(&node);
0079   }
0080
0081   virtual void VisitFunction(FunctionProto* function) {
0082     if (ProcessFunction(function))
0083       for (auto& node : *(function->mutable_node()))
0084         VisitNode(&node);
0085   }
0086
0087   virtual void VisitNode(NodeProto* node) {
0088     if (ProcessNode(node)) {
0089       for (auto& attr : *(node->mutable_attribute())) {
0090         VisitAttribute(&attr);
0091       }
0092     }
0093   }
0094
0095   virtual void VisitAttribute(AttributeProto* attr) {
0096     if (ProcessAttribute(attr)) {
0097       if (attr->has_g()) {
0098         VisitGraph(attr->mutable_g());
0099       }
0100       for (auto& graph : *(attr->mutable_graphs()))
0101         VisitGraph(&graph);
0102     }
0103   }
0104
0105   virtual bool ProcessGraph(GraphProto* graph) {
0106     ONNX_UNUSED_PARAMETER(graph);
0107     return true;
0108   }
0109
0110   virtual bool ProcessFunction(FunctionProto* function) {
0111     ONNX_UNUSED_PARAMETER(function);
0112     return true;
0113   }
0114
0115   virtual bool ProcessNode(NodeProto* node) {
0116     ONNX_UNUSED_PARAMETER(node);
0117     return true;
0118   }
0119
0120   virtual bool ProcessAttribute(AttributeProto* attr) {
0121     ONNX_UNUSED_PARAMETER(attr);
0122     return true;
0123   }
0124
0125   virtual ~MutableVisitor() {}
0126 };
0127
0128 } // namespace internal
0129 } // namespace ONNX_NAMESPACE