Update LLZK and use new APIs

circom/tests/misc/large_constant3.circom

@@ -13,7 +13,7 @@ template CompConstant(ct) {
 
 component main = CompConstant(10944121435919637611123202872628637544274182200208017171849102093287904247808);
 
-// CHECK-LABEL: module attributes {llzk.lang, llzk.main = !struct.type<@CompConstant::@CompConstant<[#felt<const 10944121435919637611123202872628637544274182200208017171849102093287904247808>]>>} {
+// CHECK-LABEL: module attributes {llzk.lang, llzk.main = !struct.type<@CompConstant::@CompConstant<[#felt<const 10944121435919637611123202872628637544274182200208017171849102093287904247808 : <"bn128">>]>>} {
 // CHECK-NEXT:    poly.template @CompConstant {
 // CHECK-NEXT:      poly.param @ct
 // CHECK-NEXT:      struct.def @CompConstant {

llzk_backend/src/codegen.rs

@@ -10,13 +10,10 @@ pub use crate::shared::LlzkConfig;
 use ansi_term::Color;
 use anyhow::anyhow;
 use anyhow::Result;
-use llzk::prelude::FeltConstAttribute;
-use llzk::prelude::IntegerAttribute;
 use llzk::prelude::LlzkContext;
 use llzk::prelude::Module;
 use llzk::prelude::OperationMutLike as _;
 use llzk::prelude::StructType;
-use llzk::prelude::Type;
 use llzk::prelude::TypeAttribute;
 use llzk::prelude::MAIN_ATTR_NAME;
 use melior::ir::Attribute;
@@ -39,7 +36,7 @@ fn to_u64_digits(n: &BigUint) -> Vec<u64> {
         .collect()
 }
 
-/// Prepares the paremeters of the main component StructType.
+/// Prepares the parameters of the main component StructType.
 ///
 /// TODO: This approach does not currently handle ArrayInLine or Call expressions that can be
 /// used as parameters to the main component. The Call case could be handled by finding the
@@ -60,18 +57,12 @@ fn prepare_main_component_params<'ctx>(
                 n.ok_or_else(|| anyhow!("main component parameter {i} is not a positive constant"))
             })?;
             Ok(match n.to_i64() {
-                Some(n) => IntegerAttribute::new(Type::index(context), n).into(),
+                Some(n) => context.index_attr(n).into(),
                 None => {
                     // Increase by one to ensure the value is kept unsigned.
                     let bitlen = n.bits() + 1;
                     let parts = to_u64_digits(&n);
-                    FeltConstAttribute::from_parts(
-                        context,
-                        bitlen.try_into().unwrap(),
-                        &parts,
-                        None,
-                    )
-                    .into()
+                    context.felt_attr_from_parts(bitlen.try_into().unwrap(), &parts).into()
                 }
             })
         })
