feat(llzk): more local invariants checks (array), refactor common checks

Veir/Dialects/LLZK/Array/OpInfo.lean

@@ -91,12 +91,62 @@ def LLZK.Array.verifyLocalInvariants {OpInfo : Type} [IsOpCode OpInfo]
   if op.getNumSuccessors ctx.raw opIn ≠ 0 then
     throw s!"{instrName}: Expected 0 successors"
   match opType with
-  | .new => requireResults 1
-  | .read | .extract => do requireAtLeastOperands 1; requireResults 1
-  | .write | .insert => do requireAtLeastOperands 2; requireResults 0
+  | .new => do
+    requireResults 1
+    if !(op.getResultTypes! ctx.raw)[0]!.val.isLLZKArrayType then
+      throw s!"{instrName}: Expected valid LLZK array.type return type"
+    -- TODO: check operands types/count
+  | .read => do
+    requireAtLeastOperands 1
+    requireResults 1
+    for (operType,idx) in (op.getOperandTypes! ctx.raw).zipIdx do
+      if idx = 0 && !operType.val.isLLZKArrayType then
+        throw s!"{instrName}: Expected operand 0 to have a valid LLZK array type"
+      if idx ≠ 0 && !(operType.val matches .indexType _) then
+        throw s!"{instrName}: Expected operand {idx} to have an index type"
+    if !(op.getResultTypes! ctx.raw)[0]!.val.isLLZKArrayElemType then
+      throw s!"{instrName}: Expected return to have a valid LLZK array element type"
+  | .extract => do
+    requireAtLeastOperands 1
+    requireResults 1
+    for (operType,idx) in (op.getOperandTypes! ctx.raw).zipIdx do
+      if idx = 0 && !operType.val.isLLZKArrayType then
+        throw s!"{instrName}: Expected operand 0 to have a valid LLZK array type"
+      if idx ≠ 0 && !(operType.val matches .indexType _) then
+        throw s!"{instrName}: Expected operand {idx} to have an index type"
+    if !(op.getResultTypes! ctx.raw)[0]!.val.isLLZKArrayType then
+      throw s!"{instrName}: Expected return to have a valid LLZK array type"
+  | .write => do
+    requireAtLeastOperands 2
+    requireResults 0
+    for (operType,idx) in (op.getOperandTypes! ctx.raw).zipIdx do
+      if idx = 0 && !operType.val.isLLZKArrayType then
+        throw s!"{instrName}: Expected operand 0 to be a valid LLZK array"
+      if idx ≠ 0 && idx ≠ (op.getNumOperands! ctx.raw)-1 && !(operType.val matches .indexType _) then
+        throw s!"{instrName}: Expected operand {idx} to have an index type"
+      if idx = (op.getNumOperands! ctx.raw)-1 && !operType.val.isLLZKArrayElemType then
+        throw s!"{instrName}: Expected operand {idx} to have a valid LLZK array element type"
+  | .insert => do
+    requireAtLeastOperands 2
+    requireResults 0
+    for (operType,idx) in (op.getOperandTypes! ctx.raw).zipIdx do
+      if idx = 0 && !operType.val.isLLZKArrayType then
+        throw s!"{instrName}: Expected operand 0 to be a valid LLZK array"
+      if idx ≠ 0 && idx ≠ (op.getNumOperands! ctx.raw)-1 && !(operType.val matches .indexType _) then
+        throw s!"{instrName}: Expected operand {idx} to have an index type"
+      if idx = (op.getNumOperands! ctx.raw)-1 && !operType.val.isLLZKArrayType then
+        throw s!"{instrName}: Expected operand {idx} to have a valid LLZK array type"
   | .len => do
+    requireResults 1
     if op.getNumOperands ctx.raw opIn ≠ 2 then
       throw s!"{instrName}: Expected 2 operands (the array and the dimension)"
+    if !(op.getOperandTypes! ctx.raw)[0]!.val.isLLZKArrayType then
+      throw s!"{instrName}: Expected operand 0 to be a valid LLZK array"
+    if !((op.getOperandTypes! ctx.raw)[1]!.val matches .indexType _) then
+      throw s!"{instrName}: Expected operand 1 to have an index type"
+    if !((op.getResultTypes! ctx.raw)[0]!.val matches .indexType _) then
+      throw s!"{instrName}: Expected return to have index type"
+
     requireResults 1
 
 instance : HasOpInfo LLZK.Array where

Veir/Dialects/LLZK/Constrain/OpInfo.lean

@@ -65,11 +65,11 @@ instance : IsOpCode LLZK.Constrain where
   fromAttrDict := LLZK.Constrain.fromAttrDict
   toAttrDict := LLZK.Constrain.toAttrDict
 
