program.rs

// SPDX-License-Identifier: CC0-1.0

use std::sync::Arc;

use super::Extractor;
use simplicity::node::{
    CoreConstructible as _, DisconnectConstructible as _, JetConstructible as _,
    WitnessConstructible as _,
};
use simplicity::types;
use simplicity::{jet::Core, Cmr, ConstructNode, FailEntropy};

/// Structure controlling the kind of programs generated by the fuzzer.
pub struct ProgramControl {
    /// Whether to attempt to insert "type bombs" into the generated program.
    ///
    /// Importantly, type bombs may have 2^n nodes for fairly large n, and will
    /// not respect `max_nodes`. So if you are trying to generate small programs
    /// you should not enable this.
    pub enable_type_bomb: bool,
    /// Whether to attempt to insert disconnect nodes into the generated program.
    pub enable_disconnect: bool,
    /// Whether to attempt to insert witness nodes into the generated program.
    pub enable_witness: bool,
    /// Whether to attempt to insert fail nodes into the generated program.
    pub enable_fail: bool,
    /// Whether to attempt to insert assertl and assertr nodes into the generated program.
    pub enable_asserts: bool,
    /// Maximum number of nodes a generated program may have. This limit may not
    /// be exactly enforced. If it is `None`, no limit is enforced.
    pub max_nodes: Option<usize>,
}

impl ProgramControl {
    fn from_u16(u: u16) -> Self {
        ProgramControl {
            enable_type_bomb: u & 0x8000 == 0x8000,
            enable_disconnect: u & 0x4000 == 0x4000,
            enable_witness: u & 0x2000 == 0x2000,
            enable_fail: u & 0x1000 == 0x1000,
            enable_asserts: u & 0x0800 == 0x0800,
            max_nodes: Some(5 * usize::from(u & 0x07ff)),
        }
    }
}

impl Extractor<'_> {
    pub fn extract_core_construct_node(
        &mut self,
        force_control: Option<ProgramControl>,
    ) -> Option<Arc<ConstructNode<Core>>> {
        type ArcNode = Arc<ConstructNode<Core>>;

        let ctx = types::Context::new();
        let mut stack: Vec<ArcNode> = vec![];

        let program_control =
            force_control.unwrap_or(ProgramControl::from_u16(self.extract_u16()?));

        let mut count = 0usize;
        for _ in 0..program_control.max_nodes.unwrap_or(usize::MAX) {
            let control = self.extract_u8()?;
            if program_control.enable_type_bomb && control & 0x80 == 0x80 {
                let mut ret = stack.pop()?;
                // Special-case: type bomb. Iterate x -> pair(x, x) on the top stack
                // item up to 128 times. Its CPU cost and target type will blow up
                // by a factor 2^128. If its target type has nonzero size this should
                // fail to construct; if it's 0 we should be able to construct it but
                // the bit machine should reject it.
                for _ in 0..control & 0x7f {
                    // FIXME should we refuse to make the type-bomb if `ret` contains any
                    //  witness or disconnect nodes? In this case the encoding of our
                    //  CommitNode won't round-trip, since we're force-sharing both children
                    //  of this `pair` but when decoding `CommitNode` we reject anything that
                    //  shares witnesses or disconnects, which at commit-time we treat as
                    //  being unique and never shared.
                    ret = ArcNode::pair(&ret, &ret).unwrap();
                }
                stack.push(ret);
            } else {
                match control {
                    // Return whatever we've got (note that this will "waste" everything else
                    // on the stack)
                    0 => {
                        if stack.len() == 1 {
                            return stack.pop();
                        } else {
                            return None;
                        }
                    }
                    // 1 through 63
                    1 => stack.push(ArcNode::unit(&ctx)),
                    2 => stack.push(ArcNode::iden(&ctx)),
                    3 => {
                        use simplicity::dag::DagLike as _;

                        let val = self.extract_value_direct()?;
                        if program_control.max_nodes.is_some() {
                            for _ in val.as_ref().pre_order_iter::<simplicity::dag::NoSharing>() {
                                count = count.checked_add(1)?;
                            }
                        }
                        if let Some(max) = program_control.max_nodes {
                            if val.compact_len() > max {
                                return None;
                            }
                        }
                        stack.push(ArcNode::scribe(&ctx, &val));
                    }
                    4 if program_control.enable_witness => stack.push(ArcNode::witness(&ctx, None)),
                    5 => {
                        let child = stack.pop()?;
                        stack.push(ArcNode::injl(&child));
                    }
                    6 => {
                        let child = stack.pop()?;
                        stack.push(ArcNode::injr(&child));
                    }
                    7 => {
                        let child = stack.pop()?;
                        stack.push(ArcNode::drop_(&child));
                    }
                    8 => {
                        let child = stack.pop()?;
                        stack.push(ArcNode::take(&child));
                    }
                    9 => {
                        let child = stack.pop()?;
                        let cmr_u8 = self.extract_u8()?;
                        let cmr = Cmr::from_byte_array([cmr_u8; 32]);
                        stack.push(ArcNode::assertl(&child, cmr).ok()?);
                    }
                    10 => {
                        let child = stack.pop()?;
                        let cmr_u8 = self.extract_u8()?;
                        let cmr = Cmr::from_byte_array([cmr_u8; 32]);
                        stack.push(ArcNode::assertr(cmr, &child).ok()?);
                    }
                    11 if program_control.enable_fail => {
                        let fail_u8 = self.extract_u8()?;
                        let fail = FailEntropy::from_byte_array([fail_u8; 64]);
                        stack.push(ArcNode::fail(&ctx, fail));
                    }
                    12 => {
                        let rchild = stack.pop()?;
                        let lchild = stack.pop()?;
                        stack.push(ArcNode::pair(&lchild, &rchild).ok()?);
                    }
                    13 => {
                        let rchild = stack.pop()?;
                        let lchild = stack.pop()?;
                        stack.push(ArcNode::case(&lchild, &rchild).ok()?);
                    }
                    14 => {
                        let rchild = stack.pop()?;
                        let lchild = stack.pop()?;
                        stack.push(ArcNode::comp(&lchild, &rchild).ok()?);
                    }
                    15 if program_control.enable_disconnect => {
                        let child = stack.pop()?;
                        stack.push(ArcNode::disconnect(&child, &None).ok()?);
                    }
                    // We assume that the above cases did not cover 64-255, so that if we
                    // right-shift by 6 we can get all 4 values.
                    _ => {
                        let extra_bits = usize::from(control >> 6);
                        let idx = (extra_bits << 8) + usize::from(self.extract_u8()?);
                        stack.push(ArcNode::jet(&ctx, Core::ALL[idx % Core::ALL.len()]));
                    }
                }
            }

            if let Some(max) = program_control.max_nodes {
                count = count.checked_add(1)?;
                if count > max {
                    return None;
                }
            }
        }

        None
    }
}