llzk 0.4.0

use llzk_sys::{
    llzkOperationIsA_Struct_MemberDefOp, llzkOperationIsA_Struct_StructDefOp,
    llzkStruct_MemberDefOpHasPublicAttr, llzkStruct_MemberDefOpSetPublicAttr,
    llzkStruct_MemberReadOpBuild, llzkStruct_StructDefOpGetBody,
    llzkStruct_StructDefOpGetBodyRegion, llzkStruct_StructDefOpGetComputeFuncOp,
    llzkStruct_StructDefOpGetConstrainFuncOp, llzkStruct_StructDefOpGetMemberDef,
    llzkStruct_StructDefOpGetMemberDefs, llzkStruct_StructDefOpGetNumMemberDefs,
    llzkStruct_StructDefOpGetNumTemplateExprOpNames,
    llzkStruct_StructDefOpGetNumTemplateParamOpNames, llzkStruct_StructDefOpGetTemplateExprOpNames,
    llzkStruct_StructDefOpGetTemplateParamOpNames, llzkStruct_StructDefOpGetType,
    llzkStruct_StructDefOpGetTypeWithParams, llzkStruct_StructDefOpHasColumns,
    llzkStruct_StructDefOpIsMainComponent,
};
use melior::ir::{
    Attribute, AttributeLike, Block, BlockLike as _, BlockRef, Identifier, Location, Operation,
    OperationRef, Region, RegionLike as _, RegionRef, Type, TypeLike, Value, ValueLike,
    attribute::{ArrayAttribute, FlatSymbolRefAttribute, StringAttribute, TypeAttribute},
    operation::{OperationBuilder, OperationLike, OperationMutLike},
};
use mlir_sys::{MlirAttribute, MlirOperation};

use crate::{
    builder::{OpBuilder, OpBuilderLike},
    dialect::function::FuncDefOpRef,
    error::Error,
    ident,
    macros::llzk_op_type,
};

use super::StructType;

//===----------------------------------------------------------------------===//
// StructDefOpLike
//===----------------------------------------------------------------------===//

