[Heavy] Fix mangle parsing for module file

Author: ricejasonf

heavy/lib/Mangle.cpp

@@ -235,7 +235,13 @@ bool consumeNameSegment(llvm::StringRef& Buf) {
 //                     but is undefined and may trigger assertions.
 llvm::StringRef Mangler::parseModulePrefix(llvm::StringRef Name) {
   llvm::StringRef Ending = Name.drop_front(getManglePrefix().size());
+  size_t PrevEndingSize = 0;
   while (Ending.size() > 0) {
+    // Check for progress.
+    if (Ending.size() == PrevEndingSize)
+      return {};
+    PrevEndingSize = Ending.size();
+
     switch (Ending[0]) {
       case 'V': case 'F': case 'A':
         // Break from the loop
@@ -244,6 +250,12 @@ llvm::StringRef Mangler::parseModulePrefix(llvm::StringRef Name) {
         // Library name
         Ending = Ending.drop_front(1);
         continue;
+      case '_': {
+        llvm::StringRef ModuleFile("__module_file");
+        if (Ending.starts_with(ModuleFile))
+          Ending.consume_front(ModuleFile);
+        continue;
+      }
       default:
         if (!consumeNameSegment(Ending))
           return {};  // Something was invalid.

heavy/lib/OpEval.cpp

@@ -408,7 +408,9 @@ class OpEvalImpl {
     llvm::StringRef ModuleName = Mangler::parseModulePrefix(MangledName);
     auto M = dyn_cast_or_null<mlir::ModuleOp>(
         TopModuleOp.lookupSymbol(ModuleName));
-    auto F = dyn_cast_or_null<heavy::FuncOp>(M.lookupSymbol(MangledName));
+    heavy::FuncOp F;
+    if (M)
+      F = dyn_cast_or_null<heavy::FuncOp>(M.lookupSymbol(MangledName));
     if (!F) {
       String* ErrMsg = Context.CreateString(
           "undefined reference to function ", MangledName);