diff --git a/src/ast/compiler/analyzer.rs b/src/ast/compiler/analyzer.rs index 1e57a14..d0155f7 100644 --- a/src/ast/compiler/analyzer.rs +++ b/src/ast/compiler/analyzer.rs @@ -123,6 +123,15 @@ impl<'a> Analyzer<'a> { (BoundKind::Call { callee: Box::new(callee_m), args: Box::new(args_m) }, p) } + BoundKind::Again { args } => { + let args_m = self.visit(Rc::new((**args).clone())); + if let Some(lambda_id) = self.lambda_stack.last() { + self.recursive_identities.insert(lambda_id.clone()); + is_recursive = true; + } + (BoundKind::Again { args: Box::new(args_m) }, Purity::Impure) + } + BoundKind::Block { exprs } => { let mut new_exprs = Vec::with_capacity(exprs.len()); let mut p = Purity::Pure; diff --git a/src/ast/compiler/binder.rs b/src/ast/compiler/binder.rs index 187d151..da3fcfe 100644 --- a/src/ast/compiler/binder.rs +++ b/src/ast/compiler/binder.rs @@ -289,6 +289,19 @@ impl Binder { )) } + UntypedKind::Again { args } => { + if self.functions.len() <= 1 { + return Err("'again' is only allowed inside a function or lambda.".to_string()); + } + let args = self.bind(args)?; + Ok(self.make_node( + node.identity.clone(), + BoundKind::Again { + args: Box::new(args), + }, + )) + } + UntypedKind::Block { exprs } => { let mut bound_exprs = Vec::new(); for expr in exprs { diff --git a/src/ast/compiler/bound_nodes.rs b/src/ast/compiler/bound_nodes.rs index 6edfa5e..c9193be 100644 --- a/src/ast/compiler/bound_nodes.rs +++ b/src/ast/compiler/bound_nodes.rs @@ -96,6 +96,10 @@ pub enum BoundKind { args: Box>, }, + Again { + args: Box>, + }, + Block { exprs: Vec>, }, @@ -207,6 +211,7 @@ where args: ab, }, ) => ca == cb && aa == ab, + (BoundKind::Again { args: aa }, BoundKind::Again { args: ab }) => aa == ab, (BoundKind::Block { exprs: ea }, BoundKind::Block { exprs: eb }) => ea == eb, (BoundKind::Tuple { elements: ea }, BoundKind::Tuple { elements: eb }) => ea == eb, (BoundKind::Record { fields: fa }, BoundKind::Record { fields: fb }) => fa == fb, @@ -254,6 +259,7 @@ impl BoundKind { format!("LAMBDA({}, Captures:{})", p_str, upvalues.len()) } BoundKind::Call { .. } => "CALL".to_string(), + BoundKind::Again { .. } => "AGAIN".to_string(), BoundKind::Block { .. } => "BLOCK".to_string(), BoundKind::Tuple { elements } => format!("TUPLE({})", elements.len()), BoundKind::Record { fields } => format!("RECORD({})", fields.len()), diff --git a/src/ast/compiler/dumper.rs b/src/ast/compiler/dumper.rs index e5d3154..f844873 100644 --- a/src/ast/compiler/dumper.rs +++ b/src/ast/compiler/dumper.rs @@ -187,6 +187,13 @@ impl Dumper { self.indent -= 1; } + BoundKind::Again { args } => { + self.log("Again", node); + self.indent += 1; + self.visit(args); + self.indent -= 1; + } + BoundKind::Block { exprs } => { self.log("Block", node); self.indent += 1; diff --git a/src/ast/compiler/optimizer.rs b/src/ast/compiler/optimizer.rs index 70c4e93..a70b93f 100644 --- a/src/ast/compiler/optimizer.rs +++ b/src/ast/compiler/optimizer.rs @@ -313,6 +313,16 @@ impl Optimizer { ) } + BoundKind::Again { args } => { + let args = self.visit_node(*args, sub, path); + ( + BoundKind::Again { + args: Box::new(args), + }, + node.ty.clone(), + ) + } + BoundKind::If { ref cond, ref then_br, diff --git a/src/ast/compiler/tco.rs b/src/ast/compiler/tco.rs index 43752b5..e67a47a 100644 --- a/src/ast/compiler/tco.rs +++ b/src/ast/compiler/tco.rs @@ -40,6 +40,15 @@ impl TCO { args: Box::new(Self::transform(Rc::new((**args).clone()), false)), }, + BoundKind::Again { args } => { + if !is_tail_position { + panic!("'again' is only allowed in tail position to avoid dead code."); + } + BoundKind::Again { + args: Box::new(Self::transform(Rc::new((**args).clone()), false)), + } + } + BoundKind::If { cond, then_br, diff --git a/src/ast/compiler/type_checker.rs b/src/ast/compiler/type_checker.rs index 1ec0568..2316160 100644 --- a/src/ast/compiler/type_checker.rs +++ b/src/ast/compiler/type_checker.rs @@ -11,6 +11,8 @@ struct TypeContext<'a> { slots: Vec, /// Types of captured variables (passed from outer scope) upvalue_types: Vec, + /// The expected parameters of the current function (for 'again' validation) + current_params_ty: Option, } impl<'a> TypeContext<'a> { @@ -23,6 +25,7 @@ impl<'a> TypeContext<'a> { _parent: parent, slots: vec![StaticType::Any; slot_count as usize], upvalue_types, + current_params_ty: None, } } @@ -324,6 +327,10 @@ impl TypeChecker { // 3. Check parameters and body let params_typed = self.check_params(params.as_ref().clone(), &StaticType::Any, &mut lambda_ctx)?; + + // Set current params type for 'again' validation + lambda_ctx.current_params_ty = Some(params_typed.ty.clone()); + let body_typed = self.check_node((*body).clone(), &mut lambda_ctx)?; let ret_ty = body_typed.ty.clone(); @@ -362,6 +369,26 @@ impl TypeChecker { ) } + BoundKind::Again { args } => { + let args_typed = self.check_node(*args, ctx)?; + + if let Some(expected_ty) = &ctx.current_params_ty + && !expected_ty.is_assignable_from(&args_typed.ty) + { + return Err(format!( + "Type mismatch in 'again' call: expected {}, but got {}", + expected_ty, args_typed.ty + )); + } + + ( + BoundKind::Again { + args: Box::new(args_typed), + }, + StaticType::Any, + ) + } + BoundKind::Tuple { elements } => { let mut typed_elements = Vec::new(); for e in elements { diff --git a/src/ast/compiler/upvalues.rs b/src/ast/compiler/upvalues.rs index fe665ab..48833dd 100644 --- a/src/ast/compiler/upvalues.rs +++ b/src/ast/compiler/upvalues.rs @@ -99,6 +99,9 @@ impl UpvalueAnalyzer { Self::visit(callee, scopes, capture_map, current_lambda.clone()); Self::visit(args, scopes, capture_map, current_lambda.clone()); } + UntypedKind::Again { args } => { + Self::visit(args, scopes, capture_map, current_lambda.clone()); + } UntypedKind::Block { exprs } => { for expr in exprs { Self::visit(expr, scopes, capture_map, current_lambda.clone()); diff --git a/src/ast/nodes.rs b/src/ast/nodes.rs index 0812573..e1e629a 100644 --- a/src/ast/nodes.rs +++ b/src/ast/nodes.rs @@ -86,6 +86,9 @@ pub enum UntypedKind { callee: Box>, args: Box>, }, + Again { + args: Box>, + }, Block { exprs: Vec>, }, diff --git a/src/ast/parser.rs b/src/ast/parser.rs index 7344cc3..363ea49 100644 --- a/src/ast/parser.rs +++ b/src/ast/parser.rs @@ -140,6 +140,7 @@ impl<'a> Parser<'a> { match sym.name.as_ref() { "if" => self.parse_if(identity), "fn" => self.parse_fn(identity), + "again" => self.parse_again(identity), "def" => self.parse_def(identity), "assign" => self.parse_assign(identity), "do" => self.parse_do(identity), @@ -154,6 +155,15 @@ impl<'a> Parser<'a> { result } + fn parse_again(&mut self, identity: Identity) -> Result, String> { + let args = Box::new(self.parse_expression()?); + Ok(Node { + identity, + kind: UntypedKind::Again { args }, + ty: (), + }) + } + fn parse_if(&mut self, identity: Identity) -> Result, String> { let cond = Box::new(self.parse_expression()?); let then_br = Box::new(self.parse_expression()?); diff --git a/src/ast/types.rs b/src/ast/types.rs index f5e2f32..f7c56db 100644 --- a/src/ast/types.rs +++ b/src/ast/types.rs @@ -249,6 +249,16 @@ impl StaticType { } elements.iter().all(|e| e.is_assignable_from(inner)) } + // Tuple to Tuple (Element-wise) + (StaticType::Tuple(elements), StaticType::Tuple(other_elements)) => { + if elements.len() != other_elements.len() { + return false; + } + elements + .iter() + .zip(other_elements.iter()) + .all(|(e, o)| e.is_assignable_from(o)) + } // A Matrix is a Vector (of Vectors/Matrices) (StaticType::Vector(inner, len), StaticType::Matrix(m_inner, shape)) => { if shape.is_empty() || shape[0] != *len { diff --git a/src/ast/vm.rs b/src/ast/vm.rs index f9ad455..1618a4b 100644 --- a/src/ast/vm.rs +++ b/src/ast/vm.rs @@ -433,6 +433,37 @@ impl VM { } } } + BoundKind::Again { args } => { + let arg_vals = match &args.kind { + BoundKind::Tuple { elements } => { + let mut vals = Vec::with_capacity(elements.len()); + let mut is_complex = false; + for e in elements { + if matches!(e.kind, BoundKind::Tuple { .. }) { + is_complex = true; + break; + } + vals.push(self.eval_internal(obs, e)?); + } + if is_complex { + self.prepare_args_internal(obs, args)? + } else { + vals + } + } + _ => self.prepare_args_internal(obs, args)?, + }; + + let frame = self.frames.last().ok_or("No call frame for 'again'")?; + if let Some(closure) = &frame.closure { + Ok(Value::TailCallRequest(Box::new(( + Rc::new(closure.as_ref().clone()) as Rc, + arg_vals, + )))) + } else { + Err("'again' called outside of a closure".to_string()) + } + } BoundKind::Tuple { elements } => { let mut vals = Vec::with_capacity(elements.len()); for e in elements { diff --git a/src/integration_test.rs b/src/integration_test.rs index 973d33f..9c6cc3d 100644 --- a/src/integration_test.rs +++ b/src/integration_test.rs @@ -263,6 +263,15 @@ mod tests { } } + #[test] + #[should_panic(expected = "'again' is only allowed in tail position to avoid dead code.")] + fn test_again_non_tail_panic() { + let env = Environment::new(); + let source = "(do (def f (fn [x] (do (again [(- x 1)]) x))) (f 5))"; + // This will trigger the TCO pass which contains the validation logic + let _ = env.run_script(source); + } + #[test] fn test_dynamic_call_destructuring_underflow() { let env = Environment::new();