/// Defines the public API of the 'struct.def' op.
pub trait StructDefOpLike<'c: 'a, 'a>: OperationLike<'c, 'a> {
    /// Returns the associated StructType to this op using the const params defined by the op.
    ///
    /// # Panics
    ///
    /// If the 'struct.def' op type is not `!struct.type`.
    fn r#type(&self) -> StructType<'c> {
        unsafe { Type::from_raw(llzkStruct_StructDefOpGetType(self.to_raw())) }
            .try_into()
            .expect("StructDefOpLike::type error")
    }

    /// Returns the name of the struct
    ///
    /// # Panics
    ///
    /// If the 'struct.def' op doesn't have an attribute named `sym_name`.
    fn name(&'a self) -> &'c str {
        self.attribute("sym_name")
            .and_then(StringAttribute::try_from)
            .map(|a| a.value())
            .unwrap()
    }

    /// Returns the single body Region of the StructDefOp.
    fn body_region(&self) -> RegionRef<'c, 'a> {
        unsafe { RegionRef::from_raw(llzkStruct_StructDefOpGetBodyRegion(self.to_raw())) }
    }

    /// Returns the single body Block within the StructDefOp's Region.
    fn body(&self) -> BlockRef<'c, 'a> {
        unsafe { BlockRef::from_raw(llzkStruct_StructDefOpGetBody(self.to_raw())) }
    }

    /// Returns the associated StructType to this op using the given const params instead of the
    /// parameters defined by the op.
    ///
    /// # Panics
    ///
    /// If the 'struct.def' op type is not `!struct.type`.
    fn type_with_params(&self, params: ArrayAttribute<'c>) -> StructType<'c> {
        unsafe {
            Type::from_raw(llzkStruct_StructDefOpGetTypeWithParams(
                self.to_raw(),
                params.to_raw(),
            ))
        }
        .try_into()
        .expect("StructDefOpLike::type error")
    }

    /// Returns the operation that defines the member with the given name, if present.
    ///
    /// # Panics
    ///
    /// If the nested symbol operation with the given name is not a `struct.member`.
    fn get_member_def(&self, name: &str) -> Option<MemberDefOpRef<'c, 'a>> {
        let raw_op = unsafe {
            llzkStruct_StructDefOpGetMemberDef(
                self.to_raw(),
                Identifier::new(self.context().to_ref(), name).to_raw(),
            )
        };
        if raw_op.ptr.is_null() {
            return None;
        }
        Some(
            unsafe { OperationRef::from_raw(raw_op) }
                .try_into()
                .expect("op of type 'struct.member'"),
        )
    }

    /// Returns the operation that defines the member with the given name, creating a new operation
    /// if not present.
    fn get_or_create_member_def<F>(&self, name: &str, f: F) -> Result<MemberDefOpRef<'c, 'a>, Error>
    where
        F: FnOnce() -> Result<MemberDefOp<'c>, Error>,
    {
        match self.get_member_def(name) {
            Some(f) => Ok(f),
            None => {
                let op = f()?;
                let region = self.region(0)?;
                let block = region
                    .first_block()
                    .unwrap_or_else(|| region.append_block(Block::new(&[])));

                let member_ref = block.append_operation(op.into());

                Ok(member_ref.try_into()?)
            }
        }
    }

    /// Fills the given array with the MemberDefOp operations inside this struct.
    ///
    /// # Panics
    ///
    /// If any of the result operations is not a `struct.member` op.
    fn get_member_defs(&self) -> Vec<MemberDefOpRef<'c, '_>> {
        let num_members =
            usize::try_from(unsafe { llzkStruct_StructDefOpGetNumMemberDefs(self.to_raw()) })
                .unwrap();
        let mut raw_ops: Vec<MlirOperation> = Vec::with_capacity(num_members);
        unsafe {
            llzkStruct_StructDefOpGetMemberDefs(self.to_raw(), raw_ops.as_mut_ptr());
            raw_ops.set_len(num_members);
        };
        raw_ops
            .into_iter()
            .map(|op| {
                unsafe { OperationRef::from_raw(op) }
                    .try_into()
                    .expect("op of type 'struct.member'")
            })
            .collect()
    }

    /// Returns true if the struct has members marked as columns.
    fn has_columns(&self) -> bool {
        unsafe { llzkStruct_StructDefOpHasColumns(self.to_raw()) }.value != 0
    }

    /// Returns a [`FuncDefOpRef`] reference to the operation that defines the witness computation
    /// of the struct.
    ///
    /// # Panics
    ///
    /// If the result operation is not a `function.def`.
    fn get_compute_func<'b>(&self) -> Option<FuncDefOpRef<'c, 'b>> {
        let raw_op = unsafe { llzkStruct_StructDefOpGetComputeFuncOp(self.to_raw()) };
        if raw_op.ptr.is_null() {
            return None;
        }
        Some(
            unsafe { OperationRef::from_raw(raw_op) }
                .try_into()
                .expect("op of type 'function.def'"),
        )
    }

    /// Returns a [`FuncDefOpRef`] reference to the operation that defines the constraints of the
    /// struct.
    ///
    /// # Panics
    ///
    /// If the result operation is not a `function.def`.
    fn get_constrain_func<'b>(&self) -> Option<FuncDefOpRef<'c, 'b>> {
        let raw_op = unsafe { llzkStruct_StructDefOpGetConstrainFuncOp(self.to_raw()) };
        if raw_op.ptr.is_null() {
            return None;
        }
        Some(
            unsafe { OperationRef::from_raw(raw_op) }
                .try_into()
                .expect("op of type 'function.def'"),
        )
    }

    /// Returns the names of all template parameters accessible by the struct,
    /// if the struct is within a template op. Otherwise, returns an empty vec.
    fn get_template_param_op_names(&self) -> Vec<FlatSymbolRefAttribute<'c>> {
        let num_attrs = usize::try_from(unsafe {
            llzkStruct_StructDefOpGetNumTemplateParamOpNames(self.to_raw())
        })
        .unwrap();
        let mut raw_attrs: Vec<MlirAttribute> = Vec::with_capacity(num_attrs);
        unsafe {
            llzkStruct_StructDefOpGetTemplateParamOpNames(self.to_raw(), raw_attrs.as_mut_ptr());
            raw_attrs.set_len(num_attrs);
        };
        raw_attrs
            .into_iter()
            .map(|attr| {
                FlatSymbolRefAttribute::try_from(unsafe { Attribute::from_raw(attr) }).unwrap()
            })
            .collect()
    }

    /// Returns the names of all template expressions accessible by the struct,
    /// if the struct is within a template op. Otherwise, returns an empty vec.
    fn get_template_expr_op_names(&self) -> Vec<FlatSymbolRefAttribute<'c>> {
        let num_attrs = usize::try_from(unsafe {
            llzkStruct_StructDefOpGetNumTemplateExprOpNames(self.to_raw())
        })
        .unwrap();
        let mut raw_attrs: Vec<MlirAttribute> = Vec::with_capacity(num_attrs);
        unsafe {
            llzkStruct_StructDefOpGetTemplateExprOpNames(self.to_raw(), raw_attrs.as_mut_ptr());
            raw_attrs.set_len(num_attrs);
        };
        raw_attrs
            .into_iter()
            .map(|attr| {
                FlatSymbolRefAttribute::try_from(unsafe { Attribute::from_raw(attr) }).unwrap()
            })
            .collect()
    }

    /// Returns a StringAttr with the fully qualified name of the struct.
    fn get_fully_qualified_name(&self) -> Attribute<'_> {
        todo!("melior does not have a SymbolRefAttribute type")
    }

    /// Returns true if the struct is the main entry point of the circuit.
    fn is_main_component(&self) -> bool {
        unsafe { llzkStruct_StructDefOpIsMainComponent(self.to_raw()) }
    }
}

