tx3_lang/

lowering.rs

//! Lowers the Tx3 language to the intermediate representation.
//!
//! This module takes an AST and performs lowering on it. It converts the AST
//! into the intermediate representation (IR) of the Tx3 language.

use std::cell::RefCell;
use std::rc::Rc;

use crate::ast;
use tx3_tir::model::core::{Type, UtxoRef};
use tx3_tir::model::v1beta0 as ir;

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("missing analyze phase for {0}")]
    MissingAnalyzePhase(String),

    #[error("symbol '{0}' expected to be '{1}'")]
    InvalidSymbol(String, &'static str),

    #[error("symbol '{0}' expected to be of type '{1}'")]
    InvalidSymbolType(String, &'static str),

    #[error("invalid ast: {0}")]
    InvalidAst(String),

    #[error("invalid property {0} on type {1:?}")]
    InvalidProperty(String, String),

    #[error("missing required field {0} for {1:?}")]
    MissingRequiredField(String, &'static str),

    #[error("failed to decode hex string {0}")]
    DecodeHexError(String),
}

#[inline]
fn hex_decode(s: &str) -> Result<Vec<u8>, Error> {
    hex::decode(s).map_err(|_| Error::DecodeHexError(s.to_string()))
}

fn expect_type_def(ident: &ast::Identifier) -> Result<&ast::TypeDef, Error> {
    let symbol = ident
        .symbol
        .as_ref()
        .ok_or(Error::MissingAnalyzePhase(ident.value.clone()))?;

    symbol
        .as_type_def()
        .ok_or(Error::InvalidSymbol(ident.value.clone(), "TypeDef"))
}

fn expect_alias_def(ident: &ast::Identifier) -> Result<&ast::AliasDef, Error> {
    let symbol = ident
        .symbol
        .as_ref()
        .ok_or(Error::MissingAnalyzePhase(ident.value.clone()))?;

    symbol
        .as_alias_def()
        .ok_or(Error::InvalidSymbol(ident.value.clone(), "AliasDef"))
}

fn expect_case_def(ident: &ast::Identifier) -> Result<&ast::VariantCase, Error> {
    let symbol = ident
        .symbol
        .as_ref()
        .ok_or(Error::MissingAnalyzePhase(ident.value.clone()))?;

    symbol
        .as_variant_case()
        .ok_or(Error::InvalidSymbol(ident.value.clone(), "VariantCase"))
}

#[allow(dead_code)]
fn expect_field_def(ident: &ast::Identifier) -> Result<&ast::RecordField, Error> {
    let symbol = ident
        .symbol
        .as_ref()
        .ok_or(Error::MissingAnalyzePhase(ident.value.clone()))?;

    symbol
        .as_field_def()
        .ok_or(Error::InvalidSymbol(ident.value.clone(), "FieldDef"))
}

fn coerce_identifier_into_asset_def(identifier: &ast::Identifier) -> Result<ast::AssetDef, Error> {
    match identifier.try_symbol()? {
        ast::Symbol::AssetDef(x) => Ok(x.as_ref().clone()),
        _ => Err(Error::InvalidSymbol(identifier.value.clone(), "AssetDef")),
    }
}

/// Reference-script UTxOs collected during lowering, drained into
/// `Tx.references`.
#[derive(Debug, Default)]
struct RefAccumulator {
    refs: Vec<ir::Expression>,
}

impl RefAccumulator {
    fn record(&mut self, r#ref: ir::Expression) {
        if !self.refs.contains(&r#ref) {
            self.refs.push(r#ref);
        }
    }
}

#[derive(Debug, Default, Clone)]
pub(crate) struct Context {
    is_asset_expr: bool,
    is_datum_expr: bool,
    is_address_expr: bool,
    // Within this subtree, a ref-backed policy's script runs, so its ref UTxO
    // is captured. Sticky across `enter_*`.
    capture_policy_ref: bool,
    // Shared across `enter_*` clones so captures from any depth reach `Tx`.
    script_refs: Rc<RefCell<RefAccumulator>>,
}

impl Context {
    pub fn enter_asset_expr(&self) -> Self {
        Self {
            is_asset_expr: true,
            is_datum_expr: false,
            is_address_expr: false,
            capture_policy_ref: self.capture_policy_ref,
            script_refs: self.script_refs.clone(),
        }
    }

    pub fn enter_datum_expr(&self) -> Self {
        Self {
            is_asset_expr: false,
            is_datum_expr: true,
            is_address_expr: false,
            capture_policy_ref: self.capture_policy_ref,
            script_refs: self.script_refs.clone(),
        }
    }

    pub fn enter_address_expr(&self) -> Self {
        Self {
            is_asset_expr: false,
            is_datum_expr: false,
            is_address_expr: true,
            capture_policy_ref: self.capture_policy_ref,
            script_refs: self.script_refs.clone(),
        }
    }

    /// Mark this subtree as one where a referenced policy's script runs, so its
    /// ref UTxO is captured. Sticky across nested `enter_*`.
    pub fn capturing_policy_refs(&self) -> Self {
        Self {
            capture_policy_ref: true,
            ..self.clone()
        }
    }

    pub fn is_address_expr(&self) -> bool {
        self.is_address_expr
    }

    pub fn is_asset_expr(&self) -> bool {
        self.is_asset_expr
    }

    pub fn is_datum_expr(&self) -> bool {
        self.is_datum_expr
    }

    pub fn captures_policy_refs(&self) -> bool {
        self.capture_policy_ref
    }

    /// Record a reference-script UTxO, deduplicating.
    pub fn record_script_ref(&self, r#ref: ir::Expression) {
        self.script_refs.borrow_mut().record(r#ref);
    }

    /// Take all recorded reference-script UTxOs, in discovery order.
    pub fn drain_script_refs(&self) -> Vec<ir::Expression> {
        std::mem::take(&mut self.script_refs.borrow_mut().refs)
    }
}

pub(crate) trait IntoLower {
    type Output;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error>;
}

impl<T> IntoLower for Option<&T>
where
    T: IntoLower,
{
    type Output = Option<T::Output>;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        self.map(|x| x.into_lower(ctx)).transpose()
    }
}

impl<T> IntoLower for Box<T>
where
    T: IntoLower,
{
    type Output = T::Output;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        self.as_ref().into_lower(ctx)
    }
}

impl IntoLower for ast::Identifier {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let symbol = self
            .symbol
            .as_ref()
            .ok_or(Error::MissingAnalyzePhase(self.value.clone()))?;

        match symbol {
            ast::Symbol::ParamVar(n, ty) => {
                Ok(ir::Param::ExpectValue(n.to_lowercase().clone(), ty.into_lower(ctx)?).into())
            }
            ast::Symbol::LocalExpr(expr) => Ok(expr.into_lower(ctx)?),
            ast::Symbol::PartyDef(x) => Ok(ir::Param::ExpectValue(
                x.name.value.to_lowercase().clone(),
                Type::Address,
            )
            .into()),
            ast::Symbol::Input(def) => {
                let inner = def.into_lower(ctx)?.utxos;

                let out = if ctx.is_asset_expr() {
                    ir::Coerce::IntoAssets(inner).into()
                } else if ctx.is_datum_expr() {
                    ir::Coerce::IntoDatum(inner).into()
                } else {
                    inner
                };

                Ok(out)
            }
            ast::Symbol::Reference(def) => def.into_lower(ctx),
            ast::Symbol::Fees => Ok(ir::Param::ExpectFees.into()),
            ast::Symbol::EnvVar(n, ty) => {
                Ok(ir::Param::ExpectValue(n.to_lowercase().clone(), ty.into_lower(ctx)?).into())
            }
            ast::Symbol::PolicyDef(x) => {
                let policy = x.into_lower(ctx)?;

                // Capture the ref UTxO only where the script runs, not in every
                // address position (an output `to` receives funds without
                // running the script). Hash-only policies yield `None`.
                if ctx.captures_policy_refs() {
                    if let Some(r#ref) = policy.script.as_utxo_ref() {
                        ctx.record_script_ref(r#ref);
                    }
                }

                if ctx.is_address_expr() {
                    Ok(ir::CompilerOp::BuildScriptAddress(policy.hash).into())
                } else {
                    Ok(policy.hash)
                }
            }
            ast::Symbol::Output(index) => Ok(ir::Expression::Number(*index as i128)),
            _ => {
                dbg!(&self);
                todo!();
            }
        }
    }
}

impl IntoLower for ast::UtxoRef {
    type Output = ir::Expression;

    fn into_lower(&self, _: &Context) -> Result<Self::Output, Error> {
        let x = ir::Expression::UtxoRefs(vec![UtxoRef {
            txid: self.txid.clone(),
            index: self.index as u32,
        }]);

        Ok(x)
    }
}

impl IntoLower for ast::StructConstructor {
    type Output = ir::StructExpr;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let type_def = expect_type_def(&self.r#type)
            .or_else(|_| {
                expect_alias_def(&self.r#type).and_then(|alias_def| {
                    alias_def.resolve_alias_chain().ok_or_else(|| {
                        Error::InvalidAst("Alias does not resolve to a TypeDef".to_string())
                    })
                })
            })
            .map_err(|_| Error::InvalidSymbol(self.r#type.value.clone(), "TypeDef or AliasDef"))?;

        let constructor = type_def
            .find_case_index(&self.case.name.value)
            .ok_or(Error::InvalidAst("case not found".to_string()))?;

        let case_def = expect_case_def(&self.case.name)?;

        let mut fields = vec![];

        for (index, field_def) in case_def.fields.iter().enumerate() {
            let value = self.case.find_field_value(&field_def.name.value);

            if let Some(value) = value {
                fields.push(value.into_lower(ctx)?);
            } else {
                let spread_target = self
                    .case
                    .spread
                    .as_ref()
                    .expect("spread must be set for missing explicit field")
                    .into_lower(ctx)?;

                fields.push(ir::Expression::EvalBuiltIn(Box::new(
                    ir::BuiltInOp::Property(spread_target, ir::Expression::Number(index as i128)),
                )));
            }
        }

        Ok(ir::StructExpr {
            constructor,
            fields,
        })
    }
}

impl IntoLower for ast::PolicyField {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        match self {
            ast::PolicyField::Hash(x) => x.into_lower(ctx),
            ast::PolicyField::Script(x) => x.into_lower(ctx),
            ast::PolicyField::Ref(x) => x.into_lower(ctx),
        }
    }
}

impl IntoLower for ast::PolicyDef {
    type Output = ir::PolicyExpr;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        match &self.value {
            ast::PolicyValue::Assign(x) => {
                let out = ir::PolicyExpr {
                    name: self.name.value.clone(),
                    hash: ir::Expression::Hash(hex_decode(&x.value)?),
                    script: ir::ScriptSource::expect_parameter(self.name.value.clone()),
                };

                Ok(out)
            }
            ast::PolicyValue::Constructor(x) => {
                let hash = x
                    .find_field("hash")
                    .ok_or(Error::InvalidAst("Missing policy hash".to_string()))?
                    .into_lower(ctx)?;

                let rf = x.find_field("ref").map(|x| x.into_lower(ctx)).transpose()?;

                let script = x
                    .find_field("script")
                    .map(|x| x.into_lower(ctx))
                    .transpose()?;

                let script = match (rf, script) {
                    (Some(rf), Some(script)) => ir::ScriptSource::new_ref(rf, script),
                    (Some(rf), None) => {
                        ir::ScriptSource::expect_ref_input(self.name.value.clone(), rf)
                    }
                    (None, Some(script)) => ir::ScriptSource::new_embedded(script),
                    (None, None) => ir::ScriptSource::expect_parameter(self.name.value.clone()),
                };

                Ok(ir::PolicyExpr {
                    name: self.name.value.clone(),
                    hash,
                    script,
                })
            }
        }
    }
}

impl IntoLower for ast::Type {
    type Output = Type;

    fn into_lower(&self, _: &Context) -> Result<Self::Output, Error> {
        match self {
            ast::Type::Undefined => Ok(Type::Undefined),
            ast::Type::Unit => Ok(Type::Unit),
            ast::Type::Int => Ok(Type::Int),
            ast::Type::Bool => Ok(Type::Bool),
            ast::Type::Bytes => Ok(Type::Bytes),
            ast::Type::Address => Ok(Type::Address),
            ast::Type::Utxo => Ok(Type::Utxo),
            ast::Type::UtxoRef => Ok(Type::UtxoRef),
            ast::Type::AnyAsset => Ok(Type::AnyAsset),
            ast::Type::List(_) => Ok(Type::List),
            ast::Type::Map(_, _) => Ok(Type::Map),
            ast::Type::Tuple(_) => Ok(Type::Tuple),
            ast::Type::Custom(x) => Ok(Type::Custom(x.value.clone())),
        }
    }
}

impl IntoLower for ast::AddOp {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let left = self.lhs.into_lower(ctx)?;
        let right = self.rhs.into_lower(ctx)?;

        Ok(ir::Expression::EvalBuiltIn(Box::new(ir::BuiltInOp::Add(
            left, right,
        ))))
    }
}

