1 files changed, 72 insertions, 6 deletions

diff --git a/backend/tools/AssemblyEmitter/AssemblyEmitterMain.cpp b/backend/tools/AssemblyEmitter/AssemblyEmitterMain.cpp
index ff339fa..8697754 100644
--- a/backend/tools/AssemblyEmitter/AssemblyEmitterMain.cpp
+++ b/backend/tools/AssemblyEmitter/AssemblyEmitterMain.cpp
@@ -10,8 +10,10 @@
 #include "rcc/Support/TextDiagnostic.h"
 #include "rcc/Target/RISCV/RISCVAssemblyPrinter.h"
 #include "rcc/Target/RISCV/RISCVCodeGen.h"
+#include "rcc/Target/RISCV/RISCVObjectFile.h"
 #include "rcc/Target/RISCV/RISCVTarget.h"
 
+#include <cstdint>
 #include <fstream>
 #include <iostream>
 #include <iterator>
@@ -19,29 +21,86 @@
 #include <string>
 #include <string_view>
 #include <utility>
+#include <vector>
 
 namespace {
 
+enum class OutputKind : std::uint8_t {
+  Assembly,
+  Object,
+};
+
 struct CommandLine final {
   std::string InputPath;
+  std::string OutputPath = "-";
+  OutputKind Kind = OutputKind::Assembly;
 };
 
+void printUsage() {
+  std::cerr << "usage: rcc-backend-asm --input <mir-file|-> "
+               "[--filetype asm|obj] [--output <file|->]\n";
+}
+
 [[nodiscard]] bool parseCommandLine(int ArgumentCount, char *Arguments[], CommandLine &Result) {
   for (int Index = 1; Index < ArgumentCount; ++Index) {
     const std::string_view Argument(Arguments[Index]);
-    if (Argument != "--input" || Index + 1 >= ArgumentCount) {
-      std::cerr << "usage: rcc-backend-asm --input <mir-file|->\n";
+    if (Argument == "--input" && Index + 1 < ArgumentCount) {
+      Result.InputPath = Arguments[++Index];
+      continue;
+    }
+    if ((Argument == "--output" || Argument == "-o") && Index + 1 < ArgumentCount) {
+      Result.OutputPath = Arguments[++Index];
+      continue;
+    }
+    if (Argument == "--filetype" && Index + 1 < ArgumentCount) {
+      const std::string_view Value(Arguments[++Index]);
+      if (Value == "asm" || Value == "assembly") {
+        Result.Kind = OutputKind::Assembly;
+      } else if (Value == "obj" || Value == "object") {
+        Result.Kind = OutputKind::Object;
+      } else {
+        printUsage();
+        return false;
+      }
+      continue;
+    }
+    {
+      printUsage();
       return false;
     }
-    Result.InputPath = Arguments[++Index];
   }
   if (Result.InputPath.empty()) {
-    std::cerr << "usage: rcc-backend-asm --input <mir-file|->\n";
+    printUsage();
+    return false;
+  }
+  return true;
+}
+
+[[nodiscard]] bool writeOutput(const CommandLine &Command, std::span<const std::uint8_t> Contents) {
+  if (Command.OutputPath == "-") {
+    std::cout.write(reinterpret_cast<const char *>(Contents.data()),
+                    static_cast<std::streamsize>(Contents.size()));
+    return static_cast<bool>(std::cout);
+  }
+  std::ofstream Output(Command.OutputPath, std::ios::binary);
+  if (!Output) {
+    std::cerr << "rcc-backend-asm: cannot open output '" << Command.OutputPath << "'\n";
+    return false;
+  }
+  Output.write(reinterpret_cast<const char *>(Contents.data()),
+               static_cast<std::streamsize>(Contents.size()));
+  if (!Output) {
+    std::cerr << "rcc-backend-asm: cannot write output '" << Command.OutputPath << "'\n";
     return false;
   }
   return true;
 }
 
+[[nodiscard]] bool writeOutput(const CommandLine &Command, std::string_view Contents) {
+  const auto *Data = reinterpret_cast<const std::uint8_t *>(Contents.data());
+  return writeOutput(Command, std::span<const std::uint8_t>(Data, Contents.size()));
+}
+
 [[nodiscard]] bool readInput(const CommandLine &Command, std::string &Contents) {
   if (Command.InputPath == "-") {
     Contents.assign(std::istreambuf_iterator<char>(std::cin), std::istreambuf_iterator<char>());
@@ -160,12 +219,19 @@ int main(int ArgumentCount, char *Arguments[]) {
     }
   }
 
+  if (Command.Kind == OutputKind::Object) {
+    auto Object = rcc::riscv::emitObject(Unit, *Target);
+    if (!Object) {
+      std::cerr << "rcc-backend-asm: object emission failed: " << Object.error().message() << '\n';
+      return 1;
+    }
+    return writeOutput(Command, *Object) ? 0 : 1;
+  }
   auto Assembly = rcc::riscv::emitAssembly(Unit, *Target);
   if (!Assembly) {
     std::cerr << "rcc-backend-asm: assembly emission failed: " << Assembly.error().message()
               << '\n';
     return 1;
   }
-  std::cout << *Assembly;
-  return 0;
+  return writeOutput(Command, *Assembly) ? 0 : 1;
 }