Handle Piecewise in TransformVisitor (for e.g. cse)

CMakeLists.txt

@@ -1,5 +1,9 @@
 cmake_minimum_required(VERSION 2.8.12)
 
+if (POLICY CMP0074)
+  cmake_policy(SET CMP0074 NEW)
+endif()
+
 if (POLICY CMP0057)
     cmake_policy(SET CMP0057 NEW) # needed for llvm >= 16
 endif ()

symengine/cse.cpp

@@ -501,7 +501,8 @@ void tree_cse(vec_pair &replacements, vec_basic &reduced_exprs,
 
     std::function<void(RCP<const Basic> & expr)> find_repeated;
     find_repeated = [&](RCP<const Basic> expr) -> void {
-        if (is_a_Number(*expr)) {
+        // Do not replace atoms
+        if (is_a_Number(*expr) or is_a<BooleanAtom>(*expr)) {
             return;
         }
 

symengine/tests/basic/test_cse.cpp

@@ -1,19 +1,25 @@
 #include "catch.hpp"
+#include "symengine/dict.h"
+#include "symengine/logic.h"
 
 #include <symengine/basic.h>
 #include <symengine/pow.h>
 #include <symengine/add.h>
 #include <symengine/mul.h>
 #include <symengine/functions.h>
+#include <symengine/logic.h>
 
 using SymEngine::add;
 using SymEngine::Basic;
+using SymEngine::boolTrue;
 using SymEngine::cse;
 using SymEngine::div;
+using SymEngine::Gt;
 using SymEngine::integer;
 using SymEngine::mul;
 using SymEngine::neg;
 using SymEngine::one;
+using SymEngine::piecewise;
 using SymEngine::pow;
 using SymEngine::RCP;
 using SymEngine::sin;
@@ -302,6 +308,30 @@ TEST_CASE("CSE: simple", "[cse]")
         REQUIRE(unified_eq(substs, {{x0, add(x, y)}, {x1, add(x0, z)}}));
         REQUIRE(unified_eq(reduced, {x0, add(i2, x0), x1, add(i3, x1)}));
     }
+    {
+        auto pw1 = piecewise(
+            {{pow(add(x, y), i2), Gt(x, y)}, {sqrt(add(x, y)), boolTrue}});
+
+        vec_pair substs;
+        vec_basic reduced;
+        cse(substs, reduced, {pw1});
+        REQUIRE(unified_eq(substs, {{x0, add(x, y)}}));
+        REQUIRE(unified_eq(reduced, {piecewise({{pow(x0, i2), Gt(x, y)},
+                                                {sqrt(x0), boolTrue}})}));
+    }
+    {
+        auto pw2 = piecewise({{pow(x, i2), Gt(add(x, y), i2)},
+                              {sqrt(y), Gt(add(x, y), i3)},
+                              {sqrt(x), boolTrue}});
+
+        vec_pair substs;
+        vec_basic reduced;
+        cse(substs, reduced, {pw2});
+        REQUIRE(unified_eq(substs, {{x0, add(x, y)}}));
+        REQUIRE(unified_eq(reduced, {piecewise({{pow(x, i2), Gt(x0, i2)},
+                                                {sqrt(y), Gt(x0, i3)},
+                                                {sqrt(x), boolTrue}})}));
+    }
 }
 
 TEST_CASE("CSE: regression test gh-1463", "[cse]")

symengine/visitor.cpp

@@ -201,6 +201,24 @@ void TransformVisitor::bvisit(const MultiArgFunction &x)
     result_ = nbarg;
 }
 
+void TransformVisitor::bvisit(const Piecewise &x)
+{
+    auto branch_cond_pairs = x.get_vec();
+    PiecewiseVec new_pairs;
+    for (const auto &branch_cond : branch_cond_pairs) {
+        auto branch = branch_cond.first;
+        auto cond = branch_cond.second;
+        auto new_branch = apply(branch);
+        auto new_cond = apply(cond);
+        if (!is_a_Boolean(*new_cond)) {
+            new_cond = Eq(new_cond, boolTrue);
+        }
+        new_pairs.push_back(
+            {new_branch, rcp_static_cast<const Boolean>(new_cond)});
+    }
+    result_ = piecewise(new_pairs);
+}
+
 void preorder_traversal_local_stop(const Basic &b, LocalStopVisitor &v)
 {
     b.accept(v);

symengine/visitor.h

@@ -262,6 +262,7 @@ class TransformVisitor : public BaseVisitor<TransformVisitor>
     }
 
     void bvisit(const MultiArgFunction &x);
+    void bvisit(const Piecewise &x);
 };
 
 template <typename Derived, typename First, typename... Rest>