-private def Attribute.isSupportedLLZKConstrainEqType (type : Attribute) : Bool :=
-  match type with
-  | .integerType intType => intType.bitwidth = 1
-  | .indexType _ | .feltType _ | .arrayType _ => true
-  | _ => false
+
+private def Attribute.isLLZKEmitEqType (attr : Attribute) : Bool :=
+  match attr with
+  | .structType _ | .stringType _ => false
+  | _ => attr.isLLZKType
 
 /-- Whether `candidate` is an element or trailing-dimensional subarray of `arrayType`. -/
 private def isLLZKSubArrayOrElementType
@@ -91,15 +91,15 @@ def LLZK.Constrain.verifyLocalInvariants {OpInfo : Type} [IsOpCode OpInfo]
     op.verifyPlainOpCounts ctx opIn 2 0
     let operandType ← op.verifyOperandTypesMatch ctx 0 1
       "constrain.eq: expected operands to have the same type"
-    if !operandType.val.isSupportedLLZKConstrainEqType then
+    if !operandType.val.isLLZKEmitEqType then
       throw s!"constrain.eq: unsupported operand type {operandType}"
   | .«in» => do
     op.verifyPlainOpCounts ctx opIn 2 0
     let lhsType := (op.getOperand! ctx.raw 0).getType! ctx.raw
     let rhsType := (op.getOperand! ctx.raw 1).getType! ctx.raw
     let .arrayType arrayType := lhsType.val
       | throw s!"constrain.in: expected first operand to have array type, got {lhsType}"
-    if !arrayType.elementType.isSupportedLLZKConstrainEqType then
+    if !arrayType.elementType.isLLZKEmitEqType then
       throw s!"constrain.in: unsupported array element type {arrayType.elementType}"
     if !isLLZKSubArrayOrElementType arrayType rhsType.val then
       throw s!"constrain.in: {rhsType} is not an element or compatible subarray of {lhsType}"

Veir/Dialects/LLZK/Struct/OpInfo.lean

@@ -107,6 +107,7 @@ instance : IsOpCode LLZK.Struct where
 def LLZK.Struct.verifyLocalInvariants {OpInfo : Type} [IsOpCode OpInfo]
     [HasDialect OpInfo LLZK.Struct] (opType : LLZK.Struct) (op : OperationPtr)
     (ctx : WfIRContext OpInfo) (opIn : op.InBounds ctx.raw) : Except String PUnit := do
+  let instrName := String.fromUTF8! (IsOpCode.name (op.getOpType ctx.raw opIn))
   match opType with
   | .«def» => do
     if op.getNumOperands ctx.raw opIn ≠ 0 then
@@ -118,17 +119,14 @@ def LLZK.Struct.verifyLocalInvariants {OpInfo : Type} [IsOpCode OpInfo]
     if op.getNumSuccessors ctx.raw opIn ≠ 0 then
       throw "struct.def: Expected 0 successors"
   | .member => do
-    let instrName := String.fromUTF8! (IsOpCode.name (op.getOpType ctx.raw opIn))
     op.verifyPlainOpCounts ctx opIn 0 0
-    (op.getProperties! ctx.raw LLZK.Struct.member).type.verifySupportedLLZKType
-      s!"{instrName}: expected 'type' to be a supported LLZK type"
+    if !(op.getProperties! ctx.raw LLZK.Struct.member).type.val.isLLZKType then
+      throw s!"{instrName}: expected 'type' to be a supported LLZK type"
   | .new => do
-    let instrName := String.fromUTF8! (IsOpCode.name (op.getOpType ctx.raw opIn))
     op.verifyPlainOpCounts ctx opIn 0 1
     if !((op.getResultTypes! ctx.raw)[0]!.val matches Attribute.structType _) then
       throw s!"{instrName}: expected result 0 to have !struct.type"
   | .readm => do
-    let instrName := String.fromUTF8! (IsOpCode.name (op.getOpType ctx.raw opIn))
     if op.getNumOperands ctx.raw opIn < 1 then
       throw s!"{instrName}: Expected at least 1 operand (the component)"
     if op.getNumResults ctx.raw opIn ≠ 1 then
@@ -139,15 +137,14 @@ def LLZK.Struct.verifyLocalInvariants {OpInfo : Type} [IsOpCode OpInfo]
       throw s!"{instrName}: Expected 0 successors"
     if !((op.getOperandTypes! ctx.raw)[0]!.val matches Attribute.structType _ ) then
       throw s!"{instrName}: expected operand 0 to have !struct.type"
-    (op.getResultTypes! ctx.raw)[0]!.verifySupportedLLZKType
-      s!"{instrName}: expected result 0 to have a supported LLZK type"
+    if !(op.getResultTypes! ctx.raw)[0]!.val.isLLZKType then
+      throw s!"{instrName}: expected result 0 to have a supported LLZK type"
   | .writem => do