/// Defines the mutable public API of the 'struct.def' op.
pub trait StructDefOpMutLike<'c: 'a, 'a>:
    StructDefOpLike<'c, 'a> + OperationMutLike<'c, 'a>
{
}

//===----------------------------------------------------------------------===//
// StructDefOp, StructDefOpRef, and StructDefOpRefMut
//===----------------------------------------------------------------------===//

llzk_op_type!(
    StructDefOp,
    llzkOperationIsA_Struct_StructDefOp,
    "struct.def"
);

impl<'a, 'c: 'a> StructDefOpLike<'c, 'a> for StructDefOp<'c> {}

impl<'a, 'c: 'a> StructDefOpLike<'c, 'a> for StructDefOpRef<'c, 'a> {}

impl<'a, 'c: 'a> StructDefOpLike<'c, 'a> for StructDefOpRefMut<'c, 'a> {}

impl<'a, 'c: 'a> StructDefOpMutLike<'c, 'a> for StructDefOp<'c> {}

impl<'a, 'c: 'a> StructDefOpMutLike<'c, 'a> for StructDefOpRefMut<'c, 'a> {}

//===----------------------------------------------------------------------===//
// MemberDefOpLike
//===----------------------------------------------------------------------===//

/// Defines the public API of the 'struct.member' op.
pub trait MemberDefOpLike<'c: 'a, 'a>: OperationLike<'c, 'a> {
    /// Returns true if the member op has a `llzk.pub` attribute.
    fn has_public_attr(&self) -> bool {
        unsafe { llzkStruct_MemberDefOpHasPublicAttr(self.to_raw()) }
    }

    /// Sets or unsets the `llzk.pub` attribute.
    fn set_public_attr(&self, value: bool) {
        unsafe {
            llzkStruct_MemberDefOpSetPublicAttr(self.to_raw(), value);
        }
    }

    /// Returns the name of the member.
    ///
    /// # Panics
    ///
    /// If the 'struct.member' op doesn't have an attribute named `sym_name`.
    fn member_name(&self) -> &'c str {
        self.attribute("sym_name")
            .and_then(StringAttribute::try_from)
            .expect("malformed 'struct.member' op")
            .value()
    }

    /// Returns the type of the member.
    ///
    /// # Panics
    ///
    /// If the 'struct.member' op doesn't have a attribute named `type`.
    fn member_type(&self) -> Type<'c> {
        self.attribute("type")
            .and_then(TypeAttribute::try_from)
            .expect("malformed 'struct.member' op")
            .value()
    }
}

//===----------------------------------------------------------------------===//
// MemberDefOp, MemberDefOpRef, MemberDefOpRefMut
//===----------------------------------------------------------------------===//

llzk_op_type!(
    MemberDefOp,
    llzkOperationIsA_Struct_MemberDefOp,
    "struct.member"
);

impl<'a, 'c: 'a> MemberDefOpLike<'c, 'a> for MemberDefOp<'c> {}

impl<'a, 'c: 'a> MemberDefOpLike<'c, 'a> for MemberDefOpRef<'c, 'a> {}

impl<'a, 'c: 'a> MemberDefOpLike<'c, 'a> for MemberDefOpRefMut<'c, 'a> {}

//===----------------------------------------------------------------------===//
// Operation factories
//===----------------------------------------------------------------------===//

