[SMT] Add basic array operations

include/circt/Dialect/SMT/SMT.td

@@ -15,6 +15,7 @@ include "circt/Dialect/SMT/SMTAttributes.td"
 include "circt/Dialect/SMT/SMTDialect.td"
 include "circt/Dialect/SMT/SMTTypes.td"
 include "circt/Dialect/SMT/SMTOps.td"
+include "circt/Dialect/SMT/SMTArrayOps.td"
 include "circt/Dialect/SMT/SMTBitVectorOps.td"
 
 #endif // CIRCT_DIALECT_SMT_SMT_TD

include/circt/Dialect/SMT/SMTArrayOps.td

@@ -0,0 +1,99 @@
+//===- SMTArrayOps.td - SMT array operations ---------------*- tablegen -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef CIRCT_DIALECT_SMT_SMTARRAYOPS_TD
+#define CIRCT_DIALECT_SMT_SMTARRAYOPS_TD
+
+include "circt/Dialect/SMT/SMTDialect.td"
+include "circt/Dialect/SMT/SMTAttributes.td"
+include "circt/Dialect/SMT/SMTTypes.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+
+class SMTArrayOp<string mnemonic, list<Trait> traits = []> :
+  SMTOp<"array." # mnemonic, traits>;
+
+def ArrayStoreOp : SMTArrayOp<"store", [
+  Pure,
+  TypesMatchWith<"summary", "array", "index", 
+                 "cast<ArrayType>($_self).getDomainType()">,
+  TypesMatchWith<"summary", "array", "value",
+                 "cast<ArrayType>($_self).getRangeType()">,
+  AllTypesMatch<["array", "result"]>,
+]> {
+  let summary = "stores a value at a given index and returns the new array";
+  let description = [{
+    This operation returns a new array which is the same as the 'array' operand
+    except that the value at the given 'index' is changed to the given 'value'.
+    The semantics are equivalent to the 'store' operator described in the
+    [SMT ArrayEx theory](https://smtlib.cs.uiowa.edu/Theories/ArraysEx.smt2) of
+    the SMT-LIB standard 2.6.
+  }];
+
+  let arguments = (ins ArrayType:$array, AnySMTType:$index, AnySMTType:$value);
+  let results = (outs ArrayType:$result);
+
+  let assemblyFormat = [{
+    $array `[` $index `]` `,` $value attr-dict `:` qualified(type($array))
+  }];
+}
+
+def ArraySelectOp : SMTArrayOp<"select", [
+  Pure,
+  TypesMatchWith<"summary", "array", "index",
+                 "cast<ArrayType>($_self).getDomainType()">,
+  TypesMatchWith<"summary", "array", "result",
+                 "cast<ArrayType>($_self).getRangeType()">,
+]> {
+  let summary = "get the value stored in the array at the given index";
+  let description = [{
+    This operation is retuns the value stored in the given array at the given
+    index. The semantics are equivalent to the `select` operator defined in the
+    [SMT ArrayEx theory](https://smtlib.cs.uiowa.edu/Theories/ArraysEx.smt2) of
+    the SMT-LIB standard 2.6.
+  }];
+
+  let arguments = (ins ArrayType:$array, AnySMTType:$index);
+  let results = (outs AnySMTType:$result);
+
+  let assemblyFormat = [{
+    $array `[` $index `]` attr-dict `:` qualified(type($array))
+  }];
+}
+
+def ArrayBroadcastOp : SMTArrayOp<"broadcast", [
+  Pure,
+  TypesMatchWith<"summary", "result", "value",
+                 "cast<ArrayType>($_self).getRangeType()">,
+]> {
+  let summary = "construct an array with the given value stored at every index";
+  let description = [{
+    This operation represents a broadcast of the 'value' operand to all indices
+    of the array. It is equivalent to
+    ```
+    %0 = smt.declare_const "array" : !smt.array<[!smt.int -> !smt.bool]>
+    %1 = smt.forall ["idx"] {
+    ^bb0(%idx: !smt.int):
+      %2 = smt.array.select %0[%idx] : !smt.array<[!smt.int -> !smt.bool]>
+      %3 = smt.eq %value, %2 : !smt.bool
+      smt.yield %3 : !smt.bool
+    }
+    smt.assert %1
+    // return %0
+    ```
+
+    In SMT-LIB, this is frequently written as
+    `((as const (Array Int Bool)) value)`.
+  }];
+
+  let arguments = (ins AnySMTType:$value);
+  let results = (outs ArrayType:$result);
+
+  let assemblyFormat = "$value attr-dict `:` qualified(type($result))";
+}
+
+#endif // CIRCT_DIALECT_SMT_SMTARRAYOPS_TD

include/circt/Dialect/SMT/SMTTypes.h

@@ -15,4 +15,13 @@
 #define GET_TYPEDEF_CLASSES
 #include "circt/Dialect/SMT/SMTTypes.h.inc"
 
+namespace circt {
+namespace smt {
+
+/// Returns whether the given type is an SMT value type.
+bool isAnySMTValueType(mlir::Type type);
+
+} // namespace smt
+} // namespace circt
+
 #endif // CIRCT_DIALECT_SMT_SMTTYPES_H

include/circt/Dialect/SMT/SMTTypes.td

@@ -35,4 +35,22 @@ def BitVectorType : SMTTypeDef<"BitVector"> {
   let genVerifyDecl = true;
 }
 
+def ArrayType : SMTTypeDef<"Array"> {
+  let mnemonic = "array";
+  let description = [{
+    This type represents the `(Array X Y)` sort, where X and Y are any
+    sort/type, as described in the
+    [SMT ArrayEx theory](https://smtlib.cs.uiowa.edu/Theories/ArraysEx.smt2) of
+    the SMT-LIB standard 2.6.
+  }];
+
+  let parameters = (ins "mlir::Type":$domainType, "mlir::Type":$rangeType);
+  let assemblyFormat = "`<` `[` $domainType `->` $rangeType `]` `>`";
+
+  let genVerifyDecl = true;
+}
+
+def AnySMTType : Type<CPred<"smt::isAnySMTValueType($_self)">,
+                      "any SMT value type">;
+
 #endif // CIRCT_DIALECT_SMT_SMTTYPES_TD

lib/Dialect/SMT/SMTTypes.cpp

@@ -26,6 +26,10 @@ void SMTDialect::registerTypes() {
       >();
 }
 
+bool smt::isAnySMTValueType(Type type) {
+  return isa<BoolType, BitVectorType, ArrayType>(type);
+}
+
 //===----------------------------------------------------------------------===//
 // BitVectorType
 //===----------------------------------------------------------------------===//
@@ -37,3 +41,17 @@ BitVectorType::verify(function_ref<InFlightDiagnostic()> emitError,
     return emitError() << "bit-vector must have at least a width of one";
   return success();
 }
+
+//===----------------------------------------------------------------------===//
+// ArrayType
+//===----------------------------------------------------------------------===//
+
+LogicalResult ArrayType::verify(function_ref<InFlightDiagnostic()> emitError,
+                                Type domainType, Type rangeType) {
+  if (!isAnySMTValueType(domainType))
+    return emitError() << "domain must be any SMT value type";
+  if (!isAnySMTValueType(rangeType))
+    return emitError() << "range must be any SMT value type";
+
+  return success();
+}

test/Dialect/SMT/array-errors.mlir

@@ -0,0 +1,13 @@
+// RUN: circt-opt %s --split-input-file --verify-diagnostics
+
+// expected-error @below {{domain must be any SMT value type}}
+func.func @array_domain_no_smt_type(%arg0: !smt.array<[i32 -> !smt.bool]>) {
+  return
+}
+
+// -----
+
+// expected-error @below {{range must be any SMT value type}}
+func.func @array_range_no_smt_type(%arg0: !smt.array<[!smt.bool -> i32]>) {
+  return
+}

test/Dialect/SMT/array.mlir

@@ -0,0 +1,14 @@
+// RUN: circt-opt %s | circt-opt | FileCheck %s
+
+// CHECK-LABEL: func @arrayOperations
+// CHECK-SAME:  ([[A0:%.+]]: !smt.bool)
+func.func @arrayOperations(%arg0: !smt.bool) {
+  // CHECK-NEXT: [[V0:%.+]] = smt.array.broadcast [[A0]] {smt.some_attr} : !smt.array<[!smt.bool -> !smt.bool]>
+  %0 = smt.array.broadcast %arg0 {smt.some_attr} : !smt.array<[!smt.bool -> !smt.bool]>
+  // CHECK-NEXT: [[V1:%.+]] = smt.array.select [[V0]][[[A0]]] {smt.some_attr} : !smt.array<[!smt.bool -> !smt.bool]>
+  %1 = smt.array.select %0[%arg0] {smt.some_attr} : !smt.array<[!smt.bool -> !smt.bool]>
+  // CHECK-NEXT: [[V2:%.+]] = smt.array.store [[V0]][[[A0]]], [[A0]] {smt.some_attr} : !smt.array<[!smt.bool -> !smt.bool]>
+  %2 = smt.array.store %0[%arg0], %arg0 {smt.some_attr} : !smt.array<[!smt.bool -> !smt.bool]>
+
+  return
+}