impl IntoLower for ast::SubOp {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let left = self.lhs.into_lower(ctx)?;
        let right = self.rhs.into_lower(ctx)?;

        Ok(ir::Expression::EvalBuiltIn(Box::new(ir::BuiltInOp::Sub(
            left, right,
        ))))
    }
}

impl IntoLower for ast::MulOp {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let left = self.lhs.into_lower(ctx)?;
        let right = self.rhs.into_lower(ctx)?;

        Ok(ir::Expression::EvalBuiltIn(Box::new(ir::BuiltInOp::Mul(
            left, right,
        ))))
    }
}

impl IntoLower for ast::DivOp {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let left = self.lhs.into_lower(ctx)?;
        let right = self.rhs.into_lower(ctx)?;

        Ok(ir::Expression::EvalBuiltIn(Box::new(ir::BuiltInOp::Div(
            left, right,
        ))))
    }
}

impl IntoLower for ast::ConcatOp {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let left = self.lhs.into_lower(ctx)?;
        let right = self.rhs.into_lower(ctx)?;

        Ok(ir::Expression::EvalBuiltIn(Box::new(
            ir::BuiltInOp::Concat(left, right),
        )))
    }
}

impl IntoLower for ast::NegateOp {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let operand = self.operand.into_lower(ctx)?;

