Fix ARM64 profile disassembly CFG parsing

src/ProfileExplorerCore/Compilers/ASM/DisassemblerSectionLoader.cs

@@ -2,7 +2,9 @@
 // Licensed under the MIT license.
 using System;
 using System.Collections.Generic;
+using System.Reflection.PortableExecutable;
 using ProfileExplorer.Core.Binary;
+using ProfileExplorer.Core.Compilers.Architecture;
 using ProfileExplorer.Core.IR;
 using ProfileExplorer.Core.Providers;
 
@@ -24,8 +26,8 @@ public sealed class DisassemblerSectionLoader : IRTextSectionLoader {
 
   public DisassemblerSectionLoader(string binaryFilePath, ICompilerInfoProvider compilerInfo, IDebugInfoProvider debugInfo,
     bool preloadFunctions = true, bool isManagedImage = false) {
-    Initialize(compilerInfo.IR, false);
     binaryFilePath_ = binaryFilePath;
+    Initialize(SelectIRInfo(compilerInfo, binaryFilePath), false);
     debugInfo_ = debugInfo;
     debugFileFinder_ = compilerInfo.DebugFileFinder;
     debugInfoProviderFactory_ = compilerInfo.DebugInfoProviderFactory;
@@ -36,6 +38,17 @@ public DisassemblerSectionLoader(string binaryFilePath, ICompilerInfoProvider co
     funcToDebugInfoMap_ = new Dictionary<IRTextFunction, FunctionDebugInfo>();
   }
 
+  private static ICompilerIRInfo SelectIRInfo(ICompilerInfoProvider compilerInfo, string binaryFilePath) {
+    var binaryInfo = PEBinaryInfoProvider.GetBinaryFileInfo(binaryFilePath);
+
+    return binaryInfo?.Architecture switch {
+      Machine.Arm64 => new ASMCompilerIRInfo(IRMode.ARM64),
+      Machine.I386  => new ASMCompilerIRInfo(IRMode.x86_64),
+      Machine.Amd64 => new ASMCompilerIRInfo(IRMode.x86_64),
+      _             => compilerInfo.IR
+    };
+  }
+
   public IDebugInfoProvider DebugInfo {
     get => debugInfo_;
     set => debugInfo_ = value;
@@ -233,4 +246,4 @@ protected override void Dispose(bool disposing) {
     disassembler_?.Dispose();
     debugInfo_?.Dispose();
   }
-}
+}

src/ProfileExplorerCoreTests/ASMParserTests.cs

@@ -0,0 +1,41 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+using ProfileExplorer.Core;
+using ProfileExplorer.Core.Compilers.ASM;
+using ProfileExplorer.Core.Compilers.Architecture;
+using ProfileExplorer.Core.IR;
+
+namespace ProfileExplorer.CoreTests;
+
+[TestClass]
+public class ASMParserTests {
+  [TestMethod]
+  public void Parse_ARM64ConditionalBranch_CreatesBasicBlocksAndCFGEdges() {
+    string text =
+      """
+      1000:    54000040  b.eq #0x1008
+      1004:    d503201f  nop
+      1008:    d65f03c0  ret
+      """;
+
+    var function = ParseARM64(text, 12);
+
+    Assert.AreEqual(3, function.BlockCount);
+    Assert.AreEqual(2, function.EntryBlock.Successors.Count);
+    CollectionAssert.Contains(function.EntryBlock.Successors, function.Blocks[1]);
+    CollectionAssert.Contains(function.EntryBlock.Successors, function.Blocks[2]);
+    Assert.IsTrue(function.EntryBlock.FirstInstruction.IsBranch);
+    Assert.IsTrue(function.Blocks[2].FirstInstruction.IsReturn);
+  }
+
+  private static FunctionIR ParseARM64(string text, long functionSize) {
+    var irInfo = new ASMCompilerIRInfo(IRMode.ARM64);
+    var parentFunction = new IRTextFunction("test");
+    var output = new IRPassOutput(0, text.Length, 1, text.Split('\n').Length);
+    var section = new IRTextSection(parentFunction, parentFunction.Name, output);
+    var parser = new ASMIRSectionParser(functionSize, irInfo, irInfo.CreateParsingErrorHandler());
+
+    return parser.ParseSection(section, text);
+  }
+}