Add tracing for head arithmetic (#374)

* Add tracing for head arithmetic

* Change display of numeric literal

* Use checked operations for evaluating a constant

* Use logical type of the predicate as the basis for evaluating constant expressions during tracing

---------

Co-authored-by: Alex Ivliev <[email protected]>

nemo/src/error.rs

@@ -17,7 +17,7 @@ pub use nemo_physical::error::ReadingError;
 pub enum Error {
     /// Currently tracing doesn't work for all language features
     #[error(
-        "Tracing is currently not supported for rules with arithmetic operations in the head."
+        "Tracing is currently not supported for some rules with arithmetic operations in the head."
     )]
     TraceUnsupportedFeature(),
     /// Error which implies a needed Rollback

nemo/src/execution/execution_engine.rs

@@ -443,10 +443,14 @@ impl<Strategy: RuleSelectionStrategy> ExecutionEngine<Strategy> {
             }
 
             // Unify the head atom with the given fact
-            // If unification is possible `compatible` remain true
-            // and `assignment` will contain the match which is responsible for the fact
+
+            // If unification is possible `compatible` remains true
             let mut compatible = true;
-            let mut assignment = HashMap::<Variable, Constant>::new();
+            // Contains the head variable and the constant it aligns with.
+            let mut assignment_constant = HashMap::<Variable, Constant>::new();
+            // For each constructor variable, contains the term which describes its calculation
+            // and the constant it should equal to based on the input fact
+            let mut assignment_constructor = HashMap::<Variable, (Constant, Term)>::new();
 
             for (ty, (head_term, fact_term)) in predicate_types
                 .iter()
@@ -468,22 +472,77 @@ impl<Strategy: RuleSelectionStrategy> ExecutionEngine<Strategy> {
                             continue;
                         }
 
