Fix dictionary lookup for modules (#14)

xoofx · Nov 29, 2024 · 225fdc1 · 225fdc1
1 parent 20b8517
commit 225fdc1
Showing 1 changed file with 28 additions and 30 deletions.
diff --git a/src/Ultra.Core/EtwConverterToFirefox.cs b/src/Ultra.Core/EtwConverterToFirefox.cs
@@ -128,12 +128,6 @@ private FirefoxProfiler.Profile Convert(List<int> processIds)
 
         foreach (var processId in processIds)
         {
-            // Reset all maps and default values before processing a new process
-            _mapModuleFileIndexToFirefox.Clear();
-            _setManagedModules.Clear();
-            _clrJitModuleIndex = ModuleFileIndex.Invalid;
-            _coreClrModuleIndex = ModuleFileIndex.Invalid;
-
             var process = _traceLog.Processes.LastProcessWithID(processId);
 
             ConvertProcess(process);
@@ -559,34 +553,37 @@ private void ConvertProcess(TraceProcess process)
     /// <param name="process">The process to load the modules.</param>
     private void LoadModules(TraceProcess process)
     {
-        _options.LogProgress?.Invoke($"Loading Modules for process {process.Name}");
+        _options.LogProgress?.Invoke($"Loading Modules for process {process.Name} ({process.ProcessID})");
+
+        _setManagedModules.Clear();
+        _clrJitModuleIndex = ModuleFileIndex.Invalid;
+        _coreClrModuleIndex = ModuleFileIndex.Invalid;
 
         var allModules = process.LoadedModules.ToList();
         for (var i = 0; i < allModules.Count; i++)
         {
             var module = allModules[i];
-            if (_mapModuleFileIndexToFirefox.ContainsKey(module.ModuleFile.ModuleFileIndex))
-            {
-                continue; // Skip in case
-            }
-
-            _options.LogStepProgress?.Invoke($"Loading Symbols [{i}/{allModules.Count}] for Module `{module.Name}`, ImageSize: {ByteSize.FromBytes(module.ModuleFile.ImageSize)}");
-
-            var lib = new FirefoxProfiler.Lib
+            if (!_mapModuleFileIndexToFirefox.ContainsKey(module.ModuleFile.ModuleFileIndex))
             {
-                Name = module.Name,
-                AddressStart = module.ImageBase,
-                AddressEnd = module.ModuleFile.ImageEnd,
-                Path = module.ModuleFile.FilePath,
-                DebugPath = module.ModuleFile.PdbName,
-                DebugName = module.ModuleFile.PdbName,
-                BreakpadId = $"0x{module.ModuleID:X16}",
-                Arch = "x64" // TODO
-            };
+                _options.LogStepProgress?.Invoke($"Loading Symbols [{i}/{allModules.Count}] for Module `{module.Name}`, ImageSize: {ByteSize.FromBytes(module.ModuleFile.ImageSize)}");
 
-            _traceLog!.CodeAddresses.LookupSymbolsForModule(_symbolReader, module.ModuleFile);
-            _mapModuleFileIndexToFirefox.Add(module.ModuleFile.ModuleFileIndex, _profile.Libs.Count);
-            _profile.Libs.Add(lib);
+                var lib = new FirefoxProfiler.Lib
+                {
+                    Name = module.Name,
+                    AddressStart = module.ImageBase,
+                    AddressEnd = module.ModuleFile.ImageEnd,
+                    Path = module.ModuleFile.FilePath,
+                    DebugPath = module.ModuleFile.PdbName,
+                    DebugName = module.ModuleFile.PdbName,
+                    BreakpadId = $"0x{module.ModuleID:X16}",
+                    Arch = "x64" // TODO
+                };
+
+                _traceLog!.CodeAddresses.LookupSymbolsForModule(_symbolReader, module.ModuleFile);
+
+                _mapModuleFileIndexToFirefox.Add(module.ModuleFile.ModuleFileIndex, _profile.Libs.Count);
+                _profile.Libs.Add(lib);
+            }
 
             var fileName = Path.GetFileName(module.FilePath);
             if (fileName.Equals("clrjit.dll", StringComparison.OrdinalIgnoreCase))
@@ -832,12 +829,13 @@ private int ConvertMethod(CodeAddressIndex codeAddressIndex, MethodIndex methodI
             funcTable.ColumnNumber.Add(null);
 
             var moduleIndex = _traceLog.CodeAddresses.ModuleFileIndex(codeAddressIndex);
-            if (moduleIndex != ModuleFileIndex.Invalid)
+            if (moduleIndex != ModuleFileIndex.Invalid && _mapModuleFileIndexToFirefox.TryGetValue(moduleIndex, out var firefoxModuleIndex))
             {
-                funcTable.Resource.Add(profileThread.ResourceTable.Length); // TODO
+                funcTable.Resource.Add(profileThread.ResourceTable.Length);
+
                 var moduleName = Path.GetFileName(_traceLog.ModuleFiles[moduleIndex].FilePath);
                 profileThread.ResourceTable.Name.Add(GetOrCreateString(moduleName, profileThread));
-                profileThread.ResourceTable.Lib.Add(_mapModuleFileIndexToFirefox[moduleIndex]);
+                profileThread.ResourceTable.Lib.Add(firefoxModuleIndex);
                 profileThread.ResourceTable.Length++;
             }
             else