xlog_logic/
expand.rs

1//! Inline expansion of user-defined functions.
2
3use crate::ast::{
4    ArithExpr, Atom, BodyLiteral, Comparison, Constraint, FuncBody, FuncDef, IsExpr, Term, Univ,
5};
6use crate::function::{is_builtin, FunctionError, FunctionRegistry};
7use std::collections::{HashMap, HashSet};
8
9const GENERATED_FUNCTION_VARIABLE_PREFIX: &str = "__XLOG_FUNCTION_";
10
11pub(crate) fn generated_function_variable_name(
12    function_name: &str,
13    source_name: &str,
14    counter: usize,
15) -> String {
16    format!(
17        "{GENERATED_FUNCTION_VARIABLE_PREFIX}{}_{}_{}",
18        function_name.to_ascii_uppercase(),
19        source_name,
20        counter
21    )
22}
23
24pub(crate) fn generated_function_variable_source<'a>(
25    generated_name: &'a str,
26    function_name: &str,
27) -> Option<&'a str> {
28    let prefix = format!(
29        "{GENERATED_FUNCTION_VARIABLE_PREFIX}{}_",
30        function_name.to_ascii_uppercase()
31    );
32    let (source_name, counter) = generated_name.strip_prefix(&prefix)?.rsplit_once('_')?;
33    (!source_name.is_empty()
34        && !counter.is_empty()
35        && counter.chars().all(|character| character.is_ascii_digit()))
36    .then_some(source_name)
37}
38
39#[derive(Debug)]
40struct ExpandedExpression {
41    generated_literals: Vec<BodyLiteral>,
42    expression: ArithExpr,
43}
44
45fn inline_trailing_predicate_result_binding(
46    mut literals: Vec<BodyLiteral>,
47    result: ArithExpr,
48) -> (Vec<BodyLiteral>, ArithExpr) {
49    let ArithExpr::Variable(result_name) = &result else {
50        return (literals, result);
51    };
52    let Some(BodyLiteral::IsExpr(binding)) = literals.last() else {
53        return (literals, result);
54    };
55    let ArithExpr::Variable(source_name) = &binding.expr else {
56        return (literals, result);
57    };
58    if !predicate_prefix_binds_variable(&literals[..literals.len() - 1], source_name)
59        || binding.target != *result_name
60        || binding
61            .expr
62            .variables()
63            .into_iter()
64            .any(|name| name == result_name)
65        || literals[..literals.len() - 1].iter().any(|literal| {
66            literal
67                .variables()
68                .into_iter()
69                .any(|name| name == result_name)
70        })
71    {
72        return (literals, result);
73    }
74
75    let Some(BodyLiteral::IsExpr(binding)) = literals.pop() else {
76        unreachable!("trailing predicate result binding was checked above")
77    };
78    (literals, binding.expr)
79}
80
81fn predicate_prefix_binds_variable(literals: &[BodyLiteral], wanted: &str) -> bool {
82    let mut bound = HashSet::new();
83    for literal in literals {
84        match literal {
85            BodyLiteral::Positive(atom) => {
86                bound.extend(atom.variables().into_iter().map(ToOwned::to_owned));
87            }
88            BodyLiteral::IsExpr(binding)
89                if binding
90                    .expr
91                    .variables()
92                    .into_iter()
93                    .all(|name| bound.contains(name)) =>
94            {
95                bound.insert(binding.target.clone());
96            }
97            BodyLiteral::Negated(_)
98            | BodyLiteral::Epistemic(_)
99            | BodyLiteral::Comparison(_)
100            | BodyLiteral::IsExpr(_)
101            | BodyLiteral::Univ(_) => {}
102        }
103    }
104    bound.contains(wanted)
105}
106
107impl ExpandedExpression {
108    fn value(expression: ArithExpr) -> Self {
109        Self {
110            generated_literals: Vec::new(),
111            expression,
112        }
113    }
114}
115
116type BinaryExpressionConstructor = fn(Box<ArithExpr>, Box<ArithExpr>) -> ArithExpr;
117
118enum ExpansionTask {
119    Expression {
120        expression: ArithExpr,
121        subst: HashMap<String, ArithExpr>,
122        in_conditional_branch: bool,
123    },
124    FinishCall {
125        name: String,
126        argument_count: usize,
127        in_conditional_branch: bool,
128    },
129    EnterFunction {
130        name: String,
131        args: Vec<ArithExpr>,
132        in_conditional_branch: bool,
133    },
134    FunctionBody {
135        function_name: String,
136        body: FuncBody,
137        subst: HashMap<String, ArithExpr>,
138        in_conditional_branch: bool,
139    },
140    LeaveFunction,
141    PrependGenerated {
142        literals: Vec<BodyLiteral>,
143    },
144    FinishBinary {
145        constructor: BinaryExpressionConstructor,
146    },
147    FinishAbs,
148    FinishCast(xlog_core::ScalarType),
149    FinishConditional(crate::ast::CompOp),
150    PredicateLiteral(PreparedPredicateLiteral),
151    FinishPredicateBinding {
152        target: String,
153    },
154    FinishPredicateBody {
155        literal_count: usize,
156        result: ArithExpr,
157    },
158}
159
160enum PreparedPredicateLiteral {
161    Literal(BodyLiteral),
162    Binding {
163        target: String,
164        expression: ArithExpr,
165        subst: HashMap<String, ArithExpr>,
166    },
167}
168
169enum TermSubstitutionTask<'a> {
170    Term(&'a Term),
171    FinishList(usize),
172    FinishCons,
173    FinishCompound {
174        functor: String,
175        argument_count: usize,
176    },
177}
178
179/// Context for inline expansion of user-defined functions.
180pub struct ExpansionContext<'a> {
181    registry: &'a FunctionRegistry,
182    depth: u32,
183    max_depth: u32,
184    fresh_counter: usize,
185}
186
187impl<'a> ExpansionContext<'a> {
188    /// Create an expansion context with the given function registry and recursion limit.
189    pub fn new(registry: &'a FunctionRegistry, max_depth: u32) -> Self {
190        Self {
191            registry,
192            depth: 0,
193            max_depth,
194            fresh_counter: 0,
195        }
196    }
197
198    /// Expand a scalar function call to its body with arguments substituted.
199    ///
200    /// Calls that produce relational literals require a surrounding rule-like
201    /// body and are rejected by this expression-only API.
202    pub fn expand_call(
203        &mut self,
204        name: &str,
205        args: &[ArithExpr],
206    ) -> Result<ArithExpr, FunctionError> {
207        if !self.registry.contains(name) {
208            return Err(FunctionError::UndefinedFunction {
209                name: name.to_string(),
210            });
211        }
212
213        let call = ArithExpr::FuncCall {
214            name: name.to_string(),
215            args: args.to_vec(),
216        };
217        let mut used_variables = call
218            .variables()
219            .into_iter()
220            .map(ToOwned::to_owned)
221            .collect();
222        let expanded =
223            self.expand_expr_for_rule(&call, &HashMap::new(), &mut used_variables, false)?;
224        if expanded.generated_literals.is_empty() {
225            Ok(expanded.expression)
226        } else {
227            Err(FunctionError::PredicateBodyRequiresRuleContext {
228                name: name.to_string(),
229            })
230        }
231    }
232
233    fn check_arity(func: &FuncDef, args: &[ArithExpr]) -> Result<(), FunctionError> {
234        if func.params.len() == args.len() {
235            return Ok(());
236        }
237        Err(FunctionError::ArityMismatch {
238            name: func.name.clone(),
239            expected: func.params.len(),
240            received: args.len(),
241        })
242    }
243
244    fn fresh_variable(
245        &mut self,
246        function_name: &str,
247        source_name: &str,
248        used_variables: &mut HashSet<String>,
249    ) -> String {
250        loop {
251            let candidate =
252                generated_function_variable_name(function_name, source_name, self.fresh_counter);
253            self.fresh_counter += 1;
254            if used_variables.insert(candidate.clone()) {
255                return candidate;
256            }
257        }
258    }
259
260    /// Freshen and substitute a predicate body before the expansion machine visits it.
261    fn prepare_predicate_func(
262        &mut self,
263        function_name: &str,
264        result: String,
265        body: Vec<BodyLiteral>,
266        mut subst: HashMap<String, ArithExpr>,
267        used_variables: &mut HashSet<String>,
268    ) -> Result<(Vec<PreparedPredicateLiteral>, ArithExpr), FunctionError> {
269        let parameter_names: HashSet<String> = subst.keys().cloned().collect();
270        let mut local_names = Vec::new();
271        let mut seen_locals = HashSet::new();
272
273        if !parameter_names.contains(&result) && seen_locals.insert(result.clone()) {
274            local_names.push(result.clone());
275        }
276        for literal in &body {
277            for variable in literal.variables() {
278                if !parameter_names.contains(variable) && seen_locals.insert(variable.to_string()) {
279                    local_names.push(variable.to_string());
280                }
281            }
282        }
283
284        for local in local_names {
285            let fresh = self.fresh_variable(function_name, &local, used_variables);
286            subst.insert(local, ArithExpr::Variable(fresh));
287        }
288
289        let substituted_body = body
290            .into_iter()
291            .map(|literal| {
292                if let BodyLiteral::IsExpr(binding) = literal {
293                    Ok(PreparedPredicateLiteral::Binding {
294                        target: self.substitute_binding_target(
295                            function_name,
296                            &binding.target,
297                            &subst,
298                        )?,
299                        expression: binding.expr,
300                        subst: subst.clone(),
301                    })
302                } else {
303                    self.substitute_literal(function_name, &literal, &subst)
304                        .map(PreparedPredicateLiteral::Literal)
305                }
306            })
307            .collect::<Result<Vec<_>, _>>()?;
308        let expression = subst
309            .get(&result)
310            .cloned()
311            .unwrap_or(ArithExpr::Variable(result));
312        Ok((substituted_body, expression))
313    }
314
315    fn substitute_literal(
316        &self,
317        function_name: &str,
318        lit: &BodyLiteral,
319        subst: &HashMap<String, ArithExpr>,
320    ) -> Result<BodyLiteral, FunctionError> {
321        Ok(match lit {
322            BodyLiteral::Positive(atom) => {
323                BodyLiteral::Positive(self.substitute_atom(function_name, atom, subst)?)
324            }
325            BodyLiteral::Negated(atom) => {
326                BodyLiteral::Negated(self.substitute_atom(function_name, atom, subst)?)
327            }
328            BodyLiteral::Epistemic(lit) => BodyLiteral::Epistemic(crate::ast::EpistemicLiteral {
329                op: lit.op,
330                negated: lit.negated,
331                atom: self.substitute_atom(function_name, &lit.atom, subst)?,
332            }),
333            BodyLiteral::Comparison(cmp) => BodyLiteral::Comparison(Comparison {
334                left: self.substitute_term(function_name, &cmp.left, subst)?,
335                op: cmp.op,
336                right: self.substitute_term(function_name, &cmp.right, subst)?,
337            }),
338            BodyLiteral::IsExpr(_) => unreachable!("predicate bindings are prepared separately"),
339            BodyLiteral::Univ(univ) => BodyLiteral::Univ(Univ {
340                term: self.substitute_term(function_name, &univ.term, subst)?,
341                parts: self.substitute_term(function_name, &univ.parts, subst)?,
342            }),
343        })
344    }
345
346    fn substitute_atom(
347        &self,
348        function_name: &str,
349        atom: &Atom,
350        subst: &HashMap<String, ArithExpr>,
351    ) -> Result<Atom, FunctionError> {
352        Ok(Atom {
353            predicate: atom.predicate.clone(),
354            terms: atom
355                .terms
356                .iter()
357                .map(|term| self.substitute_term(function_name, term, subst))
358                .collect::<Result<Vec<_>, _>>()?,
359        })
360    }
361
362    fn substitute_term(
363        &self,
364        function_name: &str,
365        term: &Term,
366        subst: &HashMap<String, ArithExpr>,
367    ) -> Result<Term, FunctionError> {
368        let mut tasks = vec![TermSubstitutionTask::Term(term)];
369        let mut values = Vec::new();
370
371        while let Some(task) = tasks.pop() {
372            match task {
373                TermSubstitutionTask::Term(term) => match term {
374                    Term::Variable(name) => values.push(match subst.get(name) {
375                        Some(ArithExpr::Variable(new_name)) => Term::Variable(new_name.clone()),
376                        Some(ArithExpr::Integer(value)) => Term::Integer(*value),
377                        Some(ArithExpr::Float(value)) => Term::Float(*value),
378                        Some(_) => {
379                            return Err(FunctionError::UnsupportedPredicateTermArgument {
380                                name: function_name.to_string(),
381                                parameter: name.clone(),
382                            });
383                        }
384                        None => Term::Variable(name.clone()),
385                    }),
386                    Term::List(items) => {
387                        tasks.push(TermSubstitutionTask::FinishList(items.len()));
388                        for item in items.iter().rev() {
389                            tasks.push(TermSubstitutionTask::Term(item));
390                        }
391                    }
392                    Term::Cons { head, tail } => {
393                        tasks.push(TermSubstitutionTask::FinishCons);
394                        tasks.push(TermSubstitutionTask::Term(tail));
395                        tasks.push(TermSubstitutionTask::Term(head));
396                    }
397                    Term::Compound { functor, args } => {
398                        tasks.push(TermSubstitutionTask::FinishCompound {
399                            functor: functor.clone(),
400                            argument_count: args.len(),
401                        });
402                        for argument in args.iter().rev() {
403                            tasks.push(TermSubstitutionTask::Term(argument));
404                        }
405                    }
406                    Term::Aggregate(aggregate) => {
407                        values.push(Term::Aggregate(crate::ast::AggExpr {
408                            op: aggregate.op,
409                            variable: self.substitute_binding_target(
410                                function_name,
411                                &aggregate.variable,
412                                subst,
413                            )?,
414                        }));
415                    }
416                    Term::Anonymous => values.push(Term::Anonymous),
417                    Term::Integer(value) => values.push(Term::Integer(*value)),
418                    Term::Float(value) => values.push(Term::Float(*value)),
419                    Term::String(value) => values.push(Term::String(value.clone())),
420                    Term::Symbol(value) => values.push(Term::Symbol(*value)),
421                    Term::PredRef(value) => values.push(Term::PredRef(value.clone())),
422                },
423                TermSubstitutionTask::FinishList(item_count) => {
424                    let start = values
425                        .len()
426                        .checked_sub(item_count)
427                        .expect("list items produce substituted terms");
428                    let items = values.split_off(start);
429                    values.push(Term::List(items));
430                }
431                TermSubstitutionTask::FinishCons => {
432                    let tail = values.pop().expect("cons tail produces a substituted term");
433                    let head = values.pop().expect("cons head produces a substituted term");
434                    values.push(Term::Cons {
435                        head: Box::new(head),
436                        tail: Box::new(tail),
437                    });
438                }
439                TermSubstitutionTask::FinishCompound {
440                    functor,
441                    argument_count,
442                } => {
443                    let start = values
444                        .len()
445                        .checked_sub(argument_count)
446                        .expect("compound arguments produce substituted terms");
447                    let args = values.split_off(start);
448                    values.push(Term::Compound { functor, args });
449                }
450            }
451        }
452
453        let substituted = values.pop().expect("term substitution produces one term");
454        debug_assert!(values.is_empty());
455        Ok(substituted)
456    }
457
458    fn substitute_binding_target(
459        &self,
460        function_name: &str,
461        variable: &str,
462        subst: &HashMap<String, ArithExpr>,
463    ) -> Result<String, FunctionError> {
464        match subst.get(variable) {
465            Some(ArithExpr::Variable(new_name)) => Ok(new_name.clone()),
466            Some(_) => Err(FunctionError::InvalidPredicateBindingTarget {
467                name: function_name.to_string(),
468                parameter: variable.to_string(),
469            }),
470            None => Ok(variable.to_string()),
471        }
472    }
473
474    fn expand_literal_for_rule(
475        &mut self,
476        literal: &BodyLiteral,
477        used_variables: &mut HashSet<String>,
478    ) -> Result<Vec<BodyLiteral>, FunctionError> {
479        let BodyLiteral::IsExpr(binding) = literal else {
480            return Ok(vec![literal.clone()]);
481        };
482
483        let expanded =
484            self.expand_expr_for_rule(&binding.expr, &HashMap::new(), used_variables, false)?;
485        let mut literals = expanded.generated_literals;
486        literals.push(BodyLiteral::IsExpr(IsExpr {
487            target: binding.target.clone(),
488            expr: expanded.expression,
489        }));
490        Ok(literals)
491    }
492
493    fn expand_expr_for_rule(
494        &mut self,
495        expression: &ArithExpr,
496        subst: &HashMap<String, ArithExpr>,
497        used_variables: &mut HashSet<String>,
498        in_conditional_branch: bool,
499    ) -> Result<ExpandedExpression, FunctionError> {
500        let saved_depth = self.depth;
501        let saved_fresh_counter = self.fresh_counter;
502        let saved_used_variables = used_variables.clone();
503        let result =
504            self.run_expansion_machine(expression, subst, used_variables, in_conditional_branch);
505        self.depth = saved_depth;
506        if result.is_err() {
507            self.fresh_counter = saved_fresh_counter;
508            *used_variables = saved_used_variables;
509        }
510        result
511    }
512
513    fn run_expansion_machine(
514        &mut self,
515        expression: &ArithExpr,
516        subst: &HashMap<String, ArithExpr>,
517        used_variables: &mut HashSet<String>,
518        in_conditional_branch: bool,
519    ) -> Result<ExpandedExpression, FunctionError> {
520        let mut tasks = vec![ExpansionTask::Expression {
521            expression: expression.clone(),
522            subst: subst.clone(),
523            in_conditional_branch,
524        }];
525        let mut values = Vec::new();
526        let mut predicate_literal_values: Vec<Vec<BodyLiteral>> = Vec::new();
527
528        while let Some(task) = tasks.pop() {
529            match task {
530                ExpansionTask::Expression {
531                    expression,
532                    subst,
533                    in_conditional_branch,
534                } => match expression {
535                    ArithExpr::Variable(name) => values.push(ExpandedExpression::value(
536                        subst
537                            .get(&name)
538                            .cloned()
539                            .unwrap_or(ArithExpr::Variable(name)),
540                    )),
541                    ArithExpr::Integer(_) | ArithExpr::Float(_) => {
542                        values.push(ExpandedExpression::value(expression));
543                    }
544                    ArithExpr::FuncCall { name, args } => {
545                        if let Some(func) = self.registry.get(&name) {
546                            Self::check_arity(func, &args)?;
547                        }
548                        tasks.push(ExpansionTask::FinishCall {
549                            name,
550                            argument_count: args.len(),
551                            in_conditional_branch,
552                        });
553                        for argument in args.into_iter().rev() {
554                            tasks.push(ExpansionTask::Expression {
555                                expression: argument,
556                                subst: subst.clone(),
557                                in_conditional_branch,
558                            });
559                        }
560                    }
561                    ArithExpr::Add(left, right) => {
562                        Self::schedule_binary(
563                            &mut tasks,
564                            *left,
565                            *right,
566                            subst,
567                            in_conditional_branch,
568                            ArithExpr::Add,
569                        );
570                    }
571                    ArithExpr::Sub(left, right) => {
572                        Self::schedule_binary(
573                            &mut tasks,
574                            *left,
575                            *right,
576                            subst,
577                            in_conditional_branch,
578                            ArithExpr::Sub,
579                        );
580                    }
581                    ArithExpr::Mul(left, right) => {
582                        Self::schedule_binary(
583                            &mut tasks,
584                            *left,
585                            *right,
586                            subst,
587                            in_conditional_branch,
588                            ArithExpr::Mul,
589                        );
590                    }
591                    ArithExpr::Div(left, right) => {
592                        Self::schedule_binary(
593                            &mut tasks,
594                            *left,
595                            *right,
596                            subst,
597                            in_conditional_branch,
598                            ArithExpr::Div,
599                        );
600                    }
601                    ArithExpr::Mod(left, right) => {
602                        Self::schedule_binary(
603                            &mut tasks,
604                            *left,
605                            *right,
606                            subst,
607                            in_conditional_branch,
608                            ArithExpr::Mod,
609                        );
610                    }
611                    ArithExpr::Min(left, right) => {
612                        Self::schedule_binary(
613                            &mut tasks,
614                            *left,
615                            *right,
616                            subst,
617                            in_conditional_branch,
618                            ArithExpr::Min,
619                        );
620                    }
621                    ArithExpr::Max(left, right) => {
622                        Self::schedule_binary(
623                            &mut tasks,
624                            *left,
625                            *right,
626                            subst,
627                            in_conditional_branch,
628                            ArithExpr::Max,
629                        );
630                    }
631                    ArithExpr::Pow(left, right) => {
632                        Self::schedule_binary(
633                            &mut tasks,
634                            *left,
635                            *right,
636                            subst,
637                            in_conditional_branch,
638                            ArithExpr::Pow,
639                        );
640                    }
641                    ArithExpr::Abs(inner) => {
642                        tasks.push(ExpansionTask::FinishAbs);
643                        tasks.push(ExpansionTask::Expression {
644                            expression: *inner,
645                            subst,
646                            in_conditional_branch,
647                        });
648                    }
649                    ArithExpr::Cast(inner, target) => {
650                        tasks.push(ExpansionTask::FinishCast(target));
651                        tasks.push(ExpansionTask::Expression {
652                            expression: *inner,
653                            subst,
654                            in_conditional_branch,
655                        });
656                    }
657                    ArithExpr::Conditional {
658                        cond_left,
659                        cond_op,
660                        cond_right,
661                        then_expr,
662                        else_expr,
663                    } => {
664                        tasks.push(ExpansionTask::FinishConditional(cond_op));
665                        tasks.push(ExpansionTask::Expression {
666                            expression: *else_expr,
667                            subst: subst.clone(),
668                            in_conditional_branch: true,
669                        });
670                        tasks.push(ExpansionTask::Expression {
671                            expression: *then_expr,
672                            subst: subst.clone(),
673                            in_conditional_branch: true,
674                        });
675                        tasks.push(ExpansionTask::Expression {
676                            expression: *cond_right,
677                            subst: subst.clone(),
678                            in_conditional_branch,
679                        });
680                        tasks.push(ExpansionTask::Expression {
681                            expression: *cond_left,
682                            subst,
683                            in_conditional_branch,
684                        });
685                    }
686                },
687                ExpansionTask::FinishCall {
688                    name,
689                    argument_count,
690                    in_conditional_branch,
691                } => {
692                    let start = values
693                        .len()
694                        .checked_sub(argument_count)
695                        .expect("function arguments produce expansion values");
696                    let argument_values = values.split_off(start);
697                    let mut generated_literals = Vec::new();
698                    let mut args = Vec::with_capacity(argument_count);
699                    for argument in argument_values {
700                        generated_literals.extend(argument.generated_literals);
701                        args.push(argument.expression);
702                    }
703                    if self.registry.contains(&name) {
704                        tasks.push(ExpansionTask::PrependGenerated {
705                            literals: generated_literals,
706                        });
707                        tasks.push(ExpansionTask::EnterFunction {
708                            name,
709                            args,
710                            in_conditional_branch,
711                        });
712                    } else if is_builtin(&name) {
713                        values.push(ExpandedExpression {
714                            generated_literals,
715                            expression: ArithExpr::FuncCall { name, args },
716                        });
717                    } else {
718                        return Err(FunctionError::UndefinedFunction { name });
719                    }
720                }
721                ExpansionTask::EnterFunction {
722                    name,
723                    args,
724                    in_conditional_branch,
725                } => {
726                    let func =
727                        self.registry.get(&name).cloned().ok_or_else(|| {
728                            FunctionError::UndefinedFunction { name: name.clone() }
729                        })?;
730                    Self::check_arity(&func, &args)?;
731                    if self.depth >= self.max_depth {
732                        return Err(FunctionError::MaxRecursionDepth {
733                            name,
734                            depth: self.max_depth,
735                        });
736                    }
737                    if in_conditional_branch && matches!(func.body, FuncBody::Predicate { .. }) {
738                        return Err(FunctionError::PredicateCallInConditionalBranch { name });
739                    }
740
741                    self.depth += 1;
742                    let subst = func
743                        .params
744                        .iter()
745                        .zip(&args)
746                        .map(|(param, argument)| (param.name.clone(), argument.clone()))
747                        .collect();
748                    tasks.push(ExpansionTask::LeaveFunction);
749                    tasks.push(ExpansionTask::FunctionBody {
750                        function_name: func.name,
751                        body: func.body,
752                        subst,
753                        in_conditional_branch,
754                    });
755                }
756                ExpansionTask::FunctionBody {
757                    function_name,
758                    body,
759                    subst,
760                    in_conditional_branch,
761                } => match body {
762                    FuncBody::Arithmetic(expression) => {
763                        tasks.push(ExpansionTask::Expression {
764                            expression,
765                            subst,
766                            in_conditional_branch,
767                        });
768                    }
769                    FuncBody::Predicate { result, body } => {
770                        if in_conditional_branch {
771                            return Err(FunctionError::PredicateCallInConditionalBranch {
772                                name: function_name,
773                            });
774                        }
775                        let (literals, result) = self.prepare_predicate_func(
776                            &function_name,
777                            result,
778                            body,
779                            subst,
780                            used_variables,
781                        )?;
782                        tasks.push(ExpansionTask::FinishPredicateBody {
783                            literal_count: literals.len(),
784                            result,
785                        });
786                        for literal in literals.into_iter().rev() {
787                            tasks.push(ExpansionTask::PredicateLiteral(literal));
788                        }
789                    }
790                    FuncBody::Conditional(conditional) => {
791                        tasks.push(ExpansionTask::FinishConditional(conditional.cond_op));
792                        tasks.push(ExpansionTask::FunctionBody {
793                            function_name: function_name.clone(),
794                            body: *conditional.else_branch,
795                            subst: subst.clone(),
796                            in_conditional_branch: true,
797                        });
798                        tasks.push(ExpansionTask::FunctionBody {
799                            function_name,
800                            body: *conditional.then_branch,
801                            subst: subst.clone(),
802                            in_conditional_branch: true,
803                        });
804                        tasks.push(ExpansionTask::Expression {
805                            expression: conditional.cond_right,
806                            subst: subst.clone(),
807                            in_conditional_branch,
808                        });
809                        tasks.push(ExpansionTask::Expression {
810                            expression: conditional.cond_left,
811                            subst,
812                            in_conditional_branch,
813                        });
814                    }
815                },
816                ExpansionTask::LeaveFunction => {
817                    debug_assert!(self.depth > 0);
818                    self.depth -= 1;
819                }
820                ExpansionTask::PrependGenerated { mut literals } => {
821                    let mut expanded = values
822                        .pop()
823                        .expect("function call produces an expansion value");
824                    literals.append(&mut expanded.generated_literals);
825                    expanded.generated_literals = literals;
826                    values.push(expanded);
827                }
828                ExpansionTask::FinishBinary { constructor } => {
829                    let right = values.pop().expect("right expression is expanded");
830                    let mut left = values.pop().expect("left expression is expanded");
831                    left.generated_literals.extend(right.generated_literals);
832                    left.expression =
833                        constructor(Box::new(left.expression), Box::new(right.expression));
834                    values.push(left);
835                }
836                ExpansionTask::FinishAbs => {
837                    let mut expanded = values.pop().expect("absolute-value operand is expanded");
838                    expanded.expression = ArithExpr::Abs(Box::new(expanded.expression));
839                    values.push(expanded);
840                }
841                ExpansionTask::FinishCast(target) => {
842                    let mut expanded = values.pop().expect("cast operand is expanded");
843                    expanded.expression = ArithExpr::Cast(Box::new(expanded.expression), target);
844                    values.push(expanded);
845                }
846                ExpansionTask::FinishConditional(cond_op) => {
847                    let else_value = values.pop().expect("else branch is expanded");
848                    let then_value = values.pop().expect("then branch is expanded");
849                    let right = values.pop().expect("condition right side is expanded");
850                    let mut left = values.pop().expect("condition left side is expanded");
851                    left.generated_literals.extend(right.generated_literals);
852                    left.generated_literals
853                        .extend(then_value.generated_literals);
854                    left.generated_literals
855                        .extend(else_value.generated_literals);
856                    left.expression = ArithExpr::Conditional {
857                        cond_left: Box::new(left.expression),
858                        cond_op,
859                        cond_right: Box::new(right.expression),
860                        then_expr: Box::new(then_value.expression),
861                        else_expr: Box::new(else_value.expression),
862                    };
863                    values.push(left);
864                }
865                ExpansionTask::PredicateLiteral(literal) => match literal {
866                    PreparedPredicateLiteral::Binding {
867                        target,
868                        expression,
869                        subst,
870                    } => {
871                        tasks.push(ExpansionTask::FinishPredicateBinding { target });
872                        tasks.push(ExpansionTask::Expression {
873                            expression,
874                            subst,
875                            in_conditional_branch: false,
876                        });
877                    }
878                    PreparedPredicateLiteral::Literal(literal) => {
879                        predicate_literal_values.push(vec![literal]);
880                    }
881                },
882                ExpansionTask::FinishPredicateBinding { target } => {
883                    let mut expanded = values
884                        .pop()
885                        .expect("predicate-body binding expression is expanded");
886                    expanded
887                        .generated_literals
888                        .push(BodyLiteral::IsExpr(IsExpr {
889                            target,
890                            expr: expanded.expression,
891                        }));
892                    predicate_literal_values.push(expanded.generated_literals);
893                }
894                ExpansionTask::FinishPredicateBody {
895                    literal_count,
896                    result,
897                } => {
898                    let start = predicate_literal_values
899                        .len()
900                        .checked_sub(literal_count)
901                        .expect("predicate literals produce expansion values");
902                    let generated_literals = predicate_literal_values
903                        .split_off(start)
904                        .into_iter()
905                        .flatten()
906                        .collect();
907                    let (generated_literals, result) =
908                        inline_trailing_predicate_result_binding(generated_literals, result);
909                    values.push(ExpandedExpression {
910                        generated_literals,
911                        expression: result,
912                    });
913                }
914            }
915        }
916
917        debug_assert!(predicate_literal_values.is_empty());
918        let expanded = values
919            .pop()
920            .expect("expression expansion produces one value");
921        debug_assert!(values.is_empty());
922        Ok(expanded)
923    }
924
925    fn schedule_binary(
926        tasks: &mut Vec<ExpansionTask>,
927        left: ArithExpr,
928        right: ArithExpr,
929        subst: HashMap<String, ArithExpr>,
930        in_conditional_branch: bool,
931        constructor: BinaryExpressionConstructor,
932    ) {
933        tasks.push(ExpansionTask::FinishBinary { constructor });
934        tasks.push(ExpansionTask::Expression {
935            expression: right,
936            subst: subst.clone(),
937            in_conditional_branch,
938        });
939        tasks.push(ExpansionTask::Expression {
940            expression: left,
941            subst,
942            in_conditional_branch,
943        });
944    }
945
946    /// Check if a function has a predicate body.
947    #[allow(dead_code)]
948    pub(crate) fn is_predicate_func(&self, name: &str) -> bool {
949        self.registry
950            .get(name)
951            .map(|f| matches!(f.body, FuncBody::Predicate { .. }))
952            .unwrap_or(false)
953    }
954}
955
956use crate::ast::{Program, Rule};
957
958/// Expand user-defined function calls in ordinary rules and constraints.
959///
960/// Scalar calls become arithmetic expressions. Predicate-bodied calls also
961/// contribute relational literals immediately before their source binding.
962pub fn expand_program_functions(
963    program: &Program,
964    max_depth: u32,
965) -> Result<Program, FunctionError> {
966    // If no functions defined, return program unchanged
967    if program.functions.is_empty() {
968        return Ok(program.clone());
969    }
970    expand_program_functions_impl(program, max_depth)
971}
972
973/// [`expand_program_functions`] taking the program by value: without function
974/// definitions (the common, fact-heavy case) the program is returned as is,
975/// without a clone.
976pub fn expand_program_functions_owned(
977    program: Program,
978    max_depth: u32,
979) -> Result<Program, FunctionError> {
980    if program.functions.is_empty() {
981        return Ok(program);
982    }
983    expand_program_functions_impl(&program, max_depth)
984}
985
986fn expand_program_functions_impl(
987    program: &Program,
988    max_depth: u32,
989) -> Result<Program, FunctionError> {
990    let mut registry = FunctionRegistry::new();
991    for function in &program.functions {
992        registry.register(function.clone())?;
993    }
994    let mut ctx = ExpansionContext::new(&registry, max_depth);
995
996    // Expand function calls in each rule
997    let expanded_rules: Result<Vec<Rule>, FunctionError> = program
998        .rules
999        .iter()
1000        .map(|rule| expand_rule_functions(&mut ctx, rule))
1001        .collect();
1002    let expanded_constraints: Result<Vec<Constraint>, FunctionError> = program
1003        .constraints
1004        .iter()
1005        .map(|constraint| expand_constraint_functions(&mut ctx, constraint))
1006        .collect();
1007
1008    Ok(Program {
1009        rules: expanded_rules?,
1010        directives: program.directives.clone(),
1011        queries: program.queries.clone(),
1012        predicates: program.predicates.clone(),
1013        constraints: expanded_constraints?,
1014        authored_constraint_source_bound: program.authored_constraint_source_bound,
1015        imports: program.imports.clone(),
1016        functions: program.functions.clone(),
1017        domains: program.domains.clone(),
1018        prob_facts: program.prob_facts.clone(),
1019        annotated_disjunctions: program.annotated_disjunctions.clone(),
1020        evidence: program.evidence.clone(),
1021        prob_queries: program.prob_queries.clone(),
1022        neural_predicates: program.neural_predicates.clone(),
1023        learnable_rules: program.learnable_rules.clone(),
1024    })
1025}
1026
1027/// Expand function calls in a single rule.
1028fn expand_rule_functions(ctx: &mut ExpansionContext, rule: &Rule) -> Result<Rule, FunctionError> {
1029    let mut used_variables: HashSet<String> = rule
1030        .head
1031        .variables()
1032        .into_iter()
1033        .chain(rule.body.iter().flat_map(BodyLiteral::variables))
1034        .map(ToOwned::to_owned)
1035        .collect();
1036    let expanded_body = expand_body_functions(ctx, &rule.body, &mut used_variables)?;
1037
1038    Ok(Rule {
1039        head: rule.head.clone(),
1040        body: expanded_body,
1041    })
1042}
1043
1044fn expand_constraint_functions(
1045    ctx: &mut ExpansionContext,
1046    constraint: &Constraint,
1047) -> Result<Constraint, FunctionError> {
1048    let mut used_variables: HashSet<String> = constraint
1049        .body
1050        .iter()
1051        .flat_map(BodyLiteral::variables)
1052        .map(ToOwned::to_owned)
1053        .collect();
1054    Ok(Constraint {
1055        authored_index: constraint.authored_index,
1056        body: expand_body_functions(ctx, &constraint.body, &mut used_variables)?,
1057    })
1058}
1059
1060fn expand_body_functions(
1061    ctx: &mut ExpansionContext,
1062    body: &[BodyLiteral],
1063    used_variables: &mut HashSet<String>,
1064) -> Result<Vec<BodyLiteral>, FunctionError> {
1065    let mut expanded_body = Vec::new();
1066    for literal in body {
1067        expanded_body.extend(ctx.expand_literal_for_rule(literal, used_variables)?);
1068    }
1069    Ok(expanded_body)
1070}
1071
1072#[cfg(test)]
1073mod tests {
1074    use super::*;
1075    use crate::ast::{FuncDef, FuncParam};
1076
1077    #[test]
1078    fn test_simple_expansion() {
1079        let mut reg = FunctionRegistry::new();
1080
1081        // func double(X) = X + X
1082        let double = FuncDef {
1083            name: "double".to_string(),
1084            params: vec![FuncParam {
1085                name: "X".to_string(),
1086                typ: None,
1087            }],
1088            return_type: None,
1089            body: FuncBody::Arithmetic(ArithExpr::Add(
1090                Box::new(ArithExpr::Variable("X".to_string())),
1091                Box::new(ArithExpr::Variable("X".to_string())),
1092            )),
1093            is_private: false,
1094        };
1095        reg.register(double).unwrap();
1096
1097        let mut ctx = ExpansionContext::new(&reg, 100);
1098
1099        // double(5) should expand to 5 + 5
1100        let result = ctx.expand_call("double", &[ArithExpr::Integer(5)]).unwrap();
1101
1102        match result {
1103            ArithExpr::Add(l, r) => {
1104                assert!(matches!(*l, ArithExpr::Integer(5)));
1105                assert!(matches!(*r, ArithExpr::Integer(5)));
1106            }
1107            _ => panic!("Expected Add expression"),
1108        }
1109    }
1110
1111    #[test]
1112    fn test_nested_expansion() {
1113        let mut reg = FunctionRegistry::new();
1114
1115        // func double(X) = X + X
1116        let double = FuncDef {
1117            name: "double".to_string(),
1118            params: vec![FuncParam {
1119                name: "X".to_string(),
1120                typ: None,
1121            }],
1122            return_type: None,
1123            body: FuncBody::Arithmetic(ArithExpr::Add(
1124                Box::new(ArithExpr::Variable("X".to_string())),
1125                Box::new(ArithExpr::Variable("X".to_string())),
1126            )),
1127            is_private: false,
1128        };
1129
1130        // func quadruple(X) = double(double(X))
1131        let quadruple = FuncDef {
1132            name: "quadruple".to_string(),
1133            params: vec![FuncParam {
1134                name: "X".to_string(),
1135                typ: None,
1136            }],
1137            return_type: None,
1138            body: FuncBody::Arithmetic(ArithExpr::FuncCall {
1139                name: "double".to_string(),
1140                args: vec![ArithExpr::FuncCall {
1141                    name: "double".to_string(),
1142                    args: vec![ArithExpr::Variable("X".to_string())],
1143                }],
1144            }),
1145            is_private: false,
1146        };
1147
1148        reg.register(double).unwrap();
1149        reg.register(quadruple).unwrap();
1150
1151        let mut ctx = ExpansionContext::new(&reg, 100);
1152
1153        // quadruple(2) should expand to (2 + 2) + (2 + 2)
1154        let result = ctx
1155            .expand_call("quadruple", &[ArithExpr::Integer(2)])
1156            .unwrap();
1157
1158        // Result should be Add(Add(2, 2), Add(2, 2))
1159        match &result {
1160            ArithExpr::Add(l, r) => {
1161                assert!(matches!(l.as_ref(), ArithExpr::Add(_, _)));
1162                assert!(matches!(r.as_ref(), ArithExpr::Add(_, _)));
1163            }
1164            _ => panic!("Expected nested Add expression, got {:?}", result),
1165        }
1166    }
1167
1168    #[test]
1169    fn test_max_recursion_depth() {
1170        let mut reg = FunctionRegistry::new();
1171
1172        // func infinite(X) = infinite(X)
1173        let infinite = FuncDef {
1174            name: "infinite".to_string(),
1175            params: vec![FuncParam {
1176                name: "X".to_string(),
1177                typ: None,
1178            }],
1179            return_type: None,
1180            body: FuncBody::Arithmetic(ArithExpr::FuncCall {
1181                name: "infinite".to_string(),
1182                args: vec![ArithExpr::Variable("X".to_string())],
1183            }),
1184            is_private: false,
1185        };
1186        reg.register(infinite).unwrap();
1187
1188        let mut ctx = ExpansionContext::new(&reg, 10);
1189
1190        let result = ctx.expand_call("infinite", &[ArithExpr::Integer(1)]);
1191        assert!(matches!(
1192            result,
1193            Err(FunctionError::MaxRecursionDepth { .. })
1194        ));
1195    }
1196
1197    #[test]
1198    fn test_undefined_function() {
1199        let reg = FunctionRegistry::new();
1200        let mut ctx = ExpansionContext::new(&reg, 100);
1201
1202        let result = ctx.expand_call("undefined", &[ArithExpr::Integer(1)]);
1203        assert!(matches!(
1204            result,
1205            Err(FunctionError::UndefinedFunction { .. })
1206        ));
1207    }
1208
1209    #[test]
1210    fn test_builtin_function_passthrough() {
1211        let mut reg = FunctionRegistry::new();
1212
1213        // func abs_x(X) = abs(X)
1214        let abs_x = FuncDef {
1215            name: "abs_x".to_string(),
1216            params: vec![FuncParam {
1217                name: "X".to_string(),
1218                typ: None,
1219            }],
1220            return_type: None,
1221            body: FuncBody::Arithmetic(ArithExpr::FuncCall {
1222                name: "abs".to_string(),
1223                args: vec![ArithExpr::Variable("X".to_string())],
1224            }),
1225            is_private: false,
1226        };
1227        reg.register(abs_x).unwrap();
1228
1229        let mut ctx = ExpansionContext::new(&reg, 100);
1230
1231        let result = ctx.expand_call("abs_x", &[ArithExpr::Integer(-5)]).unwrap();
1232
1233        // Should preserve abs call with substituted arg
1234        match result {
1235            ArithExpr::FuncCall { name, args } => {
1236                assert_eq!(name, "abs");
1237                assert_eq!(args.len(), 1);
1238                assert!(matches!(args[0], ArithExpr::Integer(-5)));
1239            }
1240            _ => panic!("Expected FuncCall for builtin"),
1241        }
1242    }
1243
1244    #[test]
1245    fn test_variable_substitution() {
1246        let mut reg = FunctionRegistry::new();
1247
1248        // func add(X, Y) = X + Y
1249        let add = FuncDef {
1250            name: "add".to_string(),
1251            params: vec![
1252                FuncParam {
1253                    name: "X".to_string(),
1254                    typ: None,
1255                },
1256                FuncParam {
1257                    name: "Y".to_string(),
1258                    typ: None,
1259                },
1260            ],
1261            return_type: None,
1262            body: FuncBody::Arithmetic(ArithExpr::Add(
1263                Box::new(ArithExpr::Variable("X".to_string())),
1264                Box::new(ArithExpr::Variable("Y".to_string())),
1265            )),
1266            is_private: false,
1267        };
1268        reg.register(add).unwrap();
1269
1270        let mut ctx = ExpansionContext::new(&reg, 100);
1271
1272        // add(3, 7) should expand to 3 + 7
1273        let result = ctx
1274            .expand_call("add", &[ArithExpr::Integer(3), ArithExpr::Integer(7)])
1275            .unwrap();
1276
1277        match result {
1278            ArithExpr::Add(l, r) => {
1279                assert!(matches!(*l, ArithExpr::Integer(3)));
1280                assert!(matches!(*r, ArithExpr::Integer(7)));
1281            }
1282            _ => panic!("Expected Add expression"),
1283        }
1284    }
1285
1286    #[test]
1287    fn test_expansion_with_variable_args() {
1288        let mut reg = FunctionRegistry::new();
1289
1290        // func double(X) = X + X
1291        let double = FuncDef {
1292            name: "double".to_string(),
1293            params: vec![FuncParam {
1294                name: "X".to_string(),
1295                typ: None,
1296            }],
1297            return_type: None,
1298            body: FuncBody::Arithmetic(ArithExpr::Add(
1299                Box::new(ArithExpr::Variable("X".to_string())),
1300                Box::new(ArithExpr::Variable("X".to_string())),
1301            )),
1302            is_private: false,
1303        };
1304        reg.register(double).unwrap();
1305
1306        let mut ctx = ExpansionContext::new(&reg, 100);
1307
1308        // double(Y) should expand to Y + Y
1309        let result = ctx
1310            .expand_call("double", &[ArithExpr::Variable("Y".to_string())])
1311            .unwrap();
1312
1313        match result {
1314            ArithExpr::Add(l, r) => {
1315                assert!(matches!(l.as_ref(), ArithExpr::Variable(n) if n == "Y"));
1316                assert!(matches!(r.as_ref(), ArithExpr::Variable(n) if n == "Y"));
1317            }
1318            _ => panic!("Expected Add expression"),
1319        }
1320    }
1321
1322    #[test]
1323    fn test_predicate_func_expansion() {
1324        // func get_parent(X) = P :- parent(X, P).
1325        // get_parent(alice) should expand to: parent(alice, P)
1326
1327        let func = FuncDef {
1328            name: "get_parent".to_string(),
1329            params: vec![FuncParam {
1330                name: "X".to_string(),
1331                typ: None,
1332            }],
1333            return_type: None,
1334            body: FuncBody::Predicate {
1335                result: "P".to_string(),
1336                body: vec![BodyLiteral::Positive(Atom {
1337                    predicate: "parent".to_string(),
1338                    terms: vec![
1339                        Term::Variable("X".to_string()),
1340                        Term::Variable("P".to_string()),
1341                    ],
1342                })],
1343            },
1344            is_private: false,
1345        };
1346
1347        let mut reg = FunctionRegistry::new();
1348        reg.register(func).unwrap();
1349
1350        let mut ctx = ExpansionContext::new(&reg, 100);
1351
1352        // Call get_parent with "alice"
1353        let args = vec![ArithExpr::Variable("alice".to_string())];
1354        let mut used = HashSet::from(["alice".to_string()]);
1355        let expanded = ctx
1356            .expand_expr_for_rule(
1357                &ArithExpr::FuncCall {
1358                    name: "get_parent".to_string(),
1359                    args,
1360                },
1361                &HashMap::new(),
1362                &mut used,
1363                false,
1364            )
1365            .unwrap();
1366        let body = expanded.generated_literals;
1367        let ArithExpr::Variable(result) = expanded.expression else {
1368            panic!("Expected variable result")
1369        };
1370
1371        assert_ne!(result, "P");
1372        assert_eq!(body.len(), 1);
1373
1374        // Check the expanded literal
1375        if let BodyLiteral::Positive(atom) = &body[0] {
1376            assert_eq!(atom.predicate, "parent");
1377            assert!(matches!(&atom.terms[0], Term::Variable(v) if v == "alice"));
1378            assert!(matches!(&atom.terms[1], Term::Variable(v) if v == &result));
1379        } else {
1380            panic!("Expected Positive literal");
1381        }
1382    }
1383
1384    #[test]
1385    fn test_predicate_func_with_constant_arg() {
1386        // func get_child(P) = C :- parent(C, P).
1387        // get_child(bob) should expand to: parent(C, bob)
1388
1389        let func = FuncDef {
1390            name: "get_child".to_string(),
1391            params: vec![FuncParam {
1392                name: "P".to_string(),
1393                typ: None,
1394            }],
1395            return_type: None,
1396            body: FuncBody::Predicate {
1397                result: "C".to_string(),
1398                body: vec![BodyLiteral::Positive(Atom {
1399                    predicate: "parent".to_string(),
1400                    terms: vec![
1401                        Term::Variable("C".to_string()),
1402                        Term::Variable("P".to_string()),
1403                    ],
1404                })],
1405            },
1406            is_private: false,
1407        };
1408
1409        let mut reg = FunctionRegistry::new();
1410        reg.register(func).unwrap();
1411
1412        let mut ctx = ExpansionContext::new(&reg, 100);
1413
1414        // Call get_child with integer constant
1415        let args = vec![ArithExpr::Integer(42)];
1416        let mut used = HashSet::new();
1417        let expanded = ctx
1418            .expand_expr_for_rule(
1419                &ArithExpr::FuncCall {
1420                    name: "get_child".to_string(),
1421                    args,
1422                },
1423                &HashMap::new(),
1424                &mut used,
1425                false,
1426            )
1427            .unwrap();
1428        let body = expanded.generated_literals;
1429        let ArithExpr::Variable(result) = expanded.expression else {
1430            panic!("Expected variable result")
1431        };
1432
1433        assert_ne!(result, "C");
1434        assert_eq!(body.len(), 1);
1435
1436        // Check the expanded literal has integer substituted
1437        if let BodyLiteral::Positive(atom) = &body[0] {
1438            assert_eq!(atom.predicate, "parent");
1439            assert!(matches!(&atom.terms[0], Term::Variable(v) if v == &result));
1440            assert!(matches!(&atom.terms[1], Term::Integer(42)));
1441        } else {
1442            panic!("Expected Positive literal");
1443        }
1444    }
1445
1446    #[test]
1447    fn test_predicate_func_multiple_body_literals() {
1448        // func get_grandparent(X) = G :- parent(X, P), parent(P, G).
1449        // get_grandparent(alice) should expand to: parent(alice, P), parent(P, G)
1450
1451        let func = FuncDef {
1452            name: "get_grandparent".to_string(),
1453            params: vec![FuncParam {
1454                name: "X".to_string(),
1455                typ: None,
1456            }],
1457            return_type: None,
1458            body: FuncBody::Predicate {
1459                result: "G".to_string(),
1460                body: vec![
1461                    BodyLiteral::Positive(Atom {
1462                        predicate: "parent".to_string(),
1463                        terms: vec![
1464                            Term::Variable("X".to_string()),
1465                            Term::Variable("P".to_string()),
1466                        ],
1467                    }),
1468                    BodyLiteral::Positive(Atom {
1469                        predicate: "parent".to_string(),
1470                        terms: vec![
1471                            Term::Variable("P".to_string()),
1472                            Term::Variable("G".to_string()),
1473                        ],
1474                    }),
1475                ],
1476            },
1477            is_private: false,
1478        };
1479
1480        let mut reg = FunctionRegistry::new();
1481        reg.register(func).unwrap();
1482
1483        let mut ctx = ExpansionContext::new(&reg, 100);
1484
1485        let args = vec![ArithExpr::Variable("alice".to_string())];
1486        let mut used = HashSet::from(["alice".to_string()]);
1487        let expanded = ctx
1488            .expand_expr_for_rule(
1489                &ArithExpr::FuncCall {
1490                    name: "get_grandparent".to_string(),
1491                    args,
1492                },
1493                &HashMap::new(),
1494                &mut used,
1495                false,
1496            )
1497            .unwrap();
1498        let body = expanded.generated_literals;
1499        let ArithExpr::Variable(result) = expanded.expression else {
1500            panic!("Expected variable result")
1501        };
1502
1503        assert_ne!(result, "G");
1504        assert_eq!(body.len(), 2);
1505
1506        // First literal: parent(alice, P)
1507        if let BodyLiteral::Positive(atom) = &body[0] {
1508            assert_eq!(atom.predicate, "parent");
1509            assert!(matches!(&atom.terms[0], Term::Variable(v) if v == "alice"));
1510            assert!(matches!(&atom.terms[1], Term::Variable(v) if v != "P"));
1511        } else {
1512            panic!("Expected Positive literal for first body");
1513        }
1514
1515        // Second literal: parent(P, G)
1516        if let BodyLiteral::Positive(atom) = &body[1] {
1517            assert_eq!(atom.predicate, "parent");
1518            assert_eq!(atom.terms[0], body[0].atom().unwrap().terms[1]);
1519            assert!(matches!(&atom.terms[1], Term::Variable(v) if v == &result));
1520        } else {
1521            panic!("Expected Positive literal for second body");
1522        }
1523    }
1524
1525    #[test]
1526    fn test_is_predicate_func() {
1527        let mut reg = FunctionRegistry::new();
1528
1529        // Arithmetic function
1530        let arith_func = FuncDef {
1531            name: "double".to_string(),
1532            params: vec![FuncParam {
1533                name: "X".to_string(),
1534                typ: None,
1535            }],
1536            return_type: None,
1537            body: FuncBody::Arithmetic(ArithExpr::Add(
1538                Box::new(ArithExpr::Variable("X".to_string())),
1539                Box::new(ArithExpr::Variable("X".to_string())),
1540            )),
1541            is_private: false,
1542        };
1543
1544        // Predicate function
1545        let pred_func = FuncDef {
1546            name: "get_parent".to_string(),
1547            params: vec![FuncParam {
1548                name: "X".to_string(),
1549                typ: None,
1550            }],
1551            return_type: None,
1552            body: FuncBody::Predicate {
1553                result: "P".to_string(),
1554                body: vec![BodyLiteral::Positive(Atom {
1555                    predicate: "parent".to_string(),
1556                    terms: vec![
1557                        Term::Variable("X".to_string()),
1558                        Term::Variable("P".to_string()),
1559                    ],
1560                })],
1561            },
1562            is_private: false,
1563        };
1564
1565        reg.register(arith_func).unwrap();
1566        reg.register(pred_func).unwrap();
1567
1568        let ctx = ExpansionContext::new(&reg, 100);
1569
1570        assert!(!ctx.is_predicate_func("double"));
1571        assert!(ctx.is_predicate_func("get_parent"));
1572        assert!(!ctx.is_predicate_func("nonexistent"));
1573    }
1574
1575    #[test]
1576    fn generated_variable_names_round_trip_underscored_identifiers() {
1577        let generated = generated_function_variable_name("get_parent", "Parent_Value", 42);
1578        assert_eq!(
1579            generated_function_variable_source(&generated, "get_parent"),
1580            Some("Parent_Value")
1581        );
1582        assert_eq!(generated, "__XLOG_FUNCTION_GET_PARENT_Parent_Value_42");
1583        assert!(generated_function_variable_source(
1584            "__XLOG_FUNCTION_GET_PARENT_Parent_Value_not_a_counter",
1585            "get_parent"
1586        )
1587        .is_none());
1588    }
1589}
xlog_logic/expand.rs

xlog_logic/
expand.rs