-                        if rule.get_constructor(variable).is_some() {
-                            // TODO: Support arbitrary operations in the head
-                            return Err(Error::TraceUnsupportedFeature());
-                        }
+                        if let Some(constructor) = rule.get_constructor(variable) {
+                            match assignment_constructor.entry(variable.clone()) {
+                                Entry::Occupied(entry) => {
+                                    let (stored_constant, _) = entry.get();
 
-                        match assignment.entry(variable.clone()) {
-                            Entry::Occupied(entry) => {
-                                if ty.ground_term_to_data_value_t(entry.get().clone())
-                                    != ty.ground_term_to_data_value_t(fact_term.clone())
-                                {
-                                    compatible = false;
-                                    break;
+                                    if stored_constant != fact_term {
+                                        compatible = false;
+                                        break;
+                                    }
+                                }
+                                Entry::Vacant(entry) => {
+                                    if constructor.term().variables().next().is_none() {
+                                        if let Term::Primitive(PrimitiveTerm::Constant(constant)) =
+                                            constructor.term()
+                                        {
+                                            if ty.ground_term_to_data_value_t(constant.clone())
+                                                != ty.ground_term_to_data_value_t(fact_term.clone())
+                                            {
+                                                compatible = false;
+                                                break;
+                                            }
+                                        } else {
+                                            if let Ok(fact_term_data) =
+                                                ty.ground_term_to_data_value_t(fact_term.clone())
+                                            {
+                                                if let Some(fact_term_storage) = fact_term_data
+                                                    .to_storage_value(
+                                                        &self.table_manager.get_dict(),
+                                                    )
+                                                {
+                                                    if let Some(constructor_value_storage) =
+                                                        constructor
+                                                            .term()
+                                                            .evaluate_constant_numeric(
+                                                                ty,
+                                                                &self.table_manager.get_dict(),
+                                                            )
+                                                    {
+                                                        if fact_term_storage
+                                                            == constructor_value_storage
+                                                        {
+                                                            continue;
+                                                        }
+                                                    }
+                                                }
+                                            }
+
+                                            compatible = false;
+                                            break;
+                                        }
+                                    } else {
+                                        entry.insert((
+                                            fact_term.clone(),
+                                            constructor.term().clone(),
+                                        ));
+                                    }
                                 }
                             }
-                            Entry::Vacant(entry) => {
-                                entry.insert(fact_term.clone());
+                        } else {
+                            match assignment_constant.entry(variable.clone()) {
+                                Entry::Occupied(entry) => {
+                                    if ty.ground_term_to_data_value_t(entry.get().clone())
+                                        != ty.ground_term_to_data_value_t(fact_term.clone())
+                                    {
+                                        compatible = false;
+                                        break;
+                                    }
+                                }
+                                Entry::Vacant(entry) => {
+                                    entry.insert(fact_term.clone());
+                                }
                             }
                         }
                     }
@@ -498,18 +557,28 @@ impl<Strategy: RuleSelectionStrategy> ExecutionEngine<Strategy> {
             // The goal of this part of the code is to apply the rule which led to the given fact
             // but with the variable binding derived from the unification above
 
-            let new_constraints: Vec<Constraint> = assignment
-                .iter()
-                .map(|(variable, term)| {
-                    Constraint::Equals(
-                        Term::Primitive(PrimitiveTerm::Variable(variable.clone())),
-                        Term::Primitive(PrimitiveTerm::Constant(term.clone())),
-                    )
-                })
-                .collect();
+            let unification_constraints: Vec<Constraint> =
+                assignment_constant
+                    .into_iter()
+                    .map(|(variable, term)| {
+                        Constraint::Equals(
+                            Term::Primitive(PrimitiveTerm::Variable(variable)),
+                            Term::Primitive(PrimitiveTerm::Constant(term)),
+                        )
+                    })
+                    .chain(assignment_constructor.into_iter().map(
+                        |(_variable, (constant, term))| {
+                            Constraint::Equals(
+                                Term::Primitive(PrimitiveTerm::Constant(constant)),
+                                term,
+                            )
+                        },
+                    ))
+                    .collect();
 
             let mut rule = self.program.rules()[rule_index].clone();
-            rule.positive_constraints_mut().extend(new_constraints);
+            rule.positive_constraints_mut()
+                .extend(unification_constraints);
             let analysis = &self.analysis.rule_analysis[rule_index];
 
             let body_execution = SeminaiveStrategy::initialize(&rule, analysis);

nemo/src/execution/planning.rs

@@ -19,4 +19,4 @@ pub mod plan_util;
 
 pub mod negation;
 
-mod arithmetic;
+pub mod arithmetic;

nemo/src/execution/planning/arithmetic.rs

@@ -17,7 +17,9 @@ use crate::{
     program_analysis::variable_order::VariableOrder,
 };
 
-pub(super) fn termtree_to_arithmetictree(
+/// Builds an [`ArithmeticTree`] with [`DataValueT`]
+/// from a given [`Term`].
+pub fn termtree_to_arithmetictree(
     term: &Term,
     order: &VariableOrder,
     logical_type: &PrimitiveType,

nemo/src/model/rule_model/numeric_literal.rs

@@ -1,7 +1,7 @@
 use nemo_physical::datatypes::Double;
 
 /// A numerical literal.
-#[derive(Debug, Eq, PartialEq, Copy, Clone, PartialOrd, Ord)]
+#[derive(Eq, PartialEq, Copy, Clone, PartialOrd, Ord)]
 pub enum NumericLiteral {
     /// An integer literal.
     Integer(i64),
@@ -11,7 +11,7 @@ pub enum NumericLiteral {
     Double(Double),
 }
 
-impl std::fmt::Display for NumericLiteral {
+impl std::fmt::Debug for NumericLiteral {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             NumericLiteral::Integer(value) => write!(f, "{value}"),
@@ -20,3 +20,15 @@ impl std::fmt::Display for NumericLiteral {
         }
     }
 }
+
+impl std::fmt::Display for NumericLiteral {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            NumericLiteral::Integer(value) => write!(f, "{value}"),
+            NumericLiteral::Decimal(left, right) => write!(f, "{left}.{right}"),
+            NumericLiteral::Double(value) => {
+                f.write_str(format!("{:.4}", f64::from(*value)).trim_end_matches(['.', '0']))
+            }
+        }
+    }
+}

nemo/src/model/rule_model/term.rs

@@ -1,6 +1,18 @@
 use std::fmt::{Debug, Display};
 
-use crate::model::{types::primitive_logical_value::LOGICAL_NULL_PREFIX, VariableAssignment};
+use nemo_physical::{
+    columnar::operations::arithmetic::expression::ArithmeticTreeLeaf,
+    datatypes::{DataValueT, StorageTypeName, StorageValueT},
+    management::database::Dict,
+};
+
+use crate::{
+    execution::planning::arithmetic::termtree_to_arithmetictree,
+    model::{
+        types::primitive_logical_value::LOGICAL_NULL_PREFIX, PrimitiveType, VariableAssignment,
+    },
+    program_analysis::variable_order::VariableOrder,
+};
 
 use super::{Aggregate, Identifier, NumericLiteral, RdfLiteral};
 
@@ -367,6 +379,50 @@ impl Term {
             Term::Function(sub) => sub.subterms.iter().any(Self::aggregate_subterm_recursive),
         }
     }
+
+    /// Evaluates a constant (numeric) term.
+    pub fn evaluate_constant_numeric(
+        &self,
+        ty: &PrimitiveType,
+        dict: &Dict,
+    ) -> Option<StorageValueT> {
+        let arithmetic_tree = termtree_to_arithmetictree(self, &VariableOrder::new(), ty);
+        let storage_type = ty.datatype_name().to_storage_type_name();
+
+        macro_rules! translate_data_type {
+            ($variant:ident, $type:ty) => {{
+                let translate_function = |l: ArithmeticTreeLeaf<DataValueT>| match l {
+                    ArithmeticTreeLeaf::Constant(t) => {
+                        if let StorageValueT::$variant(value) = t
+                            .to_storage_value(dict)
+                            .expect("We don't have string operations so this cannot fail.")
+                        {
+                            ArithmeticTreeLeaf::Constant(value)
+                        } else {
+                            panic!(
+                                "Expected a operation tree value of type {}",
+                                stringify!($src_name)
+                            );
+                        }
+                    }
+                    ArithmeticTreeLeaf::Reference(index) => ArithmeticTreeLeaf::Reference(index),
+                };
+
+                let arithmetic_tree_typed = arithmetic_tree.map(&translate_function);
+                Some(StorageValueT::$variant(
+                    arithmetic_tree_typed.evaluate(&[])?,
+                ))
+            }};
+        }
+
+        match storage_type {
+            StorageTypeName::U32 => translate_data_type!(U32, u32),
+            StorageTypeName::U64 => translate_data_type!(U64, u64),
+            StorageTypeName::I64 => translate_data_type!(I64, i64),
+            StorageTypeName::Float => translate_data_type!(Float, f32),
+            StorageTypeName::Double => translate_data_type!(Double, f64),
+        }
+    }
 }
 
 impl From<PrimitiveTerm> for Term {