@@ -106,7 +97,8 @@ fn new_llzk_module<'ctx>(
 /// Generate LLZK IR from the given `ProgramArchive` and write it to a file with the given filename.
 #[allow(clippy::result_unit_err)]
 pub fn generate_llzk(program: &impl ProgramLike, config: LlzkConfig) -> Result<(), ()> {
-    let ctx = LlzkContext::new();
+    let mut ctx = LlzkContext::new();
+    ctx.set_field(config.prime_str.as_str());
     let module = new_llzk_module(&ctx, program, &config.prime).map_err(|err| {
         if config.verbose {
             eprintln!("{} {err:?}", Color::Red.paint("Failed to generate LLZK IR:"));

llzk_backend/src/lvalue.rs

@@ -12,7 +12,6 @@ use anyhow::Result;
 use llzk::dialect::llzk::nondet;
 use llzk::dialect::pod;
 use llzk::dialect::r#struct;
-use llzk::prelude::IntegerAttribute;
 use llzk::prelude::Location;
 use llzk::prelude::PodType;
 use llzk::prelude::StructType;
@@ -204,12 +203,7 @@ impl<'ast> Lvalue<'ast> {
         // Constant true used as starting point for concatenating the conditions for a particular
         // index set together with a fold.
         let true_value = {
-            let attr = IntegerAttribute::new(codegen.bool_type().into(), 1);
-            block_gen.as_mut().append_op_unnamed_result(arith::constant(
-                codegen.context,
-                attr.into(),
-                location,
-            ))
+            block_gen.as_mut().append_op_unnamed_result(codegen.new_bool_const_op(true, location))
         }?;
 
         let entries = mixed_subcmp_layout

llzk_backend/src/shared.rs

@@ -28,7 +28,6 @@ use llzk::dialect::array::ArrayCtor;
 use llzk::dialect::felt;
 use llzk::dialect::pod;
 use llzk::dialect::poly;
-use llzk::dialect::poly::TVarType;
 use llzk::prelude::is_felt_type;
 use llzk::prelude::melior_dialects::arith;
 use llzk::prelude::replace_uses_of_with;
@@ -40,7 +39,6 @@ use llzk::prelude::Block;
 use llzk::prelude::BlockLike;
 use llzk::prelude::BlockRef;
 use llzk::prelude::BoolAttribute;
-use llzk::prelude::FeltConstAttribute;
 use llzk::prelude::FeltType;
 use llzk::prelude::FlatSymbolRefAttribute;
 use llzk::prelude::FuncDefOp;
@@ -58,7 +56,6 @@ use llzk::prelude::OperationLike;
 use llzk::prelude::OperationMutLike;
 use llzk::prelude::OperationRef;
 use llzk::prelude::PassManager;
-use llzk::prelude::PodRecordAttribute;
 use llzk::prelude::PodType;
 use llzk::prelude::RecordValue;
 use llzk::prelude::Region;
@@ -782,7 +779,7 @@ impl<'ast: 'r, 'ctx: 'r, 'r, P: ProgramLike> LlzkCodegen<'ast, 'ctx, 'r, P> {
     /// Get the unknown location.
     #[inline]
     pub fn location_unknown(&self) -> Location<'ctx> {
-        Location::unknown(self.context)
+        self.context.unknown_location()
     }
 
     /// Convert circom location information to MLIR location.
@@ -837,34 +834,28 @@ impl<'ast: 'r, 'ctx: 'r, 'r, P: ProgramLike> LlzkCodegen<'ast, 'ctx, 'r, P> {
         Ok(dialect::llzk::nondet(location, result_type))
     }
 
-    /// Get the integer type of the given bitwidth.
-    #[inline]
-    pub fn int_type(&self, bits: u32) -> IntegerType<'ctx> {
-        IntegerType::new(self.context, bits)
-    }
-
     /// Get the boolean type (`i1`).
     #[inline]
-    pub fn bool_type(&self) -> IntegerType<'ctx> {
-        self.int_type(1)
+    pub fn bool_type(&self) -> Type<'ctx> {
+        self.context.bool_type()
     }
 
     /// Get the index type.
     #[inline]
     pub fn index_type(&self) -> Type<'ctx> {
-        Type::index(self.context)
+        self.context.index_type()
     }
 
     /// Get the felt type.
     #[inline]
     pub fn felt_type(&self) -> FeltType<'ctx> {
-        FeltType::with_field(self.context, &self.config.prime_str)
+        self.context.felt_type()
     }
 
     /// Get the polymorphic type variable type for the given name.
     #[inline]
     pub fn tvar_type(&self, name: &str) -> Type<'ctx> {
-        TVarType::new(self.context, StringRef::new(name)).into()
+        self.context.tvar_type(name).into()
     }
 
     /// Get the struct type for the given struct name and parameters.
@@ -885,29 +876,31 @@ impl<'ast: 'r, 'ctx: 'r, 'r, P: ProgramLike> LlzkCodegen<'ast, 'ctx, 'r, P> {
         self.struct_type_with_params(name, &[])
     }
 