-    let instrName := String.fromUTF8! (IsOpCode.name (op.getOpType ctx.raw opIn))
     op.verifyPlainOpCounts ctx opIn 2 0
     if !((op.getOperandTypes! ctx.raw)[0]!.val matches Attribute.structType _ ) then
       throw s!"{instrName}: expected operand 0 to have !struct.type"
-    (op.getOperandTypes! ctx.raw)[1]!.verifySupportedLLZKType
-      s!"{instrName}: expected operand 1 to have a supported LLZK type"
+    if !(op.getOperandTypes! ctx.raw)[1]!.val.isLLZKType then
+      throw s!"{instrName}: expected operand 1 to have a supported LLZK type"
 
 instance : HasOpInfo LLZK.Struct where
   verifyLocalInvariants := LLZK.Struct.verifyLocalInvariants

Veir/Dialects/LLZK/Struct/Properties.lean

@@ -2,6 +2,7 @@ module
 
 public import Veir.IR.Attribute
 public import Veir.IR.OpInfo
+public import Veir.Verifier.Basic
 public import Veir.Dialects.Builtin.Properties
 
 namespace Veir
@@ -31,14 +32,6 @@ structure StructMemberProperties where
   signal : Bool
 deriving Inhabited, Repr, Hashable, DecidableEq
 
-/-- Verify that a type belongs to the subset of LLZK types currently represented by VeIR.
-https://github.com/project-llzk/llzk-lib/blob/265d68f678ab15018e3f6253b85557fbaeac9c0d/lib/Util/TypeHelper.cpp#L482-L511 -/
-def TypeAttr.verifySupportedLLZKType (ty : TypeAttr) (errMsg : String) : Except String PUnit :=
-  match ty.val with
-  | .integerType intType => if intType.bitwidth = 1 then pure () else throw errMsg
-  | .indexType _ | .feltType _ | .structType _ | .stringType _ | .arrayType _ => pure ()
-  | _ => throw errMsg
-
 def StructMemberProperties.fromAttrDict (opName : String)
     (attrDict : Std.HashMap ByteArray Attribute) :
     Except String StructMemberProperties := do
@@ -48,12 +41,10 @@ def StructMemberProperties.fromAttrDict (opName : String)
     | throw s!"{opName}: expected 'sym_name' to be a string attribute, got {symAttr}"
   let typeAttr ← match attrDict["type".toUTF8]? with
     | some attr =>
-      if _ : attr.isType = true then
-        attr.asType.verifySupportedLLZKType
-          s!"{opName}: expected 'type' to be a supported LLZK type"
+      if _ : attr.isType = true && attr.isLLZKType then
         pure (attr.asType)
       else
-        throw s!"{opName}: expected 'type' to be a type attribute"
+        throw s!"{opName}: expected 'type' to be a supported LLZK type"
     | _ =>
       throw s!"{opName}: missing 'type' property"
   let column ← getUnitAttr "column" attrDict

Veir/Verifier/Basic.lean

@@ -38,6 +38,25 @@ def Attribute.isKnownNonZero (attr : Attribute) : Bool :=
   | .floatAttr fltAttr => fltAttr.value != 0.0
   | _ => false
 
+/-- Verify that a type belongs to the subset of LLZK types currently represented by VeIR.
+  TypeVarType, PodType, and NoneType are currently missing.
+  https://github.com/project-llzk/llzk-lib/blob/265d68f678ab15018e3f6253b85557fbaeac9c0d/lib/Util/TypeHelper.cpp#L482-L511 -/
+def Attribute.isLLZKType (attr : Attribute)  : Bool :=
+  match attr with
+  | .integerType intType => decide (intType.bitwidth = 1)
+  | .indexType _ | .feltType _ | .structType _ | .stringType _ | .arrayType _ => true
+  | _ => false
+
+def Attribute.isLLZKArrayElemType (attr : Attribute)  : Bool :=
+  match attr with
+  | .arrayType _ => false
+  | _ => attr.isLLZKType
+
+def Attribute.isLLZKArrayType (attr : Attribute) : Bool :=
+  match attr with
+  | .arrayType arrT  => arrT.elementType.isLLZKArrayElemType
+  | _ => false
+
 /--
   Verify the result, region, and successor counts of a terminator: one that
   produces no results, has no regions, and transfers control to `successors`
@@ -189,9 +208,9 @@ def TypeAttr.verifyI1
 
 def TypeAttr.verifyI64
     (ty : TypeAttr) (errMsg : String) : Except String PUnit :=
-  match ty.val with
-  | .integerType { bitwidth := 64 } => pure ()
-  | _ => throw errMsg
+  match ty.val with
+  | .integerType { bitwidth := 64 } => pure ()
+  | _ => throw errMsg
 
 /--
   Verify the operand and result counts of a "plain" operation: one that has no