        Ok(ir::Expression::EvalBuiltIn(Box::new(
            ir::BuiltInOp::Negate(operand),
        )))
    }
}

impl IntoLower for ast::FnCall {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        // A callee that resolves to a function definition is either a built-in
        // (lowered to a dedicated compiler op) or a user-defined function
        // (inlined below).
        if let Some(fn_def) = self.callee.symbol.as_ref().and_then(|s| s.as_fn_def()) {
            if let Some(builtin) = fn_def.builtin {
                return crate::builtins::resolve(builtin).lower_call(&self.args, ctx);
            }

            // Inline a user-defined function: lower its analyzed body, then
            // substitute each parameter with the lowered call argument.
            let body = fn_def.body.as_ref().ok_or_else(|| {
                Error::InvalidAst(format!(
                    "function '{}' has neither a body nor a built-in kind",
                    fn_def.name.value
                ))
            })?;

            let lowered_body = body.result.into_lower(ctx)?;

            let mut subs = std::collections::HashMap::new();
            for (param, arg) in fn_def.parameters.parameters.iter().zip(&self.args) {
                subs.insert(param.name.value.to_lowercase(), arg.into_lower(ctx)?);
            }

            use tx3_tir::Node;
            let mut visitor = ParamSubstituter { subs: &subs };
            return lowered_body
                .apply(&mut visitor)
                .map_err(|e| Error::InvalidAst(e.to_string()));
        }

