From dfe6b6e0c93f2f469366fa22b172d4fd68fcfee9 Mon Sep 17 00:00:00 2001
From: Mike Urbach <mike.urbach@sifive.com>
Date: Thu, 8 Aug 2024 15:11:12 -0700
Subject: [PATCH] [OM] Add AnyType C API and Python bindings.

We want to expose this type through the Python bindings for isinstance
queries, etc., so add the necessary boilerplate.
---
 include/circt-c/Dialect/OM.h                    | 6 ++++++
 integration_test/Bindings/Python/dialects/om.py | 6 +++++-
 lib/Bindings/Python/OMModule.cpp                | 3 +++
 lib/Bindings/Python/dialects/om.py              | 2 +-
 lib/CAPI/Dialect/OM.cpp                         | 6 ++++++
 5 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/include/circt-c/Dialect/OM.h b/include/circt-c/Dialect/OM.h
index 2248113d98c1..2085edfb9edb 100644
--- a/include/circt-c/Dialect/OM.h
+++ b/include/circt-c/Dialect/OM.h
@@ -25,6 +25,12 @@ MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(OM, om);
 // Type API.
 //===----------------------------------------------------------------------===//
 
+/// Is the Type an AnyType.
+MLIR_CAPI_EXPORTED bool omTypeIsAAnyType(MlirType type);
+
+/// Get the TypeID for an AnyType.
+MLIR_CAPI_EXPORTED MlirTypeID omAnyTypeGetTypeID(void);
+
 /// Is the Type a ClassType.
 MLIR_CAPI_EXPORTED bool omTypeIsAClassType(MlirType type);
 
diff --git a/integration_test/Bindings/Python/dialects/om.py b/integration_test/Bindings/Python/dialects/om.py
index e7afceb1ddc3..bb863a691417 100644
--- a/integration_test/Bindings/Python/dialects/om.py
+++ b/integration_test/Bindings/Python/dialects/om.py
@@ -3,7 +3,7 @@
 
 import circt
 from circt.dialects import om
-from circt.ir import Context, InsertionPoint, Location, Module, IntegerAttr, IntegerType
+from circt.ir import Context, InsertionPoint, Location, Module, IntegerAttr, IntegerType, Type
 from circt.support import var_to_attribute
 
 from dataclasses import dataclass
@@ -307,3 +307,7 @@
       IntegerAttr.get(IntegerType.get_unsigned(64), -42))
   # CHECK: 18446744073709551574
   print(str(int_attr6))
+
+  # Test AnyType
+  any_type = Type.parse("!om.any")
+  assert isinstance(any_type, om.AnyType)
diff --git a/lib/Bindings/Python/OMModule.cpp b/lib/Bindings/Python/OMModule.cpp
index 16fa03873ea1..eea989b73201 100644
--- a/lib/Bindings/Python/OMModule.cpp
+++ b/lib/Bindings/Python/OMModule.cpp
@@ -612,6 +612,9 @@ void circt::python::populateDialectOMSubmodule(py::module &m) {
       .def("__len__", &omMapAttrGetNumElements);
   PyMapAttrIterator::bind(m);
 
+  // Add the AnyType class definition.
+  mlir_type_subclass(m, "AnyType", omTypeIsAAnyType, omAnyTypeGetTypeID);
+
   // Add the ClassType class definition.
   mlir_type_subclass(m, "ClassType", omTypeIsAClassType, omClassTypeGetTypeID)
       .def_property_readonly("name", [](MlirType type) {
diff --git a/lib/Bindings/Python/dialects/om.py b/lib/Bindings/Python/dialects/om.py
index bdf51eed84e1..123c6d104df8 100644
--- a/lib/Bindings/Python/dialects/om.py
+++ b/lib/Bindings/Python/dialects/om.py
@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 from ._om_ops_gen import *
-from .._mlir_libs._circt._om import Evaluator as BaseEvaluator, Object as BaseObject, List as BaseList, Tuple as BaseTuple, Map as BaseMap, BasePath as BaseBasePath, BasePathType, Path, PathType, ClassType, ReferenceAttr, ListAttr, MapAttr, OMIntegerAttr
+from .._mlir_libs._circt._om import AnyType, Evaluator as BaseEvaluator, Object as BaseObject, List as BaseList, Tuple as BaseTuple, Map as BaseMap, BasePath as BaseBasePath, BasePathType, Path, PathType, ClassType, ReferenceAttr, ListAttr, MapAttr, OMIntegerAttr
 
 from ..ir import Attribute, Diagnostic, DiagnosticSeverity, Module, StringAttr, IntegerAttr, IntegerType
 from ..support import attribute_to_var, var_to_attribute
diff --git a/lib/CAPI/Dialect/OM.cpp b/lib/CAPI/Dialect/OM.cpp
index 1263974d5ab8..fd40fb5acac0 100644
--- a/lib/CAPI/Dialect/OM.cpp
+++ b/lib/CAPI/Dialect/OM.cpp
@@ -29,6 +29,12 @@ MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(OM, om, OMDialect)
 // Type API.
 //===----------------------------------------------------------------------===//
 
+/// Is the Type an AnyType.
+bool omTypeIsAAnyType(MlirType type) { return isa<AnyType>(unwrap(type)); }
+
+/// Get the TypeID for an AnyType.
+MlirTypeID omAnyTypeGetTypeID(void) { return wrap(AnyType::getTypeID()); }
+
 /// Is the Type a ClassType.
 bool omTypeIsAClassType(MlirType type) { return isa<ClassType>(unwrap(type)); }