/// Creates a 'struct.def' op
pub fn def<'c, I>(
    location: Location<'c>,
    name: &str,
    region_ops: I,
) -> Result<StructDefOp<'c>, Error>
where
    I: IntoIterator<Item = Result<Operation<'c>, Error>>,
{
    let ctx = location.context();
    let region = Region::new();
    let block = Block::new(&[]);
    region_ops
        .into_iter()
        .try_for_each(|op| -> Result<(), Error> {
            block.append_operation(op?);
            Ok(())
        })?;
    region.append_block(block);
    let name: Attribute = StringAttribute::new(unsafe { ctx.to_ref() }, name).into();
    let attrs = [(ident!(ctx, "sym_name"), name)];

    OperationBuilder::new("struct.def", location)
        .add_attributes(&attrs)
        .add_regions([region])
        .build()
        .map_err(Into::into)
        .and_then(TryInto::try_into)
}

/// Return `true` iff the given op is `struct.def`.
#[inline]
pub fn is_struct_def<'c: 'a, 'a>(op: &impl OperationLike<'c, 'a>) -> bool {
    crate::operation::isa(op, "struct.def")
}

/// Creates a 'struct.member' op
pub fn member<'c, T>(
    location: Location<'c>,
    name: &str,
    r#type: T,
    is_column: bool,
    is_public: bool,
) -> Result<MemberDefOp<'c>, Error>
where
    T: Into<Type<'c>>,
{
    let ctx = location.context();
    let r#type = TypeAttribute::new(r#type.into());
    let mut builder = OperationBuilder::new("struct.member", location).add_attributes(&[
        (
            ident!(ctx, "sym_name"),
            StringAttribute::new(unsafe { ctx.to_ref() }, name).into(),
        ),
        (ident!(ctx, "type"), r#type.into()),
    ]);

    builder = if is_column {
        builder.add_attributes(&[(
            ident!(ctx, "column"),
            Attribute::unit(unsafe { ctx.to_ref() }),
        )])
    } else {
        builder
    };

    builder
        .build()
        .map_err(Into::into)
        .and_then(TryInto::try_into)
        .inspect(|op: &MemberDefOp<'c>| op.set_public_attr(is_public))
}

/// Return `true` iff the given op is `struct.member`.
#[inline]
pub fn is_struct_member<'c: 'a, 'a>(op: &impl OperationLike<'c, 'a>) -> bool {
    crate::operation::isa(op, "struct.member")
}

/// Creates a 'struct.readm' op
pub fn readm<'c>(
    builder: &OpBuilder<'c>,
    location: Location<'c>,
    result_type: Type<'c>,
    component: Value<'c, '_>,
    member_name: &str,
) -> Result<Operation<'c>, Error> {
    unsafe {
        let raw = llzkStruct_MemberReadOpBuild(
            builder.to_raw(),
            location.to_raw(),
            result_type.to_raw(),
            component.to_raw(),
            Identifier::new(result_type.context().to_ref(), member_name).to_raw(),
        );
        if raw.ptr.is_null() {
            Err(Error::BuildMethodFailed("readm"))
        } else {
            Ok(Operation::from_raw(raw))
        }
    }
}

/// Creates a 'struct.readm' op.
///
/// This factory method is not implemented yet.
pub fn readm_with_offset<'c>() -> Operation<'c> {
    todo!()
}

/// Return `true` iff the given op is `struct.readm`.
#[inline]
pub fn is_struct_readm<'c: 'a, 'a>(op: &impl OperationLike<'c, 'a>) -> bool {
    crate::operation::isa(op, "struct.readm")
}

/// Creates a 'struct.writem' op.
pub fn writem<'c>(
    location: Location<'c>,
    component: Value<'c, '_>,
    member_name: &str,
    value: Value<'c, '_>,
) -> Result<Operation<'c>, Error> {
    let context = location.context();
    let member_name = FlatSymbolRefAttribute::new(unsafe { context.to_ref() }, member_name);
    let attrs = [(ident!(context, "member_name"), member_name.into())];
    OperationBuilder::new("struct.writem", location)
        .add_operands(&[component, value])
        .add_attributes(&attrs)
        .build()
        .map_err(Into::into)
}

/// Return `true` iff the given op is `struct.writem`.
#[inline]
pub fn is_struct_writem<'c: 'a, 'a>(op: &impl OperationLike<'c, 'a>) -> bool {
    crate::operation::isa(op, "struct.writem")
}

/// Creates a 'struct.new' op
pub fn new<'c>(location: Location<'c>, r#type: StructType<'c>) -> Operation<'c> {
    OperationBuilder::new("struct.new", location)
        .add_results(&[r#type.into()])
        .build()
        .expect("valid operation")
}

/// Return `true` iff the given op is `struct.new`.
#[inline]
pub fn is_struct_new<'c: 'a, 'a>(op: &impl OperationLike<'c, 'a>) -> bool {
    crate::operation::isa(op, "struct.new")
}