Use native AddIncludeFile method (#18)

raviqqe · web-flow · commit 6495e35af1f3 · 2025-01-07T13:08:30.000Z
This PR changes the `add_source_file` method in Rust to use `AddIncludeFile` in C++ which resolves given file paths from include directories. This PR depends on #19. # References - mlir-rs/melior#654 - https://llvm.org/doxygen/classllvm_1_1SourceMgr.html
diff --git a/cc/include/TableGen.h b/cc/include/TableGen.h
@@ -26,10 +26,10 @@ extern "C" {
 #endif
 
 typedef enum {
-    TABLEGEN_DK_ERROR,
-    TABLEGEN_DK_WARNING,
-    TABLEGEN_DK_REMARK,
-    TABLEGEN_DK_NOTE,
+  TABLEGEN_DK_ERROR,
+  TABLEGEN_DK_WARNING,
+  TABLEGEN_DK_REMARK,
+  TABLEGEN_DK_NOTE,
 } TableGenDiagKind;
 
 typedef enum {
@@ -54,9 +54,8 @@ typedef void (*TableGenStringCallback)(TableGenStringRef, void *);
 TableGenParserRef tableGenGet();
 void tableGenFree(TableGenParserRef tg_ref);
 TableGenBool tableGenAddSource(TableGenParserRef tg_ref, const char *source);
-TableGenBool tableGenAddSourceFile(TableGenParserRef tg_ref,
-                                   TableGenStringRef source);
-void tableGenAddIncludePath(TableGenParserRef tg_ref,
+void tableGenAddSourceFile(TableGenParserRef tg_ref, TableGenStringRef source);
+void tableGenAddIncludeDirectory(TableGenParserRef tg_ref,
                             TableGenStringRef include);
 
 /// NOTE: TableGen currently relies on global state within a given parser
@@ -162,14 +161,17 @@ TableGenBool tableGenIntInitGetValue(TableGenTypedInitRef ti, int64_t *integer);
 TableGenStringRef tableGenStringInitGetValue(TableGenTypedInitRef ti);
 char *tableGenStringInitGetValueNewString(TableGenTypedInitRef ti);
 TableGenRecordRef tableGenDefInitGetValue(TableGenTypedInitRef ti);
-void tableGenInitPrint(TableGenTypedInitRef ti,
-                         TableGenStringCallback callback, void *userData);
+void tableGenInitPrint(TableGenTypedInitRef ti, TableGenStringCallback callback,
+                       void *userData);
 void tableGenInitDump(TableGenTypedInitRef ti);
-TableGenBool tableGenPrintError(TableGenParserRef ref, TableGenSourceLocationRef loc_ref, TableGenDiagKind dk,
-                        TableGenStringRef message,
-                        TableGenStringCallback callback, void *userData);
+TableGenBool tableGenPrintError(TableGenParserRef ref,
+                                TableGenSourceLocationRef loc_ref,
+                                TableGenDiagKind dk, TableGenStringRef message,
+                                TableGenStringCallback callback,
+                                void *userData);
 TableGenSourceLocationRef tableGenSourceLocationNull();
-TableGenSourceLocationRef tableGenSourceLocationClone(TableGenSourceLocationRef loc_ref);
+TableGenSourceLocationRef
+tableGenSourceLocationClone(TableGenSourceLocationRef loc_ref);
 
 // Memory
 void tableGenSourceLocationFree(TableGenSourceLocationRef loc_ref);
diff --git a/cc/lib/TableGen.cpp b/cc/lib/TableGen.cpp
@@ -19,6 +19,14 @@ using ctablegen::tableGenFromRecType;
 RecordKeeper *ctablegen::TableGenParser::parse() {
   auto recordKeeper = new RecordKeeper;
   sourceMgr.setIncludeDirs(includeDirs);
+
+  for (const auto &file : files) {
+    std::string full_path;
+    if (!sourceMgr.AddIncludeFile(file, SMLoc(), full_path)) {
+      return nullptr;
+    }
+  }
+
   bool result = TableGenParseFile(sourceMgr, *recordKeeper);
   if (!result) {
     return recordKeeper;
@@ -27,7 +35,7 @@ RecordKeeper *ctablegen::TableGenParser::parse() {
   return nullptr;
 }
 
-void ctablegen::TableGenParser::addIncludePath(const StringRef include) {
+void ctablegen::TableGenParser::addIncludeDirectory(const StringRef include) {
   includeDirs.push_back(std::string(include));
 }
 
@@ -43,16 +51,8 @@ bool ctablegen::TableGenParser::addSource(const char *source) {
   return true;
 }
 
-bool ctablegen::TableGenParser::addSourceFile(const StringRef source) {
-  ErrorOr<std::unique_ptr<MemoryBuffer>> FileOrErr =
-      MemoryBuffer::getFile(source);
-
-  if (std::error_code EC = FileOrErr.getError()) {
-    return false;
-  }
-
-  sourceMgr.AddNewSourceBuffer(std::move(*FileOrErr), SMLoc());
-  return true;
+void ctablegen::TableGenParser::addSourceFile(const StringRef file) {
+  files.push_back(std::string(file));
 }
 
 TableGenParserRef tableGenGet() {
@@ -61,18 +61,18 @@ TableGenParserRef tableGenGet() {
 
 void tableGenFree(TableGenParserRef tg_ref) { delete unwrap(tg_ref); }
 
-TableGenBool tableGenAddSourceFile(TableGenParserRef tg_ref,
-                                   TableGenStringRef source) {
-  return unwrap(tg_ref)->addSourceFile(StringRef(source.data, source.len));
+void tableGenAddSourceFile(TableGenParserRef tg_ref, TableGenStringRef source) {
+  unwrap(tg_ref)->addSourceFile(StringRef(source.data, source.len));
 }
 
 TableGenBool tableGenAddSource(TableGenParserRef tg_ref, const char *source) {
   return unwrap(tg_ref)->addSource(source);
 }
 
-void tableGenAddIncludePath(TableGenParserRef tg_ref,
+void tableGenAddIncludeDirectory(TableGenParserRef tg_ref,
                             TableGenStringRef include) {
-  return unwrap(tg_ref)->addIncludePath(StringRef(include.data, include.len));
+  return unwrap(tg_ref)->addIncludeDirectory(
+      StringRef(include.data, include.len));
 }
 
 TableGenRecordKeeperRef tableGenParse(TableGenParserRef tg_ref) {
diff --git a/cc/lib/TableGen.hpp b/cc/lib/TableGen.hpp
@@ -38,13 +38,15 @@ class TableGenParser {
 public:
   TableGenParser() {}
   bool addSource(const char *source);
-  bool addSourceFile(const StringRef source);
-  void addIncludePath(const StringRef include);
+  void addSourceFile(const StringRef source);
+  void addIncludeDirectory(const StringRef include);
   RecordKeeper *parse();
 
   SourceMgr sourceMgr;
+
 private:
   std::vector<std::string> includeDirs;
+  std::vector<std::string> files;
 };
 
 // Utility
@@ -60,7 +62,7 @@ class CallbackOstream : public llvm::raw_ostream {
         opaqueData(opaqueData), pos(0u) {}
 
   void write_impl(const char *ptr, size_t size) override {
-    TableGenStringRef string = TableGenStringRef { .data = ptr, .len = size };
+    TableGenStringRef string = TableGenStringRef{.data = ptr, .len = size};
     callback(string, opaqueData);
     pos += size;
   }
@@ -75,7 +77,8 @@ class CallbackOstream : public llvm::raw_ostream {
 
 } // namespace ctablegen
 
-DEFINE_SIMPLE_CONVERSION_FUNCTIONS(ctablegen::TableGenParser, TableGenParserRef);
+DEFINE_SIMPLE_CONVERSION_FUNCTIONS(ctablegen::TableGenParser,
+                                   TableGenParserRef);
 DEFINE_SIMPLE_CONVERSION_FUNCTIONS(RecordKeeper, TableGenRecordKeeperRef);
 
 DEFINE_SIMPLE_CONVERSION_FUNCTIONS(ctablegen::RecordMap, TableGenRecordMapRef);
diff --git a/src/lib.rs b/src/lib.rs
@@ -62,7 +62,25 @@
 //! # fn main() -> Result<(), Box<dyn std::error::Error>> {
 //! let keeper: RecordKeeper = TableGenParser::new()
 //!     .add_source(r#"include "mlir/IR/OpBase.td""#)?
-//!     .add_include_path(&format!("{}/include", std::env::var("TABLEGEN_190_PREFIX")?))
+//!     .add_include_directory(&format!("{}/include", std::env::var("TABLEGEN_190_PREFIX")?))
+//!     .parse()?;
+//! let i32_def = keeper.def("I32").expect("has I32 def");
+//! assert!(i32_def.subclass_of("I"));
+//! assert_eq!(i32_def.int_value("bitwidth"), Ok(32));
+//! # Ok(())
+//! # }
+//! ```
+//!
+//! You can also pass an included filename directly.
+//!
+//! ```rust
+//! use tblgen_alt::{TableGenParser, RecordKeeper};
+//! use std::path::Path;
+//!
+//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
+//! let keeper: RecordKeeper = TableGenParser::new()
+//!     .add_source_file("mlir/IR/OpBase.td")
+//!     .add_include_directory(&format!("{}/include", std::env::var("TABLEGEN_190_PREFIX")?))
 //!     .parse()?;
 //! let i32_def = keeper.def("I32").expect("has I32 def");
 //! assert!(i32_def.subclass_of("I"));
@@ -108,8 +126,8 @@ pub use record::RecordValue;
 pub use record_keeper::RecordKeeper;
 
 use raw::{
-    tableGenAddIncludePath, tableGenAddSource, tableGenAddSourceFile, tableGenFree, tableGenGet,
-    tableGenParse, TableGenParserRef,
+    tableGenAddIncludeDirectory, tableGenAddSource, tableGenAddSourceFile, tableGenFree,
+    tableGenGet, tableGenParse, TableGenParserRef,
 };
 use string_ref::StringRef;
 
@@ -144,18 +162,15 @@ impl<'s> TableGenParser<'s> {
     }
 
     /// Adds the given path to the list of included directories.
-    pub fn add_include_path(self, include: &str) -> Self {
-        unsafe { tableGenAddIncludePath(self.raw, StringRef::from(include).to_raw()) }
+    pub fn add_include_directory(self, include: &str) -> Self {
+        unsafe { tableGenAddIncludeDirectory(self.raw, StringRef::from(include).to_raw()) }
         self
     }
 
     /// Reads TableGen source code from the file at the given path.
-    pub fn add_source_file(self, source: &str) -> Result<Self, Error> {
-        if unsafe { tableGenAddSourceFile(self.raw, StringRef::from(source).to_raw()) > 0 } {
-            Ok(self)
-        } else {
-            Err(TableGenError::InvalidSource.into())
-        }
+    pub fn add_source_file(self, source: &str) -> Self {
+        unsafe { tableGenAddSourceFile(self.raw, StringRef::from(source).to_raw()) }
+        self
     }
 
     /// Adds the given TableGen source string.
