Merge pull request #36 from DanexCodr/copilot/extend-natural-arrays-nd-support

Add ND-style natural array indexing, shape-aware array broadcasting, and array literal methods

Author: DanexCodr

src/main/cod/src/main/test/Broadcast.cod

@@ -1,5 +1,40 @@
 unit test
 
 Broadcaster {
+    share main() {
+        out("=== Array Broadcasting (not unit broadcast) ===")
+        
+        a := [1 to 5]
+        b := [1 to 5]
+        c := a * b
+        out("c[0] = " + c[0] + " (expect 1)")
+        out("c[4] = " + c[4] + " (expect 25)")
+        
+        matrix := [[1, 2, 3], [4, 5, 6]]
+        vector := [10, 20, 30]
+        sum := matrix + vector
+        out("sum[0,0] = " + sum[0,0] + " (expect 11)")
+        out("sum[1,2] = " + sum[1,2] + " (expect 36)")
+        
+        matrix[1,2] = 99
+        out("matrix[1,2] = " + matrix[1,2] + " (expect 99)")
+        
+        naturalGrid := [[1 to 3], [4 to 6]]
+        naturalGrid[0,1] = 777
+        out("naturalGrid[0,1] = " + naturalGrid[0,1] + " (expect 777)")
+        
+        nd := [1 to 2, 1 to 3]
+        out("nd[1,2] = " + nd[1,2] + " (expect 3)")
+        
+        mapped := [1, 2, 3].map("+", 10)
+        filtered := [1, 2, 3, 4].filter(">=", 3)
+        reduced := [1, 2, 3, 4].reduce("+", 0)
+        
+        out("mapped[0] = " + mapped[0] + " (expect 11)")
+        out("mapped[2] = " + mapped[2] + " (expect 13)")
+        out("filtered[0] = " + filtered[0] + " (expect 3)")
+        out("filtered[1] = " + filtered[1] + " (expect 4)")
+        out("reduced = " + reduced + " (expect 10)")
+    }
 
 }

src/main/java/cod/interpreter/InterpreterVisitor.java

@@ -1229,6 +1229,21 @@ public Object visit(PropertyAccessNode node) {
                 }
             }
 
