2 files changed, 22 insertions, 0 deletions
diff --git a/gcc/rust/util/rust-hir-map.cc b/gcc/rust/util/rust-hir-map.cc
index c62c111..05b26d6 100644
--- a/gcc/rust/util/rust-hir-map.cc
+++ b/gcc/rust/util/rust-hir-map.cc
@@ -17,9 +17,11 @@
 // <http://www.gnu.org/licenses/>.
 
 #include "rust-hir-map.h"
+#include "optional.h"
 #include "rust-ast-full.h"
 #include "rust-diagnostics.h"
 #include "rust-hir-full.h"
+#include "rust-item.h"
 #include "rust-macro-builtins.h"
 #include "rust-mapping-common.h"
 #include "rust-attribute-values.h"
@@ -1159,6 +1161,23 @@ Mappings::lookup_module_children (NodeId module)
 }
 
 void
+Mappings::insert_ast_module (AST::Module *module)
+{
+  rust_assert (modules.find (module->get_node_id ()) == modules.end ());
+  modules[module->get_node_id ()] = module;
+}
+
+tl::optional<AST::Module *>
+Mappings::lookup_ast_module (NodeId id)
+{
+  auto it = modules.find (id);
+  if (it == modules.end ())
+    return tl::nullopt;
+
+  return {it->second};
+}
+
+void
 Mappings::insert_module_child_item (NodeId module,
 				    Resolver::CanonicalPath child)
 {
diff --git a/gcc/rust/util/rust-hir-map.h b/gcc/rust/util/rust-hir-map.h
index 927c012..92af067 100644
--- a/gcc/rust/util/rust-hir-map.h
+++ b/gcc/rust/util/rust-hir-map.h
@@ -328,6 +328,8 @@ public:
   void insert_visibility (NodeId id, Privacy::ModuleVisibility visibility);
   bool lookup_visibility (NodeId id, Privacy::ModuleVisibility &def);
 
+  void insert_ast_module (AST::Module *);
+  tl::optional<AST::Module *> lookup_ast_module (NodeId id);
   void insert_module_child (NodeId module, NodeId child);
   tl::optional<std::vector<NodeId> &> lookup_module_children (NodeId module);
 
@@ -427,6 +429,7 @@ private:
   std::map<NodeId, std::vector<NodeId>> module_child_map;
   std::map<NodeId, std::vector<Resolver::CanonicalPath>> module_child_items;
   std::map<NodeId, NodeId> child_to_parent_module_map;
+  std::map<NodeId, AST::Module *> modules;
 
   // AST mappings
   std::map<NodeId, AST::Item *> ast_item_mappings;