initial terms utils

loda-lang · Jun 19, 2024 · c407fe5 · c407fe5
1 parent 5eb2475
commit c407fe5
Show file tree

Hide file tree

Showing 5 changed files with 48 additions and 21 deletions.
diff --git a/src/form/expression_util.cpp b/src/form/expression_util.cpp
@@ -280,6 +280,14 @@ bool ExpressionUtil::isSimpleFunction(const Expression& e, bool strict) {
   }
 }
 
+bool ExpressionUtil::isInitialTerm(const Expression& e) {
+  if (e.type != Expression::Type::FUNCTION || e.children.size() != 1) {
+    return false;
+  }
+  const auto arg = e.children.front();
+  return arg.type == Expression::Type::CONSTANT;
+}
+
 bool ExpressionUtil::canBeNegative(const Expression& e) {
   switch (e.type) {
     case Expression::Type::CONSTANT:

diff --git a/src/form/expression_util.hpp b/src/form/expression_util.hpp
@@ -20,6 +20,8 @@ class ExpressionUtil {
 
   static bool isSimpleFunction(const Expression& e, bool strict = true);
 
+  static bool isInitialTerm(const Expression& e);
+
   static bool canBeNegative(const Expression& e);
 
   static void collectNames(const Expression& e, Expression::Type type,

diff --git a/src/form/formula_util.cpp b/src/form/formula_util.cpp
@@ -177,3 +177,35 @@ void FormulaUtil::resolveSimpleRecursions(Formula& formula) {
     formula.entries[func] = sum;
   }
 }
+
+void FormulaUtil::convertInitialTermsToIf(Formula& formula) {
+  auto it = formula.entries.begin();
+  while (it != formula.entries.end()) {
+    auto left = it->first;
+    auto general = ExpressionUtil::newFunction(left.name);
+    auto general_it = formula.entries.find(general);
+    if (ExpressionUtil::isInitialTerm(left) &&
+        general_it != formula.entries.end()) {
+      general_it->second =
+          Expression(Expression::Type::IF, "",
+                     {left.children.front(), it->second, general_it->second});
+      it = formula.entries.erase(it);
+    } else {
+      it++;
+    }
+  }
+}
+
+Formula FormulaUtil::extractInitialTerms(Formula& formula) {
+  Formula initial_terms;
+  auto it = formula.entries.begin();
+  while (it != formula.entries.end()) {
+    if (ExpressionUtil::isInitialTerm(it->first)) {
+      initial_terms.entries[it->first] = it->second;
+      it = formula.entries.erase(it);
+    } else {
+      it++;
+    }
+  }
+  return initial_terms;
+}
diff --git a/src/form/formula_util.hpp b/src/form/formula_util.hpp
@@ -9,4 +9,8 @@ class FormulaUtil {
   static void resolveSimpleFunctions(Formula& formula);
 
   static void resolveSimpleRecursions(Formula& formula);
+
+  static void convertInitialTermsToIf(Formula& formula);
+
+  static Formula extractInitialTerms(Formula& formula);
 };
diff --git a/src/form/pari.cpp b/src/form/pari.cpp
@@ -4,28 +4,9 @@
 #include <sstream>
 
 #include "form/expression_util.hpp"
+#include "form/formula_util.hpp"
 #include "sys/log.hpp"
 
-void convertInitialTermsToIf(Formula& formula) {
-  auto it = formula.entries.begin();
-  while (it != formula.entries.end()) {
-    auto left = it->first;
-    Expression general(Expression::Type::FUNCTION, left.name,
-                       {Expression(Expression::Type::PARAMETER, "n")});
-    auto general_it = formula.entries.find(general);
-    if (left.type == Expression::Type::FUNCTION && left.children.size() == 1 &&
-        left.children.front().type == Expression::Type::CONSTANT &&
-        general_it != formula.entries.end()) {
-      general_it->second =
-          Expression(Expression::Type::IF, "",
-                     {left.children.front(), it->second, general_it->second});
-      it = formula.entries.erase(it);
-    } else {
-      it++;
-    }
-  }
-}
-
 bool convertExprToPari(Expression& expr, const Formula& f, bool as_vector) {
   // convert bottom-up!
   for (auto& c : expr.children) {
@@ -99,7 +80,7 @@ bool Pari::convertToPari(Formula& f, bool as_vector) {
   }
   f = tmp;
   addLocalVars(f);
-  convertInitialTermsToIf(f);
+  FormulaUtil::convertInitialTermsToIf(f);
   return true;
 }