        // Otherwise the callee must name an asset; treat the call as an asset
        // constructor.
        match coerce_identifier_into_asset_def(&self.callee) {
            Ok(asset_def) => {
                let policy = asset_def.policy.into_lower(ctx)?;
                let asset_name = asset_def.asset_name.into_lower(ctx)?;
                let amount = self.args[0].into_lower(ctx)?;

                Ok(ir::Expression::Assets(vec![ir::AssetExpr {
                    policy,
                    asset_name,
                    amount,
                }]))
            }
            Err(_) => Err(Error::InvalidAst(format!(
                "unknown function: {}",
                self.callee.value
            ))),
        }
    }
}

/// TIR visitor used by function inlining to replace each `EvalParam` standing
/// for a function parameter with the lowered call argument.
struct ParamSubstituter<'a> {
    subs: &'a std::collections::HashMap<String, ir::Expression>,
}

impl tx3_tir::Visitor for ParamSubstituter<'_> {
    fn reduce(
        &mut self,
        expr: ir::Expression,
    ) -> Result<ir::Expression, tx3_tir::reduce::Error> {
        if let ir::Expression::EvalParam(ref param) = expr {
            if let ir::Param::ExpectValue(name, _) = param.as_ref() {
                if let Some(replacement) = self.subs.get(name) {
                    return Ok(replacement.clone());
                }
            }
        }
        Ok(expr)
    }
}

impl IntoLower for ast::PropertyOp {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let object = self.operand.into_lower(ctx)?;

        let ty = self
            .operand
            .target_type()
            .ok_or(Error::MissingAnalyzePhase(format!("{0:?}", self.operand)))?;

