Add support for struct type & corresponding operations (#53)

We are likely missing something like `struct_inject` for assigning a field for struct values, but this could be done later.

Lots of small fixes and refactorings while there.

---------

Signed-off-by: Anton Korobeynikov <anton@korobeynikov.info>

include/p4mlir/Dialect/P4HIR/CMakeLists.txt

@@ -14,3 +14,11 @@ mlir_tablegen(P4HIR_Attrs.cpp.inc -gen-attrdef-defs -attrdefs-dialect=p4hir)
 
 add_public_tablegen_target(P4MLIR_P4HIR_IncGen)
 add_dependencies(mlir-headers P4MLIR_P4HIR_IncGen)
+
+set(LLVM_TARGET_DEFINITIONS P4HIR_TypeInterfaces.td)
+mlir_tablegen(P4HIR_TypeInterfaces.h.inc -gen-type-interface-decls)
+mlir_tablegen(P4HIR_TypeInterfaces.cpp.inc -gen-type-interface-defs)
+
+add_public_tablegen_target(P4MLIR_P4HIR_TypeInterfacesIncGen)
+add_dependencies(mlir-headers P4MLIR_P4HIR_TypeInterfacesIncGen)
+

include/p4mlir/Dialect/P4HIR/P4HIR.td

@@ -4,5 +4,6 @@
 include "p4mlir/Dialect/P4HIR/P4HIR_Dialect.td"
 include "p4mlir/Dialect/P4HIR/P4HIR_Ops.td"
 include "p4mlir/Dialect/P4HIR/P4HIR_Types.td"
+include "p4mlir/Dialect/P4HIR/P4HIR_TypeInterfaces.td"
 
 #endif // P4MLIR_DIALECT_P4HIR_P4HIR_TD

include/p4mlir/Dialect/P4HIR/P4HIR_Attrs.td

@@ -59,6 +59,35 @@ def P4HIR_IntAttr : P4HIR_Attr<"Int", "int", [TypedAttrInterface]> {
   let hasCustomAssemblyFormat = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// AggAttr
+//===----------------------------------------------------------------------===//
+
+def P4HIR_AggAttr : P4HIR_Attr<"Agg", "aggregate", [TypedAttrInterface]> {
+  let summary = "An Attribute containing an aggregate value";
+  let description = [{
+    An aggregate attribute is a literal attribute that represents an aggregate
+    value of the specified type. For nested aggregates, embedded arrays are
+    used.
+  }];
+  let parameters = (ins AttributeSelfTypeParameter<"">:$type,
+                        "mlir::ArrayAttr":$fields);
+
+  let builders = [
+    AttrBuilderWithInferredContext<(ins "mlir::Type":$type,
+                                        "mlir::ArrayAttr":$members), [{
+      return $_get(type.getContext(), type, members);
+    }]>
+  ];
+  // let genVerifyDecl = 1;
+  //let hasCustomAssemblyFormat = 1;
+  let assemblyFormat = [{
+    `<` $fields `>`
+  }];
+
+}
+
+
 //===----------------------------------------------------------------------===//
 // ParamDirAttr
 //===----------------------------------------------------------------------===//

include/p4mlir/Dialect/P4HIR/P4HIR_Ops.td

@@ -853,4 +853,98 @@ def CallOp : P4HIR_Op<"call",
   }];
 }
 
