From 584a68e523c74e18de7c57dcfb7a38e43daaabfd Mon Sep 17 00:00:00 2001
From: Dalia Shaaban <144673861+dshaaban01@users.noreply.github.com>
Date: Fri, 27 Dec 2024 18:51:17 +0100
Subject: [PATCH] squashed commit

---
 .../substrait-mlir/Dialect/Substrait/IR/CMakeLists.txt   | 7 +++++++
 include/substrait-mlir/Dialect/Substrait/IR/Substrait.h  | 3 +++
 .../Dialect/Substrait/IR/SubstraitTypes.td               | 6 ++++++
 lib/Dialect/Substrait/IR/Substrait.cpp                   | 9 ++++++++-
 4 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/include/substrait-mlir/Dialect/Substrait/IR/CMakeLists.txt b/include/substrait-mlir/Dialect/Substrait/IR/CMakeLists.txt
index 3d8b935..bb01665 100644
--- a/include/substrait-mlir/Dialect/Substrait/IR/CMakeLists.txt
+++ b/include/substrait-mlir/Dialect/Substrait/IR/CMakeLists.txt
@@ -8,6 +8,13 @@ mlir_tablegen(SubstraitEnums.cpp.inc -gen-enum-defs)
 add_public_tablegen_target(MLIRSubstraitEnumsIncGen)
 add_dependencies(MLIRSubstraitDialect MLIRSubstraitEnumsIncGen)
 
+# Add custom type attributes
+set(LLVM_TARGET_DEFINITIONS SubstraitTypes.td)
+mlir_tablegen(SubstraitOpsAttrs.h.inc --gen-attrdef-decls)
+mlir_tablegen(SubstraitOpsAttrs.cpp.inc --gen-attrdef-defs)
+add_public_tablegen_target(MLIRSubstraitAttrsIncGen)
+add_dependencies(MLIRSubstraitDialect MLIRSubstraitAttrsIncGen)
+
 set(LLVM_TARGET_DEFINITIONS SubstraitInterfaces.td)
 mlir_tablegen(SubstraitOpInterfaces.h.inc -gen-op-interface-decls)
 mlir_tablegen(SubstraitOpInterfaces.cpp.inc -gen-op-interface-defs)
diff --git a/include/substrait-mlir/Dialect/Substrait/IR/Substrait.h b/include/substrait-mlir/Dialect/Substrait/IR/Substrait.h
index 7b7413d..710e853 100644
--- a/include/substrait-mlir/Dialect/Substrait/IR/Substrait.h
+++ b/include/substrait-mlir/Dialect/Substrait/IR/Substrait.h
@@ -25,6 +25,9 @@
 #define GET_TYPEDEF_CLASSES
 #include "substrait-mlir/Dialect/Substrait/IR/SubstraitOpsTypes.h.inc" // IWYU: export
 
+#define GET_ATTRDEF_CLASSES
+#include "substrait-mlir/Dialect/Substrait/IR/SubstraitOpsAttrs.h.inc" // IWYU: export
+
 #define GET_OP_CLASSES
 #include "substrait-mlir/Dialect/Substrait/IR/SubstraitOps.h.inc" // IWYU: export
 
diff --git a/include/substrait-mlir/Dialect/Substrait/IR/SubstraitTypes.td b/include/substrait-mlir/Dialect/Substrait/IR/SubstraitTypes.td
index b2e2cae..49416ee 100644
--- a/include/substrait-mlir/Dialect/Substrait/IR/SubstraitTypes.td
+++ b/include/substrait-mlir/Dialect/Substrait/IR/SubstraitTypes.td
@@ -19,6 +19,12 @@ class Substrait_Type<string name, string typeMnemonic, list<Trait> traits = []>
   let mnemonic = typeMnemonic;
 }
 
+// Base class for Substrait dialect attribute types.
+class Substrait_Attr<string name, string typeMnemonic, list<Trait> traits = []>
+    : AttrDef<Substrait_Dialect, name, traits> {
+  let mnemonic = typeMnemonic;
+}
+
 /// Currently supported atomic types. These correspond directly to the types in
 /// https://github.com/substrait-io/substrait/blob/main/proto/substrait/type.proto.
 // TODO(ingomueller): Add the other low-hanging fruits here.
diff --git a/lib/Dialect/Substrait/IR/Substrait.cpp b/lib/Dialect/Substrait/IR/Substrait.cpp
index b15ed8d..c3f9880 100644
--- a/lib/Dialect/Substrait/IR/Substrait.cpp
+++ b/lib/Dialect/Substrait/IR/Substrait.cpp
@@ -32,6 +32,10 @@ void SubstraitDialect::initialize() {
 #define GET_TYPEDEF_LIST
 #include "substrait-mlir/Dialect/Substrait/IR/SubstraitOpsTypes.cpp.inc"
       >();
+  addAttributes<
+#define GET_ATTRDEF_LIST
+#include "substrait-mlir/Dialect/Substrait/IR/SubstraitOpsAttrs.cpp.inc"
+      >();
 }
 
 #include "substrait-mlir/Dialect/Substrait/IR/SubstraitEnums.cpp.inc"
@@ -432,8 +436,11 @@ LogicalResult ProjectOp::verifyRegions() {
 } // namespace mlir
 
 //===----------------------------------------------------------------------===//
-// Substrait types
+// Substrait types and attributes
 //===----------------------------------------------------------------------===//
 
 #define GET_TYPEDEF_CLASSES
 #include "substrait-mlir/Dialect/Substrait/IR/SubstraitOpsTypes.cpp.inc"
+
+#define GET_ATTRDEF_CLASSES
+#include "substrait-mlir/Dialect/Substrait/IR/SubstraitOpsAttrs.cpp.inc"
\ No newline at end of file