doc.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
#ifndef TVM_SCRIPT_PRINTER_DOC_H_
#define TVM_SCRIPT_PRINTER_DOC_H_
 
#include <tvm/ffi/reflection/access_path.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/expr.h>
#include <tvm/node/node.h>
#include <tvm/runtime/data_type.h>
#include <tvm/runtime/device_api.h>
 
#include <string>
 
namespace tvm {
namespace script {
namespace printer {
 
using AccessPath = ffi::reflection::AccessPath;
 
// Forward declaration
class Doc;
 
String DocToPythonScript(Doc doc, const PrinterConfig& cfg);
 
class DocNode : public Object {
 public:
  mutable Array<ffi::reflection::AccessPath> source_paths;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<DocNode>().def_rw("source_paths", &DocNode::source_paths);
  }
 
  static constexpr const char* _type_key = "script.printer.Doc";
  static constexpr bool _type_mutable = true;
 
  TVM_DECLARE_BASE_OBJECT_INFO(DocNode, Object);
 
 public:
  virtual ~DocNode() = default;
};
 
class Doc : public ObjectRef {
 protected:
  Doc() = default;
 
 public:
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(Doc, ObjectRef, DocNode);
};
 
class ExprDoc;
 
class ExprDocNode : public DocNode {
 public:
  ExprDoc Attr(String attr) const;
 
  ExprDoc operator[](Array<Doc> indices) const;
 
  ExprDoc Call(Array<ExprDoc, void> args) const;
 
  ExprDoc Call(Array<ExprDoc, void> args,  //
               Array<String> kwargs_keys,  //
               Array<ExprDoc, void> kwargs_values) const;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ExprDocNode>();
  }
 
  static constexpr const char* _type_key = "script.printer.ExprDoc";
 
  TVM_DECLARE_BASE_OBJECT_INFO(ExprDocNode, DocNode);
};
 
class ExprDoc : public Doc {
 protected:
  ExprDoc() = default;
 
 public:
  ExprDoc operator[](Array<Doc> indices) const;
 
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ExprDoc, Doc, ExprDocNode);
};
 
class StmtDocNode : public DocNode {
 public:
  mutable Optional<String> comment{std::nullopt};
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<StmtDocNode>().def_rw("comment", &StmtDocNode::comment);
  }
 
  static constexpr const char* _type_key = "script.printer.StmtDoc";
 
  TVM_DECLARE_BASE_OBJECT_INFO(StmtDocNode, DocNode);
};
 
class StmtDoc : public Doc {
 protected:
  StmtDoc() = default;
 
 public:
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(StmtDoc, Doc, StmtDocNode);
};
 
class StmtBlockDocNode : public DocNode {
 public:
  Array<StmtDoc> stmts;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<StmtBlockDocNode>().def_ro("stmts", &StmtBlockDocNode::stmts);
  }
 
  static constexpr const char* _type_key = "script.printer.StmtBlockDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(StmtBlockDocNode, DocNode);
};
 
class StmtBlockDoc : public Doc {
 public:
  explicit StmtBlockDoc(Array<StmtDoc> stmts);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(StmtBlockDoc, Doc, StmtBlockDocNode);
};
 
class LiteralDocNode : public ExprDocNode {
 public:
  ffi::Any value;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<LiteralDocNode>().def_ro("value", &LiteralDocNode::value);
  }
 
  static constexpr const char* _type_key = "script.printer.LiteralDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(LiteralDocNode, ExprDocNode);
};
 
class LiteralDoc : public ExprDoc {
 protected:
  explicit LiteralDoc(ffi::Any value, const Optional<AccessPath>& object_path);
 
