Handle RootSquare functions with RootPattern

core/src/main/java/it/cavallium/warppi/math/rules/dsl/patterns/RootPattern.java

@@ -2,16 +2,21 @@ package it.cavallium.warppi.math.rules.dsl.patterns;
 
 import it.cavallium.warppi.math.Function;
 import it.cavallium.warppi.math.MathContext;
+import it.cavallium.warppi.math.functions.Number;
 import it.cavallium.warppi.math.functions.Root;
+import it.cavallium.warppi.math.functions.RootSquare;
 import it.cavallium.warppi.math.rules.dsl.Pattern;
 import it.cavallium.warppi.math.rules.dsl.PatternUtils;
 import it.cavallium.warppi.math.rules.dsl.VisitorPattern;
 
+import java.math.BigDecimal;
 import java.util.Map;
 import java.util.Optional;
 
 /**
  * Matches and generates a root of degree and radicand patterns.
+ * <p>
+ * Also matches and generates functions of type <code>RootSquare</code>.
  */
 public class RootPattern extends VisitorPattern {
     private final Pattern degree;
@@ -27,12 +32,21 @@ public class RootPattern extends VisitorPattern {
         return PatternUtils.matchFunctionOperatorParameters(root, degree, radicand);
     }
 
+    @Override
+    public Optional<Map<String, Function>> visit(RootSquare rootSquare) {
+        return PatternUtils.matchFunctionOperatorParameters(rootSquare, degree, radicand);
+    }
+
     @Override
     public Function replace(final MathContext mathContext, final Map<String, Function> subFunctions) {
-        return new Root(
+        final Function newDegree = degree.replace(mathContext, subFunctions);
-                mathContext,
+        final Function newRadicand = radicand.replace(mathContext, subFunctions);
-                degree.replace(mathContext, subFunctions),
+
-                radicand.replace(mathContext, subFunctions)
+        if (newDegree instanceof Number
-        );
+                && ((Number) newDegree).getTerm().compareTo(new BigDecimal(2)) == 0) {
+            return new RootSquare(mathContext, newRadicand);
+        } else {
+            return new Root(mathContext, newDegree, newRadicand);
+        }
     }
 }

core/src/test/java/it/cavallium/warppi/math/rules/dsl/PatternTest.java

@@ -160,6 +160,34 @@ public class PatternTest {
         assertEquals(shouldMatch, pattern.replace(mathContext, subFunctions.get()));
     }
 
+    @Test
+    public void rootPatternForRootSquare() {
+        final Pattern pattern = new RootPattern(
+                new SubFunctionPattern("x"),
+                new SubFunctionPattern("y")
+        );
+
+        final Function root = new Root(
+                mathContext,
+                new Number(mathContext, 2),
+                new Number(mathContext, 1)
+        );
+        final Optional<Map<String, Function>> rootSubFunctions = pattern.match(root);
+        assertTrue(rootSubFunctions.isPresent());
+
+        final Function rootSquare = new RootSquare(
+                mathContext,
+                new Number(mathContext, 1)
+        );
+        final Optional<Map<String, Function>> rootSquareSubFunctions = pattern.match(rootSquare);
+        assertTrue(rootSquareSubFunctions.isPresent());
+        assertEquals(rootSubFunctions.get(), rootSquareSubFunctions.get());
+
+        final Function replacement = pattern.replace(mathContext, rootSubFunctions.get());
+        assertTrue(replacement instanceof RootSquare);
+        assertEquals(rootSquare, replacement);
+    }
+
     @Test
     public void otherBinaryPatterns() {
         final Number one = new Number(mathContext, 1);