        let prop_index =
            ty.property_index(*self.property.clone())
                .ok_or(Error::InvalidProperty(
                    format!("{:?}", self.property),
                    ty.to_string(),
                ))?;

        Ok(ir::Expression::EvalBuiltIn(Box::new(
            ir::BuiltInOp::Property(object, prop_index.into_lower(ctx)?),
        )))
    }
}

impl IntoLower for ast::ListConstructor {
    type Output = Vec<ir::Expression>;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let elements = self
            .elements
            .iter()
            .map(|x| x.into_lower(ctx))
            .collect::<Result<Vec<_>, _>>()?;

        Ok(elements)
    }
}

impl IntoLower for ast::TupleConstructor {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let elements = self
            .elements
            .iter()
            .map(|x| x.into_lower(ctx))
            .collect::<Result<Vec<_>, _>>()?;

        Ok(ir::Expression::Tuple(elements))
    }
}

impl IntoLower for ast::MapConstructor {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let pairs = self
            .fields
            .iter()
            .map(|field| {
                let key = field.key.into_lower(ctx)?;
                let value = field.value.into_lower(ctx)?;
                Ok((key, value))
            })
            .collect::<Result<Vec<_>, _>>()?;

        Ok(ir::Expression::Map(pairs))
    }
}

impl IntoLower for ast::DataExpr {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let out = match self {
            ast::DataExpr::None => ir::Expression::None,
            ast::DataExpr::Number(x) => Self::Output::Number(*x as i128),
            ast::DataExpr::Bool(x) => ir::Expression::Bool(*x),
            ast::DataExpr::String(x) => ir::Expression::String(x.value.clone()),
            ast::DataExpr::HexString(x) => ir::Expression::Bytes(hex_decode(&x.value)?),
            ast::DataExpr::StructConstructor(x) => ir::Expression::Struct(x.into_lower(ctx)?),
            ast::DataExpr::ListConstructor(x) => ir::Expression::List(x.into_lower(ctx)?),
            ast::DataExpr::MapConstructor(x) => x.into_lower(ctx)?,
            ast::DataExpr::TupleConstructor(x) => x.into_lower(ctx)?,
            ast::DataExpr::AnyAssetConstructor(x) => x.into_lower(ctx)?,
            ast::DataExpr::Unit => ir::Expression::Struct(ir::StructExpr::unit()),
            ast::DataExpr::Identifier(x) => x.into_lower(ctx)?,
            ast::DataExpr::AddOp(x) => x.into_lower(ctx)?,
            ast::DataExpr::SubOp(x) => x.into_lower(ctx)?,
            ast::DataExpr::MulOp(x) => x.into_lower(ctx)?,
            ast::DataExpr::DivOp(x) => x.into_lower(ctx)?,
            ast::DataExpr::ConcatOp(x) => x.into_lower(ctx)?,
            ast::DataExpr::NegateOp(x) => x.into_lower(ctx)?,
            ast::DataExpr::PropertyOp(x) => x.into_lower(ctx)?,
            ast::DataExpr::UtxoRef(x) => x.into_lower(ctx)?,
            ast::DataExpr::FnCall(x) => x.into_lower(ctx)?,
        };

        Ok(out)
    }
}

impl IntoLower for ast::AnyAssetConstructor {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let ctx = &ctx.enter_datum_expr();
        let policy = self.policy.into_lower(ctx)?;

        let ctx = &ctx.enter_datum_expr();
        let asset_name = self.asset_name.into_lower(ctx)?;

        let ctx = &ctx.enter_datum_expr();
        let amount = self.amount.into_lower(ctx)?;

        Ok(ir::Expression::Assets(vec![ir::AssetExpr {
            policy,
            asset_name,
            amount,
        }]))
    }
}

impl IntoLower for ast::InputBlockField {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        match self {
            ast::InputBlockField::From(x) => {
                // Spending from a script address runs its script.
                let ctx = ctx.enter_address_expr().capturing_policy_refs();
                x.into_lower(&ctx)
            }
            ast::InputBlockField::DatumIs(_) => todo!(),
            ast::InputBlockField::MinAmount(x) => {
                let ctx = ctx.enter_asset_expr();
                x.into_lower(&ctx)
            }
            ast::InputBlockField::Redeemer(x) => {
                let ctx = ctx.enter_datum_expr();
                x.into_lower(&ctx)
            }
            ast::InputBlockField::Ref(x) => x.into_lower(ctx),
        }
    }
}

