#pragma once #include #include namespace torch { namespace jit { namespace script { // Simple data structure for containing a type T in nested control blocks // Should only be used after initial compilation where type checking and // loads and stores are emitted template struct MiniEnvironment { MiniEnvironment(Block* b, std::shared_ptr next = nullptr) : next(std::move(next)) {} std::shared_ptr> next; T findInThisFrame(const std::string& name) { auto it = table.find(name); if (it != table.end()) { return it->second; } return nullptr; } T findInAnyFrame(const std::string& name) { for (auto runner = this; runner; runner = runner->next.get()) { if (auto r = runner->findInThisFrame(name)) { return r; } } return nullptr; } void setVar(const std::string& name, T value) { table[name] = value; } std::vector definedVariables() { std::vector result; for (auto& kv : table) { result.push_back(kv.first); } std::sort(result.begin(), result.end()); return result; } private: std::unordered_map table; }; } // namespace script } // namespace jit } // namespace torch