+def StructOp : P4HIR_Op<"struct",
+  [Pure,
+   DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>]> {
+  let summary = "Create a struct from constituent parts.";
+  // FIXME: Better constraint type
+  let arguments = (ins Variadic<AnyP4Type>:$input);
+  let results = (outs StructType:$result);
+  let hasCustomAssemblyFormat = 1;
+  let hasVerifier = 1;
+}
+
+def StructExtractOp : P4HIR_Op<"struct_extract",
+    [Pure,
+     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>
+  ]> {
+  let summary = "Extract a named field from a struct.";
+  let description = [{
+    ```
+    %result = p4hir.struct_extract %input["field"] : !p4hir.struct<field: type>
+    TODO: Support the nested extractions
+    %result2 = p4hir.struct_extract %input["field1", "field2"] : !p4hir.struct<field1: !p4hir.struct<field2: type>>
+
+    ```
+  }];
+
+  let arguments = (ins StructType:$input, I32Attr:$fieldIndex);
+  // FIXME: Better constraint type
+  let results = (outs AnyP4Type:$result);
+  let hasCustomAssemblyFormat = 1;
+  let hasVerifier = 1;
+
+  let builders = [
+    OpBuilder<(ins "mlir::Value":$input, "StructType::FieldInfo":$field)>,
+    OpBuilder<(ins "mlir::Value":$input, "mlir::StringAttr":$fieldName)>,
+    OpBuilder<(ins "mlir::Value":$input, "llvm::StringRef":$fieldName), [{
+      build($_builder, $_state, input, $_builder.getStringAttr(fieldName));
+    }]>
+  ];
+
+  let extraClassDeclaration = [{
+    /// Return the name attribute of the accessed field.
+    mlir::StringAttr getFieldNameAttr() {
+      StructType type = getInput().getType();
+      return type.getElements()[getFieldIndex()].name;
+    }
+
+    /// Return the name of the accessed field.
+    llvm::StringRef getFieldName() {
+      return getFieldNameAttr().getValue();
+    }
+  }];
+}
+
+def StructExtractRefOp : P4HIR_Op<"struct_extract_ref",
+    [Pure,
+    DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmResultNames"]>
+  ]> {
+  let summary = "Project from a struct reference to a reference to a named struct field";
+  let description = [{
+    ```
+    %result = p4hir.struct_extract_ref %input["field"] : <!p4hir.struct<field: type>>
+    TODO: Support the nested extractions
+    %result2 = p4hir.struct_extract_ref %input["field1", "field2"] : <!p4hir.struct<field1: !p4hir.struct<field2: type>>>
+
+    ```
+  }];
+
+  let arguments = (ins StructRefType:$input, I32Attr:$fieldIndex);
+  let results = (outs ReferenceType:$result);
+  let hasCustomAssemblyFormat = 1;
+  let hasVerifier = 1;
+
+  let builders = [
+    OpBuilder<(ins "mlir::Value":$input, "StructType::FieldInfo":$field)>,
+    OpBuilder<(ins "mlir::Value":$input, "mlir::StringAttr":$fieldName)>,
+    OpBuilder<(ins "mlir::Value":$input, "llvm::StringRef":$fieldName), [{
+      build($_builder, $_state, input, $_builder.getStringAttr(fieldName));
+    }]>
+  ];
+
+  let extraClassDeclaration = [{
+    /// Return the name attribute of the accessed field.
+    mlir::StringAttr getFieldNameAttr() {
+      auto type = mlir::cast<StructType>(mlir::cast<ReferenceType>(getInput().getType()).getObjectType());
+      return type.getElements()[getFieldIndex()].name;
+    }
+
+    /// Return the name of the accessed field.
+    llvm::StringRef getFieldName() {
+      return getFieldNameAttr().getValue();
+    }
+  }];
+}
+
 #endif // P4MLIR_DIALECT_P4HIR_P4HIR_OPS_TD

include/p4mlir/Dialect/P4HIR/P4HIR_TypeInterfaces.h

@@ -0,0 +1,27 @@
+#ifndef P4MLIR_DIALECT_P4HIR_P4HIR_TYPEINTERFACES_H
+#define P4MLIR_DIALECT_P4HIR_P4HIR_TYPEINTERFACES_H
+
+#include "mlir/IR/Types.h"
+
+namespace P4::P4MLIR::P4HIR {
+namespace FieldIdImpl {
+unsigned getMaxFieldID(::mlir::Type);
+
+std::pair<::mlir::Type, unsigned> getSubTypeByFieldID(::mlir::Type, unsigned fieldID);
+
+::mlir::Type getFinalTypeByFieldID(::mlir::Type type, unsigned fieldID);
+
+std::pair<unsigned, bool> projectToChildFieldID(::mlir::Type, unsigned fieldID, unsigned index);
+
+std::pair<unsigned, unsigned> getIndexAndSubfieldID(::mlir::Type type, unsigned fieldID);
+
+unsigned getFieldID(::mlir::Type type, unsigned index);
+
+unsigned getIndexForFieldID(::mlir::Type type, unsigned fieldID);
+
+}  // namespace FieldIdImpl
+}  // namespace P4::P4MLIR::P4HIR
+
+#include "p4mlir/Dialect/P4HIR/P4HIR_TypeInterfaces.h.inc"
+
+#endif  // P4MLIR_DIALECT_P4HIR_P4HIR_TYPEINTERFACES_H