impl IntoLower for ast::InputBlock {
    type Output = ir::Input;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let from_field = self.find("from");

        let address = from_field.map(|x| x.into_lower(ctx)).transpose()?;

        let min_amount = self
            .find("min_amount")
            .map(|x| x.into_lower(ctx))
            .transpose()?;

        let r#ref = self.find("ref").map(|x| x.into_lower(ctx)).transpose()?;

        let redeemer = self
            .find("redeemer")
            .map(|x| x.into_lower(ctx))
            .transpose()?
            .unwrap_or(ir::Expression::None);

        let query = ir::InputQuery {
            address: address.unwrap_or(ir::Expression::None),
            min_amount: min_amount.unwrap_or(ir::Expression::None),
            r#ref: r#ref.unwrap_or(ir::Expression::None),
            many: self.many,
            collateral: false,
        };

        let param = ir::Param::ExpectInput(self.name.to_lowercase().clone(), query);

        let input = ir::Input {
            name: self.name.to_lowercase().clone(),
            utxos: param.into(),
            redeemer,
        };

        Ok(input)
    }
}

impl IntoLower for ast::OutputBlockField {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        match self {
            ast::OutputBlockField::To(x) => {
                let ctx = ctx.enter_address_expr();
                x.into_lower(&ctx)
            }
            ast::OutputBlockField::Amount(x) => {
                let ctx = ctx.enter_asset_expr();
                x.into_lower(&ctx)
            }
            ast::OutputBlockField::Datum(x) => {
                let ctx = ctx.enter_datum_expr();
                x.into_lower(&ctx)
            }
        }
    }
}

impl IntoLower for ast::OutputBlock {
    type Output = ir::Output;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let address = self.find("to").into_lower(ctx)?.unwrap_or_default();
        let datum = self.find("datum").into_lower(ctx)?.unwrap_or_default();
        let amount = self.find("amount").into_lower(ctx)?.unwrap_or_default();

        Ok(ir::Output {
            address,
            datum,
            amount,
            optional: self.optional,
        })
    }
}

impl IntoLower for ast::ValidityBlockField {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        match self {
            ast::ValidityBlockField::SinceSlot(x) => x.into_lower(ctx),
            ast::ValidityBlockField::UntilSlot(x) => x.into_lower(ctx),
        }
    }
}

impl IntoLower for ast::ValidityBlock {
    type Output = ir::Validity;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let since = self.find("since_slot").into_lower(ctx)?.unwrap_or_default();
        let until = self.find("until_slot").into_lower(ctx)?.unwrap_or_default();

        Ok(ir::Validity { since, until })
    }
}

impl IntoLower for ast::MintBlockField {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        match self {
            // Minting/burning runs the asset's policy script.
            ast::MintBlockField::Amount(x) => x.into_lower(&ctx.capturing_policy_refs()),
            ast::MintBlockField::Redeemer(x) => x.into_lower(ctx),
        }
    }
}

impl IntoLower for ast::MintBlock {
    type Output = ir::Mint;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let amount = self.find("amount").into_lower(ctx)?.unwrap_or_default();
        let redeemer = self.find("redeemer").into_lower(ctx)?.unwrap_or_default();

        Ok(ir::Mint { amount, redeemer })
    }
}

impl IntoLower for ast::MetadataBlockField {
    type Output = ir::Metadata;
    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        Ok(ir::Metadata {
            key: self.key.into_lower(ctx)?,
            value: self.value.into_lower(ctx)?,
        })
    }
}

impl IntoLower for ast::MetadataBlock {
    type Output = Vec<ir::Metadata>;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let fields = self
            .fields
            .iter()
            .map(|metadata_field| metadata_field.into_lower(ctx))
            .collect::<Result<Vec<_>, _>>()?;

        Ok(fields)
    }
}

impl IntoLower for ast::ChainSpecificBlock {
    type Output = ir::AdHocDirective;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        match self {
            ast::ChainSpecificBlock::Cardano(x) => x.into_lower(ctx),
        }
    }
}

impl IntoLower for ast::ReferenceBlock {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let r#ref = self.r#ref.into_lower(ctx)?;

        let query = ir::InputQuery {
            address: ir::Expression::None,
            min_amount: ir::Expression::None,
            r#ref,
            many: false,
            collateral: false,
        };

