chore: reorganize IntegerType testing in export file (#53)

substrait-io · Jan 9, 2025 · fe77d82 · fe77d82
1 parent 7234dc3
commit fe77d82
Showing 1 changed file with 14 additions and 5 deletions.
diff --git a/lib/Target/SubstraitPB/Export.cpp b/lib/Target/SubstraitPB/Export.cpp
@@ -89,12 +89,11 @@ class SubstraitExporter {
   std::unique_ptr<SymbolTable> symbolTable;     // Symbol table cache.
 };
 
-FailureOr<std::unique_ptr<proto::Type>>
-SubstraitExporter::exportType(Location loc, mlir::Type mlirType) {
-  MLIRContext *context = mlirType.getContext();
+std::unique_ptr<proto::Type> exportIntegerType(mlir::Type mlirType,
+                                               MLIRContext *context) {
+  // Function that handles `IntegerType`'s.
 
-  // TODO (daliashaaban): Reorganize, test isa<IntegerType>(...) first, then
-  // handle cases. Handle SI1.
+  // Handle SI1.
   auto si1 = IntegerType::get(context, 1, IntegerType::Signed);
   if (mlirType == si1) {
     // TODO(ingomueller): support other nullability modes.
@@ -158,6 +157,16 @@ SubstraitExporter::exportType(Location loc, mlir::Type mlirType) {
     type->set_allocated_i64(i64Type.release());
     return std::move(type);
   }
+}
+
+FailureOr<std::unique_ptr<proto::Type>>
+SubstraitExporter::exportType(Location loc, mlir::Type mlirType) {
+  MLIRContext *context = mlirType.getContext();
+
+  // Handle `IntegerType`'s.
+  if (mlirType.isa<IntegerType>()) {
+    return exportIntegerType(mlirType, context);
+  }
 
   // TODO (daliashaaban): Reorganize, test isa<FloatType>(...) first, then
   // handle cases. Handle FP32.