doc.cc 3.09 KB
Newer Older
1 2 3 4 5 6 7 8
/*
 * 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
Yinghai Lu committed
9
 *
10
 *   http://www.apache.org/licenses/LICENSE-2.0
Yinghai Lu committed
11
 *
12 13 14 15 16 17 18 19
 * 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.
 */

20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52
/*!
 * \file src/tvm/relay/doc.cc
 * \brief Doc ADT used for pretty printing.
 * Based on Section 1 of https://homepages.inf.ed.ac.uk/wadler/papers/prettier/prettier.pdf.
 */
#include <memory>
#include <vector>
#include "doc.h"

namespace tvm {
namespace relay {

// Text constructor
DocAtom Text(const std::string& str) {
  return std::make_shared<TextNode>(str);
}

// Line constructor
DocAtom Line(int indent = 0) {
  return std::make_shared<LineNode>(indent);
}

Doc::Doc(const std::string& str) {
  if (str == "\n") {
    this->stream_ = {Line()};
  } else {
    this->stream_ = {Text(str)};
  }
}

// DSL function implementations

Doc& Doc::operator<<(const Doc& right) {
Yinghai Lu committed
53
  CHECK(this != &right);
54 55 56 57 58 59 60 61
  this->stream_.insert(this->stream_.end(), right.stream_.begin(), right.stream_.end());
  return *this;
}

Doc& Doc::operator<<(const std::string& right) {
  return *this << Doc(right);
}

62 63 64 65 66
Doc& Doc::operator<<(const DocAtom& right) {
  this->stream_.push_back(right);
  return *this;
}

67 68 69 70 71 72 73
Doc Indent(int indent, const Doc& doc) {
  Doc ret;
  for (auto atom : doc.stream_) {
    if (auto text = std::dynamic_pointer_cast<TextNode>(atom)) {
      ret.stream_.push_back(text);
    } else if (auto line = std::dynamic_pointer_cast<LineNode>(atom)) {
      ret.stream_.push_back(Line(indent + line->indent));
Yinghai Lu committed
74
    } else {CHECK(false);}
75 76 77 78 79 80 81 82 83 84 85
  }
  return ret;
}

std::string Doc::str() {
  std::ostringstream os;
  for (auto atom : this->stream_) {
    if (auto text = std::dynamic_pointer_cast<TextNode>(atom)) {
      os << text->str;
    } else if (auto line = std::dynamic_pointer_cast<LineNode>(atom)) {
      os << "\n" << std::string(line->indent, ' ');
Yinghai Lu committed
86
    } else {CHECK(false);}
87 88 89 90
  }
  return os.str();
}

91
Doc PrintSep(const std::vector<Doc>& vec, const Doc& sep) {
92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110
  Doc seq;
  if (vec.size() != 0) {
    seq = vec[0];
    for (size_t i = 1; i < vec.size(); i++) {
      seq << sep << vec[i];
    }
  }
  return seq;
}

Doc PrintBool(bool value) {
  if (value) {
    return Doc("True");
  } else {
    return Doc("False");
  }
}

Doc PrintDType(DataType dtype) {
111
  return Doc(runtime::TVMType2String(dtype));
112 113 114 115 116 117 118 119
}

Doc PrintString(const std::string& value) {
  // TODO(M.K.): add escape.
  Doc doc;
  return doc << "\"" << value << "\"";
}

120 121 122 123 124
Doc PrintNewLine(int ident) {
  Doc doc;
  return doc << Line(ident);
}

125 126
}  // namespace relay
}  // namespace tvm