diff --git a/src/ast.rs b/src/ast.rs index e3a084cb..1fce5956 100644 --- a/src/ast.rs +++ b/src/ast.rs @@ -303,6 +303,8 @@ pub enum CallName { Panic, /// [`dbg!`]. Debug, + /// Padding that inflates the program without changing what it computes. + Padding(NonZeroUsize), /// Cast from the given source type. TypeCast(ResolvedType), /// A custom function that was defined previously. @@ -332,6 +334,7 @@ impl PartialEq for CallName { (Self::Assert, Self::Assert) => true, (Self::Panic, Self::Panic) => true, (Self::Debug, Self::Debug) => true, + (Self::Padding(a), Self::Padding(b)) => a == b, (Self::TypeCast(a), Self::TypeCast(b)) => a == b, (Self::Custom(a), Self::Custom(b)) => a == b, (Self::Fold(a, b), Self::Fold(c, d)) => a == c && b == d, @@ -2175,6 +2178,12 @@ impl AbstractSyntaxTree for Call { scope.track_call(from, TrackedCallName::Debug(arg_ty)); args } + CallName::Padding(_) => { + let args_tys = []; + check_argument_types(from.args(), &args_tys).with_span(from)?; + check_output_type(&ResolvedType::unit(), ty).with_span(from)?; + analyze_arguments(from.args(), &args_tys, scope)? + } CallName::TypeCast(source) => { // Casts prove structural equality, but enums are nominal: // every enum must map to itself at its structural position @@ -2310,6 +2319,19 @@ impl AbstractSyntaxTree for CallName { parse::CallName::Assert => Ok(Self::Assert), parse::CallName::Panic => Ok(Self::Panic), parse::CallName::Debug => Ok(Self::Debug), + parse::CallName::Padding(size) => { + // The text parser rejects oversized padding, but a `parse::Program` + // can also be built directly (fuzzing), so re-check it here. + if size.get() > parse::MAX_PADDING_SIZE { + Err(Error::PaddingSizeTooLarge { + size: size.get(), + max: parse::MAX_PADDING_SIZE, + }) + .with_span(from) + } else { + Ok(Self::Padding(*size)) + } + } parse::CallName::TypeCast(target) => { scope.resolve(target).map(Self::TypeCast).with_span(from) } @@ -3646,3 +3668,91 @@ mod literal_tests { } } } + +#[cfg(test)] +mod padding_tests { + use crate::{error, parse::ParseFromStr}; + + use super::*; + + /// Parse `input` and return the single `padding::()` call it consists of. + fn parse_padding(input: &str) -> parse::Call { + let parsed = parse::Expression::parse_from_str(input).expect("Failed to parse"); + + match parsed.inner() { + parse::ExpressionInner::Single(single) => match single.inner() { + parse::SingleExpressionInner::Call(call) => match call.name() { + parse::CallName::Padding(_) => call.clone(), + name => panic!("Expected a padding call, found `{name}`"), + }, + _ => panic!("Expected a call expression"), + }, + _ => panic!("Expected a single expression"), + } + } + + #[test] + fn test_ast_padding() { + let input = "padding::<10>()"; + + let parsed_call = &parse_padding(input); + + let mut scope = Scope::default(); + let ast_padding = Call::analyze(parsed_call, &ResolvedType::unit(), &mut scope) + .expect("Failed to analyze Padding expression"); + + assert_eq!( + ast_padding.args().len(), + 0, + "Args did not analyse correctly" + ); + assert_eq!( + ast_padding.name(), + &CallName::Padding(NonZeroUsize::new(10).unwrap()), + "Call name was not padding" + ); + } + + #[test] + fn test_ast_padding_rejects_size_above_max() { + // The text parser caps the size, but a parse tree can be built directly + // (as the fuzz targets do), so analysis has to reject it too. + let oversized = parse::CallName::Padding( + NonZeroUsize::new(parse::MAX_PADDING_SIZE + 1).expect("non-zero"), + ); + let call = parse::Call::new_unparsed(oversized, Arc::from([]), Span::new(0, 0..0)); + + let mut scope = Scope::default(); + let res = Call::analyze(&call, &ResolvedType::unit(), &mut scope); + + assert!( + matches!( + res.unwrap_err().error(), + error::Error::PaddingSizeTooLarge { size, max } + if *size == parse::MAX_PADDING_SIZE + 1 && *max == parse::MAX_PADDING_SIZE + ), + "oversized padding was accepted but should have been rejected" + ); + } + + #[test] + fn test_ast_padding_should_fail_with_args() { + let input = "padding::<10>(1)"; + + let parsed_call = &parse_padding(input); + + let mut scope = Scope::default(); + let res = Call::analyze(parsed_call, &ResolvedType::unit(), &mut scope); + + assert!( + matches!( + res.unwrap_err().error(), + error::Error::InvalidNumberOfArguments { + expected: 0, + found: 1 + } + ), + "padding parsed correctly but should have failed" + ); + } +} diff --git a/src/compile/mod.rs b/src/compile/mod.rs index ed16801a..ddfca893 100644 --- a/src/compile/mod.rs +++ b/src/compile/mod.rs @@ -448,6 +448,16 @@ impl Call { let iden = ProgNode::iden(scope.ctx()); scope.with_debug_symbol(args, &iden, self) } + CallName::Padding(size) => { + // Each layer composes one more `unit` onto the chain. Used for increasing + // the size of the program to increase the budget. + let unit = ProgNode::unit(scope.ctx()); + let mut padded = PairBuilder::unit(scope.ctx()); + for _ in 0..size.get() { + padded = padded.comp(&unit).with_span(self)?; + } + Ok(padded) + } CallName::TypeCast(..) => { // A cast converts between two structurally equal types. // Structural equality of SimplicityHL types A and B means @@ -759,3 +769,77 @@ impl EnumMatch { input.comp(&dispatch).with_span(self) } } + +#[cfg(test)] +mod padding_tests { + use crate::{ + ast::{self, ElementsJetHinter}, + parse::{self, ParseFromStr}, + }; + + use super::*; + + fn compile_program(input: &str) -> Result, Diagnostic> { + let parse_program = parse::Program::parse_from_str(input).expect("Failed to parse"); + let ast_program = ast::Program::analyze(&parse_program, Box::new(ElementsJetHinter)) + .expect("Failed to analyze"); + ast_program.compile(Arguments::default(), false, Box::new(ElementsJetHinter)) + } + + /// Compile `fn main() { padding::(); }` and return the encoded program length. + fn encoded_len(size: usize) -> usize { + let program = compile_program(&format!("fn main() {{ padding::<{size}>(); }}")) + .expect("padding expression should compile"); + named::forget_names(&program).to_vec_without_witness().len() + } + + #[test] + fn test_padding_compiles() { + let input_program = r#" + fn main() { + padding::<20>(); + }"#; + + compile_program(input_program).expect("padding expression should compile"); + } + + /// Count the nodes of the compiled program under the same maximal sharing + /// that the encoder applies. + fn shared_node_count(size: usize) -> usize { + let program = compile_program(&format!("fn main() {{ padding::<{size}>(); }}")) + .expect("padding expression should compile"); + let program = named::forget_names(&program); + simplicity::dag::DagLike::post_order_iter::< + simplicity::dag::MaxSharing, + >(&*program) + .count() + } + + #[test] + fn test_padding_layers_are_not_shared() { + // Each layer's `comp` has a distinct left child (the chain built so far), so it + // has a distinct CMR. Maximal sharing -- what the encoder applies -- therefore + // cannot collapse the chain. If it could, padding would be a no-op. + let base = shared_node_count(10); + let extra = shared_node_count(110); + + assert_eq!( + extra - base, + 100, + "expected one fresh node per padding layer, got {} across 100 extra layers", + extra - base + ); + } + + #[test] + fn test_padding_grows_the_program() { + // The whole point of padding is to make the program bigger. + let small = encoded_len(1); + let large = encoded_len(20); + + assert!( + large > small, + "padding::<20> encoded to {large} bytes, no larger than padding::<1> at {small}" + ); + } +} diff --git a/src/error.rs b/src/error.rs index 5384ca5e..77d46c26 100644 --- a/src/error.rs +++ b/src/error.rs @@ -895,6 +895,11 @@ pub enum Error { declared: ResolvedType, assigned: ResolvedType, }, + PaddingSizeZero, + PaddingSizeTooLarge { + size: usize, + max: usize, + }, } #[rustfmt::skip] @@ -1156,6 +1161,11 @@ impl fmt::Display for Error { f, "Parameter `{name}` was declared with type `{declared}` but its assigned argument is of type `{assigned}`" ), + Error::PaddingSizeZero => write!(f, "Padding size cannot be zero"), + Error::PaddingSizeTooLarge { size, max } => write!( + f, + "Expected a padding size of at most {max}, found {size}" + ), } } } diff --git a/src/lexer.rs b/src/lexer.rs index ffd15bee..f9188363 100644 --- a/src/lexer.rs +++ b/src/lexer.rs @@ -537,6 +537,80 @@ pub fn is_keyword(s: &str) -> bool { mod original_lexer { use super::*; + /// Helper function to get the variant name of a token + fn variant_name(token: &Token) -> &'static str { + match token { + Token::Pub => "Pub", + Token::Use => "Use", + Token::As => "As", + Token::Fn => "Fn", + Token::Let => "Let", + Token::Type => "Type", + Token::Mod => "Mod", + Token::Const => "Const", + Token::Match => "Match", + Token::Enum => "Enum", + Token::Crate => "Crate", + Token::Simc => "Simc", + Token::Arrow => "Arrow", + Token::DoubleColon => "DoubleColon", + Token::Colon => "Colon", + Token::Semi => "Semi", + Token::Comma => "Comma", + Token::Eq => "Eq", + Token::FatArrow => "FatArrow", + Token::LParen => "LParen", + Token::RParen => "RParen", + Token::LBracket => "LBracket", + Token::RBracket => "RBracket", + Token::LBrace => "LBrace", + Token::RBrace => "RBrace", + Token::LAngle => "LAngle", + Token::RAngle => "RAngle", + Token::DecLiteral(_) => "DecLiteral", + Token::HexLiteral(_) => "HexLiteral", + Token::BinLiteral(_) => "BinLiteral", + Token::Bool(_) => "Bool", + Token::Ident(_) => "Ident", + Token::Jet(_) => "Jet", + Token::Witness(_) => "Witness", + Token::Param(_) => "Param", + Token::Macro(_) => "Macro", + } + } + + /// Macro to assert that a sequence of tokens matches the expected variant types + macro_rules! assert_tokens_match { + ($tokens:expr, $($expected:ident),* $(,)?) => { + { + let tokens = $tokens.as_ref().expect("Expected Some tokens"); + let expected_variants = vec![$( stringify!($expected) ),*]; + + assert_eq!( + tokens.len(), + expected_variants.len(), + "Expected {} tokens, got {}.\nTokens: {:?}", + expected_variants.len(), + tokens.len(), + tokens + ); + + for (idx, ((token, _span), expected_variant)) in tokens.iter().zip(expected_variants.iter()).enumerate() { + let actual_variant = variant_name(token); + assert_eq!( + actual_variant, + *expected_variant, + "Token at index {} does not match: expected {}, got {} (token: {:?})", + idx, + expected_variant, + actual_variant, + token + ); + } + } + }; + } + fn lex<'src>( input: &'src str, ) -> (Option>>, Vec>) { @@ -809,6 +883,28 @@ mod original_lexer { assert!(lex_errs.is_empty()); } + + #[test] + fn test_lexer_padding_detection() { + let expr = "padding::<10>()"; + + let (tokens, lex_errs) = lexer().parse(expr).into_output_errors(); + + // let _ = tokens.unwrap(); + + assert!(lex_errs.is_empty()); + + assert_tokens_match!( + tokens, + Ident, + DoubleColon, + LAngle, + DecLiteral, + RAngle, + LParen, + RParen + ); + } } #[cfg(feature = "fmt")] diff --git a/src/parse.rs b/src/parse.rs index 547af341..94740d83 100644 --- a/src/parse.rs +++ b/src/parse.rs @@ -475,6 +475,15 @@ pub struct Call { } impl Call { + /// Construct a call directly, bypassing the parser. + /// + /// Mirrors what `arbitrary` does for the fuzz targets, which build a + /// [`Program`] without going through the text parser. + #[cfg(test)] + pub(crate) fn new_unparsed(name: CallName, args: Arc<[Expression]>, span: Span) -> Self { + Self { name, args, span } + } + /// Access the name of the call. pub fn name(&self) -> &CallName { &self.name @@ -495,6 +504,16 @@ impl_eq_hash!(Call; name, args); impl_require_feature!(Call {recurse: name, args; }); +/// Maximum number of layers that `padding::()` may request. +/// +/// Padding compiles to a chain of nodes whose depth is the layer count, and the rest +/// of the compiler recurses over that depth, so an unbounded count overflows the stack. +/// +/// At this limit a program compiles on the 8 MB main-thread stack in both debug and +/// release, and on the 2 MB stack that spawned threads get by default in release. A +/// debug build on a 2 MB thread stack is the one measured case that still runs out. +pub const MAX_PADDING_SIZE: usize = 8192; + /// Name of a call. #[derive(Clone, Debug, Eq, PartialEq, Hash)] #[cfg_attr(feature = "arbitrary", derive(arbitrary::Arbitrary))] @@ -525,6 +544,9 @@ pub enum CallName { ArrayFold(FunctionName, NonZeroUsize), /// Loop over the given function a bounded number of times until it returns success. ForWhile(FunctionName), + /// Padding that inflates the program, raising the transaction weight so that + /// CPU-heavy programs fit in the budget. At most [`MAX_PADDING_SIZE`] layers. + Padding(NonZeroUsize), } impl_require_feature!(CallName { @@ -542,6 +564,7 @@ impl_require_feature!(CallName { Fold(_, _), ArrayFold(_, _), ForWhile(_), + Padding(_), }); /// A type alias. @@ -1609,6 +1632,7 @@ impl fmt::Display for CallName { CallName::Fold(name, bound) => write!(f, "fold::<{name}, {bound}>"), CallName::ArrayFold(name, size) => write!(f, "array_fold::<{name}, {size}>"), CallName::ForWhile(name) => write!(f, "for_while::<{name}>"), + CallName::Padding(size) => write!(f, "padding::<{size}>"), } } } @@ -2446,6 +2470,41 @@ impl ChumskyParse for CallName { Token::Macro("dbg!") => CallName::Debug, }; + let padding = just(Token::Ident("padding")) + .ignore_then(turbofish_start.clone()) + .then(select! { Token::DecLiteral(s) => s }.labelled("size")) + .then_ignore(generics_close.clone()) + .validate(|((), size_str), e, emit| { + let size = match size_str.as_inner().parse::() { + Ok(0) => { + emit.emit(Error::PaddingSizeZero.with_span(e.span())); + NonZeroUsize::new(1).unwrap() + } + Ok(n) if n > MAX_PADDING_SIZE => { + emit.emit( + Error::PaddingSizeTooLarge { + size: n, + max: MAX_PADDING_SIZE, + } + .with_span(e.span()), + ); + NonZeroUsize::new(1).unwrap() + } + Ok(n) => NonZeroUsize::new(n).unwrap(), + Err(_) => { + emit.emit( + Error::CannotParse { + msg: format!("Invalid number: {size_str}"), + } + .with_span(e.span()), + ); + NonZeroUsize::new(1).unwrap() + } + }; + + CallName::Padding(size) + }); + let jet = select! { Token::Jet(s) => JetName::from_str_unchecked(s) }.map(CallName::Jet); let custom_func = FunctionName::parser().map(CallName::Custom); @@ -2460,6 +2519,8 @@ impl ChumskyParse for CallName { for_while, simple_builtins, jet, + padding, + // Note: Add built-in functions before this, otherwise they will not be matched. custom_func, )) } @@ -4327,3 +4388,80 @@ fn main() { ); } } + +#[cfg(test)] +mod padding_tests { + use super::*; + + /// Parse `input` and return the single `padding::()` call it consists of. + fn parse_padding(input: &str) -> Result { + let parsed = Expression::parse_from_str(input).map_err(|_| "Failed to parse")?; + + match parsed.inner() { + ExpressionInner::Single(single) => match single.inner() { + SingleExpressionInner::Call(call) => match call.name() { + CallName::Padding(_) => Ok(call.clone()), + _ => Err("Expected a padding call"), + }, + _ => Err("Expected a call expression"), + }, + _ => Err("Expected a single expression"), + } + } + + #[test] + fn test_parse_padding() { + let input = "padding::<10>()"; + + parse_padding(input).unwrap(); + } + + #[test] + fn test_parse_padding_accepts_max_size() { + let input = format!("padding::<{MAX_PADDING_SIZE}>()"); + + parse_padding(&input).unwrap(); + } + + #[test] + fn test_parse_padding_rejects_size_above_max() { + let input = format!("padding::<{}>()", MAX_PADDING_SIZE + 1); + + let err = Expression::parse_from_str(&input) + .expect_err("padding above the maximum size should be rejected"); + + assert!( + matches!( + err.error(), + Error::PaddingSizeTooLarge { + size, + max: MAX_PADDING_SIZE + } if *size == MAX_PADDING_SIZE + 1 + ), + "expected a PaddingSizeTooLarge diagnostic, got {err:?}" + ); + } + + #[test] + fn test_parse_padding_rejects_zero() { + let err = Expression::parse_from_str("padding::<0>()") + .expect_err("padding of size zero should be rejected"); + + assert!( + matches!(err.error(), Error::PaddingSizeZero), + "expected a PaddingSizeZero diagnostic, got {err:?}" + ); + } + + #[test] + fn test_parse_padding_should_fail_with_multiple_generics() { + let input = "padding::<10, 22>()"; + + let parsed_call = parse_padding(input); + + assert!( + parsed_call.is_err(), + "padding parsed correctly but should have failed" + ); + } +}