include/p4mlir/Dialect/P4HIR/P4HIR_TypeInterfaces.td

@@ -0,0 +1,72 @@
+#ifndef P4MLIR_DIALECT_P4HIR_P4HIR_TYPEINTERFACES_TD
+#define P4MLIR_DIALECT_P4HIR_P4HIR_TYPEINTERFACES_TD
+
+include "mlir/IR/OpBase.td"
+
+def FieldIDTypeInterface : TypeInterface<"FieldIDTypeInterface"> {
+  let description = [{
+    Common methods for types which can be indexed by a FieldID.
+    FieldID is a depth-first numbering of the elements of a type.  For example:
+    ```
+    struct a  /* 0 */ {
+      int b; /* 1 */
+      struct c /* 2 */ {
+        int d; /* 3 */
+      }
+    }
+
+    int e; /* 0 */
+    ```
+  }];
+
+  let methods = [
+    InterfaceMethod<"Get the maximum field ID for this type",
+      "unsigned", "getMaxFieldID">,
+
+    InterfaceMethod<[{
+      Get the sub-type of a type for a field ID, and the subfield's ID. Strip
+      off a single layer of this type and return the sub-type and a field ID
+      targeting the same field, but rebased on the sub-type.
+
+      The resultant type *may* not be a FieldIDTypeInterface if the resulting 
+      fieldID is zero.  This means that leaf types may be ground without 
+      implementing an interface.  An empty aggregate will also appear as a 
+      zero.
+    }],
+      "std::pair<::mlir::Type, unsigned>", "getSubTypeByFieldID", (ins "unsigned":$fieldID)>,
+
+    InterfaceMethod<[{
+      Returns the effective field id when treating the index field as the
+      root of the type.  Essentially maps a fieldID to a fieldID after a
+      subfield op. Returns the new id and whether the id is in the given
+      child.
+    }],
+      "std::pair<unsigned, bool>", "projectToChildFieldID", (ins "unsigned":$fieldID, "unsigned":$index)>,
+
+    InterfaceMethod<[{
+      Returns the index (e.g. struct or vector element) for a given FieldID.
+      This returns the containing index in the case that the fieldID points to a
+      child field of a field.
+    }],
+      "unsigned", "getIndexForFieldID", (ins "unsigned":$fieldID)>,
+
+    InterfaceMethod<[{
+      Return the fieldID of a given index (e.g. struct or vector element).
+      Field IDs start at 1, and are assigned
+      to each field in a recursive depth-first walk of all
+      elements. A field ID of 0 is used to reference the type itself.
+    }],
+      "unsigned", "getFieldID", (ins "unsigned":$index)>,
+
+    InterfaceMethod<[{
+      Find the index of the element that contains the given fieldID.
+      As well, rebase the fieldID to the element.
+    }],
+      "std::pair<unsigned, unsigned>", "getIndexAndSubfieldID", (ins "unsigned":$fieldID)>,
+
+  ];
+
+  let cppNamespace = "::P4::P4MLIR::P4HIR";
+}
+
+#endif // P4MLIR_DIALECT_P4HIR_P4HIR_TYPEINTERFACES_TD

include/p4mlir/Dialect/P4HIR/P4HIR_Types.h

@@ -6,9 +6,22 @@
 #pragma GCC diagnostic ignored "-Wunused-parameter"
 
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Interfaces/MemorySlotInterfaces.h"
 #include "p4mlir/Dialect/P4HIR/P4HIR_OpsEnums.h"
+#include "p4mlir/Dialect/P4HIR/P4HIR_TypeInterfaces.h"
+
+namespace P4::P4MLIR::P4HIR {
+
+namespace detail {
+/// Struct defining a field. Used in structs.
+struct FieldInfo {
+    mlir::StringAttr name;
+    mlir::Type type;
+};
+}  // namespace detail
+}  // namespace P4::P4MLIR::P4HIR
 
 #define GET_TYPEDEF_CLASSES
 #include "p4mlir/Dialect/P4HIR/P4HIR_Types.h.inc"
 
-#endif // P4MLIR_DIALECT_P4HIR_P4HIR_TYPES_H
+#endif  // P4MLIR_DIALECT_P4HIR_P4HIR_TYPES_H

