Harden rail fence cipher decoding against resource exhaustion

src/main/java/kyu3/RailFenceCipher.java

@@ -1,7 +1,6 @@
 package kyu3;
 
 
-import java.util.LinkedList;
 import java.util.Map;
 import java.util.TreeMap;
 
@@ -40,56 +39,36 @@ static String decode(String s, int n) {
         // First computes the exact size of each rail, splits ciphertext into
         // contiguous rail segments, then reconstructs plaintext by replaying the
         // same rail traversal cycle.
+        if (n < 2) {
+            throw new IllegalArgumentException("Number of rails must be at least 2");
+        }
+        if (s.isEmpty() || n >= s.length()) {
+            return s;
+        }
+
         int[] size = new int[n];
         Counter counter = new Counter(n - 1);
         for (int i = 0; i < s.length(); i++) {
             int tik = counter.tik();
             size[tik]++;
         }
-        Map<Integer, StringBuilder> map = new TreeMap<>();
-        StringBuilder result = new StringBuilder();
+        int[] position = new int[n];
+        int start = 0;
         for (int i = 0; i < n; i++) {
-            StringBuilder stringBuilder = new StringBuilder();
-            if (i == 0) {
-                stringBuilder.append(s, 0, size[0]);
-            } else {
-                int start = sumArr(size, i);
-                int end = start + size[i];
-                stringBuilder.append(s, start, end);
-            }
-            map.put(i, stringBuilder);
-        }
-
-        Map<Integer, LinkedList<String>> map2 = new TreeMap<>();
-
-        for (Map.Entry<Integer, StringBuilder> e : map.entrySet()
-        ) {
-            LinkedList<String> strings = new LinkedList<>();
-            for (char c : e.getValue().toString().toCharArray()) {
-                strings.add(String.valueOf(c));
-            }
-            map2.put(e.getKey(), strings);
+            position[i] = start;
+            start += size[i];
         }
 
+        StringBuilder result = new StringBuilder(s.length());
         Counter counterRes = new Counter(n - 1);
         for (int i = 0; i < s.length(); i++) {
             int tik = counterRes.tik();
-            result.append(map2.get(tik).poll());
+            result.append(s.charAt(position[tik]++));
         }
 
         return result.toString();
     }
 
-    private static int sumArr(int[] arr, int i) {
-        // Sum widths of all previous rails to determine the start offset
-        // of the current rail segment.
-        int result = 0;
-        for (int j = 0; j < i; j++) {
-            result += arr[j];
-        }
-        return result;
-    }
-
     /**
      * Create two functions to encode and then decode a string using the Rail Fence Cipher. This cipher is used to encode
      * a string by placing each character successively in a diagonal along a set of "rails". First start off moving diagonally

src/test/java/kyu3/RailFenceCipherTest.java

@@ -20,6 +20,12 @@ void shouldEncodeAndDecodeKnownExample() {
         assertEquals("WEAREDISCOVEREDFLEEATONCE", RailFenceCipher.decode(encoded, 3));
     }
 
+    @Test
+    void shouldDecodeWithoutAllocatingForUnusedRails() {
+        assertEquals("", RailFenceCipher.decode("", Integer.MAX_VALUE));
+        assertEquals("abc", RailFenceCipher.decode("abc", Integer.MAX_VALUE));
+    }
+
     @ParameterizedTest
     @MethodSource("roundTripCases")
     void shouldRoundTripPlainTextForDifferentRails(String text, int rails) {