split up functions in reshape to make more readable

Author: konstibob

src/main/java/dev/zarr/zarrjava/v3/codec/core/ReshapeCodec.java

@@ -131,15 +131,47 @@ public ArrayMetadata.CoreArrayMetadata resolveArrayMetadata() throws ZarrExcepti
 
     /**
      * Resolves the concrete output shape {@code B_shape} from the {@code shape} configuration and the
-     * given input shape {@code A_shape}, validating all invariants mandated by the specification.
+     * given input  shape {@code A_shape}, validating all invariants mandated by the specification.
      */
     int[] resolveOutputShape(int[] inputShape) throws ZarrException {
-        int ndim = inputShape.length;
-        long inputTotal = 1L;
-        for (int size : inputShape) {
-            inputTotal *= size;
+        long inputTotal = product(inputShape);
+
+        // read the 'shape' config into a concrete (still -1-placeholder) output shape.
+        ParsedShape parsed = parseConfiguredShape(inputShape);
+        
+        // reject configs that reorder input dimensions (use transport for them).
+        checkNoReordering(parsed.flatInputDims);
+        
+        // fill in the single -1 entry, if present, so the element count matches.
+        resolveInferredDimension(parsed.outputShape, parsed.minusOnePos, inputTotal);
+        
+        // The fundamental reshape law, prod(output) == prod(input).
+        checkElementCountPreserved(parsed.outputShape, inputTotal);
+        
+        // Step 5: every merge entry must combine adjacent input dims sitting in the same flat position.
+        checkMergesAligned(parsed.outputShape, parsed.inputDimsPerOutput, inputShape);
+        
+        // Step 6: narrow long -> int for the array API.
+        return toIntShape(parsed.outputShape);
+    }
+
+    /**
+     * Product of all entries of {@code shape}, computed in {@code long} to avoid overflow.
+     */
+    private static long product(int[] shape) {
+        long total = 1L;
+        for (int size : shape) {
+            total *= size;
         }
+        return total;
+    }
 
+    /**
+     * Step 1. Reads {@code configuration.shape} into the intermediate {@link ParsedShape}: each entry
+     * becomes a literal size, a {@code -1} placeholder, or the product of the merged input dimensions.
+     */
+    private ParsedShape parseConfiguredShape(int[] inputShape) throws ZarrException {
+        int ndim = inputShape.length;
         int n = configuration.shape.length;
         if (n == 0) {
             throw new ZarrException("reshape codec: 'shape' must not be empty.");
@@ -183,34 +215,49 @@ int[] resolveOutputShape(int[] inputShape) throws ZarrException {
                 outputShape[i] = product;
             }
         }
+        return new ParsedShape(outputShape, inputDimsPerOutput, minusOnePos, flatInputDims);
+    }
 
-        // The flattened list of input dimensions must be strictly monotonically increasing. This rules
-        // out configurations that would (incorrectly) suggest a transpose.
+    /**
+     * Step 2. The flattened list of referenced input dimensions must be strictly increasing. This rules
+     * out configurations that would (incorrectly) suggest a transpose.
+     */
+    private static void checkNoReordering(List<Integer> flatInputDims) throws ZarrException {
         for (int j = 1; j < flatInputDims.size(); j++) {
             if (flatInputDims.get(j) <= flatInputDims.get(j - 1)) {
                 throw new ZarrException(
                         "reshape codec: the flattened list of input dimensions must be strictly increasing, but got "
                                 + flatInputDims + ".");
             }
         }
+    }
 
-        // Resolve the automatic dimension (-1).
-        if (minusOnePos != -1) {
-            long known = 1L;
-            for (int i = 0; i < n; i++) {
-                if (i != minusOnePos) {
-                    known *= outputShape[i];
-                }
-            }
-            if (known == 0 || inputTotal % known != 0) {
-                throw new ZarrException(
-                        "reshape codec: cannot infer the -1 dimension because prod(output shape) would not equal "
-                                + "prod(input shape) (" + inputTotal + ").");
+    /**
+     * Step 3. Resolves the automatic dimension ({@code -1}) so that {@code prod(output) == prod(input)}.
+     */
+    private static void resolveInferredDimension(long[] outputShape, int minusOnePos, long inputTotal)
+            throws ZarrException {
+        if (minusOnePos == -1) {
+            return;
+        }
+        long known = 1L;
+        for (int i = 0; i < outputShape.length; i++) {
+            if (i != minusOnePos) {
+                known *= outputShape[i];
             }
-            outputShape[minusOnePos] = inputTotal / known;
         }
+        if (known == 0 || inputTotal % known != 0) {
+            throw new ZarrException(
+                    "reshape codec: cannot infer the -1 dimension because prod(output shape) would not equal "
+                            + "prod(input shape) (" + inputTotal + ").");
+        }
+        outputShape[minusOnePos] = inputTotal / known;
+    }
 
-        // Invariant: prod(B_shape) == prod(A_shape).
+    /**
+     * Step 4. Invariant: {@code prod(B_shape) == prod(A_shape)}. A reshape never changes the element count.
+     */
+    private static void checkElementCountPreserved(long[] outputShape, long inputTotal) throws ZarrException {
         long outputTotal = 1L;
         for (long size : outputShape) {
             outputTotal *= size;
@@ -220,9 +267,18 @@ int[] resolveOutputShape(int[] inputShape) throws ZarrException {
                     "reshape codec: prod(output shape)=" + outputTotal + " does not equal prod(input shape)="
                             + inputTotal + ".");
         }
+    }
 
-        // When an output dimension is specified by input_dims, verify that the coordinates along those
-        // input dimensions actually correspond to the raveled index along that output dimension.
+    /**
+     * Step 5. For each output dimension specified by input dimensions, verify that the merged input
+     * dimensions actually correspond to the raveled index along that output dimension: the flat range
+     * before and after the merged block must match the flat range before and after the output dimension.
+     * This is a pure check &mdash; it never modifies {@code outputShape}.
+     */
+    private static void checkMergesAligned(long[] outputShape, int[][] inputDimsPerOutput, int[] inputShape)
+            throws ZarrException {
+        int n = outputShape.length;
+        int ndim = inputShape.length;
         for (int i = 0; i < n; i++) {
             int[] inputDims = inputDimsPerOutput[i];
             if (inputDims == null || inputDims.length == 0) {
@@ -252,9 +308,15 @@ int[] resolveOutputShape(int[] inputShape) throws ZarrException {
                                 + " vs " + inputSuffix + ").");
             }
         }
+    }
 
-        int[] result = new int[n];
-        for (int i = 0; i < n; i++) {
+    /**
+     * Step 6. Narrows the resolved {@code long} shape to {@code int[]} (the array API type), rejecting
+     * any dimension that would not fit.
+     */
+    private static int[] toIntShape(long[] outputShape) throws ZarrException {
+        int[] result = new int[outputShape.length];
+        for (int i = 0; i < outputShape.length; i++) {
             if (outputShape[i] > Integer.MAX_VALUE) {
                 throw new ZarrException("reshape codec: output dimension " + i + " exceeds Integer.MAX_VALUE.");
             }
@@ -263,6 +325,29 @@ int[] resolveOutputShape(int[] inputShape) throws ZarrException {
         return result;
     }
 
+    /**
+     * Intermediate result of {@link #parseConfiguredShape}: the (possibly still {@code -1}-containing)
+     * output shape plus the bookkeeping the later validation steps need.
+     */
+    private static final class ParsedShape {
+        /** Output shape; the {@code -1} entry (if any) is still a placeholder until step 3. */
+        final long[] outputShape;
+        /** Per output dimension, the merged input dims (or {@code null} for a literal/{@code -1} entry). */
+        final int[][] inputDimsPerOutput;
+        /** Index of the single {@code -1} entry, or {@code -1} if there is none. */
+        final int minusOnePos;
+        /** All referenced input dimensions, flattened in reading order (used by the no-reorder check). */
+        final List<Integer> flatInputDims;
+
+        ParsedShape(long[] outputShape, int[][] inputDimsPerOutput, int minusOnePos,
+                    List<Integer> flatInputDims) {
+            this.outputShape = outputShape;
+            this.inputDimsPerOutput = inputDimsPerOutput;
+            this.minusOnePos = minusOnePos;
+            this.flatInputDims = flatInputDims;
+        }
+    }
+
     private static int[] asInputDims(Object element) throws ZarrException {
         if (element instanceof int[]) {
             return (int[]) element;
@@ -291,7 +376,8 @@ private static int asInteger(Object element) throws ZarrException {
             return ((Number) element).intValue();
         }
         throw new ZarrException(
-                "reshape codec: 'shape' entries must be integers or arrays of integers, but got " + element + ".");
+            "reshape codec: 'shape' entries must be integers or arrays of integers, but got " + element + "."
+        );
     }
 
     public static final class Configuration {