        let inner = ir::Param::ExpectInput(self.name.to_lowercase(), query).into();

        let out = if ctx.is_asset_expr() {
            ir::Coerce::IntoAssets(inner).into()
        } else if ctx.is_datum_expr() {
            ir::Coerce::IntoDatum(inner).into()
        } else {
            inner
        };

        Ok(out)
    }
}

impl IntoLower for ast::CollateralBlockField {
    type Output = ir::Expression;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        match self {
            ast::CollateralBlockField::From(x) => x.into_lower(ctx),
            ast::CollateralBlockField::MinAmount(x) => x.into_lower(ctx),
            ast::CollateralBlockField::Ref(x) => x.into_lower(ctx),
        }
    }
}

impl IntoLower for ast::CollateralBlock {
    type Output = ir::Collateral;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        let from = self.find("from").map(|x| x.into_lower(ctx)).transpose()?;

        let min_amount = self
            .find("min_amount")
            .map(|x| x.into_lower(ctx))
            .transpose()?;

        let r#ref = self.find("ref").map(|x| x.into_lower(ctx)).transpose()?;

        let query = ir::InputQuery {
            address: from.unwrap_or(ir::Expression::None),
            min_amount: min_amount.unwrap_or(ir::Expression::None),
            r#ref: r#ref.unwrap_or(ir::Expression::None),
            many: false,
            collateral: true,
        };

        let param = ir::Param::ExpectInput("collateral".to_string(), query);

        let collateral = ir::Collateral {
            utxos: param.into(),
        };

        Ok(collateral)
    }
}

impl IntoLower for ast::SignersBlock {
    type Output = ir::Signers;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        Ok(ir::Signers {
            signers: self
                .signers
                .iter()
                .map(|x| x.into_lower(ctx))
                .collect::<Result<Vec<_>, _>>()?,
        })
    }
}

impl IntoLower for ast::TxDef {
    type Output = ir::Tx;

    fn into_lower(&self, ctx: &Context) -> Result<Self::Output, Error> {
        // Seed with explicit `reference` blocks first so they dedup against,
        // and precede, refs that the body derives from ref-backed policies.
        for reference in self.references.iter() {
            let r#ref = reference.r#ref.into_lower(ctx)?;
            ctx.record_script_ref(r#ref);
        }

        let inputs = self
            .inputs
            .iter()
            .map(|x| x.into_lower(ctx))
            .collect::<Result<Vec<_>, _>>()?;
        let outputs = self
            .outputs
            .iter()
            .map(|x| x.into_lower(ctx))
            .collect::<Result<Vec<_>, _>>()?;
        let validity = self
            .validity
            .as_ref()
            .map(|x| x.into_lower(ctx))
            .transpose()?;
        let mints = self
            .mints
            .iter()
            .map(|x| x.into_lower(ctx))
            .collect::<Result<Vec<_>, _>>()?;
        let burns = self
            .burns
            .iter()
            .map(|x| x.into_lower(ctx))
            .collect::<Result<Vec<_>, _>>()?;
        let adhoc = self
            .adhoc
            .iter()
            .map(|x| x.into_lower(ctx))
            .collect::<Result<Vec<_>, _>>()?;
        let collateral = self
            .collateral
            .iter()
            .map(|x| x.into_lower(ctx))
            .collect::<Result<Vec<_>, _>>()?;
        let signers = self
            .signers
            .as_ref()
            .map(|x| x.into_lower(ctx))
            .transpose()?;
        let metadata = self
            .metadata
            .as_ref()
            .map(|x| x.into_lower(ctx))
            .transpose()?
            .unwrap_or(vec![]);

        let ir = ir::Tx {
            references: ctx.drain_script_refs(),
            inputs,
            outputs,
            validity,
            mints,
            burns,
            adhoc,
            fees: ir::Param::ExpectFees.into(),
            collateral,
            signers,
            metadata,
        };

        Ok(ir)
    }
}

pub fn lower_tx(ast: &ast::TxDef) -> Result<ir::Tx, Error> {
    let ctx = &Context::default();

    let tx = ast.into_lower(ctx)?;

    Ok(tx)
}

/// Lowers the Tx3 language to the intermediate representation.
///
/// This function takes an AST and converts it into the intermediate
/// representation (IR) of the Tx3 language.
///
/// # Arguments
///
/// * `ast` - The AST to lower
///
/// # Returns
///
/// * `Result<ir::Program, Error>` - The lowered intermediate representation
pub fn lower(ast: &ast::Program, template: &str) -> Result<ir::Tx, Error> {
    let tx = ast
        .txs
        .iter()
        .find(|x| x.name.value == template)
        .ok_or(Error::InvalidAst("tx not found".to_string()))?;

    lower_tx(tx)
}

#[cfg(test)]
mod tests {
    use assert_json_diff::assert_json_eq;
    use paste::paste;