+            if (node.right instanceof MethodCallNode) {
+                MethodCallNode literalMethod = (MethodCallNode) node.right;
+                String methodName = literalMethod.name;
+                if (literalRegistry.hasMethod(leftObj, methodName)) {
+                    List<Object> evaluatedArgs = new ArrayList<Object>();
+                    if (literalMethod.arguments != null) {
+                        for (ExprNode arg : literalMethod.arguments) {
+                            Object argValue = dispatch(arg);
+                            evaluatedArgs.add(typeSystem.unwrap(argValue));
+                        }
+                    }
+                    return literalRegistry.handleMethod(leftObj, methodName, evaluatedArgs, ctx);
+                }
+            }
+            
             if (leftObj instanceof NaturalArray) {
                 NaturalArray natural = (NaturalArray) leftObj;
                 if (natural.hasPendingUpdates()) {
@@ -1395,23 +1410,42 @@ public Object visit(TypeCastNode node) {
 
     @SuppressWarnings("unchecked")
     @Override
-public Object visit(MethodCallNode node) {
+    public Object visit(MethodCallNode node) {
     if (node == null) {
         throw new InternalError("visit(MethodCallNode) called with null node");
     }
 
-    try {
-        // Handle super calls first
-        if (node.isSuperCall) {
-            return handleSuperMethodCall(node);
-        }
-        
-        // Evaluate all arguments first
-        List<Object> evaluatedArgs = new ArrayList<Object>();
-        for (ExprNode arg : node.arguments) {
-            Object argValue = dispatch(arg);
-            evaluatedArgs.add(typeSystem.unwrap(argValue));
-        }
+        try {
+            // Handle super calls first
+            if (node.isSuperCall) {
+                return handleSuperMethodCall(node);
+            }
+            
+            ExecutionContext ctx = getCurrentContext();
+            if (ctx != null && node.qualifiedName != null && node.qualifiedName.contains(".")) {
+                String[] parts = node.qualifiedName.split("\\.");
+                if (parts.length == 2) {
+                    String receiverName = parts[0];
+                    String methodName = parts[1];
+                    Object receiverValue = ctx.getVariable(receiverName);
+                    receiverValue = typeSystem.unwrap(receiverValue);
+                    if (literalRegistry.hasMethod(receiverValue, methodName)) {
+                        List<Object> evaluatedArgs = new ArrayList<Object>();
+                        for (ExprNode arg : node.arguments) {
+                            Object argValue = dispatch(arg);
+                            evaluatedArgs.add(typeSystem.unwrap(argValue));
+                        }
+                        return literalRegistry.handleMethod(receiverValue, methodName, evaluatedArgs, ctx);
+                    }
+                }
+            }
+            
+            // Evaluate all arguments first
+            List<Object> evaluatedArgs = new ArrayList<Object>();
+            for (ExprNode arg : node.arguments) {
+                Object argValue = dispatch(arg);
+                evaluatedArgs.add(typeSystem.unwrap(argValue));
+            }
 
         // ========== CHECK GLOBAL FUNCTIONS FIRST ==========
         // This must come BEFORE any other resolution to ensure out(), in(), etc.
@@ -1425,9 +1459,7 @@ public Object visit(MethodCallNode node) {
         // ========== END GLOBAL CHECK ==========
 
         // Try to find method in current class hierarchy
-        ExecutionContext ctx = getCurrentContext();
         MethodNode method = null;
-
         if (ctx.currentClass != null) {
             method = interpreter
                 .getConstructorResolver()
@@ -1693,6 +1725,10 @@ public Object visit(ArrayNode node) {
                     return new NaturalArray(range, this, getCurrentContext());
                 }
             }
+            
+            if (node.elements.size() > 1 && allElementsAreRanges(node.elements)) {
+                return buildDimensionArray(node.elements, 0);
+            }
 
             // Regular array literal handling
             List<Object> result = new ArrayList<Object>();
@@ -1728,6 +1764,35 @@ public Object visit(ArrayNode node) {
             throw new InternalError("Array creation failed", e);
         }
     }
+    
+    private boolean allElementsAreRanges(List<ExprNode> elements) {
+        if (elements == null || elements.isEmpty()) return false;
+        for (ExprNode element : elements) {
+            if (!(element instanceof RangeNode)) {
+                return false;
+            }
+        }
+        return true;
+    }
+    
+    private Object buildDimensionArray(List<ExprNode> ranges, int dimension) {
+        RangeNode currentRange = (RangeNode) ranges.get(dimension);
+        NaturalArray currentNatural = new NaturalArray(currentRange, this, getCurrentContext());
+        if (dimension == ranges.size() - 1) {
+            return currentNatural;
+        }
+        
+        long length = currentNatural.size();
+        if (length > Integer.MAX_VALUE) {
+            throw new ProgramError("Dimension size too large for nested ND array literal: " + length + " (max " + Integer.MAX_VALUE + ")");
+        }
+        
+        List<Object> result = new ArrayList<Object>((int) length);
+        for (int i = 0; i < (int) length; i++) {
+            result.add(buildDimensionArray(ranges, dimension + 1));
+        }
+        return result;
+    }
 
     @SuppressWarnings("unchecked")
     @Override
@@ -1750,6 +1815,10 @@ public Object visit(IndexAccessNode node) {
 
             Object indexObj = dispatch(node.index);
             indexObj = typeSystem.unwrap(indexObj);
+            
+            if (indexObj instanceof List) {
+                return applyTupleIndices(arrayObj, (List<?>) indexObj);
+            }
 
             if (indexObj instanceof RangeSpec) {
                 return applyRangeIndex(arrayObj, (RangeSpec) indexObj);
@@ -1984,6 +2053,40 @@ private Object applyMultiRangeIndex(Object array, MultiRangeSpec multiRange) {
             (array != null ? array.getClass().getSimpleName() : "null"));
     }
 
+    @SuppressWarnings("unchecked")
+    private Object applyTupleIndices(Object array, List<?> indices) {
+        Object current = array;
+        for (Object rawIndex : indices) {
+            Object indexObj = typeSystem.unwrap(rawIndex);
+            if (indexObj instanceof RangeSpec) {
+                current = applyRangeIndex(current, (RangeSpec) indexObj);
+                continue;
+            }
+            if (indexObj instanceof MultiRangeSpec) {
+                current = applyMultiRangeIndex(current, (MultiRangeSpec) indexObj);
+                continue;
+            }
+            if (current instanceof NaturalArray) {
+                NaturalArray natural = (NaturalArray) current;
+                long idx = expressionHandler.toLongIndex(indexObj);
+                current = natural.needsConversion() ? natural.get(idx, true) : natural.get(idx);
+                continue;
+            }
+            if (current instanceof List) {
+                List<Object> list = (List<Object>) current;
+                int idx = expressionHandler.toIntIndex(indexObj);
+                if (idx < 0 || idx >= list.size()) {
+                    throw new ProgramError("Index out of bounds: " + idx + " for array of size " + list.size());
+                }
+                current = list.get(idx);
+                continue;
+            }
+            throw new ProgramError("Invalid array access during multidimensional indexing: expected NaturalArray or List, got "
+                + (current != null ? current.getClass().getSimpleName() : "null"));
+        }
+        return current;
+    }
+
     private List<Object> getListRange(List<Object> list, RangeSpec range) {
         try {
             long start, end;

src/main/java/cod/interpreter/handler/AssignmentHandler.java

@@ -174,6 +174,10 @@ private Object handleIndexAssignment(IndexAccessNode indexAccess, Object newValu
             Object indexObj = dispatcher.dispatch(indexAccess.index);
             indexObj = typeSystem.unwrap(indexObj);
 
+            if (indexObj instanceof List) {
+                return assignTupleIndex(arrayObj, (List<?>) indexObj, newValue);
+            }
+            
             if (indexObj instanceof RangeSpec) {
                 return assignRange(arrayObj, (RangeSpec) indexObj, newValue);
             }
@@ -204,6 +208,130 @@ private Object handleIndexAssignment(IndexAccessNode indexAccess, Object newValu
             throw new InternalError("Index assignment failed", e);
         }
     }
+
+    @SuppressWarnings("unchecked")
+    private Object assignTupleIndex(Object arrayObj, List<?> tupleIndices, Object newValue) {
+        if (tupleIndices == null || tupleIndices.isEmpty()) {
+            throw new ProgramError("Invalid multidimensional assignment: empty index tuple");
+        }
+        
+        Object current = arrayObj;
+        for (int i = 0; i < tupleIndices.size() - 1; i++) {
+            Object idxObj = typeSystem.unwrap(tupleIndices.get(i));
+            if (idxObj instanceof RangeSpec) {
+                current = extractRange(current, (RangeSpec) idxObj);
+                continue;
+            }
+            if (idxObj instanceof MultiRangeSpec) {
+                current = extractMultiRange(current, (MultiRangeSpec) idxObj);
+                continue;
+            }
+            if (current instanceof NaturalArray) {
+                NaturalArray natural = (NaturalArray) current;
+                long idx = expressionHandler.toLongIndex(idxObj);
+                current = natural.get(idx);
+                continue;
+            }
+            if (current instanceof List) {
+                List<Object> list = (List<Object>) current;
+                int idx = expressionHandler.toIntIndex(idxObj);
+                if (idx < 0 || idx >= list.size()) {
+                    throw new ProgramError("Index out of bounds: " + idx + " for array of size " + list.size());
+                }
+                current = list.get(idx);
+                continue;
+            }
+            throw new ProgramError("Invalid multidimensional assignment path: expected NaturalArray or List, got "
+                + (current != null ? current.getClass().getSimpleName() : "null"));
+        }
+        
+        Object lastIdxObj = typeSystem.unwrap(tupleIndices.get(tupleIndices.size() - 1));
+        if (lastIdxObj instanceof RangeSpec) {
+            return assignRange(current, (RangeSpec) lastIdxObj, newValue);
+        }
+        if (lastIdxObj instanceof MultiRangeSpec) {
+            return assignMultiRange(current, (MultiRangeSpec) lastIdxObj, newValue);
+        }
+        if (current instanceof NaturalArray) {
+            NaturalArray natural = (NaturalArray) current;
+            long idx = expressionHandler.toLongIndex(lastIdxObj);
+            natural.set(idx, newValue);
+            return newValue;
+        }
+        if (current instanceof List) {
+            List<Object> list = (List<Object>) current;
+            int idx = expressionHandler.toIntIndex(lastIdxObj);
+            list.set(idx, newValue);
+            return newValue;
+        }
+        
+        throw new ProgramError("Invalid array assignment target in multidimensional assignment: " +
+            (current != null ? current.getClass().getSimpleName() : "null"));
+    }
+
+    @SuppressWarnings("unchecked")
+    private Object extractRange(Object array, RangeSpec range) {
+        if (array instanceof NaturalArray) {
+            return ((NaturalArray) array).getRange(range);
+        }
+        if (array instanceof List) {
+            List<Object> list = (List<Object>) array;
+            List<Object> result = new ArrayList<Object>();
+            long start = expressionHandler.toLongIndex(range.start);
+            long end = expressionHandler.toLongIndex(range.end);
+            long step = expressionHandler.calculateStep(range);
+            start = normalizeListIndex(start, list.size());
+            end = normalizeListIndex(end, list.size());
+            
+            if (start < 0 || start >= list.size()) {
+                throw new ProgramError("Range start index out of bounds: " + start + " for array of size " + list.size());
+            }
+            if (end < 0 || end >= list.size()) {
+                throw new ProgramError("Range end index out of bounds: " + end + " for array of size " + list.size());
+            }
+            if (step > 0) {
+                for (long i = start; i <= end && i < list.size(); i += step) {
+                    result.add(list.get((int) i));
+                }
+            } else if (step < 0) {
+                for (long i = start; i >= end && i >= 0; i += step) {
+                    result.add(list.get((int) i));
+                }
+            } else {
+                throw new ProgramError("Range step cannot be zero");
+            }
+            return result;
+        }
+        throw new ProgramError("Cannot apply range index during multidimensional assignment to " +
+            (array != null ? array.getClass().getSimpleName() : "null"));
+    }
+
+    @SuppressWarnings("unchecked")
+    private Object extractMultiRange(Object array, MultiRangeSpec multiRange) {
+        if (array instanceof NaturalArray) {
+            return ((NaturalArray) array).getMultiRange(multiRange);
+        }
+        if (array instanceof List) {
+            List<Object> list = (List<Object>) array;
+            List<Object> result = new ArrayList<Object>();
+            for (RangeSpec range : multiRange.ranges) {
+                Object sub = extractRange(list, range);
+                if (sub instanceof List) {
+                    result.addAll((List<Object>) sub);
+                }
+            }
+            return result;
+        }
+        throw new ProgramError("Cannot apply multi-range index during multidimensional assignment to " +
+            (array != null ? array.getClass().getSimpleName() : "null"));
+    }
+    
+    private long normalizeListIndex(long index, int size) {
+        if (index < 0) {
+            return size + index;
+        }
+        return index;
+    }
 
     private Object handleVariableAssignment(ExprNode target, Object newValue, ExecutionContext ctx) {
         try {