Fix character loss in RailFenceCipher for inputs containing newlines

Author: SEPURI-SAI-KRISHNA

src/main/java/com/thealgorithms/ciphers/RailFenceCipher.java

@@ -1,7 +1,5 @@
 package com.thealgorithms.ciphers;
 
-import java.util.Arrays;
-
 /**
  * The rail fence cipher (also called a zigzag cipher) is a classical type of transposition cipher.
  * It derives its name from the manner in which encryption is performed, in analogy to a fence built with horizontal rails.
@@ -14,28 +12,27 @@ public class RailFenceCipher {
     // Encrypts the input string using the rail fence cipher method with the given number of rails.
     public String encrypt(String str, int rails) {
 
+        checkInput(str, rails);
+
         // Base case of single rail or rails are more than the number of characters in the string
         if (rails == 1 || rails >= str.length()) {
             return str;
         }
 
-        // Boolean flag to determine if the movement is downward or upward in the rail matrix.
+        // Boolean flag to determine if the movement is downward or upward in the rail pattern.
         boolean down = true;
-        // Create a 2D array to represent the rails (rows) and the length of the string (columns).
-        char[][] strRail = new char[rails][str.length()];
-
-        // Initialize all positions in the rail matrix with a placeholder character ('\n').
+        // Collect the characters of every rail separately. Using one buffer per rail (instead of a
+        // rails x length matrix with a placeholder character) keeps every character of the input,
+        // including characters that would otherwise be indistinguishable from the placeholder.
+        StringBuilder[] railBuffers = new StringBuilder[rails];
         for (int i = 0; i < rails; i++) {
-            Arrays.fill(strRail[i], '\n');
+            railBuffers[i] = new StringBuilder();
         }
 
-        int row = 0; // Start at the first row
-        int col = 0; // Start at the first column
+        int row = 0; // Start at the first rail
 
-        int i = 0;
-
-        // Fill the rail matrix with characters from the string based on the rail pattern.
-        while (col < str.length()) {
+        // Distribute the characters of the string over the rails following the zigzag pattern.
+        for (int i = 0; i < str.length(); i++) {
             // Change direction to down when at the first row.
             if (row == 0) {
                 down = true;
@@ -45,33 +42,28 @@ else if (row == rails - 1) {
                 down = false;
             }
 
-            // Place the character in the current position of the rail matrix.
-            strRail[row][col] = str.charAt(i);
-            col++; // Move to the next column.
+            // Append the character to the rail it belongs to.
+            railBuffers[row].append(str.charAt(i));
             // Move to the next row based on the direction.
             if (down) {
                 row++;
             } else {
                 row--;
             }
-
-            i++;
         }
 
-        // Construct the encrypted string by reading characters row by row.
-        StringBuilder encryptedString = new StringBuilder();
-        for (char[] chRow : strRail) {
-            for (char ch : chRow) {
-                if (ch != '\n') {
-                    encryptedString.append(ch);
-                }
-            }
+        // Construct the encrypted string by reading the rails top to bottom.
+        StringBuilder encryptedString = new StringBuilder(str.length());
+        for (StringBuilder railBuffer : railBuffers) {
+            encryptedString.append(railBuffer);
         }
         return encryptedString.toString();
     }
     // Decrypts the input string using the rail fence cipher method with the given number of rails.
     public String decrypt(String str, int rails) {
 
+        checkInput(str, rails);
+
         // Base case of single rail or rails are more than the number of characters in the string
         if (rails == 1 || rails >= str.length()) {
             return str;
@@ -144,4 +136,14 @@ else if (row == rails - 1) {
 
         return decryptedString.toString();
     }
+
+    // Rejects inputs the zigzag pattern is not defined for.
+    private static void checkInput(String str, int rails) {
+        if (str == null) {
+            throw new IllegalArgumentException("Input string must not be null");
+        }
+        if (rails <= 0) {
+            throw new IllegalArgumentException("Number of rails must be positive, but was " + rails);
+        }
+    }
 }

src/test/java/com/thealgorithms/ciphers/RailFenceCipherTest.java

@@ -0,0 +1,70 @@
+package com.thealgorithms.ciphers;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+class RailFenceCipherTest {
+
+    private final RailFenceCipher railFenceCipher = new RailFenceCipher();
+
+    @Test
+    void testEncrypt() {
+        assertEquals("WECRLTEERDSOEEFEAOCAIVDEN", railFenceCipher.encrypt("WEAREDISCOVEREDFLEEATONCE", 3));
+    }
+
+    @Test
+    void testDecrypt() {
+        assertEquals("WEAREDISCOVEREDFLEEATONCE", railFenceCipher.decrypt("WECRLTEERDSOEEFEAOCAIVDEN", 3));
+    }
+
+    @ParameterizedTest
+    @CsvSource({"HELLOWORLD, 2", "HELLOWORLD, 3", "HELLOWORLD, 4", "ATTACKATDAWN, 5", "abcdefghij, 6"})
+    void testRoundTrip(String message, int rails) {
+        assertEquals(message, railFenceCipher.decrypt(railFenceCipher.encrypt(message, rails), rails));
+    }
+
+    /**
+     * Every character of the input must survive encryption, including the ones that used to collide
+     * with the placeholder that marked unused cells of the rail matrix.
+     */
+    @ParameterizedTest
+    @ValueSource(strings = {"ab\ncdef", "line1\nline2\nline3", "\n\n\n\n\n", "a\nb", "tabs\tand\nnewlines\r\n"})
+    void testControlCharactersArePreserved(String message) {
+        for (int rails = 2; rails <= 5; rails++) {
+            String encrypted = railFenceCipher.encrypt(message, rails);
+            assertEquals(message.length(), encrypted.length(), "characters were dropped with " + rails + " rails");
+            assertEquals(message, railFenceCipher.decrypt(encrypted, rails), "round trip failed with " + rails + " rails");
+        }
+    }
+
+    @Test
+    void testEncryptWithNewlineMatchesReferencePattern() {
+        // Rails of "ab\ncdef" with 3 rails: {a, d} / {b, c, e} / {\n, f}
+        assertEquals("adbce\nf", railFenceCipher.encrypt("ab\ncdef", 3));
+    }
+
+    @ParameterizedTest
+    @CsvSource({"HELLO, 1", "HELLO, 5", "HELLO, 9", "'', 1", "'', 4"})
+    void testDegenerateRailCountsReturnInput(String message, int rails) {
+        assertEquals(message, railFenceCipher.encrypt(message, rails));
+        assertEquals(message, railFenceCipher.decrypt(message, rails));
+    }
+
+    @ParameterizedTest
+    @ValueSource(ints = {0, -1, -7})
+    void testNonPositiveRailCountThrows(int rails) {
+        assertThrows(IllegalArgumentException.class, () -> railFenceCipher.encrypt("HELLO", rails));
+        assertThrows(IllegalArgumentException.class, () -> railFenceCipher.decrypt("HELLO", rails));
+    }
+
+    @Test
+    void testNullInputThrows() {
+        assertThrows(IllegalArgumentException.class, () -> railFenceCipher.encrypt(null, 3));
+        assertThrows(IllegalArgumentException.class, () -> railFenceCipher.decrypt(null, 3));
+    }
+}