1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19#include "./utils.h"
20
21namespace tvm {
22namespace script {
23namespace printer {
24
25TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
26 .set_dispatch<String>("", [](String s, ObjectPath p, IRDocsifier d) -> Doc {
27 if (HasMultipleLines(s)) {
28 return d->AddMetadata(s);
29 }
30 return LiteralDoc::Str(s, p);
31 });
32
33TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
34 .set_dispatch<Array<ObjectRef>>( //
35 "", [](Array<ObjectRef> array, ObjectPath p, IRDocsifier d) -> Doc {
36 int n = array.size();
37 Array<ExprDoc> results;
38 results.reserve(n);
39 for (int i = 0; i < n; ++i) {
40 results.push_back(d->AsDoc<ExprDoc>(array[i], p->ArrayIndex(i)));
41 }
42 return ListDoc(results);
43 });
44
45TVM_STATIC_IR_FUNCTOR(IRDocsifier, vtable)
46 .set_dispatch<Map<ObjectRef, ObjectRef>>( //
47 "", [](Map<ObjectRef, ObjectRef> dict, ObjectPath p, IRDocsifier d) -> Doc {
48 using POO = std::pair<ObjectRef, ObjectRef>;
49 std::vector<POO> items{dict.begin(), dict.end()};
50 bool is_str_map = true;
51 for (const auto& kv : items) {
52 if (!kv.first.as<runtime::StringObj>()) {
53 is_str_map = false;
54 break;
55 }
56 }
57 if (is_str_map) {
58 std::sort(items.begin(), items.end(), [](const POO& lhs, const POO& rhs) {
59 return Downcast<String>(lhs.first) < Downcast<String>(rhs.first);
60 });
61 } else {
62 std::sort(items.begin(), items.end(), [](const POO& lhs, const POO& rhs) {
63 return lhs.first.get() < rhs.first.get();
64 });
65 }
66 int n = dict.size();
67 Array<ExprDoc> ks;
68 Array<ExprDoc> vs;
69 ks.reserve(n);
70 vs.reserve(n);
71 for (int i = 0; i < n; ++i) {
72 ks.push_back(d->AsDoc<ExprDoc>(items[i].first, p->MissingMapEntry()));
73 vs.push_back(d->AsDoc<ExprDoc>(items[i].second, p->MapValue(items[i].first)));
74 }
75 return DictDoc(ks, vs);
76 });
77
78} // namespace printer
79} // namespace script
80} // namespace tvm
81