-    /// Get a pod struct type with the given records.
+    /// Get a pod type with the given records.
     #[inline]
     pub fn pod_type(&self, records: &[(&str, Type<'ctx>)]) -> PodType<'ctx> {
-        let records = records
-            .iter()
-            .map(|(name, r#type)| PodRecordAttribute::new(name, *r#type))
-            .collect::<Vec<_>>();
-        PodType::new(self.context, &records)
+        self.context.pod_type(records)
     }
 
-    /// Create an index attribute.
+    /// Returns an [`IntegerAttribute`] with index type and the given value.
     #[inline]
     pub fn index_attr<T>(&self, integer: T) -> IntegerAttribute<'ctx>
     where
         T: Into<i64>,
     {
-        IntegerAttribute::new(self.index_type(), integer.into())
+        self.context.index_attr(integer)
     }
 
-    /// Creates a [`FlatSymbolRefAttribute`] from the given string.
+    /// Returns a [`BoolAttribute`] with the given value.
+    #[inline]
+    pub fn bool_attr(&self, val: bool) -> BoolAttribute<'_> {
+        self.context.bool_attr(val)
+    }
+
+    /// Returns a [`FlatSymbolRefAttribute`] created from the given string.
     #[inline]
     pub fn flat_sym(&self, sym: impl AsRef<str>) -> FlatSymbolRefAttribute<'ctx> {
-        FlatSymbolRefAttribute::new(self.context, sym.as_ref())
+        self.context.flat_sym_attr(sym)
     }
 
     /// Creates a [`SymbolRefAttribute`] from the given string as "@str::@str"
@@ -949,18 +942,7 @@ impl<'ast: 'r, 'ctx: 'r, 'r, P: ProgramLike> LlzkCodegen<'ast, 'ctx, 'r, P> {
     /// Create an LLZK operation that produces a boolean constant value.
     #[inline]
     pub fn new_bool_const_op(&self, val: bool, location: Location<'ctx>) -> Operation<'ctx> {
-        arith::constant(self.context, BoolAttribute::new(self.context, val).into(), location)
-    }
-
-    /// Create an LLZK operation that produces an integer constant value.
-    #[inline]
-    pub fn new_int_const_op(
-        &self,
-        ty: Type<'ctx>,
-        val: i64,
-        location: Location<'ctx>,
-    ) -> Operation<'ctx> {
-        arith::constant(self.context, IntegerAttribute::new(ty, val).into(), location)
+        self.context.new_bool_const_op(val, location)
     }
 
     /// Create an LLZK operation that produces an index constant value.
@@ -969,7 +951,7 @@ impl<'ast: 'r, 'ctx: 'r, 'r, P: ProgramLike> LlzkCodegen<'ast, 'ctx, 'r, P> {
     where
         T: Into<i64>,
     {
-        arith::constant(self.context, self.index_attr(val).into(), location)
+        self.context.new_index_const_op(val, location)
     }
 
     /// Create an LLZK `array.new` operation with the given type and constructor info.
@@ -993,31 +975,12 @@ impl<'ast: 'r, 'ctx: 'r, 'r, P: ProgramLike> LlzkCodegen<'ast, 'ctx, 'r, P> {
         // ASSERT: The circom parser always produces non-negative constants. These can be negated
         // via PrefixOp but negative BigInt constants are never created directly.
         assert_ne!(val.sign(), num_bigint_dig::Sign::Minus, "Felt constants must be non-negative");
-        let attr = FeltConstAttribute::parse(
-            self.context,
-            // use required bits +1 to ensure unsigned representation
-            u32::try_from(val.bits())? + 1,
-            val.to_string().as_str(),
-            Some(&self.config.prime_str),
-        );
+        // Increase by one to ensure the value is kept unsigned.
+        let bitlen = val.bits() + 1;
+        let attr = self.context.felt_attr_from_str(bitlen.try_into().unwrap(), val.to_string());
         felt::constant(location, attr).map_err(Into::into)
     }
 
-    /// Creates a new constant op to create a constant of the given type.
-    /// Assumes `type` will be a felt or integral/index type.
-    pub fn new_const_op(
-        &self,
-        location: Location<'ctx>,
-        r#type: Type<'ctx>,
-        val: i64,
-    ) -> Result<Operation<'ctx>> {
-        if is_felt_type(r#type) {
-            self.new_felt_const_op(&BigInt::from(val), location)
-        } else {
-            Ok(self.new_int_const_op(r#type, val, location))
-        }
-    }
-
     /// Run the given pass pipeline on the given operation.
     pub fn run_pass_pipeline_on<'c: 'a, 'a>(
         &self,

llzk_backend/src/write_chain.rs

@@ -20,12 +20,10 @@ use crate::template::TemplateContext;
 use anyhow::Result;
 use llzk::dialect::r#struct;
 use llzk::prelude::FuncDefOpLike as _;
-use llzk::prelude::IntegerAttribute;
 use llzk::prelude::Location;
 use llzk::prelude::PodType;
 use llzk::prelude::Value;
 use llzk::prelude::ValueLike as _;
-use melior::dialect::arith;
 use program_structure::ast::Access;
 use program_structure::ast::Expression;
 use std::convert::TryFrom as _;
@@ -253,11 +251,7 @@ impl<'ast> WriteChain<'ast> {
             location,
             InfoProviders { subcmp_info, signal_write_info },
         )?;
-        let true_value = {
-            let attr = IntegerAttribute::new(codegen.bool_type().into(), 1);
-            fc.append_op_unnamed_result(arith::constant(codegen.context, attr.into(), location))
-        }?;
-
+        let true_value = fc.append_op_unnamed_result(codegen.new_bool_const_op(true, location))?;
         let memory_value = *fc.block_ctx.get_named_value(root_var)?;
         let memory_type = PodType::try_from(memory_value.r#type())?;
         let inputs_name = crate::subcmp::names::inputs(root_var);