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    #[cfg(test)]
947    fn is_predicate_func(&self, name: &str) -> bool {
948        self.registry
949            .get(name)
950            .is_some_and(|function| matches!(function.body, FuncBody::Predicate { .. }))
951    }
952}
953
954use crate::ast::{Program, Rule};
955
956/// Expand user-defined function calls in ordinary rules and constraints.
957///
958/// Scalar calls become arithmetic expressions. Predicate-bodied calls also
959/// contribute relational literals immediately before their source binding.
960pub fn expand_program_functions(
961    program: &Program,
962    max_depth: u32,
963) -> Result<Program, FunctionError> {
964    // If no functions defined, return program unchanged
965    if program.functions.is_empty() {
966        return Ok(program.clone());
967    }
968    expand_program_functions_impl(program, max_depth)
969}
970
971/// [`expand_program_functions`] taking the program by value: without function
972/// definitions (the common, fact-heavy case) the program is returned as is,
973/// without a clone.
974pub fn expand_program_functions_owned(
975    program: Program,
976    max_depth: u32,
977) -> Result<Program, FunctionError> {
978    if program.functions.is_empty() {
979        return Ok(program);
980    }
981    expand_program_functions_impl(&program, max_depth)
982}
983
984fn expand_program_functions_impl(
985    program: &Program,
986    max_depth: u32,
987) -> Result<Program, FunctionError> {
988    let mut registry = FunctionRegistry::new();
989    for function in &program.functions {
990        registry.register(function.clone())?;
991    }
992    let mut ctx = ExpansionContext::new(&registry, max_depth);
993
994    // Expand function calls in each rule
995    let expanded_rules: Result<Vec<Rule>, FunctionError> = program
996        .rules
997        .iter()
998        .map(|rule| expand_rule_functions(&mut ctx, rule))
999        .collect();
1000    let expanded_constraints: Result<Vec<Constraint>, FunctionError> = program
1001        .constraints
1002        .iter()
1003        .map(|constraint| expand_constraint_functions(&mut ctx, constraint))
1004        .collect();
1005
1006    Ok(Program {
1007        rules: expanded_rules?,
1008        directives: program.directives.clone(),
1009        queries: program.queries.clone(),
1010        predicates: program.predicates.clone(),
1011        constraints: expanded_constraints?,
1012        authored_constraint_source_bound: program.authored_constraint_source_bound,
1013        imports: program.imports.clone(),
1014        functions: program.functions.clone(),
1015        domains: program.domains.clone(),
1016        prob_facts: program.prob_facts.clone(),
1017        annotated_disjunctions: program.annotated_disjunctions.clone(),
1018        evidence: program.evidence.clone(),
1019        prob_queries: program.prob_queries.clone(),
1020        neural_predicates: program.neural_predicates.clone(),
1021        learnable_rules: program.learnable_rules.clone(),
1022    })
1023}
1024
1025/// Expand function calls in a single rule.
1026fn expand_rule_functions(ctx: &mut ExpansionContext, rule: &Rule) -> Result<Rule, FunctionError> {
1027    let mut used_variables: HashSet<String> = rule
1028        .head
1029        .variables()
1030        .into_iter()
1031        .chain(rule.body.iter().flat_map(BodyLiteral::variables))
1032        .map(ToOwned::to_owned)
1033        .collect();
1034    let expanded_body = expand_body_functions(ctx, &rule.body, &mut used_variables)?;
1035
1036    Ok(Rule {
1037        head: rule.head.clone(),
1038        body: expanded_body,
1039    })
1040}
1041
1042fn expand_constraint_functions(
1043    ctx: &mut ExpansionContext,
1044    constraint: &Constraint,
1045) -> Result<Constraint, FunctionError> {
1046    let mut used_variables: HashSet<String> = constraint
1047        .body
1048        .iter()
1049        .flat_map(BodyLiteral::variables)
1050        .map(ToOwned::to_owned)
1051        .collect();
1052    Ok(Constraint {
1053        authored_index: constraint.authored_index,
1054        body: expand_body_functions(ctx, &constraint.body, &mut used_variables)?,
1055    })
1056}
1057
1058fn expand_body_functions(
1059    ctx: &mut ExpansionContext,
1060    body: &[BodyLiteral],
1061    used_variables: &mut HashSet<String>,
1062) -> Result<Vec<BodyLiteral>, FunctionError> {
1063    let mut expanded_body = Vec::new();
1064    for literal in body {
1065        expanded_body.extend(ctx.expand_literal_for_rule(literal, used_variables)?);
1066    }
1067    Ok(expanded_body)
1068}
1069
1070#[cfg(test)]
1071mod tests {
1072    use super::*;
1073    use crate::ast::{FuncDef, FuncParam};
1074
1075    #[test]
1076    fn test_simple_expansion() {
1077        let mut reg = FunctionRegistry::new();
1078
1079        // func double(X) = X + X
1080        let double = FuncDef {
1081            name: "double".to_string(),
1082            params: vec![FuncParam {
1083                name: "X".to_string(),
1084                typ: None,
1085            }],
1086            return_type: None,
1087            body: FuncBody::Arithmetic(ArithExpr::Add(
1088                Box::new(ArithExpr::Variable("X".to_string())),
1089                Box::new(ArithExpr::Variable("X".to_string())),
1090            )),
1091            is_private: false,
1092        };
1093        reg.register(double).unwrap();
1094
1095        let mut ctx = ExpansionContext::new(&reg, 100);
1096
1097        // double(5) should expand to 5 + 5
1098        let result = ctx.expand_call("double", &[ArithExpr::Integer(5)]).unwrap();
1099
1100        match result {
1101            ArithExpr::Add(l, r) => {
1102                assert!(matches!(*l, ArithExpr::Integer(5)));
1103                assert!(matches!(*r, ArithExpr::Integer(5)));
1104            }
1105            _ => panic!("Expected Add expression"),
1106        }
1107    }
1108
1109    #[test]
1110    fn test_nested_expansion() {
1111        let mut reg = FunctionRegistry::new();
1112
1113        // func double(X) = X + X
1114        let double = FuncDef {
1115            name: "double".to_string(),
1116            params: vec![FuncParam {
1117                name: "X".to_string(),
1118                typ: None,
1119            }],
1120            return_type: None,
1121            body: FuncBody::Arithmetic(ArithExpr::Add(
1122                Box::new(ArithExpr::Variable("X".to_string())),
1123                Box::new(ArithExpr::Variable("X".to_string())),
1124            )),
1125            is_private: false,
1126        };
1127
1128        // func quadruple(X) = double(double(X))
1129        let quadruple = FuncDef {
1130            name: "quadruple".to_string(),
1131            params: vec![FuncParam {
1132                name: "X".to_string(),
1133                typ: None,
1134            }],
1135            return_type: None,
1136            body: FuncBody::Arithmetic(ArithExpr::FuncCall {
1137                name: "double".to_string(),
1138                args: vec![ArithExpr::FuncCall {
1139                    name: "double".to_string(),
1140                    args: vec![ArithExpr::Variable("X".to_string())],
1141                }],
1142            }),
1143            is_private: false,
1144        };
1145
1146        reg.register(double).unwrap();
1147        reg.register(quadruple).unwrap();
1148
1149        let mut ctx = ExpansionContext::new(&reg, 100);
1150
1151        // quadruple(2) should expand to (2 + 2) + (2 + 2)
1152        let result = ctx
1153            .expand_call("quadruple", &[ArithExpr::Integer(2)])
1154            .unwrap();
1155
1156        // Result should be Add(Add(2, 2), Add(2, 2))
1157        match &result {
1158            ArithExpr::Add(l, r) => {
1159                assert!(matches!(l.as_ref(), ArithExpr::Add(_, _)));
1160                assert!(matches!(r.as_ref(), ArithExpr::Add(_, _)));
1161            }
1162            _ => panic!("Expected nested Add expression, got {:?}", result),
1163        }
1164    }
1165
1166    #[test]
1167    fn test_max_recursion_depth() {
1168        let mut reg = FunctionRegistry::new();
1169
1170        // func infinite(X) = infinite(X)
1171        let infinite = FuncDef {
1172            name: "infinite".to_string(),
1173            params: vec![FuncParam {
1174                name: "X".to_string(),
1175                typ: None,
1176            }],
1177            return_type: None,
1178            body: FuncBody::Arithmetic(ArithExpr::FuncCall {
1179                name: "infinite".to_string(),
1180                args: vec![ArithExpr::Variable("X".to_string())],
1181            }),
1182            is_private: false,
1183        };
1184        reg.register(infinite).unwrap();
1185
1186        let mut ctx = ExpansionContext::new(&reg, 10);
1187
1188        let result = ctx.expand_call("infinite", &[ArithExpr::Integer(1)]);
1189        assert!(matches!(
1190            result,
1191            Err(FunctionError::MaxRecursionDepth { .. })
1192        ));
1193    }
1194
1195    #[test]
1196    fn test_undefined_function() {
1197        let reg = FunctionRegistry::new();
1198        let mut ctx = ExpansionContext::new(&reg, 100);
1199
1200        let result = ctx.expand_call("undefined", &[ArithExpr::Integer(1)]);
1201        assert!(matches!(
1202            result,
1203            Err(FunctionError::UndefinedFunction { .. })
1204        ));
1205    }
1206
1207    #[test]
1208    fn test_builtin_function_passthrough() {
1209        let mut reg = FunctionRegistry::new();
1210
1211        // func abs_x(X) = abs(X)
1212        let abs_x = FuncDef {
1213            name: "abs_x".to_string(),
1214            params: vec![FuncParam {
1215                name: "X".to_string(),
1216                typ: None,
1217            }],
1218            return_type: None,
1219            body: FuncBody::Arithmetic(ArithExpr::FuncCall {
1220                name: "abs".to_string(),
1221                args: vec![ArithExpr::Variable("X".to_string())],
1222            }),
1223            is_private: false,
1224        };
1225        reg.register(abs_x).unwrap();
1226
1227        let mut ctx = ExpansionContext::new(&reg, 100);
1228
1229        let result = ctx.expand_call("abs_x", &[ArithExpr::Integer(-5)]).unwrap();
1230
1231        // Should preserve abs call with substituted arg
1232        match result {
1233            ArithExpr::FuncCall { name, args } => {
1234                assert_eq!(name, "abs");
1235                assert_eq!(args.len(), 1);
1236                assert!(matches!(args[0], ArithExpr::Integer(-5)));
1237            }
1238            _ => panic!("Expected FuncCall for builtin"),
1239        }
1240    }
1241
1242    #[test]
1243    fn test_variable_substitution() {
1244        let mut reg = FunctionRegistry::new();
1245
1246        // func add(X, Y) = X + Y
1247        let add = FuncDef {
1248            name: "add".to_string(),
1249            params: vec![
1250                FuncParam {
1251                    name: "X".to_string(),
1252                    typ: None,
1253                },
1254                FuncParam {
1255                    name: "Y".to_string(),
1256                    typ: None,
1257                },
1258            ],
1259            return_type: None,
1260            body: FuncBody::Arithmetic(ArithExpr::Add(
1261                Box::new(ArithExpr::Variable("X".to_string())),
1262                Box::new(ArithExpr::Variable("Y".to_string())),
1263            )),
1264            is_private: false,
1265        };
1266        reg.register(add).unwrap();
1267
1268        let mut ctx = ExpansionContext::new(&reg, 100);
1269
1270        // add(3, 7) should expand to 3 + 7
1271        let result = ctx
1272            .expand_call("add", &[ArithExpr::Integer(3), ArithExpr::Integer(7)])
1273            .unwrap();
1274
1275        match result {
1276            ArithExpr::Add(l, r) => {
1277                assert!(matches!(*l, ArithExpr::Integer(3)));
1278                assert!(matches!(*r, ArithExpr::Integer(7)));
1279            }
1280            _ => panic!("Expected Add expression"),
1281        }
1282    }
1283
1284    #[test]
1285    fn test_expansion_with_variable_args() {
1286        let mut reg = FunctionRegistry::new();
1287
1288        // func double(X) = X + X
1289        let double = FuncDef {
1290            name: "double".to_string(),
1291            params: vec![FuncParam {
1292                name: "X".to_string(),
1293                typ: None,
1294            }],
1295            return_type: None,
1296            body: FuncBody::Arithmetic(ArithExpr::Add(
1297                Box::new(ArithExpr::Variable("X".to_string())),
1298                Box::new(ArithExpr::Variable("X".to_string())),
1299            )),
1300            is_private: false,
1301        };
1302        reg.register(double).unwrap();
1303
1304        let mut ctx = ExpansionContext::new(&reg, 100);
1305
1306        // double(Y) should expand to Y + Y
1307        let result = ctx
1308            .expand_call("double", &[ArithExpr::Variable("Y".to_string())])
1309            .unwrap();
1310
1311        match result {
1312            ArithExpr::Add(l, r) => {
1313                assert!(matches!(l.as_ref(), ArithExpr::Variable(n) if n == "Y"));
1314                assert!(matches!(r.as_ref(), ArithExpr::Variable(n) if n == "Y"));
1315            }
1316            _ => panic!("Expected Add expression"),
1317        }
1318    }
1319
1320    #[test]
1321    fn test_predicate_func_expansion() {
1322        // func get_parent(X) = P :- parent(X, P).
1323        // get_parent(alice) should expand to: parent(alice, P)
1324
1325        let func = FuncDef {
1326            name: "get_parent".to_string(),
1327            params: vec![FuncParam {
1328                name: "X".to_string(),
1329                typ: None,
1330            }],
1331            return_type: None,
1332            body: FuncBody::Predicate {
1333                result: "P".to_string(),
1334                body: vec![BodyLiteral::Positive(Atom {
1335                    predicate: "parent".to_string(),
1336                    terms: vec![
1337                        Term::Variable("X".to_string()),
1338                        Term::Variable("P".to_string()),
1339                    ],
1340                })],
1341            },
1342            is_private: false,
1343        };
1344
1345        let mut reg = FunctionRegistry::new();
1346        reg.register(func).unwrap();
1347
1348        let mut ctx = ExpansionContext::new(&reg, 100);
1349
1350        // Call get_parent with "alice"
1351        let args = vec![ArithExpr::Variable("alice".to_string())];
1352        let mut used = HashSet::from(["alice".to_string()]);
1353        let expanded = ctx
1354            .expand_expr_for_rule(
1355                &ArithExpr::FuncCall {
1356                    name: "get_parent".to_string(),
1357                    args,
1358                },
1359                &HashMap::new(),
1360                &mut used,
1361                false,
1362            )
1363            .unwrap();
1364        let body = expanded.generated_literals;
1365        let ArithExpr::Variable(result) = expanded.expression else {
1366            panic!("Expected variable result")
1367        };
1368
1369        assert_ne!(result, "P");
1370        assert_eq!(body.len(), 1);
1371
1372        // Check the expanded literal
1373        if let BodyLiteral::Positive(atom) = &body[0] {
1374            assert_eq!(atom.predicate, "parent");
1375            assert!(matches!(&atom.terms[0], Term::Variable(v) if v == "alice"));
1376            assert!(matches!(&atom.terms[1], Term::Variable(v) if v == &result));
1377        } else {
1378            panic!("Expected Positive literal");
1379        }
1380    }
1381
1382    #[test]
1383    fn test_predicate_func_with_constant_arg() {
1384        // func get_child(P) = C :- parent(C, P).
1385        // get_child(bob) should expand to: parent(C, bob)
1386
1387        let func = FuncDef {
1388            name: "get_child".to_string(),
1389            params: vec![FuncParam {
1390                name: "P".to_string(),
1391                typ: None,
1392            }],
1393            return_type: None,
1394            body: FuncBody::Predicate {
1395                result: "C".to_string(),
1396                body: vec![BodyLiteral::Positive(Atom {
1397                    predicate: "parent".to_string(),
1398                    terms: vec![
1399                        Term::Variable("C".to_string()),
1400                        Term::Variable("P".to_string()),
1401                    ],
1402                })],
1403            },
1404            is_private: false,
1405        };
1406
1407        let mut reg = FunctionRegistry::new();
1408        reg.register(func).unwrap();
1409
1410        let mut ctx = ExpansionContext::new(&reg, 100);
1411
1412        // Call get_child with integer constant
1413        let args = vec![ArithExpr::Integer(42)];
1414        let mut used = HashSet::new();
1415        let expanded = ctx
1416            .expand_expr_for_rule(
1417                &ArithExpr::FuncCall {
1418                    name: "get_child".to_string(),
1419                    args,
1420                },
1421                &HashMap::new(),
1422                &mut used,
1423                false,
1424            )
1425            .unwrap();
1426        let body = expanded.generated_literals;
1427        let ArithExpr::Variable(result) = expanded.expression else {
1428            panic!("Expected variable result")
1429        };
1430
1431        assert_ne!(result, "C");
1432        assert_eq!(body.len(), 1);
1433
1434        // Check the expanded literal has integer substituted
1435        if let BodyLiteral::Positive(atom) = &body[0] {
1436            assert_eq!(atom.predicate, "parent");
1437            assert!(matches!(&atom.terms[0], Term::Variable(v) if v == &result));
1438            assert!(matches!(&atom.terms[1], Term::Integer(42)));
1439        } else {
1440            panic!("Expected Positive literal");
1441        }
1442    }
1443
1444    #[test]
1445    fn test_predicate_func_multiple_body_literals() {
1446        // func get_grandparent(X) = G :- parent(X, P), parent(P, G).
1447        // get_grandparent(alice) should expand to: parent(alice, P), parent(P, G)
1448
1449        let func = FuncDef {
1450            name: "get_grandparent".to_string(),
1451            params: vec![FuncParam {
1452                name: "X".to_string(),
1453                typ: None,
1454            }],
1455            return_type: None,
1456            body: FuncBody::Predicate {
1457                result: "G".to_string(),
1458                body: vec![
1459                    BodyLiteral::Positive(Atom {
1460                        predicate: "parent".to_string(),
1461                        terms: vec![
1462                            Term::Variable("X".to_string()),
1463                            Term::Variable("P".to_string()),
1464                        ],
1465                    }),
1466                    BodyLiteral::Positive(Atom {
1467                        predicate: "parent".to_string(),
1468                        terms: vec![
1469                            Term::Variable("P".to_string()),
1470                            Term::Variable("G".to_string()),
1471                        ],
1472                    }),
1473                ],
1474            },
1475            is_private: false,
1476        };
1477
1478        let mut reg = FunctionRegistry::new();
1479        reg.register(func).unwrap();
1480
1481        let mut ctx = ExpansionContext::new(&reg, 100);
1482
1483        let args = vec![ArithExpr::Variable("alice".to_string())];
1484        let mut used = HashSet::from(["alice".to_string()]);
1485        let expanded = ctx
1486            .expand_expr_for_rule(
1487                &ArithExpr::FuncCall {
1488                    name: "get_grandparent".to_string(),
1489                    args,
1490                },
1491                &HashMap::new(),
1492                &mut used,
1493                false,
1494            )
1495            .unwrap();
1496        let body = expanded.generated_literals;
1497        let ArithExpr::Variable(result) = expanded.expression else {
1498            panic!("Expected variable result")
1499        };
1500
1501        assert_ne!(result, "G");
1502        assert_eq!(body.len(), 2);
1503
1504        // First literal: parent(alice, P)
1505        if let BodyLiteral::Positive(atom) = &body[0] {
1506            assert_eq!(atom.predicate, "parent");
1507            assert!(matches!(&atom.terms[0], Term::Variable(v) if v == "alice"));
1508            assert!(matches!(&atom.terms[1], Term::Variable(v) if v != "P"));
1509        } else {
1510            panic!("Expected Positive literal for first body");
1511        }
1512
1513        // Second literal: parent(P, G)
1514        if let BodyLiteral::Positive(atom) = &body[1] {
1515            assert_eq!(atom.predicate, "parent");
1516            assert_eq!(atom.terms[0], body[0].atom().unwrap().terms[1]);
1517            assert!(matches!(&atom.terms[1], Term::Variable(v) if v == &result));
1518        } else {
1519            panic!("Expected Positive literal for second body");
1520        }
1521    }
1522
1523    #[test]
1524    fn test_is_predicate_func() {
1525        let mut reg = FunctionRegistry::new();
1526
1527        // Arithmetic function
1528        let arith_func = FuncDef {
1529            name: "double".to_string(),
1530            params: vec![FuncParam {
1531                name: "X".to_string(),
1532                typ: None,
1533            }],
1534            return_type: None,
1535            body: FuncBody::Arithmetic(ArithExpr::Add(
1536                Box::new(ArithExpr::Variable("X".to_string())),
1537                Box::new(ArithExpr::Variable("X".to_string())),
1538            )),
1539            is_private: false,
1540        };
1541
1542        // Predicate function
1543        let pred_func = FuncDef {
1544            name: "get_parent".to_string(),
1545            params: vec![FuncParam {
1546                name: "X".to_string(),
1547                typ: None,
1548            }],
1549            return_type: None,
1550            body: FuncBody::Predicate {
1551                result: "P".to_string(),
1552                body: vec![BodyLiteral::Positive(Atom {
1553                    predicate: "parent".to_string(),
1554                    terms: vec![
1555                        Term::Variable("X".to_string()),
1556                        Term::Variable("P".to_string()),
1557                    ],
1558                })],
1559            },
1560            is_private: false,
1561        };
1562
1563        reg.register(arith_func).unwrap();
1564        reg.register(pred_func).unwrap();
1565
1566        let ctx = ExpansionContext::new(&reg, 100);
1567
1568        assert!(!ctx.is_predicate_func("double"));
1569        assert!(ctx.is_predicate_func("get_parent"));
1570        assert!(!ctx.is_predicate_func("nonexistent"));
1571    }
1572
1573    #[test]
1574    fn generated_variable_names_round_trip_underscored_identifiers() {
1575        let generated = generated_function_variable_name("get_parent", "Parent_Value", 42);
1576        assert_eq!(
1577            generated_function_variable_source(&generated, "get_parent"),
1578            Some("Parent_Value")
1579        );
1580        assert_eq!(generated, "__XLOG_FUNCTION_GET_PARENT_Parent_Value_42");
1581        assert!(generated_function_variable_source(
1582            "__XLOG_FUNCTION_GET_PARENT_Parent_Value_not_a_counter",
1583            "get_parent"
1584        )
1585        .is_none());
1586    }
1587}
xlog_logic/expand.rs

xlog_logic/
expand.rs