 public:
  static LiteralDoc None(const Optional<AccessPath>& p) { return LiteralDoc(ffi::Any(nullptr), p); }
  static LiteralDoc Int(int64_t v, const Optional<AccessPath>& p) {
    return LiteralDoc(IntImm(DataType::Int(64), v), p);
  }
  static LiteralDoc Boolean(bool v, const Optional<AccessPath>& p) {
    return LiteralDoc(IntImm(DataType::Bool(), v), p);
  }
  static LiteralDoc Float(double v, const Optional<AccessPath>& p) {
    return LiteralDoc(FloatImm(DataType::Float(64), v), p);
  }
  static LiteralDoc Str(const String& v, const Optional<AccessPath>& p) { return LiteralDoc(v, p); }
  static LiteralDoc DataType(const runtime::DataType& v, const Optional<AccessPath>& p) {
    std::string dtype = v.is_void() ? "void" : runtime::DLDataTypeToString(v);
    return LiteralDoc::Str(dtype, p);
  }
  static LiteralDoc Device(const DLDevice& v, const Optional<AccessPath>& p) {
    std::ostringstream os;
    runtime::operator<<(os, v);
    return LiteralDoc::Str(os.str(), p);
  }
 
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(LiteralDoc, ExprDoc, LiteralDocNode);
};
 
class IdDocNode : public ExprDocNode {
 public:
  String name;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<IdDocNode>().def_ro("name", &IdDocNode::name);
  }
 
  static constexpr const char* _type_key = "script.printer.IdDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(IdDocNode, ExprDocNode);
};
 
class IdDoc : public ExprDoc {
 public:
  explicit IdDoc(String name);
  explicit IdDoc(std::nullptr_t) : ExprDoc(nullptr) {}
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(IdDoc, ExprDoc, IdDocNode);
};
 
class AttrAccessDocNode : public ExprDocNode {
 public:
  ExprDoc value{nullptr};
  String name;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AttrAccessDocNode>()
        .def_ro("value", &AttrAccessDocNode::value)
        .def_ro("name", &AttrAccessDocNode::name);
  }
 
  static constexpr const char* _type_key = "script.printer.AttrAccessDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(AttrAccessDocNode, ExprDocNode);
};
 
class AttrAccessDoc : public ExprDoc {
 public:
  explicit AttrAccessDoc(ExprDoc value, String name);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(AttrAccessDoc, ExprDoc, AttrAccessDocNode);
};
 
class IndexDocNode : public ExprDocNode {
 public:
  ExprDoc value{nullptr};
  Array<Doc> indices;  // Each element is union of: Slice / ExprDoc
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<IndexDocNode>()
        .def_ro("value", &IndexDocNode::value)
        .def_ro("indices", &IndexDocNode::indices);
  }
 
  static constexpr const char* _type_key = "script.printer.IndexDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(IndexDocNode, ExprDocNode);
};
 
class IndexDoc : public ExprDoc {
 public:
  explicit IndexDoc(ExprDoc value, Array<Doc> indices);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(IndexDoc, ExprDoc, IndexDocNode);
};
 
class CallDocNode : public ExprDocNode {
 public:
  ExprDoc callee{nullptr};
  Array<ExprDoc> args;
  Array<String> kwargs_keys;
  Array<ExprDoc> kwargs_values;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<CallDocNode>()
        .def_ro("callee", &CallDocNode::callee)
        .def_ro("args", &CallDocNode::args)
        .def_ro("kwargs_keys", &CallDocNode::kwargs_keys)
        .def_ro("kwargs_values", &CallDocNode::kwargs_values);
  }
 
  static constexpr const char* _type_key = "script.printer.CallDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(CallDocNode, ExprDocNode);
};
 
class CallDoc : public ExprDoc {
 public:
  CallDoc(ExprDoc callee, Array<ExprDoc> args, Array<String> kwargs_keys,
          Array<ExprDoc> kwargs_values);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(CallDoc, ExprDoc, CallDocNode);
};
 
class OperationDocNode : public ExprDocNode {
 public:
  enum class Kind : int32_t {
    // Unary operators
    kUnaryStart = 0,
    kUSub = 1,    // -x
    kInvert = 2,  // ~x
    kNot = 3,     // not x
    kUnaryEnd = 4,
 
    // Binary operators
    kBinaryStart = 5,
    kAdd = 6,        // +
    kSub = 7,        // -
    kMult = 8,       // *
    kDiv = 9,        // /
    kFloorDiv = 10,  // // in Python
    kMod = 11,       // % in Python
    kPow = 12,       // ** in Python
    kLShift = 13,    // <<
    kRShift = 14,    // >>
    kBitAnd = 15,    // &
    kBitOr = 16,     // |
    kBitXor = 17,    // ^
    kLt = 18,        // <
    kLtE = 19,       // <=
    kEq = 20,        // ==
    kNotEq = 21,     // !=
    kGt = 22,        // >
    kGtE = 23,       // >=
    kAnd = 24,       // and
    kOr = 25,        // or
    kBinaryEnd = 26,
 