include/p4mlir/Dialect/P4HIR/P4HIR_Types.td

@@ -3,8 +3,10 @@
 
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/EnumAttr.td"
+include "mlir/Interfaces/MemorySlotInterfaces.td"
 
 include "p4mlir/Dialect/P4HIR/P4HIR_Dialect.td"
+include "p4mlir/Dialect/P4HIR/P4HIR_TypeInterfaces.td"
 
 //===----------------------------------------------------------------------===//
 // P4HIR type definitions.
@@ -122,6 +124,7 @@ def VoidType : P4HIR_Type<"Void", "void"> {
     llvm::StringRef getAlias() const { return "void"; };
   }];
 }
+
 //===----------------------------------------------------------------------===//
 // ReferenceType
 //===----------------------------------------------------------------------===//
@@ -152,6 +155,10 @@ def ReferenceType : P4HIR_Type<"Reference", "ref"> {
   let skipDefaultBuilders = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// FuncType
+//===----------------------------------------------------------------------===//
+
 def FuncType : P4HIR_Type<"Func", "func"> {
   let summary = "P4 function-like type (actions, methods, functions)";
   let description = [{
@@ -212,15 +219,57 @@ def FuncType : P4HIR_Type<"Func", "func"> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// StructType
+//===----------------------------------------------------------------------===//
+
+// A packed struct. Declares the P4HIR::StructType in C++.
+def StructType : P4HIR_Type<"Struct", "struct", [
+  DeclareTypeInterfaceMethods<DestructurableTypeInterface>,
+  DeclareTypeInterfaceMethods<FieldIDTypeInterface>
+]> {
+  let summary = "struct type";
+  let description = [{
+    Represents a structure of name, value pairs.
+    !p4hir.struct<"name", fieldName1: Type1, fieldName2: Type2>
+  }];
+
+  let hasCustomAssemblyFormat = 1;
+  let genVerifyDecl = 1;
+
+  let parameters = (
+    ins StringRefParameter<"struct name">:$name, ArrayRefParameter<
+      "P4HIR::StructType::FieldInfo", "struct fields">:$elements
+  );
+
+  let extraClassDeclaration = [{
+    using FieldInfo = P4HIR::detail::FieldInfo;
+    mlir::Type getFieldType(mlir::StringRef fieldName);
+    void getInnerTypes(mlir::SmallVectorImpl<mlir::Type>&);
+    std::optional<unsigned> getFieldIndex(mlir::StringRef fieldName);
+    std::optional<unsigned> getFieldIndex(mlir::StringAttr fieldName);
+  }];
+}
+
 //===----------------------------------------------------------------------===//
 // P4HIR type constraints.
 //===----------------------------------------------------------------------===//
 
-def AnyP4Type : AnyTypeOf<[BitsType, BooleanType, InfIntType,
+def AnyP4Type : AnyTypeOf<[BitsType, BooleanType, InfIntType, StructType,
                            DontcareType, ErrorType, UnknownType]> {}
 def AnyIntP4Type : AnyTypeOf<[BitsType, InfIntType]> {}
 def CallResultP4Type : AnyTypeOf<[BitsType, BooleanType, InfIntType, VoidType]> {}
-def LoadableP4Type : AnyTypeOf<[BitsType, BooleanType, InfIntType]> {}
+def LoadableP4Type : AnyTypeOf<[BitsType, BooleanType, InfIntType, StructType]> {}
+
+/// A ref type with the specified constraints on the nested type.
+class SpecificRefType<Type type> : ConfinedType<ReferenceType,
+  [SubstLeaves<"$_self", "mlir::cast<P4HIR::ReferenceType>($_self).getObjectType()",
+               type.predicate>],
+  "ref of " # type.summary, "P4HIR::ReferenceType"
+> {
+  Type objectType = type;
+}
 
+def StructRefType : SpecificRefType<StructType>;
 
 #endif // P4MLIR_DIALECT_P4HIR_P4HIR_TYPES_TD

lib/Dialect/P4HIR/CMakeLists.txt

@@ -2,6 +2,7 @@ add_mlir_dialect_library(P4MLIR_P4HIR
   P4HIR_Ops.cpp
   P4HIR_Types.cpp
   P4HIR_Attrs.cpp
+  P4HIR_TypeInterfaces.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/include/p4mlir/Dialect/P4HIR