    use super::*;
    use crate::parsing::{self};

    fn make_snapshot_if_missing(example: &str, name: &str, tx: &ir::Tx) {
        let manifest_dir = env!("CARGO_MANIFEST_DIR");

        let path = format!("{}/../../examples/{}.{}.tir", manifest_dir, example, name);

        if !std::fs::exists(&path).unwrap() {
            let ir = serde_json::to_string_pretty(tx).unwrap();
            std::fs::write(&path, ir).unwrap();
        }
    }

    /// Lowers every tx in an example, snapshot-checks each, and returns the
    /// lowered TIRs keyed by tx name so callers can assert extra invariants.
    fn test_lowering_example(example: &str) -> std::collections::BTreeMap<String, ir::Tx> {
        let manifest_dir = env!("CARGO_MANIFEST_DIR");
        let mut program = parsing::parse_well_known_example(example);

        crate::analyzing::analyze(&mut program).ok().unwrap();

        let mut lowered = std::collections::BTreeMap::new();

        for tx in program.txs.iter() {
            let tir = lower(&program, &tx.name.value).unwrap();

            make_snapshot_if_missing(example, &tx.name.value, &tir);

            let tir_file = format!(
                "{}/../../examples/{}.{}.tir",
                manifest_dir, example, tx.name.value
            );

            let expected = std::fs::read_to_string(tir_file).unwrap();
            let expected: ir::Tx = serde_json::from_str(&expected).unwrap();

            assert_json_eq!(tir, expected);

            lowered.insert(tx.name.value.clone(), tir);
        }

        lowered
    }

    #[macro_export]
    macro_rules! test_lowering {
        ($name:ident) => {
            paste! {
                #[test]
                fn [<test_example_ $name>]() {
                    test_lowering_example(stringify!($name));
                }
            }
        };
        // Variant with extra assertions: the block receives the lowered TIRs
        // keyed by tx name (a `BTreeMap<String, ir::Tx>`) bound to `$txs`.
        ($name:ident, |$txs:ident| $checks:block) => {
            paste! {
                #[test]
                fn [<test_example_ $name>]() {
                    let $txs = test_lowering_example(stringify!($name));
                    $checks
                }
            }
        };
    }

    test_lowering!(lang_tour);

    test_lowering!(tuples);

    test_lowering!(transfer);

    test_lowering!(swap);

    test_lowering!(asteria);

    test_lowering!(vesting);

    test_lowering!(faucet);

    test_lowering!(input_datum);

    test_lowering!(env_vars);

    test_lowering!(local_vars);

    test_lowering!(cardano_witness);

    test_lowering!(reference_script);

    test_lowering!(policy_reference_script, |txs| {
        // ref-backed policy as `from`
        assert_eq!(txs["spend"].references.len(), 1);
        // same ref across two inputs is deduped
        assert_eq!(txs["spend_two"].references.len(), 1);
        // hash-only policy: nothing to reference
        assert!(txs["spend_hash_only"].references.is_empty());
        // ref-backed mint
        assert_eq!(txs["mint_token"].references.len(), 1);
        // ref-backed burn
        assert_eq!(txs["burn_token"].references.len(), 1);
        // hash-only mint: nothing to reference
        assert!(txs["mint_hash_only"].references.is_empty());
        // output recipient only: script does not run, no reference input
        assert!(txs["send_to_policy"].references.is_empty());
        // ref-backed withdrawal stake credential
        assert_eq!(txs["withdraw"].references.len(), 1);
    });

    test_lowering!(withdrawal);

    test_lowering!(map);

    test_lowering!(burn);

    test_lowering!(min_utxo);

    test_lowering!(tip_slot);

    test_lowering!(posix_time);

    test_lowering!(donation);

    test_lowering!(list_concat);

    test_lowering!(buidler_fest_2026);

    test_lowering!(functions);

    test_lowering!(nested_functions);

    test_lowering!(param_field_shadow);

    test_lowering!(oracle_reference_datum);
}