    // Special
    kSpecialStart = 27,
    kIfThenElse = 28,  // <operands[1]> if <operands[0]> else <operands[2]>
    kSpecialEnd = 29
  };
 
  Kind kind;
  Array<ExprDoc> operands;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<OperationDocNode>()
        .def_ro("kind", &OperationDocNode::kind)
        .def_ro("operands", &OperationDocNode::operands);
  }
 
  static constexpr const char* _type_key = "script.printer.OperationDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(OperationDocNode, ExprDocNode);
};
 
class OperationDoc : public ExprDoc {
 public:
  explicit OperationDoc(OperationDocNode::Kind kind, Array<ExprDoc> operands);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(OperationDoc, ExprDoc, OperationDocNode);
};
 
class LambdaDocNode : public ExprDocNode {
 public:
  Array<IdDoc> args;
  ExprDoc body{nullptr};
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<LambdaDocNode>()
        .def_ro("args", &LambdaDocNode::args)
        .def_ro("body", &LambdaDocNode::body);
  }
 
  static constexpr const char* _type_key = "script.printer.LambdaDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(LambdaDocNode, ExprDocNode);
};
 
class LambdaDoc : public ExprDoc {
 public:
  explicit LambdaDoc(Array<IdDoc> args, ExprDoc body);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(LambdaDoc, ExprDoc, LambdaDocNode);
};
 
class TupleDocNode : public ExprDocNode {
 public:
  Array<ExprDoc> elements;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<TupleDocNode>().def_ro("elements", &TupleDocNode::elements);
  }
 
  static constexpr const char* _type_key = "script.printer.TupleDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(TupleDocNode, ExprDocNode);
};
 
class TupleDoc : public ExprDoc {
 public:
  TupleDoc() : TupleDoc(ffi::make_object<TupleDocNode>()) {}
  explicit TupleDoc(Array<ExprDoc> elements);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(TupleDoc, ExprDoc, TupleDocNode);
};
 
class ListDocNode : public ExprDocNode {
 public:
  Array<ExprDoc> elements;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ListDocNode>().def_ro("elements", &ListDocNode::elements);
  }
 
  static constexpr const char* _type_key = "script.printer.ListDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(ListDocNode, ExprDocNode);
};
 
class ListDoc : public ExprDoc {
 public:
  ListDoc() : ListDoc(ffi::make_object<ListDocNode>()) {}
  explicit ListDoc(Array<ExprDoc> elements);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ListDoc, ExprDoc, ListDocNode);
};
 
class DictDocNode : public ExprDocNode {
 public:
  Array<ExprDoc> keys;
  Array<ExprDoc> values;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<DictDocNode>()
        .def_ro("keys", &DictDocNode::keys)
        .def_ro("values", &DictDocNode::values);
  }
 
  static constexpr const char* _type_key = "script.printer.DictDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(DictDocNode, ExprDocNode);
};
 
class DictDoc : public ExprDoc {
 public:
  DictDoc() : DictDoc(ffi::make_object<DictDocNode>()) {}
  explicit DictDoc(Array<ExprDoc> keys, Array<ExprDoc> values);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(DictDoc, ExprDoc, DictDocNode);
};
 
class SliceDocNode : public DocNode {
 public:
  Optional<ExprDoc> start;
  Optional<ExprDoc> stop;
  Optional<ExprDoc> step;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<SliceDocNode>()
        .def_ro("start", &SliceDocNode::start)
        .def_ro("stop", &SliceDocNode::stop)
        .def_ro("step", &SliceDocNode::step);
  }
 
  static constexpr const char* _type_key = "script.printer.SliceDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(SliceDocNode, DocNode);
};
 
class SliceDoc : public Doc {
 public:
  explicit SliceDoc(Optional<ExprDoc> start, Optional<ExprDoc> stop, Optional<ExprDoc> step);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(SliceDoc, Doc, SliceDocNode);
};
 
