var_table.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_VAR_TABLE_H_
#define TVM_SCRIPT_PRINTER_VAR_TABLE_H_

#include <tvm/node/node.h>
#include <tvm/node/object_path.h>
#include <tvm/script/printer/doc.h>
#include <tvm/script/printer/frame.h>
#include <tvm/script/printer/traced_object.h>

#include <unordered_map>
#include <unordered_set>

namespace tvm {
namespace script {
namespace printer {

class VarTableNode : public Object {
 public:
  void VisitAttrs(AttrVisitor*) {}

  IdDoc Define(const ObjectRef& obj, const String& name_hint, const ObjectPath& object_path,
               const Frame& frame);

  IdDoc Define(const ObjectRef& obj, const TracedObject<String>& name_hint, const Frame& frame) {
    return Define(obj, name_hint.Get(), name_hint.GetPath(), frame);
  }

  using DocFactory = std::function<ExprDoc()>;

  void DefineByDoc(const ObjectRef& obj, DocFactory doc_factory, const Frame& frame);

  Optional<ExprDoc> GetVarDoc(const ObjectRef& obj, const ObjectPath& object_path) const;

  template <typename TObjectRef>
  Optional<ExprDoc> GetVarDoc(const TracedObject<TObjectRef> obj) const {
    return GetVarDoc(obj.Get(), obj.GetPath());
  }

  bool IsVarDefined(const ObjectRef& obj) const;

  static constexpr const char* _type_key = "script.printer.VarTable";
  TVM_DECLARE_FINAL_OBJECT_INFO(VarTableNode, Object);

 private:
  void RemoveVar(const ObjectRef& obj);

  struct VariableInfo {
    DocFactory doc_factory;
    Optional<String> name;
  };
  std::unordered_map<ObjectRef, VariableInfo, ObjectPtrHash, ObjectPtrEqual> obj2info;
  std::unordered_set<String> defined_names;
};

class VarTable : public ObjectRef {
 public:
  VarTable();
  TVM_DEFINE_MUTABLE_NOTNULLABLE_OBJECT_REF_METHODS(VarTable, ObjectRef, VarTableNode);
};

}  // namespace printer
}  // namespace script
}  // namespace tvm

#endif  // TVM_SCRIPT_PRINTER_VAR_TABLE_H_