class AssignDocNode : public StmtDocNode {
 public:
  ExprDoc lhs{nullptr};
  Optional<ExprDoc> rhs;
  Optional<ExprDoc> annotation;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AssignDocNode>()
        .def_ro("lhs", &AssignDocNode::lhs)
        .def_ro("rhs", &AssignDocNode::rhs)
        .def_ro("annotation", &AssignDocNode::annotation);
  }
 
  static constexpr const char* _type_key = "script.printer.AssignDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(AssignDocNode, StmtDocNode);
};
 
class AssignDoc : public StmtDoc {
 public:
  explicit AssignDoc(ExprDoc lhs, Optional<ExprDoc> rhs, Optional<ExprDoc> annotation);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(AssignDoc, StmtDoc, AssignDocNode);
};
 
class IfDocNode : public StmtDocNode {
 public:
  ExprDoc predicate{nullptr};
  Array<StmtDoc> then_branch;
  Array<StmtDoc> else_branch;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<IfDocNode>()
        .def_ro("predicate", &IfDocNode::predicate)
        .def_ro("then_branch", &IfDocNode::then_branch)
        .def_ro("else_branch", &IfDocNode::else_branch);
  }
 
  static constexpr const char* _type_key = "script.printer.IfDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(IfDocNode, StmtDocNode);
};
 
class IfDoc : public StmtDoc {
 public:
  explicit IfDoc(ExprDoc predicate, Array<StmtDoc> then_branch, Array<StmtDoc> else_branch);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(IfDoc, StmtDoc, IfDocNode);
};
 
class WhileDocNode : public StmtDocNode {
 public:
  ExprDoc predicate{nullptr};
  Array<StmtDoc> body;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<WhileDocNode>()
        .def_ro("predicate", &WhileDocNode::predicate)
        .def_ro("body", &WhileDocNode::body);
  }
 
  static constexpr const char* _type_key = "script.printer.WhileDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(WhileDocNode, StmtDocNode);
};
 
class WhileDoc : public StmtDoc {
 public:
  explicit WhileDoc(ExprDoc predicate, Array<StmtDoc> body);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(WhileDoc, StmtDoc, WhileDocNode);
};
 
class ForDocNode : public StmtDocNode {
 public:
  ExprDoc lhs{nullptr};
  ExprDoc rhs{nullptr};
  Array<StmtDoc> body;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ForDocNode>()
        .def_ro("lhs", &ForDocNode::lhs)
        .def_ro("rhs", &ForDocNode::rhs)
        .def_ro("body", &ForDocNode::body);
  }
 
  static constexpr const char* _type_key = "script.printer.ForDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(ForDocNode, StmtDocNode);
};
 
class ForDoc : public StmtDoc {
 public:
  explicit ForDoc(ExprDoc lhs, ExprDoc rhs, Array<StmtDoc> body);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ForDoc, StmtDoc, ForDocNode);
};
 
class ScopeDocNode : public StmtDocNode {
 public:
  Optional<ExprDoc> lhs{std::nullopt};
  ExprDoc rhs{nullptr};
  Array<StmtDoc> body;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ScopeDocNode>()
        .def_ro("lhs", &ScopeDocNode::lhs)
        .def_ro("rhs", &ScopeDocNode::rhs)
        .def_ro("body", &ScopeDocNode::body);
  }
 
  static constexpr const char* _type_key = "script.printer.ScopeDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(ScopeDocNode, StmtDocNode);
};
 
class ScopeDoc : public StmtDoc {
 public:
  explicit ScopeDoc(Optional<ExprDoc> lhs, ExprDoc rhs, Array<StmtDoc> body);
 
  explicit ScopeDoc(ExprDoc rhs, Array<StmtDoc> body);
 
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ScopeDoc, StmtDoc, ScopeDocNode);
};
 
class ExprStmtDocNode : public StmtDocNode {
 public:
  ExprDoc expr{nullptr};
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ExprStmtDocNode>().def_ro("expr", &ExprStmtDocNode::expr);
  }
 
  static constexpr const char* _type_key = "script.printer.ExprStmtDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(ExprStmtDocNode, StmtDocNode);
};
 
class ExprStmtDoc : public StmtDoc {
 public:
  explicit ExprStmtDoc(ExprDoc expr);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ExprStmtDoc, StmtDoc, ExprStmtDocNode);
};
 
class AssertDocNode : public StmtDocNode {
 public:
  ExprDoc test{nullptr};
  Optional<ExprDoc> msg{std::nullopt};
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<AssertDocNode>()
        .def_ro("test", &AssertDocNode::test)
        .def_ro("msg", &AssertDocNode::msg);
  }
 
  static constexpr const char* _type_key = "script.printer.AssertDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(AssertDocNode, StmtDocNode);
};
 
class AssertDoc : public StmtDoc {
 public:
  explicit AssertDoc(ExprDoc test, Optional<ExprDoc> msg = std::nullopt);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(AssertDoc, StmtDoc, AssertDocNode);
};
 
class ReturnDocNode : public StmtDocNode {
 public:
  ExprDoc value{nullptr};
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ReturnDocNode>().def_ro("value", &ReturnDocNode::value);
  }
 
  static constexpr const char* _type_key = "script.printer.ReturnDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(ReturnDocNode, StmtDocNode);
};
 
class ReturnDoc : public StmtDoc {
 public:
  explicit ReturnDoc(ExprDoc value);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ReturnDoc, StmtDoc, ReturnDocNode);
};
 
class FunctionDocNode : public StmtDocNode {
 public:
  IdDoc name{nullptr};
  Array<AssignDoc> args;
  Array<ExprDoc> decorators;
  Optional<ExprDoc> return_type{std::nullopt};
  Array<StmtDoc> body;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<FunctionDocNode>()
        .def_ro("name", &FunctionDocNode::name)
        .def_ro("args", &FunctionDocNode::args)
        .def_ro("decorators", &FunctionDocNode::decorators)
        .def_ro("return_type", &FunctionDocNode::return_type)
        .def_ro("body", &FunctionDocNode::body);
  }
 
  static constexpr const char* _type_key = "script.printer.FunctionDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(FunctionDocNode, StmtDocNode);
};
 
class FunctionDoc : public StmtDoc {
 public:
  explicit FunctionDoc(IdDoc name, Array<AssignDoc> args, Array<ExprDoc> decorators,
                       Optional<ExprDoc> return_type, Array<StmtDoc> body);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(FunctionDoc, StmtDoc, FunctionDocNode);
};
 
class ClassDocNode : public StmtDocNode {
 public:
  IdDoc name{nullptr};
  Array<ExprDoc> decorators;
  Array<StmtDoc> body;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ClassDocNode>()
        .def_ro("name", &ClassDocNode::name)
        .def_ro("decorators", &ClassDocNode::decorators)
        .def_ro("body", &ClassDocNode::body);
  }
 
  static constexpr const char* _type_key = "script.printer.ClassDoc";
 
  TVM_DECLARE_FINAL_OBJECT_INFO(ClassDocNode, StmtDocNode);
};
 
class ClassDoc : public StmtDoc {
 public:
  explicit ClassDoc(IdDoc name, Array<ExprDoc> decorators, Array<StmtDoc> body);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(ClassDoc, StmtDoc, ClassDocNode);
};
 
class CommentDocNode : public StmtDocNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<CommentDocNode>();
  }
 
  static constexpr const char* _type_key = "script.printer.CommentDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(CommentDocNode, StmtDocNode);
};
 
class CommentDoc : public StmtDoc {
 public:
  explicit CommentDoc(String comment);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(CommentDoc, StmtDoc, CommentDocNode);
};
 
class DocStringDocNode : public StmtDocNode {
 public:
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<DocStringDocNode>();
  }
 
  static constexpr const char* _type_key = "script.printer.DocStringDoc";
  TVM_DECLARE_FINAL_OBJECT_INFO(DocStringDocNode, StmtDocNode);
};
 
class DocStringDoc : public StmtDoc {
 public:
  explicit DocStringDoc(String docs);
  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(DocStringDoc, StmtDoc, DocStringDocNode);
};
 
}  // namespace printer
}  // namespace script
}  // namespace tvm
 
#endif  // TVM_SCRIPT_PRINTER_DOC_H_